diff --git a/docs/fsdp_model_roles.md b/docs/fsdp_model_roles.md index 74db273c4..372ee4c10 100644 --- a/docs/fsdp_model_roles.md +++ b/docs/fsdp_model_roles.md @@ -36,10 +36,10 @@ The existing loss receives the reference prediction, and custom prepare/loss hoo access `ctx.reference_models` by component name. A third source, the teacher (`--teacher-load`, with `--teacher-lora-adapter-path` for its -adapters), is reserved for a DMD2-style loss that miles does not ship yet. Until then it -has no built-in consumer: `--teacher-load` requires a custom prepare/loss/train-output hook, -which reads it from `ctx.teacher_models`, and teacher inference should run inside -`torch.no_grad()`. +adapters), is loaded by `--loss-type dmd2`, without adapters, into its critic group, which +scores the student's samples with it. Otherwise it has no built-in consumer: `--teacher-load` +requires a custom prepare/loss/train-output hook, which reads it from `ctx.teacher_models`, and +teacher inference should run inside `torch.no_grad()`. Each source may load pretrained PEFT LoRA adapters, one plain PEFT directory (what `save_pretrained` writes) per `--update-weight-target-module` entry, in the same order: diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index 04eb8f82a..f089a0c61 100644 --- a/docs/user-guide/cli-reference.md +++ b/docs/user-guide/cli-reference.md @@ -100,7 +100,7 @@ See [Dtype Control](../advanced/dtype-control.md). | Flag | Default | What | |---|---|---| -| `--loss-type` | `policy_loss` | `policy_loss` (Flow-GRPO), `nft`, `sft_loss`, `custom_loss`. | +| `--loss-type` | `policy_loss` | `policy_loss` (Flow-GRPO), `nft`, `sft_loss`, `custom_loss`, `dmd2`. | | `--advantage-estimator` | `grpo` | Only `grpo` today. | | `--diffusion-clip-range` | `1e-4` | PPO ratio clip. | | `--diffusion-adv-clip-max` | `5.0` | Advantage clamp. | @@ -260,7 +260,7 @@ See [Dtype Control](../advanced/dtype-control.md). | Flag | Type | Default | Notes | |---|---|---|---| -| `--loss-type` | enum | `policy_loss` | `policy_loss` / `nft` / `sft_loss` / `custom_loss`. | +| `--loss-type` | enum | `policy_loss` | `policy_loss` / `nft` / `sft_loss` / `custom_loss` / `dmd2`. | | `--advantage-estimator` | enum | `grpo` | | | `--disable-grpo-std-normalization` | flag | off | Subtract the group mean only; skip dividing by the group std. Std normalization is also skipped when `n_samples_per_prompt == 1`, since a one-sample group has no std. With `--globalize-reward-mean` the batch mean is subtracted instead, and the batch-wide std from `--globalize-reward-std` is skipped too. | | `--globalize-reward-mean` | flag | off | Leave **off** for flow_grpo parity. | @@ -285,6 +285,12 @@ See [Dtype Control](../advanced/dtype-control.md). | `--diffusion-nft-timestep-strategy-path` | str | `…nft.drop_random_steps` | `(args, num_steps, generator) -> list[int]` returning the denoising steps each sample trains on, in training order; `generator` is seeded per sample. Built in: `drop_random_steps` (random order), `drop_final_steps` (schedule order, noisiest first). The final clean output (σ = 0) is never a training step. | | `--diffusion-nft-num-dropped-timesteps` | int | `0` | Steps the built-in strategies drop: random ones per sample for `drop_random_steps`, the final (lowest-noise) ones for `drop_final_steps`. | | `--no-diffusion-nft-adaptive-weight` | flag | off | | +| `--dmd2-critic-steps-per-rollout` | int | `5` | Critic optimizer steps per rollout under `dmd2`, before the actor's single step. Each DP shard must split into this many steps. The last step's samples are scored by the critic before that step's update. | +| `--dmd2-teacher-guidance-scale` | float | – | Required by `dmd2`: the teacher's CFG scale when scoring the student's samples; `1` runs it unguided. | +| `--dmd2-score-sigma-range` | float float | `0.02 0.98` | `MIN MAX` of the unshifted sigma at which `dmd2` scores the student's samples: the last critic step draws grid points in this range and reads them through the visual `--fsdp-flow-shift`. The default is the reference DMD2 range. | +| `--dmd2-critic-loss-space` | str | `velocity` | `velocity` or `x0`: the space of the critic's flow-matching loss. `x0` regresses x0 = x_σ − σ·v, which weights the velocity MSE by σ². | +| `--critic-load` | str | `--teacher-load` | `dmd2` only: HF pipeline checkpoint the critic starts from. Its checkpoints go to the sibling `<--save>_critic` dir. | +| `--critic-lr` | float | `--lr` | `dmd2` only: the critic's learning rate. | ### Reward diff --git a/docs/user-guide/customization.md b/docs/user-guide/customization.md index 47516d3aa..313c015b6 100644 --- a/docs/user-guide/customization.md +++ b/docs/user-guide/customization.md @@ -33,7 +33,7 @@ signatures they expect, and the defaults they replace — in the same spirit as | **Model family** | `--train-pipeline-config-path` | Family `TrainPipelineConfig` | | | `--model-backend-path` | Model load / FSDP backend | -`--loss-type` (`policy_loss` / `nft` / `sft_loss` / `custom_loss`) is the +`--loss-type` (`policy_loss` / `nft` / `sft_loss` / `custom_loss` / `dmd2`) is the shortcut that auto-fills several of the train-side paths; see § Training. *** @@ -234,11 +234,16 @@ def buffer_filter( | `nft` | `data_conversion_hub.nft.expand_samples_to_train_pairs` | `loss_hub.nft.prepare_nft_batch` | `nft_loss_formula` | | `sft_loss` | none; convert `sft_rollout.convert_samples_to_train_data` | `loss_hub.sft.prepare_sft_batch` | `sft_loss_formula` | | `custom_loss` | your expand (or convert) | your prepare | your formula | +| `dmd2` | `data_conversion_hub.dmd2.expand_samples_to_train_pairs` | your prepare | your formula | `sft_loss` auto-fills nothing: it requires the convert, rollout-log, prepare and loss paths, plus `--train-only` and a rollout function that runs no engines (e.g. `miles.rollout.sft_rollout.generate_rollout`; see `scripts/run_diffusion_sft_wan22.py`). `custom_loss` behaves like `policy_loss`: -any hook left unset falls back to the Flow-GRPO built-in. +any hook left unset falls back to the Flow-GRPO built-in. `dmd2` also trains a `critic` group +before the actor on each rollout: it starts from `--critic-load` (default `--teacher-load`), +learns the student's samples with the SFT loss at `--critic-lr` (default `--lr`). On its last step +it keeps its own x0 predictions from that step's forward, before the update, and after training hands +them with the teacher's x0 at the same noisy inputs to the actor as `ctx.external_data["critic"]`. DiT forward always stays in the FSDP actor; the loss hook only computes the objective. @@ -308,7 +313,9 @@ def train_output(ctx, train_pairs) -> Any: Runs without grad after this trainer group's train steps on a rollout. Every trainer group that trains later on the same rollout reads the return value of its paired rank (same rank, same -train pairs) as `ctx.external_data[]`. +train pairs) as `ctx.external_data[]`. Prepare and loss hooks see the current optimizer step +as `ctx.optim_step_idx` and can leave inputs for this hook in `ctx.train_output_inputs`, which +starts empty on each train call. ### `--sde-step-backend-path` diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index dcfbb7dd4..6fb4916b5 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -316,11 +316,15 @@ def update_weights(self, previous_ema: bool = False) -> None: # type: ignore[ov self.weight_updater.update_weights() clear_memory() + def _metric_role_tag(self) -> str: + # Upstream miles tags the train metrics of non-actor trainers, e.g. train/critic-loss. + return "" if self.role == "actor" else f"{self.role}-" + def _log_metrics(self, rollout_id: int, log_dict: dict[str, float], step: int) -> None: """Emit already-reduced metrics; every DP group computed the same values.""" if dist.get_rank() != 0: return - log_dict["train/lr"] = float(self.optimizer.param_groups[0]["lr"]) + log_dict[f"train/{self._metric_role_tag()}lr"] = float(self.optimizer.param_groups[0]["lr"]) log_dict["train/epoch"] = float(rollout_id) log_dict["rollout/step"] = compute_rollout_step(self.args, rollout_id) log_dict["train/step"] = float(step) @@ -443,6 +447,7 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An microbatch_id = 0 for optim_step_idx, microbatch_ranges in enumerate(microbatch_schedule): self.optimizer.zero_grad(set_to_none=True) + loss_ctx.optim_step_idx = optim_step_idx old_log_prob_from_new = self.args.diffusion_recompute_old_log_prob and optim_step_idx == 0 @@ -489,7 +494,7 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An self.optimizer.zero_grad(set_to_none=True) self.global_step += 1 - reduced = {f"train/{key}": value for key, value in metrics.reduce().items()} + reduced = {f"train/{self._metric_role_tag()}{key}": value for key, value in metrics.reduce().items()} self._log_metrics(rollout_id, reduced, step=self.global_step) if self.custom_train_output_func is not None: diff --git a/miles/backends/fsdp_utils/loss_hub/dmd2.py b/miles/backends/fsdp_utils/loss_hub/dmd2.py new file mode 100644 index 000000000..001690a78 --- /dev/null +++ b/miles/backends/fsdp_utils/loss_hub/dmd2.py @@ -0,0 +1,155 @@ +"""DMD2 critic-group hooks: the critic (fake score) learns the student's samples and scores those of its last step.""" + +from __future__ import annotations + +import dataclasses + +import torch + +from miles.backends.fsdp_utils.ema import reshard_model +from miles.backends.fsdp_utils.input_dtype_policy import apply_input_dtype_policy +from miles.backends.fsdp_utils.loss_hub.sft import prepare_sft_batch, sft_loss_formula +from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch +from miles.utils.metric_buffer import MetricBuffer + + +def prepare_critic_batch( + ctx: DiffusionLossContext, + batch: list[dict], + *, + pad_to_len: int | None = None, +) -> PreparedBatch: + """SFT flow matching on the student's x0 predictions, conditioned on their rollout prompts. + + The last step's samples are scored at their noisy inputs, so its grid points come from the scoring range. + """ + config = ctx.train_pipeline_config + sft_batch = [ + { + "latent": {"visual": pair["x0"]}, + "cond_kwargs": config.prepare_cond_kwargs(pair["denoising_env"].pos_cond_kwargs, ctx.device), + } + for pair in batch + ] + grid_index_range = None + if _is_scoring_step(ctx): + min_sigma, max_sigma = ctx.args.dmd2_score_sigma_range + num_train_timesteps = config.num_train_timesteps + # Grid index i carries the unshifted sigma 1 - i / num_train_timesteps. + grid_index_range = (round(num_train_timesteps * (1 - max_sigma)), round(num_train_timesteps * (1 - min_sigma))) + return prepare_sft_batch(ctx, sft_batch, pad_to_len=pad_to_len, grid_index_range=grid_index_range) + + +def critic_loss_formula( + ctx: DiffusionLossContext, + batch: list[dict], + prepared: PreparedBatch, + *, + new_pred: dict[str, torch.Tensor], + old_pred: dict[str, torch.Tensor] | None, + ref_pred: dict[str, torch.Tensor] | None, + metrics: MetricBuffer, + write_old_log_prob: bool = False, + old_log_prob_from_new: bool = False, +) -> torch.Tensor: + """The SFT loss in --dmd2-critic-loss-space. + + On the last step it also keeps the critic's x0 at each noisy input, before that step's update. + """ + latents = prepared.latents["visual"] + sigmas = prepared.extras["sigmas"]["visual"].view(len(batch), *([1] * (latents.ndim - 1))) + if _is_scoring_step(ctx): + x0_fake = latents - sigmas * new_pred["visual"].detach().float() + ctx.train_output_inputs.append( + (batch, prepared.component_name, latents, prepared.timesteps["visual"], sigmas, x0_fake) + ) + if ctx.args.dmd2_critic_loss_space == "x0": + # x0 = x_sigma - sigma v, so the MSE on x0 weights the velocity MSE by sigma^2. + x0_target = latents - sigmas * prepared.extras["target"]["visual"] + prepared = dataclasses.replace(prepared, extras={**prepared.extras, "target": {"visual": x0_target}}) + new_pred = {"visual": latents - sigmas * new_pred["visual"].float()} + return sft_loss_formula( + ctx, + batch, + prepared, + new_pred=new_pred, + old_pred=old_pred, + ref_pred=ref_pred, + metrics=metrics, + write_old_log_prob=write_old_log_prob, + old_log_prob_from_new=old_log_prob_from_new, + ) + + +def score_last_step_samples(ctx: DiffusionLossContext, train_pairs: list[dict]) -> dict[int, dict]: + """The teacher's x0 at each scored sample's noisy input, next to the critic's x0 kept there, by sample.""" + predictions = {} + for batch, component_name, latents, timesteps, sigmas, x0_fake in ctx.train_output_inputs: + teacher = ctx.teacher_models[component_name] + real_pred = _predict_velocity(ctx, teacher, batch, latents, timesteps, ctx.args.dmd2_teacher_guidance_scale) + for pair, x0_real, fake in zip(batch, latents - sigmas * real_pred, x0_fake, strict=True): + predictions[pair["sample_index"]] = {"x0_real": x0_real.cpu(), "x0_fake": fake.cpu()} + + # A no-grad forward leaves FSDP's root parameters gathered; sleep must move the shards. + for model in ctx.teacher_models.values(): + reshard_model(model) + return predictions + + +def _is_scoring_step(ctx: DiffusionLossContext) -> bool: + return ctx.optim_step_idx == ctx.args.num_steps_per_rollout - 1 + + +def _predict_velocity( + ctx: DiffusionLossContext, + model: torch.nn.Module, + batch: list[dict], + latents: torch.Tensor, + timesteps: torch.Tensor, + guidance_scale: float, +) -> torch.Tensor: + """One velocity prediction, with CFG over the rollout's negative cond when ``guidance_scale > 1``.""" + config = ctx.train_pipeline_config + use_cfg = guidance_scale > 1.0 + envs = [pair["denoising_env"] for pair in batch] + if use_cfg and any(env.neg_cond_kwargs is None for env in envs): + raise ValueError( + "--dmd2-teacher-guidance-scale > 1 needs the rollout's negative cond, " + "which the engine returns only for requests that run CFG" + ) + pos_list = [config.prepare_cond_kwargs(env.pos_cond_kwargs, ctx.device) for env in envs] + neg_list = [config.prepare_cond_kwargs(env.neg_cond_kwargs, ctx.device) for env in envs] if use_cfg else [] + cfg_batching = use_cfg and config.cfg_batching + pos_cond = neg_cond = joint_cond = None + if cfg_batching: + joint_cond = config.collate_cond_for_sample_batch(pos_list + neg_list, ctx.device) + else: + pos_cond = config.collate_cond_for_sample_batch(pos_list, ctx.device) + if use_cfg: + neg_cond = config.collate_cond_for_sample_batch(neg_list, ctx.device) + + latents_in, timesteps_in, (pos_in, neg_in, joint_in) = apply_input_dtype_policy( + config.input_dtype_policy, + latents={"visual": latents}, + timesteps={"visual": config.process_timestep_as_input(timesteps)}, + conds=(pos_cond, neg_cond, joint_cond), + default_dtype=ctx.forward_dtype, + ) + with torch.autocast( + "cuda", + dtype=ctx.forward_dtype, + enabled=ctx.forward_dtype != torch.float32 and config.enable_autocast, + ): + pred = config.compute_noise_pred( + model=model, + latents_input=latents_in, + timesteps_input=timesteps_in, + pos_cond=pos_in, + neg_cond=neg_in, + joint_cond=joint_in, + use_cfg=use_cfg, + cfg_batching=cfg_batching, + guidance_scale=guidance_scale, + true_cfg_scale=None, + ) + return pred["visual"].float() diff --git a/miles/backends/fsdp_utils/loss_hub/sft.py b/miles/backends/fsdp_utils/loss_hub/sft.py index 5c120c228..b17acfae7 100644 --- a/miles/backends/fsdp_utils/loss_hub/sft.py +++ b/miles/backends/fsdp_utils/loss_hub/sft.py @@ -39,11 +39,13 @@ def prepare_sft_batch( batch: list[dict], *, pad_to_len: int | None = None, + grid_index_range: tuple[int, int] | None = None, ) -> PreparedBatch: """Corrupt every latent stream at one shared grid index, each through its own shifted grid; CFG-free cached cond. - The visual grid picks the DiT component and the index is drawn uniformly from the points it serves. Sharing the - index pairs the streams' sigmas as inference steps them, one step index read through each stream's shift. + The visual grid picks the DiT component and the index is drawn uniformly from the points it serves, within + ``grid_index_range`` (inclusive) when given. Sharing the index pairs the streams' sigmas as inference steps them, + one step index read through each stream's shift. """ device = ctx.device config = ctx.train_pipeline_config @@ -57,6 +59,8 @@ def prepare_sft_batch( for stream_name in stream_names } component_name, model, visual_pool = select_component(ctx, noise_schedules["visual"], device) + if grid_index_range is not None: + visual_pool = visual_pool[(visual_pool >= grid_index_range[0]) & (visual_pool <= grid_index_range[1])] # The visual stream keeps the single-stream seed, so single-stream runs draw as before; other streams add their # name, so adding a stream never changes another stream's noise. diff --git a/miles/backends/fsdp_utils/loss_hub/types.py b/miles/backends/fsdp_utils/loss_hub/types.py index 2d41b3537..fec342ea8 100644 --- a/miles/backends/fsdp_utils/loss_hub/types.py +++ b/miles/backends/fsdp_utils/loss_hub/types.py @@ -22,10 +22,13 @@ class DiffusionLossContext: device: torch.device rollout_id: int = 0 microbatch_id: int = 0 + optim_step_idx: int = 0 dp_rank: int = 0 reference_models: dict[str, torch.nn.Module] = field(default_factory=dict) teacher_models: dict[str, torch.nn.Module] = field(default_factory=dict) external_data: dict[str, Any] = field(default_factory=dict) + # Filled by loss hooks during one train call and read by the train-output hook. + train_output_inputs: list[Any] = field(default_factory=list) @dataclass(frozen=True) diff --git a/miles/ray/data_conversion_hub/dmd2.py b/miles/ray/data_conversion_hub/dmd2.py new file mode 100644 index 000000000..cc0b1cee3 --- /dev/null +++ b/miles/ray/data_conversion_hub/dmd2.py @@ -0,0 +1,43 @@ +"""Convert DMD2 rollout samples, each stopped at its step i, into one train pair per sample.""" + +from argparse import Namespace +from typing import Any + +from miles.utils.types import Sample + + +def expand_samples_to_train_pairs( + args: Namespace, + samples: list[Sample], + rewards: list[float], + raw_rewards: list[float], +) -> dict[str, Any]: + """One pair per sample: the student's input ``x_i`` at its stop step ``i`` and its x0 prediction there.""" + train_data: list[dict[str, Any]] = [] + for sample in samples: + traj = sample.dit_trajectory + if traj is None or traj.latent_step_indices is None or sample.denoising_env is None: + raise ValueError( + f"sample {sample.index} missing dit_trajectory.latent_step_indices or denoising_env; " + "DMD2 rollouts need --diffusion-step-strategy-path miles.rollout.step_strategy_hub.stop_at_random_step" + ) + step_index = int(traj.latent_step_indices[0]) + # The engine sends step i straight to sigma 0, so the grid ends right after sigma_i. + if traj.latent_step_indices.tolist() != [step_index, step_index + 1] or len(traj.sigmas) != step_index + 2: + raise ValueError( + f"sample {sample.index} returned steps {traj.latent_step_indices.tolist()} on a " + f"{len(traj.sigmas)}-point sigma grid; DMD2 needs a rollout stopped at its first returned step" + ) + train_data.append( + { + "latent": traj.latents[0], + "x0": traj.latents[1], + "timestep": float(traj.timesteps[step_index]), + "sigma": float(traj.sigmas[step_index]), + "step_index": step_index, + "denoising_env": sample.denoising_env, + "sample_index": sample.index, + "prompt": sample.prompt, + } + ) + return {"train_data": train_data} diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 133596006..e1b5999e9 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -943,11 +943,12 @@ def add_algo_arguments(parser): parser.add_argument( "--loss-type", type=str, - choices=["policy_loss", "sft_loss", "custom_loss", "nft"], + choices=["policy_loss", "sft_loss", "custom_loss", "nft", "dmd2"], default="policy_loss", help=( "Train objective shortcut. Diffusion default policy_loss = Flow-GRPO prepare+formula. " "nft auto-fills prepare / loss-formula paths and selects the NFT train-data converter. " + "dmd2 selects the DMD2 train-data converter and trains a critic group next to the actor. " "DiT forward always stays in the FSDP actor." ), ) @@ -1019,6 +1020,47 @@ def add_algo_arguments(parser): help="Sampling steps the built-in NFT timestep strategies drop: random ones per sample for " "drop_random_steps, the final (lowest-noise) ones for drop_final_steps.", ) + parser.add_argument( + "--dmd2-critic-steps-per-rollout", + type=int, + default=5, + help="DMD2 critic optimizer steps per rollout, taken before the actor's single step (TTUR).", + ) + parser.add_argument( + "--dmd2-teacher-guidance-scale", + type=float, + default=None, + help="CFG scale of the frozen teacher when DMD2 scores the student's samples; 1 runs it unguided.", + ) + parser.add_argument( + "--dmd2-score-sigma-range", + type=float, + nargs=2, + metavar=("MIN", "MAX"), + default=[0.02, 0.98], + help="Unshifted sigma range in which DMD2 scores the student's samples: the grid points between MIN " + "and MAX, read through the visual --fsdp-flow-shift. Defaults to the reference DMD2's range.", + ) + parser.add_argument( + "--dmd2-critic-loss-space", + type=str, + choices=["velocity", "x0"], + default="velocity", + help="Space of the DMD2 critic's flow-matching loss: velocity, or x0 = x_sigma - sigma * v, " + "which weights the velocity MSE by sigma^2.", + ) + parser.add_argument( + "--critic-load", + type=str, + default=None, + help="HF pipeline checkpoint the DMD2 critic starts from. Defaults to --teacher-load.", + ) + parser.add_argument( + "--critic-lr", + type=float, + default=None, + help="Learning rate of the DMD2 critic. Defaults to --lr.", + ) parser.add_argument( "--advantage-estimator", type=str, @@ -1774,6 +1816,16 @@ def set_default_diffusion_args(args) -> None: if args.custom_loss_function_path is None: args.custom_loss_function_path = "miles.backends.fsdp_utils.loss_hub.nft.nft_loss_formula" + if args.loss_type == "dmd2": + if args.custom_expand_samples_to_train_pairs_path is None: + args.custom_expand_samples_to_train_pairs_path = ( + "miles.ray.data_conversion_hub.dmd2.expand_samples_to_train_pairs" + ) + if args.critic_load is None: + args.critic_load = args.teacher_load + if args.critic_lr is None: + args.critic_lr = args.lr + if args.ref_mode is None: args.ref_mode = "lora_base" if args.diffusion_kl_beta > 0 else "none" @@ -1940,6 +1992,57 @@ def validate_train_sde_backend(args) -> None: ) +def validate_dmd2_args(args) -> None: + if args.loss_type != "dmd2": + if args.critic_load is not None or args.critic_lr is not None: + raise ValueError("--critic-load and --critic-lr configure the critic group of --loss-type dmd2") + return + if args.teacher_load is None: + raise ValueError("--loss-type dmd2 requires --teacher-load to score the student's samples") + if args.dmd2_teacher_guidance_scale is None or args.dmd2_teacher_guidance_scale < 1: + raise ValueError( + f"--loss-type dmd2 requires --dmd2-teacher-guidance-scale >= 1, got {args.dmd2_teacher_guidance_scale}" + ) + if args.fsdp_flow_shift is None or "visual" not in args.fsdp_flow_shift: + raise ValueError("--loss-type dmd2 requires a visual --fsdp-flow-shift for the critic's sigma grid") + min_sigma, max_sigma = args.dmd2_score_sigma_range + if not 0 <= min_sigma < max_sigma <= 1: + raise ValueError(f"--dmd2-score-sigma-range needs 0 <= MIN < MAX <= 1, got {min_sigma} {max_sigma}") + if args.diffusion_guidance_scale != 1.0 or args.diffusion_true_cfg_scale is not None: + raise ValueError( + "--loss-type dmd2 samples the student unguided; set --diffusion-guidance-scale 1.0 " + "and drop --diffusion-true-cfg-scale" + ) + if args.use_lora or args.teacher_lora_adapter_paths is not None: + raise ValueError( + "--loss-type dmd2 trains the full student and critic; drop --use-lora and --teacher-lora-adapter-path" + ) + if args.ref_mode != "none": + raise ValueError("--loss-type dmd2 uses no reference model; drop --ref-mode and --diffusion-kl-beta") + if args.custom_prepare_train_batch_path is None or args.custom_loss_function_path is None: + raise ValueError( + "--loss-type dmd2 has no built-in generator loss yet; pass --custom-prepare-train-batch-path " + "and --custom-loss-function-path" + ) + if args.num_steps_per_rollout != 1: + raise ValueError( + f"--loss-type dmd2 takes one actor step per rollout, got --num-steps-per-rollout " + f"{args.num_steps_per_rollout}; the critic scores each rollout once, for that step" + ) + dp_size = args.actor_num_gpus_per_node * args.actor_num_nodes // args.sequence_parallel_size + critic_steps = args.dmd2_critic_steps_per_rollout + if ( + critic_steps < 1 + or args.global_batch_size % (critic_steps * dp_size) != 0 + or (args.global_batch_size // (critic_steps * dp_size)) % args.micro_batch_size != 0 + ): + raise ValueError( + f"--dmd2-critic-steps-per-rollout {critic_steps} must be >= 1 and split each of the {dp_size} DP " + f"shards of the {args.global_batch_size} samples per rollout into whole --micro-batch-size " + f"{args.micro_batch_size} micro-batches" + ) + + def miles_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args) @@ -2277,6 +2380,8 @@ def miles_validate_args(args): samples_per_rollout = args.rollout_batch_size * args.n_samples_per_prompt args.num_steps_per_rollout = max(1, samples_per_rollout // args.global_batch_size) + validate_dmd2_args(args) + if args.n_samples_per_prompt == 1: args.grpo_std_normalization = False logger.info("n_samples_per_prompt is set to 1, grpo_std_normalization will be set to False.") diff --git a/miles/utils/trainer_config.py b/miles/utils/trainer_config.py index d0f19408a..54bce42be 100644 --- a/miles/utils/trainer_config.py +++ b/miles/utils/trainer_config.py @@ -11,7 +11,27 @@ class TrainerConfig: def compute_trainer_configs(args: Namespace) -> list[TrainerConfig]: - return [TrainerConfig(role="actor")] + if args.loss_type != "dmd2": + return [TrainerConfig(role="actor")] + critic_steps = args.dmd2_critic_steps_per_rollout + critic_overrides = { + "hf_checkpoint": args.critic_load, + "lr": args.critic_lr, + # A sibling dir, as upstream's --critic-save: sharing one would clobber the actor's iteration tracker. + "load": None if args.load is None else args.load.rstrip("/") + "_critic", + "save": None if args.save is None else args.save.rstrip("/") + "_critic", + "num_steps_per_rollout": critic_steps, + "global_batch_size": args.global_batch_size // critic_steps, + # Only the actor's EMA has a reader (its weight sync); the critic's would be built, stepped and saved for nothing. + "use_ema": False, + "custom_prepare_train_batch_path": "miles.backends.fsdp_utils.loss_hub.dmd2.prepare_critic_batch", + "custom_loss_function_path": "miles.backends.fsdp_utils.loss_hub.dmd2.critic_loss_formula", + "custom_train_output_function_path": "miles.backends.fsdp_utils.loss_hub.dmd2.score_last_step_samples", + } + return [ + TrainerConfig(role="critic", overrides=critic_overrides), + TrainerConfig(role="actor", overrides={"teacher_load": None}), + ] def compute_trainer_args(args: Namespace, trainer: TrainerConfig) -> Namespace: diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py b/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py new file mode 100644 index 000000000..32228c3a2 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py @@ -0,0 +1,310 @@ +"""DMD2 critic-group hooks: the critic (fake score) and the frozen teacher share this group; the actor holds the student. + + rollout pair: x0 = the student's x0 prediction at its stop step i, denoising_env = the rollout cond + │ + prepare_critic_batch ──► prepare_sft_batch on {"latent": {"visual": x0}, "cond_kwargs": prepared pos cond} + (the critic's K steps are the SFT loss on the student's samples; the last step, + whose samples get scored, draws grid points whose unshifted sigma lies in + --dmd2-score-sigma-range, default [0.02, 0.98]) + │ + critic_loss_formula ──► sft_loss_formula on v, or on x0 = x_s - sigma v under --dmd2-critic-loss-space x0; + on the last step it also keeps, per micro-batch, + x_s = that step's noisy input, sigma, and x0_fake = x_s - sigma v_fake from the + training forward, i.e. the critic before the last update + │ + score_last_step_samples, without grad after the critic's steps, per kept micro-batch: + teacher, CFG --dmd2-teacher-guidance-scale over the rollout's neg cond ──► x0_real = x_s - sigma v_real + ──► {sample_index: {"x0_real", "x0_fake"}}; the teacher is resharded so sleep moves its shards + +What each test pins: + test_critic_batch_is_the_sft_batch_on_the_student_x0 earlier steps: inputs and targets == SFT on (x0, pos cond) + test_last_step_draws_from_the_scoring_range last step: the same SFT draw over grid points [20, 980] + test_score_sigma_range_bounds_the_unshifted_sigma a custom range bounds the unshifted sigma, not 1 - sigma + test_critic_loss_is_the_sft_loss loss and metrics == sft_loss_formula on every step + test_x0_critic_loss_weights_the_velocity_mse x0 loss == x0 MSE == sigma^2 * velocity MSE, input intact + test_only_the_last_step_keeps_scoring_inputs earlier steps keep nothing, so only the last step is scored + test_scores_match_separate_forwards_at_the_kept_input fused scores == critic and teacher forwards run apart at x_s + test_teacher_runs_cfg_and_the_critic_runs_unguided CFG combine for the teacher only + test_teacher_cfg_needs_the_rollout_negative_cond rejected instead of a forward without text + test_scoring_reshards_the_teacher teacher resharded after its no-grad forwards +""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[]) + +from argparse import Namespace + +import pytest +import torch +import torch.nn as nn + +import miles.backends.fsdp_utils.loss_hub.dmd2 as dmd2 +from miles.backends.fsdp_utils.configs.sd3 import SD3TrainPipelineConfig +from miles.backends.fsdp_utils.loss_hub.sft import prepare_sft_batch, sft_loss_formula +from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, FlowNoiseSchedule +from miles.utils.types import CondKwargs, DenoisingEnv + +FLOW_SHIFT = 3.0 +LATENT_SHAPE = (2, 4, 4) +CRITIC_STEPS = 3 + + +class _CondVelocity(nn.Module): + """v = scale * mean(text embedding) per sample, so a prediction reveals which cond it saw.""" + + def __init__(self, scale): + super().__init__() + self.scale = scale + + def forward(self, hidden_states, timestep, return_dict, encoder_hidden_states, pooled_projections): + text = encoder_hidden_states.mean(dim=(1, 2)).view(-1, 1, 1, 1) + return (self.scale * text.expand_as(hidden_states),) + + +class _NoisyInputVelocity(nn.Module): + """v = x_s + timestep, so x0 = x_s - sigma v depends on both inputs.""" + + def forward(self, hidden_states, timestep, return_dict, **cond): + return (hidden_states + timestep.view(-1, 1, 1, 1),) + + +class _Metrics: + def __init__(self): + self.seen = {} + + def emit_mean(self, key, *, total, count): + prev_total, prev_count = self.seen.get(key, (0.0, 0)) + self.seen[key] = (prev_total + float(total), prev_count + count) + + +def _cond(value): + return CondKwargs(encoder_hidden_states=[torch.full((1, 3, 8), value)], pooled_projections=[torch.zeros(1, 8)]) + + +def _pairs(num_pairs, with_neg=True): + return [ + { + "x0": torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(n)), + "denoising_env": DenoisingEnv( + pos_cond_kwargs=_cond(1.0 + n), neg_cond_kwargs=_cond(-1.0) if with_neg else None + ), + "sample_index": 100 + n, + } + for n in range(num_pairs) + ] + + +def _ctx(teacher_guidance_scale=3.0, critic_scale=1.0, teacher_scale=1.0): + return DiffusionLossContext( + models={"transformer": _CondVelocity(critic_scale)}, + teacher_models={"transformer": _CondVelocity(teacher_scale)}, + train_pipeline_config=SD3TrainPipelineConfig(), + sde_backend=None, + args=Namespace( + seed=7, + micro_batch_size=2, + num_steps_per_rollout=CRITIC_STEPS, + fsdp_flow_shift={"visual": FLOW_SHIFT}, + fsdp_supervised_streams=["visual"], + log_loss_sigma_bucket=0, + dmd2_teacher_guidance_scale=teacher_guidance_scale, + dmd2_score_sigma_range=[0.02, 0.98], + dmd2_critic_loss_space="velocity", + ), + forward_dtype=torch.float32, + device=torch.device("cpu"), + rollout_id=4, + dp_rank=1, + ) + + +def _critic_step(ctx, pairs, optim_step_idx, metrics=None): + """One critic optimizer step as _train_core runs it: prepare, the critic's training forward, the loss hook.""" + ctx.optim_step_idx = optim_step_idx + steps = [] + for start in range(0, len(pairs), ctx.args.micro_batch_size): + batch = pairs[start : start + ctx.args.micro_batch_size] + ctx.microbatch_id = optim_step_idx * 10 + start + prepared = dmd2.prepare_critic_batch(ctx, batch) + new_pred = ctx.train_pipeline_config.compute_noise_pred( + model=prepared.model, + latents_input=prepared.latents, + timesteps_input=prepared.timesteps_for_model, + pos_cond=prepared.pos_cond, + neg_cond=None, + joint_cond=None, + use_cfg=prepared.use_cfg, + cfg_batching=prepared.cfg_batching, + guidance_scale=prepared.guidance_scale, + true_cfg_scale=None, + ) + loss = dmd2.critic_loss_formula( + ctx, batch, prepared, new_pred=new_pred, old_pred=None, ref_pred=None, metrics=metrics or _Metrics() + ) + steps.append((batch, prepared, new_pred, loss)) + return steps + + +def _sft_pairs(ctx, pairs): + config = ctx.train_pipeline_config + return [ + { + "latent": {"visual": pair["x0"]}, + "cond_kwargs": config.prepare_cond_kwargs(pair["denoising_env"].pos_cond_kwargs, ctx.device), + } + for pair in pairs + ] + + +def test_critic_batch_is_the_sft_batch_on_the_student_x0(): + ctx = _ctx() + pairs = _pairs(3) + + for optim_step_idx in range(CRITIC_STEPS - 1): + ctx.optim_step_idx = optim_step_idx + critic_batch = dmd2.prepare_critic_batch(ctx, pairs) + sft_batch = prepare_sft_batch(ctx, _sft_pairs(ctx, pairs)) + + assert critic_batch.model is ctx.models["transformer"] + for field in ("latents", "timesteps", "timesteps_for_model"): + assert torch.equal(getattr(critic_batch, field)["visual"], getattr(sft_batch, field)["visual"]) + assert torch.equal(critic_batch.extras["target"]["visual"], sft_batch.extras["target"]["visual"]) + assert torch.equal(critic_batch.pos_cond["encoder_hidden_states"], sft_batch.pos_cond["encoder_hidden_states"]) + assert not critic_batch.use_cfg + + +def test_last_step_draws_from_the_scoring_range(): + ctx = _ctx() + pairs = _pairs(256) + ctx.optim_step_idx = CRITIC_STEPS - 1 + + critic_batch = dmd2.prepare_critic_batch(ctx, pairs) + sft_batch = prepare_sft_batch(ctx, _sft_pairs(ctx, pairs), grid_index_range=(20, 980)) + + assert torch.equal(critic_batch.latents["visual"], sft_batch.latents["visual"]) + assert torch.equal(critic_batch.extras["target"]["visual"], sft_batch.extras["target"]["visual"]) + grid = FlowNoiseSchedule.shifted(FLOW_SHIFT, 1000).sigmas + sigmas = critic_batch.extras["sigmas"]["visual"] + assert set(sigmas.tolist()) <= set(grid[20:981].tolist()) + # The same draw without the range leaves it, so the restriction is exercised. + assert not set(prepare_sft_batch(ctx, _sft_pairs(ctx, pairs)).extras["sigmas"]["visual"].tolist()) <= set( + grid[20:981].tolist() + ) + + +def test_score_sigma_range_bounds_the_unshifted_sigma(): + ctx = _ctx() + ctx.args.dmd2_score_sigma_range = [0.1, 0.5] + pairs = _pairs(256) + ctx.optim_step_idx = CRITIC_STEPS - 1 + + critic_batch = dmd2.prepare_critic_batch(ctx, pairs) + # Grid index i carries the unshifted sigma 1 - i / 1000, so [0.1, 0.5] is grid points [500, 900]. + sft_batch = prepare_sft_batch(ctx, _sft_pairs(ctx, pairs), grid_index_range=(500, 900)) + + sigmas = critic_batch.extras["sigmas"]["visual"] + assert torch.equal(sigmas, sft_batch.extras["sigmas"]["visual"]) + unshifted = sigmas / (FLOW_SHIFT - (FLOW_SHIFT - 1) * sigmas) + assert unshifted.min() >= 0.1 - 1e-6 and unshifted.max() <= 0.5 + 1e-6 + + +def test_critic_loss_is_the_sft_loss(): + for optim_step_idx in range(CRITIC_STEPS): + ctx = _ctx() + metrics = _Metrics() + sft_metrics = _Metrics() + for batch, prepared, new_pred, loss in _critic_step(ctx, _pairs(3), optim_step_idx, metrics): + sft_loss = sft_loss_formula( + ctx, batch, prepared, new_pred=new_pred, old_pred=None, ref_pred=None, metrics=sft_metrics + ) + assert torch.equal(loss, sft_loss) + assert metrics.seen == sft_metrics.seen + + +def test_x0_critic_loss_weights_the_velocity_mse(): + for optim_step_idx in range(CRITIC_STEPS): + ctx = _ctx() + ctx.models["transformer"] = _NoisyInputVelocity() + ctx.args.dmd2_critic_loss_space = "x0" + for batch, prepared, new_pred, loss in _critic_step(ctx, _pairs(3), optim_step_idx): + x_s = prepared.latents["visual"] + sigma = prepared.extras["sigmas"]["visual"].view(-1, 1, 1, 1) + x0 = torch.stack([pair["x0"] for pair in batch]) + x0_mse = (x_s - sigma * new_pred["visual"] - x0).square().flatten(1).mean(dim=1) + velocity_mse = (new_pred["visual"] - prepared.extras["target"]["visual"]).square().flatten(1).mean(dim=1) + torch.testing.assert_close(loss, x0_mse.sum()) + torch.testing.assert_close(loss, (sigma.flatten() ** 2 * velocity_mse).sum()) + + +def test_only_the_last_step_keeps_scoring_inputs(): + ctx = _ctx() + pairs = _pairs(9) + + for optim_step_idx in range(CRITIC_STEPS - 1): + _critic_step(ctx, pairs[optim_step_idx * 3 : optim_step_idx * 3 + 3], optim_step_idx) + assert ctx.train_output_inputs == [] + + _critic_step(ctx, pairs[6:], CRITIC_STEPS - 1) + scores = dmd2.score_last_step_samples(ctx, pairs) + + assert list(scores) == [106, 107, 108] + + +def test_scores_match_separate_forwards_at_the_kept_input(): + ctx = _ctx(teacher_guidance_scale=3.0) + ctx.models["transformer"] = _NoisyInputVelocity() + pairs = _pairs(5) + + steps = _critic_step(ctx, pairs, CRITIC_STEPS - 1) + scores = dmd2.score_last_step_samples(ctx, pairs) + + assert list(scores) == [100, 101, 102, 103, 104] + for batch, prepared, _, _ in steps: + x_s = prepared.latents["visual"] + timesteps = prepared.timesteps["visual"] + sigma = prepared.extras["sigmas"]["visual"].view(-1, 1, 1, 1) + fake = dmd2._predict_velocity(ctx, ctx.models["transformer"], batch, x_s, timesteps, 1.0) + real = dmd2._predict_velocity(ctx, ctx.teacher_models["transformer"], batch, x_s, timesteps, 3.0) + for n, pair in enumerate(batch): + assert torch.equal(scores[pair["sample_index"]]["x0_fake"], x_s[n] - sigma[n] * fake[n]) + assert torch.equal(scores[pair["sample_index"]]["x0_real"], x_s[n] - sigma[n] * real[n]) + + +def test_teacher_runs_cfg_and_the_critic_runs_unguided(): + ctx = _ctx(teacher_guidance_scale=3.0, teacher_scale=1.0, critic_scale=2.0) + pairs = _pairs(2) + + [(_, prepared, _, _)] = _critic_step(ctx, pairs, CRITIC_STEPS - 1) + scores = dmd2.score_last_step_samples(ctx, pairs) + + x_s = prepared.latents["visual"] + sigma = prepared.extras["sigmas"]["visual"].view(-1, 1, 1, 1) + for n, pair in enumerate(pairs): + pos, neg = 1.0 + n, -1.0 + torch.testing.assert_close( + scores[pair["sample_index"]]["x0_real"], x_s[n] - sigma[n] * (neg + 3.0 * (pos - neg)) + ) + torch.testing.assert_close(scores[pair["sample_index"]]["x0_fake"], x_s[n] - sigma[n] * 2.0 * pos) + + +def test_teacher_cfg_needs_the_rollout_negative_cond(): + ctx = _ctx(teacher_guidance_scale=3.0) + _critic_step(ctx, _pairs(2, with_neg=False), CRITIC_STEPS - 1) + with pytest.raises(ValueError, match="needs the rollout's negative cond"): + dmd2.score_last_step_samples(ctx, _pairs(2, with_neg=False)) + # Unguided, the teacher reads only the positive cond. + ctx = _ctx(teacher_guidance_scale=1.0) + _critic_step(ctx, _pairs(2, with_neg=False), CRITIC_STEPS - 1) + dmd2.score_last_step_samples(ctx, _pairs(2, with_neg=False)) + + +def test_scoring_reshards_the_teacher(monkeypatch): + resharded = [] + monkeypatch.setattr(dmd2, "reshard_model", resharded.append) + ctx = _ctx() + _critic_step(ctx, _pairs(2), CRITIC_STEPS - 1) + + dmd2.score_last_step_samples(ctx, _pairs(2)) + + assert resharded == [ctx.teacher_models["transformer"]] diff --git a/tests/fast/rollout/test_dmd2_convert.py b/tests/fast/rollout/test_dmd2_convert.py new file mode 100644 index 000000000..3d953a144 --- /dev/null +++ b/tests/fast/rollout/test_dmd2_convert.py @@ -0,0 +1,68 @@ +"""DMD2 rollout samples -> train pairs. + + stop_at_random_step request (stop = i, return [i, i+1]) ──► engine runs steps 0..i, step i goes to sigma 0 + dit_trajectory.latents [x_i, x0_i] (x0_i = the student's x0 prediction at step i) + dit_trajectory.latent_step_indices [i, i+1] + dit_trajectory.sigmas [sigma_0 .. sigma_i, 0] + │ + expand_samples_to_train_pairs ──► one pair per sample: + latent = x_i, x0 = x0_i, timestep = timesteps[i], sigma = sigma_i, step_index = i, + denoising_env, sample_index, prompt; rewards are ignored + A trajectory that did not stop at its first returned step is rejected. +""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) + +import pytest +import torch + +from miles.ray.data_conversion_hub.dmd2 import expand_samples_to_train_pairs +from miles.utils.types import DenoisingEnv, DiTTrajectory, Sample + +FULL_SIGMAS = [1.0, 0.75, 0.5, 0.25, 0.0] +FULL_TIMESTEPS = [1000.0, 750.0, 500.0, 250.0] + + +def _stopped_sample(index, step): + return Sample( + index=index, + prompt=f"prompt {index}", + denoising_env=DenoisingEnv(), + dit_trajectory=DiTTrajectory( + latents=torch.stack([torch.full((2, 3), float(index)), torch.full((2, 3), -float(index))]), + timesteps=torch.tensor(FULL_TIMESTEPS[: step + 1]), + sigmas=torch.tensor(FULL_SIGMAS[: step + 1] + [0.0]), + latent_step_indices=torch.tensor([step, step + 1]), + ), + ) + + +def test_each_stopped_sample_becomes_one_pair(): + samples = [_stopped_sample(10, 0), _stopped_sample(11, 2)] + + pairs = expand_samples_to_train_pairs(None, samples, [0.0, 0.0], [0.0, 0.0])["train_data"] + + assert [pair["sample_index"] for pair in pairs] == [10, 11] + assert [pair["step_index"] for pair in pairs] == [0, 2] + assert [pair["sigma"] for pair in pairs] == [1.0, 0.5] + assert [pair["timestep"] for pair in pairs] == [1000.0, 500.0] + assert torch.equal(pairs[1]["latent"], torch.full((2, 3), 11.0)) + assert torch.equal(pairs[1]["x0"], torch.full((2, 3), -11.0)) + assert pairs[1]["denoising_env"] is samples[1].denoising_env + assert pairs[1]["prompt"] == "prompt 11" + + +def test_full_trajectory_is_rejected(): + sample = _stopped_sample(10, 1) + sample.dit_trajectory.latent_step_indices = None + with pytest.raises(ValueError, match="stop_at_random_step"): + expand_samples_to_train_pairs(None, [sample], [0.0], [0.0]) + + +def test_trajectory_that_runs_past_its_returned_step_is_rejected(): + sample = _stopped_sample(10, 1) + sample.dit_trajectory.sigmas = torch.tensor(FULL_SIGMAS) + with pytest.raises(ValueError, match="stopped at its first returned step"): + expand_samples_to_train_pairs(None, [sample], [0.0], [0.0]) diff --git a/tests/fast/utils/test_dmd2_arguments.py b/tests/fast/utils/test_dmd2_arguments.py new file mode 100644 index 000000000..6e49daf89 --- /dev/null +++ b/tests/fast/utils/test_dmd2_arguments.py @@ -0,0 +1,120 @@ +"""--loss-type dmd2 validation: reject what the critic group and the one-step actor cannot run, never fill args. + + defaults ──► --critic-load = --teacher-load, --critic-lr = --lr, unless the user set them + critic group ──► scores with --teacher-load (no adapters), learns on a visual --fsdp-flow-shift grid, + scores in a 0 <= MIN < MAX <= 1 --dmd2-score-sigma-range, + takes --dmd2-critic-steps-per-rollout steps over each DP shard + other losses ──► --critic-load / --critic-lr have no consumer and are rejected + teacher score ──► CFG --dmd2-teacher-guidance-scale >= 1 + actor ──► unguided rollouts, one step per rollout, generator hooks given until a built-in lands +""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-cpu", labels=[]) + +from argparse import Namespace + +import pytest + +from miles.utils.arguments import set_default_diffusion_args, validate_dmd2_args + + +def _args(**overrides): + values = dict( + loss_type="dmd2", + critic_load="teacher", + critic_lr=1e-6, + teacher_load="teacher", + teacher_lora_adapter_paths=None, + dmd2_teacher_guidance_scale=4.5, + dmd2_critic_steps_per_rollout=5, + dmd2_score_sigma_range=[0.02, 0.98], + fsdp_flow_shift={"visual": 3.0}, + diffusion_guidance_scale=1.0, + diffusion_true_cfg_scale=None, + use_lora=False, + ref_mode="none", + custom_prepare_train_batch_path="my.prepare_generator", + custom_loss_function_path="my.generator_loss", + num_steps_per_rollout=1, + global_batch_size=40, + micro_batch_size=1, + actor_num_gpus_per_node=4, + actor_num_nodes=1, + sequence_parallel_size=2, + ) + values.update(overrides) + return Namespace(**values) + + +@pytest.mark.parametrize( + ("user", "expected"), + [({}, ("teacher", 2e-6)), ({"critic_load": "critic-init", "critic_lr": 8e-6}, ("critic-init", 8e-6))], +) +def test_dmd2_critic_defaults_to_the_teacher_and_the_actor_lr(user, expected): + args = Namespace( + sglang_tp_size=1, + sglang_sp_degree=None, + sglang_enable_cfg_parallel=False, + rollout_num_gpus_per_engine=1, + loss_type="dmd2", + custom_expand_samples_to_train_pairs_path=None, + custom_prepare_train_batch_path=None, + custom_loss_function_path=None, + ref_mode=None, + diffusion_kl_beta=0.0, + teacher_load="teacher", + lr=2e-6, + **{"critic_load": None, "critic_lr": None, **user}, + ) + set_default_diffusion_args(args) + assert (args.critic_load, args.critic_lr) == expected + + +@pytest.mark.parametrize("flag", ["critic_load", "critic_lr"]) +def test_critic_flags_without_dmd2_are_rejected(flag): + with pytest.raises(ValueError, match="configure the critic group of --loss-type dmd2"): + validate_dmd2_args(Namespace(**{"loss_type": "nft", "critic_load": None, "critic_lr": None, flag: "set"})) + validate_dmd2_args(Namespace(loss_type="nft", critic_load=None, critic_lr=None)) + + +def test_valid_dmd2_args_pass_unchanged(): + args = _args() + original = vars(args).copy() + validate_dmd2_args(args) + assert vars(args) == original + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"teacher_load": None}, "requires --teacher-load"), + ({"dmd2_teacher_guidance_scale": None}, "--dmd2-teacher-guidance-scale >= 1"), + ({"dmd2_teacher_guidance_scale": 0.5}, "--dmd2-teacher-guidance-scale >= 1"), + ({"fsdp_flow_shift": None}, "visual --fsdp-flow-shift"), + ({"fsdp_flow_shift": {"audio": 3.0}}, "visual --fsdp-flow-shift"), + ({"dmd2_score_sigma_range": [0.5, 0.5]}, "0 <= MIN < MAX <= 1"), + ({"dmd2_score_sigma_range": [0.9, 0.1]}, "0 <= MIN < MAX <= 1"), + ({"dmd2_score_sigma_range": [-0.1, 0.5]}, "0 <= MIN < MAX <= 1"), + ({"dmd2_score_sigma_range": [0.5, 1.5]}, "0 <= MIN < MAX <= 1"), + ({"diffusion_guidance_scale": 4.5}, "samples the student unguided"), + ({"diffusion_true_cfg_scale": 4.5}, "samples the student unguided"), + ({"use_lora": True}, "drop --use-lora"), + ({"teacher_lora_adapter_paths": ["teacher-adapter"]}, "drop --use-lora and --teacher-lora-adapter-path"), + ({"ref_mode": "ref"}, "uses no reference model"), + ({"custom_loss_function_path": None}, "no built-in generator loss yet"), + ({"num_steps_per_rollout": 2}, "one actor step per rollout"), + ({"dmd2_critic_steps_per_rollout": 0}, "must be >= 1"), + # 40 samples over 2 DP shards = 20 per shard, which 3 critic steps cannot split. + ({"dmd2_critic_steps_per_rollout": 3}, "split each of the 2 DP shards"), + # 20 per shard over 5 critic steps = 4 per step, which micro-batches of 8 cannot split. + ({"micro_batch_size": 8}, "whole --micro-batch-size 8 micro-batches"), + ], +) +def test_unsupported_dmd2_args_are_rejected(overrides, message): + args = _args(**overrides) + original = vars(args).copy() + with pytest.raises(ValueError, match=message): + validate_dmd2_args(args) + assert vars(args) == original diff --git a/tests/fast/utils/test_trainer_config.py b/tests/fast/utils/test_trainer_config.py index e1959b808..62694d714 100644 --- a/tests/fast/utils/test_trainer_config.py +++ b/tests/fast/utils/test_trainer_config.py @@ -27,3 +27,41 @@ def test_trainer_args_apply_overrides_to_a_deep_copy(): def test_trainer_args_reject_unknown_overrides(): with pytest.raises(AssertionError, match="overrides 'lrr'"): compute_trainer_args(Namespace(lr=1e-5), TrainerConfig(role="critic", overrides={"lrr": 8e-6})) + + +@pytest.mark.parametrize( + ("load", "save", "critic_dir"), [("/ckpt/run/", "/ckpt/run", "/ckpt/run_critic"), (None, None, None)] +) +def test_dmd2_critic_trains_first_from_its_own_checkpoint(load, save, critic_dir): + args = Namespace( + loss_type="dmd2", + hf_checkpoint="student", + teacher_load="teacher", + critic_load="critic-init", + lr=2e-6, + critic_lr=8e-6, + load=load, + save=save, + global_batch_size=40, + num_steps_per_rollout=1, + dmd2_critic_steps_per_rollout=5, + use_ema=True, + custom_prepare_train_batch_path="my.prepare_generator", + custom_loss_function_path="my.generator_loss", + custom_train_output_function_path=None, + ) + original = vars(args).copy() + + trainer_configs = compute_trainer_configs(args) + + assert [config.role for config in trainer_configs] == ["critic", "actor"] + critic, actor = (compute_trainer_args(args, config) for config in trainer_configs) + assert (critic.hf_checkpoint, critic.teacher_load, critic.lr) == ("critic-init", "teacher", 8e-6) + assert (critic.load, critic.save) == (critic_dir, critic_dir) + assert (critic.num_steps_per_rollout, critic.global_batch_size) == (5, 8) + assert (critic.use_ema, actor.use_ema) == (False, True) + assert critic.custom_prepare_train_batch_path.endswith("loss_hub.dmd2.prepare_critic_batch") + assert critic.custom_loss_function_path.endswith("loss_hub.dmd2.critic_loss_formula") + assert critic.custom_train_output_function_path.endswith("loss_hub.dmd2.score_last_step_samples") + assert vars(actor) == {**original, "teacher_load": None} + assert vars(args) == original