[DMD2 1/7] feat(ray): train several colocated trainer groups in the sync driver

[zhihengy]

First of a stacked series that adds DMD2 distillation. Lets the sync driver run several colocated trainer groups, keyed by role, instead of the single actor group. Every run still has exactly one group (the actor); the critic group arrives in [DMD2 4/7].

Critic-related components are included in 4th pr of this series

What

• create_training_models builds one RayTrainGroup per entry of compute_trainer_args(args) (role → args) and returns them as a dict keyed by role. For now compute_trainer_args returns only {"actor": args}.
• All groups share the actor placement group and world size. Each train actor takes TRAIN_ACTOR_GPU / N of a bundle, so the rollout engine and reward shares do not change; with one group it is still 0.7.
• Groups initialize one after another and must restore the same rollout id.
• train_diffusion.py trains every group on each rollout in declared order, then saves and offloads all of them. Only the actor pushes weights to the rollout engines.
• train_sft.py and train_diffusion_async.py take the actor group from the dict and are otherwise unchanged.

Why

DMD2 trains a fake-score critic next to the few-step generator. Like upstream miles' PPO critic, the critic runs as its own colocated trainer group rather than as a second trainable model inside the actor. That way the actor's optimizer, checkpoint and weight-sync paths stay as they are. Keying the groups by role leaves room for more trainers later.

Upstream precedent: miles builds its trainers from a config list, creates the actor and then the optional critic, and asserts that both restored the same rollout (miles/ray/placement_group.py). Its colocated trainer workers take a 0.4 GPU share each, next to a 0.2 engine share.

Tests

• tests/fast/rollout/test_placement_group.py, with 1, 2 and 3 trainers:
  • every group gets the actor placement group, its own args and role, and TRAIN_ACTOR_GPU / N (the shares stay within 0.7 after Ray's 0.0001 truncation);
  • groups initialize in declared order, and only the actor gets the rollout manager;
  • trainers that restored different rollouts are rejected.
• tests/fast/test_train_diffusion.py: two groups train each rollout in declared order, the second only after the first finished. Both save and offload, and only the actor pushes weights.
• tests/fast/test_train_sft.py: updated for the dict return.
• With a single trainer the driver makes the same calls in the same order as before. The e2e suite was not run locally.

Checklist

• pre-commit hooks (black, isort, ruff at the pinned versions) pass on the changed files
• Added/updated tests for new behaviour (trainer group creation, driver order)
• pytest tests/fast is green (515 passed, 6 skipped; 510 on main)
• If launch flags changed, python3 train_diffusion.py --help still parses (no flags changed)
• If a public flag was added, it appears in the CLI reference docs (no flags added)
• 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. The critic (fake score) and the frozen teacher live in their own colocated trainer group; the actor holds only the student. Merge in order:

1. [DMD2 1/7] this PR: 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] 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)

Todo：

1. More precise offload train