[zhihengy]

Third of the DMD2 stacked series, on top of #329. A trainer group can now hand a per-rank result to the groups that train after it on the same rollout. For DMD2, the critic group uses this to pass the teacher's and the critic's predictions to the actor.

What

• New --custom-train-output-function-path: fn(ctx, train_pairs) -> Any, run without grad after a rollout's train steps. FSDPTrainRayActor.train returns its value (None without the hook).
• train_diffusion.py keeps each group's per-rank train refs for the current rollout and passes the earlier groups' refs to the next group.
• RayTrainGroup.async_train(..., external_data): rank i receives rank i's ref from each earlier group, the pairing upstream miles uses for its critic values. Both groups share the world size and read the same rollout shard, so the paired ranks hold the same train pairs.
• The receiving rank resolves the refs and exposes them to its prepare/loss hooks as ctx.external_data[<role>]. Without earlier groups this is {}.
• --teacher-load also counts the new hook as a consumer of ctx.teacher_models.

Why

Upstream miles hands the critic's values to the actor the same way: critic.train returns one output per worker, and actor.train(..., external_data=values) gives worker i the output of critic worker i. Here the hand-off goes through Ray object refs, so the payload never passes through the driver and is freed once the rollout's refs go out of scope. The producer side is a hook, like the existing prepare/loss hooks, so the trainer stays role-agnostic.

Tests

• tests/fast/rollout/test_placement_group.py: each rank receives the same-rank output of every earlier group, and an empty dict without them.
• tests/fast/test_train_diffusion.py: the second group receives the first group's refs from the same rollout, and the first group receives nothing.
• tests/fast/utils/test_reference_arguments.py: the train-output hook alone counts as a consumer of --teacher-load.
• tests/fast/rollout/test_async_training.py, tests/fast/backends/fsdp_utils/test_ema_optimizer.py: updated for the new train signature.

Checklist

• pre-commit hooks (black, isort, ruff at the pinned versions) pass on the changed files
• Added/updated tests for new behaviour (rank pairing, driver hand-off)
• pytest tests/fast is green (520 passed, 6 skipped; 518 on [DMD2 2/10] feat(ray): train several colocated trainer groups in the sync driver #329)
• If launch flags changed, python3 train_diffusion.py --help still parses (the new flag is listed)
• If a public flag was added, it appears in the CLI reference docs (--custom-train-output-function-path in cli-reference.md and customization.md)
• 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] this PR: 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] [DMD2 6/10] 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 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): train_dmd2.py and --dmd2-critic-scoring original, the reference's separate scoring #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.

[Rockdu]

This operation can pull the whole training output into our controller, we should return ray.put(train_output) in train in this case

[Rockdu]

Another from agent:
[Nit] docs/fsdp_model_roles.md:40 is stale after this PR: it still says --teacher-load requires a custom prepare/loss hook, but the validator now also accepts --custom-train-output-function-path alone (arguments.py:1735-1743), as cli-reference.md already says. Suggest requires a custom prepare/loss/train-output hook. (The file is outside this diff, so no inline suggestion.)

[Rockdu]

CC cmt

[P1] This test execs the real train() source, and main's train() (#361) now reads self.args.use_output_store before _train_core, so the fake args here raise AttributeError once the PR is rebased (the file is new, so the merge won't flag it).
test_ema_optimizer.py on main got use_output_store=False for the same reason; please add it here too, plus a materialize_stored_tensors stub in the exec namespace.