[DMD2 5/7] feat(dmd2): train the student with the distribution matching gradient from the critic group's scores

[zhihengy]

Fifth of the DMD2 stacked series, on top of #333. --loss-type dmd2 now trains the actor. The student recomputes its x0 at the rollout's step i, then follows the teacher-minus-critic scores that the critic group handed over.

What

• prepare_generator_batch: the student's unguided forward at (x_i, t_i) with the pair's positive cond.

• dmd2_loss_formula:

  • x0 = x_i - σ_i·v;
  • per pair, g = nan_to_num((x̂0_fake - x̂0_real) / mean|x0 - x̂0_real|);
  • loss = 0.5·mean((x0 - sg(x0 - g))²), so ∂loss/∂x0 = g / numel.

  This is the reference gradient: tianweiy/DMD2 compute_distribution_matching_loss, with the same batch-mean MSE once the actor divides by the pair count.

• --loss-type dmd2 fills these two hooks the way nft fills its own. [DMD2 4/7]'s requirement for user hooks is dropped.

• Metrics:

  • dmd2_grad_abs_mean;
  • x0_mean_abs_diff, x0_max_abs_diff and x0_rel_max, which compare the trainer's x0 with the rollout's x̂0_i (same step, same weights) through the existing record_rollout_train_abs_diff.

• Single DiT component only, as with NFT.

Why

• The numerator comes from the critic group, which scores a re-noised rollout x̂0_i. The normalizer and the gradient path use the trainer's x0. The two x0s differ only by rollout/train numerics, because the actor takes one step per rollout on synced weights. The x0_* metrics measure that gap on every step.
• No log-prob or ratio is involved, so the loss doesn't depend on the rollout's sampler numerics beyond x_i itself.

Tests

• tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py:
  • The surrogate's gradient equals the reference implementation's on the same x0 and predictions.
  • When the critic matches the teacher, the loss and gradient are zero.
  • 1-D Gaussian toy with analytic posterior means for the teacher and a perfect critic: SGD on the student's mean moves monotonically toward the teacher's mean, which catches sign and normalization bugs.
  • Generator batch fields; multiple components are rejected; the x0 diff metric is recorded.
• tests/fast/utils/test_dmd2_arguments.py: dmd2 fills the converter, prepare and loss paths, but never over a user's own.
• The GPU e2e was not run. SD3.5-Medium is gated on this devbox, and teacher CFG needs the sglang negative-cond change noted in [DMD2 4/7]. The e2e standard will come with the recipe.

Checklist

• pre-commit hooks (black, isort, ruff at the pinned versions) pass on the changed files
• Added/updated tests for new behaviour (reference-gradient parity, Gaussian toy, defaults)
• pytest tests/fast is green (556 passed, 6 skipped; 549 on [DMD2 4/7] feat(dmd2): train a fake-score critic group from the teacher and score the student's samples #333)
• If launch flags changed, python3 train_diffusion.py --help still parses (no flags added)
• If a public flag was added, it appears in the CLI reference docs (no flags added; the customization.md table lists the dmd2 hooks)
• 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] [DMD2 4/7] feat(dmd2): train a fake-score critic group from the teacher and score the student's samples #333: DMD2 critic group: critic training and teacher/critic scoring
5. [DMD2 5/7] this PR: --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)