fix(dmd2): restore multistep backward simulation

[primorLee]

Summary

• replace the multistep DMD2 path that forward-noised dataset latents with inference-time student backward simulation
• sample one selected student step across the batch and synchronize it across distributed ranks
• run preceding student steps under torch.no_grad(), using stochastic re-noising for SDE sampling or the FlowMatch Euler update for ODE sampling
• add CPU regression tests for real-sample independence, gradient isolation, timestep order, and the ODE update

Why

For dmd2_student_sample_steps > 1, _generate_student_data currently constructs the selected-step input as:

_forward_process(real_data, eps_student, sigma)

and invokes the student only once. Consequently, step k is trained on a forward-noised real latent instead of the distribution produced by student steps 0..k-1 at inference. The dmd2_student_sample_type option is also ignored.

DMD2 introduces backward simulation specifically to remove this multistep training/inference input mismatch. Its official implementation samples one step, rolls preceding student steps forward without gradients, and trains the selected step on the resulting sample:

• paper: https://proceedings.neurips.cc/paper_files/paper/2024/hash/54dcf25318f9de5a7a01f0a4125c541e-Abstract-Conference.html
• reference implementation: https://github.com/tianweiy/DMD2/blob/8d8fa55633d47cfb81bbc7a892e7248f9518763f/main/sd_unified_model.py#L127-L162

The new SDE path re-noises each preceding x0 prediction at the next student sigma. The ODE path follows the repository's FlowMatch scheduler equation:

x_next = x + ((x - x0) / sigma) * (sigma_next - sigma)

The one-step path remains a single student prediction from schedule-scaled noise.

Validation

python -m unittest discover -s tests -p test_dmd2_backward_simulation.py -v

Ran 3 tests in 0.009s
OK

The same 3 tests pass in a clean Linux Python 3.10 / CPU PyTorch container (Ran 3 tests in 0.144s). Syntax compilation and targeted Ruff correctness checks also pass.

A baseline differential check against origin/main reproduces the bug: with a fixed RNG seed and selected step, changing only real_data from 0 to 100 changes the generated output by 66.699997. With this patch, the generated rollout and selected-step input are identical because real data supplies only shape/device/dtype.

Scope / limitation

This is a code-level correction of the DMD2 sampling distribution. I do not have an A100 end-to-end rerun of the FLUX.2 Klein 4B recipe, so I am opening this as a draft and am not claiming that it resolves every source of the loss oscillation reported in #1617. A before/after training-curve and sample-quality rerun is still needed.

Addresses #1617.

[yjy415]

Thanks for this PR — the implementation and tests are both very thorough. We ran an actual training comparison on the FLUX.2-klein-base-4B DMD2 configuration, but we did not observe a clear quality improvement from the backward simulation version compared with the current implementation.

In addition, our DMD2 implementation follows FastGen, whose multi-step training path likewise applies forward noising to real latents and then trains the selected student step, rather than using the student's preceding steps for backward simulation.

For these reasons, we will not be merging this PR for now. Thanks again for the contribution and for the discussion!