[zhihengy]

Sixth of the DMD2 stacked series, on top of #331. --loss-type dmd2 adds a critic trainer group that trains before the actor on every rollout. The critic starts from the teacher by default, learns the student's samples, and scores its last step's samples with the frozen teacher and itself, before that step's update. The actor's generator loss arrives in [DMD2 7/10]; until then --loss-type dmd2 takes the actor's prepare and loss hooks from the user.

What

• --loss-type dmd2 selects data_conversion_hub.dmd2.expand_samples_to_train_pairs. Each sample comes from a rollout that stop_at_random_step ([DMD2 5/10] feat(rollout): stop the rollout at a random step for DMD2 backward simulation #331) stopped at step i, and becomes one pair: latent = x_i, x0 = x̂0_i, timestep, sigma = σ_i, step_index = i, the cond and the sample index. Rewards are ignored.

• Under dmd2, compute_trainer_configs (miles/utils/trainer_config.py, [DMD2 1/10] refactor(ray): build each trainer's args from a trainer config list #335) lists a critic trainer before the actor; other loss types still get only the actor. The critic's overrides:

  • --hf-checkpoint is the new --critic-load, which defaults to --teacher-load;
  • --lr is the new --critic-lr, which defaults to --lr;
  • --load / --save point to sibling <dir>_critic dirs, the upstream --critic-save default;
  • num_steps_per_rollout is --dmd2-critic-steps-per-rollout (K, default 5), over the same DP shard as the actor;
  • prepare, loss and train-output hooks come from loss_hub.dmd2.

  The actor's only override drops --teacher-load, so only the critic group holds the teacher.

• Critic steps (prepare_critic_batch, critic_loss_formula) use the SFT flow-matching loss on the student's x̂0_i with the rollout's positive cond. The prepare hook builds SFT-shaped pairs and calls prepare_sft_batch, and the loss is sft_loss_formula, so the critic's inputs, targets, seeds and loss are exactly SFT's. The one exception is the last of the K steps: its σ is drawn only from the grid points whose unshifted σ lies in --dmd2-score-sigma-range (default 0.02–0.98, the reference scoring range), through a new optional grid_index_range of prepare_sft_batch. Unset, SFT is unchanged.

• --dmd2-critic-loss-space (default velocity) picks the space of the critic's loss. Under x0, critic_loss_formula hands sft_loss_formula the x0 prediction x_s − σ·v and the x0 target instead of the velocity pair, so the loss is the x0 MSE, i.e. σ² times the velocity MSE.

• Scoring reuses the last step's training forward.

  • On that step, critic_loss_formula keeps each pair's noisy input x_s, its σ and timestep, and x0_fake = x_s − σ·v_fake, all detached. This is the critic after K−1 steps, before the step that trains on these samples.
  • score_last_step_samples, the train-output hook from [DMD2 3/10] feat(ray): hand earlier trainer groups' outputs to later groups as external_data #330, runs after the K steps, when no critic graph is alive. The teacher predicts at the kept x_s with CFG --dmd2-teacher-guidance-scale over the rollout's negative cond, giving x0_real.
  • Both x0 predictions, keyed by sample index, reach the paired actor rank as ctx.external_data["critic"]. Only the last step's samples are scored.
  • The teacher is resharded after its no-grad forwards, as EMAOptimizer.use_weights does, so sleep moves shards.
  • Two DiffusionLossContext fields carry this: optim_step_idx, set by the train loop, and train_output_inputs, which loss hooks fill for the train-output hook and which starts empty on each train call. Both are documented in customization.md.

• Train metrics of a non-actor group carry its role: train/critic-loss, train/critic-grad_norm, train/critic-lr. Upstream miles' log_train_step does the same. The actor's keys don't change.

• validate_dmd2_args in arguments.py:

  • requires --teacher-load, --dmd2-teacher-guidance-scale >= 1, a visual --fsdp-flow-shift and 0 <= MIN < MAX <= 1 in --dmd2-score-sigma-range;
  • requires unguided student rollouts and one actor step per rollout;
  • rejects LoRA, teacher adapters and a reference model;
  • requires K to split each DP shard evenly.

  It only rejects; defaults stay in set_default_diffusion_args. --critic-load and --critic-lr without --loss-type dmd2 are rejected because nothing would read them.

• Docs: cli-reference.md, customization.md, fsdp_model_roles.md.

Why

• One rollout of K·B samples stands for the reference's K iterations: DMD2's two-time-scale rule takes K = 5 critic updates per generator update, each on a fresh batch. The critic trains on the first (K−1)·B samples, scores the last B, then trains on them, and the actor ([DMD2 7/10]) trains on those B only. This is tianweiy/DMD2's order: the batch the generator steps on is scored by a critic that has seen K−1 batches of the current generator but not this one. The generator's weights don't change within a cycle, so drawing the K batches in one rollout is the same as drawing them one by one, and it needs one engine switch instead of K.
• Scoring a sample the critic has already trained on would pull x0_fake toward the sample itself and bias g, so the scored batch is held out until the scoring forward.
• The reference draws the scoring σ and noise apart from the critic's training draw and runs a second critic forward. Here the last step's training draw, limited to the scoring range, is the scoring draw. Each estimate keeps its own distribution, so the actor's update is distributed as in the reference. Only the two updates of one cycle share noise, and the critic's last step loses the σ tails outside the scoring range, where scoring never looks. This saves one critic forward per scored sample, and the teacher runs after training, so peak memory doesn't grow.
• The critic reuses the actor class: it is an FSDPTrainRayActor with SFT hooks and the teacher as its frozen model, so the only change to the class is the metric names. Upstream miles builds its PPO critic the same way. _compute_critic_overrides overrides load, save and lr; the critic trains first, and the actor receives its outputs as external_data.
• A separate critic init and learning rate follow the other DMD2 trainers. tianweiy/DMD2 has --guidance_lr next to --generator_lr. FastVideo's DMD2 gives each of student, teacher and critic its own init_from, and the critic its own fake_score_learning_rate (8e-6 against the student's 2e-6 in its Wan recipe). Upstream miles' PPO critic has --critic-load / --critic-lr, defaulting to --load / --lr. Here the student, teacher and critic start from --hf-checkpoint, --teacher-load and --critic-load.
• The critic ships two x0 predictions rather than the gradient. The actor computes g against its own recomputed x0 in [DMD2 7/10], so g's normalizer uses the trainer's x0, as in the reference.
• The scoring range and the critic's loss space are flags because validated DMD2 recipes differ on both. tianweiy/DMD2 scores in 2–98% (--min_step_percent / --max_step_percent), the default here. FastGen draws its noising time in sample_t_cfg.min_t..max_t before the shift, the convention --dmd2-score-sigma-range follows, and picks the critic's loss with fake_score_pred_type (flow or x0). FastVideo's MiniMax H3 4-step recipe uses 0.001–0.999 and an x0-space critic (fake_score_loss_space: x0).

Known gap

When guidance is 1 the engine doesn't encode the negative prompt, so denoising_env.neg_cond_kwargs is empty. Teacher CFG above 1 then raises an explicit error. Returning the negative cond without running CFG needs a small sglang-miles change, which will be a follow-up. Until then the teacher can only run unguided.

Tests

• tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py:
  • Before the last step, the critic batch equals prepare_sft_batch on (x̂0_i, positive cond). On the last step it equals the same draw over grid points 2–98%. A custom --dmd2-score-sigma-range bounds the unshifted σ, checked with an asymmetric range.
  • The critic loss and metrics equal sft_loss_formula on every step. Under --dmd2-critic-loss-space x0 the loss equals the x0 MSE and σ² times the velocity MSE, and the prepared batch is left as is.
  • Only the last step keeps scoring inputs, so only its samples are scored.
  • The scores equal separate critic and teacher forwards at the kept x_s, which is the reference computation at the shared draw.
  • The teacher applies CFG over the negative cond, and the critic's x0 comes from its unguided training forward.
  • A missing negative cond is rejected when the teacher runs CFG.
  • The teacher is resharded.
• tests/fast/rollout/test_dmd2_convert.py: one pair per stopped sample. Full trajectories and trajectories that run past their first returned step are rejected.
• tests/fast/utils/test_trainer_config.py:
  • The critic trains first from --critic-load at --critic-lr, with sibling checkpoint dirs and K steps.
  • The actor's copy differs only by dropping the teacher, and the input args are left untouched.
• tests/fast/utils/test_dmd2_arguments.py: valid args pass unchanged, and each unsupported combination is rejected, including an empty, reversed or out-of-[0, 1] --dmd2-score-sigma-range. --critic-load / --critic-lr default to --teacher-load / --lr unless set, and are rejected without dmd2.
• The GPU e2e was not run: DMD2 can't run end to end before [DMD2 7/10].

Checklist

• pre-commit hooks (black, isort, ruff at the pinned versions) pass on the changed files
• Added/updated tests for new behaviour (converter, critic hooks, scoring range, x0 critic loss, critic trainer args, validation)
• pytest tests/fast: 550 passed, 7 skipped; 4 tests fail in this environment and fail the same way without this change (test_hybrid_shard_mesh::test_mesh_axes and test_metric_buffer_dist::test_metric_buffer_reduction, whose torchrun children import another checkout, and two test_sft_h3_ref2va_encoder cases)
• If launch flags changed, python3 train_diffusion.py --help still parses (dmd2 and the six new flags are listed)
• If a public flag was added, it appears in the CLI reference docs (--dmd2-critic-steps-per-rollout, --dmd2-teacher-guidance-scale, --dmd2-score-sigma-range, --dmd2-critic-loss-space, --critic-load, --critic-lr)
• If an example was added, it has a real walkthrough (no examples added)

Stack

DMD2: distill a flow-matching teacher into a few-step student with distribution matching. Merge in order:

1. [DMD2 1/10] [DMD2 1/10] refactor(ray): build each trainer's args from a trainer config list #335: per-trainer args from a trainer config list
2. [DMD2 2/10] [DMD2 2/10] feat(ray): train several colocated trainer groups in the sync driver #329: several colocated trainer groups in the sync driver
3. [DMD2 3/10] [DMD2 3/10] feat(ray): hand earlier trainer groups' outputs to later groups as external_data #330: hand earlier trainer groups' outputs to later groups as external_data
4. [DMD2 4/10] [DMD2 4/10] feat(fsdp): add a train-input hook that preprocesses train pairs before micro-batching #338: a train-input hook that preprocesses a trainer group's train pairs before micro-batching
5. [DMD2 5/10] [DMD2 5/10] feat(rollout): stop the rollout at a random step for DMD2 backward simulation #331: few-step rollout and the DMD2 step strategy (needs [Feat] diffusion-rl: stop a rollout at a requested step sgl-project/sglang#45230)
6. [DMD2 6/10] this PR: DMD2 critic group: critic training and teacher/critic scoring of its last step
7. [DMD2 7/10] [DMD2 7/10] feat(dmd2): train the student with the distribution matching gradient from the critic group's scores #334: --loss-type dmd2 on the actor
8. [DMD2 8/10] [DMD2 8/10] perf(ray): sleep each trainer group right after its train #342: sleep each trainer group right after its train
9. [DMD2 9/10] [DMD2 9/10] feat(dmd2): add the MiniMax H3 4-step DMD2 recipe with eval media and DMD2 metrics #343: MiniMax H3 4-step DMD2 recipe, eval media and DMD2 metrics
10. [DMD2 10/10] [DMD2 10/10] feat(dmd2): add --dmd2-critic-scoring original, the reference's separate scoring before the critic's last step #351: --dmd2-critic-scoring original, the reference's separate scoring before the critic's last step

A GAN head and other follow-ups come after this series.

[zhihengy]

Not very elegant to leave the whole logic of training config re-ride of dmd2 here, but since dmd2 is the only multi-trainer algo supported I will just leave it here. Need a brief refactor in the future, won't be a major blocker.

[Rockdu]

optim_step_idx and train_output_inputs change during training; better to pass them as hook arguments / return values instead of mutating the ctx, which should stay rollout-level and read-only.