[DMD2 4/7] feat(dmd2): train a fake-score critic group from the teacher and score the student's samples

[zhihengy]

Fourth 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, then scores them with the frozen teacher and itself. The actor's generator loss arrives in [DMD2 5/7]; 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 3/7] 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.

• compute_trainer_args returns {"critic": ..., "actor": ...} under dmd2; other loss types still get {"actor": args}. The critic's args are a copy with these 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 and loss_hub.sft.

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

• Critic steps (prepare_critic_batch) use the SFT flow-matching loss on the student's x̂0_i with the rollout's positive cond. The hook builds SFT-shaped pairs and calls prepare_sft_batch, so the critic's inputs, targets and seeds are exactly SFT's.

• Scoring (predict_teacher_and_critic_x0) is the train-output hook from [DMD2 2/7] feat(ray): hand earlier trainer groups' outputs to later groups as external_data #330 and runs after the K steps.

  • Each pair's x̂0_i is re-noised at one σ from the --fsdp-flow-shift grid, limited to 2–98% of the train timesteps (the reference range). σ and the noise are stable_hash-seeded per rollout, micro-batch and DP rank.
  • The teacher predicts with CFG --dmd2-teacher-guidance-scale over the rollout's negative cond. The critic predicts unguided.
  • Both x0 predictions, keyed by sample index, reach the paired actor rank as ctx.external_data["critic"].
  • Both models are resharded after the no-grad forwards, as EMAOptimizer.use_weights does, so sleep moves shards.

• 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 and a visual --fsdp-flow-shift;
  • 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

• The critic trains before the actor so it catches up with the student that produced this rollout, then scores for the actor's step. This is DMD2's two-time-scale rule: 5 critic updates per generator update. Both groups train on the same samples, and the actor uses all of them in its one step.
• 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 5/7], so g's normalizer uses the trainer's x0, as in the reference.

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:
  • The critic batch equals prepare_sft_batch on (x̂0_i, positive cond).
  • Scoring: x_s and σ are reproduced from the seed, and σ stays within 2–98% of the grid.
  • The teacher applies CFG over the negative cond and the critic runs unguided.
  • A missing negative cond is rejected when the teacher runs CFG.
  • The DP rank is part of the seed.
  • Every model 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/rollout/test_placement_group.py:
  • Only dmd2 adds a trainer.
  • 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. --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 5/7].

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, critic trainer args, validation)
• pytest tests/fast is green (549 passed, 6 skipped; 518 on [DMD2 3/7] feat(rollout): stop the rollout at a random step for DMD2 backward simulation #331)
• If launch flags changed, python3 train_diffusion.py --help still parses (dmd2 and the four 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, --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/7] [DMD2 1/7] feat(ray): train several colocated trainer groups in the sync driver #329: several colocated trainer groups in the sync driver
2. [DMD2 2/7] [DMD2 2/7] 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
3. [DMD2 3/7] [DMD2 3/7] feat(rollout): stop the rollout at a random step for DMD2 backward simulation #331: few-step rollout and the DMD2 step strategy (needs [sglang-miles] [Diffusion] Stop a rollout at a requested step sgl-project/sglang#45230)
4. [DMD2 4/7] this PR: DMD2 critic group: critic training and teacher/critic scoring
5. [DMD2 5/7] [DMD2 5/7] feat(dmd2): train the student with the distribution matching gradient from the critic group's scores #334: --loss-type dmd2 on the actor
6. [DMD2 6/7] SD3.5-Medium 4-step recipe and eval (not opened yet)
7. [DMD2 7/7] GAN head (phase 2, not opened yet)