[DMD2 2/7] feat(ray): hand earlier trainer groups' outputs to later groups as external_data

[zhihengy]

Second 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 (517 passed, 6 skipped)
• 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/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] this PR: hand earlier trainer groups' outputs to later groups as external_data
3. [DMD2 3/7] few-step rollout and the DMD2 step strategy (not opened yet)
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] [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)