[zhihengy]

First of a stacked series that adds DMD2 distillation. Gives each trainer its own args: a trainer config list names the run's trainers and what each one overrides, and compute_trainer_args builds that trainer's args from it. Every run still has exactly one trainer, the actor, with no overrides, so nothing changes.

What

• New miles/utils/trainer_config.py:
  • TrainerConfig(role, overrides);
  • compute_trainer_configs(args) lists the run's trainers. For now it returns only the actor, without overrides;
  • compute_trainer_args(args, trainer) deep-copies args and applies the trainer's overrides. Overriding an argument the parser doesn't know is rejected.
• create_training_models initializes the actor group with compute_trainer_args(args, actor_config) instead of args. The copy holds the same values, so the trainer sees the same arguments.

Why

DMD2's critic runs as a second trainer group on the same FSDPTrainRayActor class, and that class reads the plain argument names (hf_checkpoint, lr, load / save, num_steps_per_rollout, the hook paths). One Namespace can't hold the actor's and the critic's values at once, so each trainer gets its own copy of the args with its overrides, and the trainer class stays role-agnostic. The per-trainer choices live in a config module next to eval_config.py, which already resolves per-dataset overrides, rather than in the Ray placement code.

Upstream precedent: miles keeps the trainer list (MegatronTrainerConfig(trainer_id, model_id, role, overrides), with the critic synthesized from --use-critic) and compute_trainer_args(args, trainer) (a deep copy whose overrides must be known arguments) in miles/backends/megatron_utils/megatron_config.py. Its placement_group.py only calls them.

Tests

• tests/fast/utils/test_trainer_config.py:
  • the actor is the only trainer by default;
  • overrides go to a deep copy, and the input args are left untouched;
  • an unknown override is rejected.
• A deep copy of fully parsed and validated args pickles to the same bytes as the original (checked locally with the test_sft_encoder_args parser fixture).
• 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 configs, trainer args)
• pytest tests/fast is green (513 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/10] this PR: 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] [DMD2 3/10] 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
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): add --dmd2-critic-scoring original, the reference's separate scoring before the critic's last step #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]

LGTM