[worker, trainer, cfg, tests, doc] feat: multi-teacher routing and standalone teacher pool for diffusion OPD - #427

What does this PR do?

Brings verl's multi-teacher on-policy distillation design (DistillationConfig, AsyncTeacherLLMServerManager, teacher_pool) to the diffusion OPD runtime merged in #300 and #325.

• Config. distillation.teacher_models accepts several named entries, each with a routing key and, for a standalone pool, a world_size. A single-teacher config is unchanged and resolves to the key default. distillation.teacher_key (default data_source) names the batch column used for routing. distillation.n_gpus_per_node and nnodes define an optional teacher_pool; nnodes: 0 (the default) keeps the teachers colocated with the actor, exactly as before.
• Worker. ActorRolloutRefWorker builds one frozen TrainingWorker per routing key (self.teachers), accepts role="teacher" for standalone groups, and infer_teacher_batch picks the teacher from a teacher_key carried in the TensorDict metadata. The call is blocking=False.
• Trainer. The new DiffusionTeacherManager is the diffusion counterpart of verl's AsyncTeacherLLMServerManager. It resolves routing keys with verl's exact error messages, sends each key's rows to its teacher group (padded to the group's DP size × its infer micro batch size so every rank's shard splits into micro-batches, all groups in flight before the first .get()), and scatters teacher_prev_sample_mean back in the original row order. The loss side is untouched.
• Placement. With nnodes > 0, main_diffusion registers teacher_pool / Role.TeacherModel, and the trainer splits the pool into one sub-pool per teacher (split_resource_pool, by world_size), each with its own worker group, so the teachers score concurrently. Teacher scoring remains a serial stage of the step, so a standalone pool buys memory isolation and distance from the rollout engine's sleep/wake cycle, not speed. Note: the default trainer.ray_master_port_range hands every spawned worker group the same rendezvous port, so multi-worker-group runs currently need trainer.ray_master_port_range=null (the standalone smoke sets it); the port-range fix, split out of this PR per review, makes the default work by slicing the range per group.

Not in scope, left for follow-ups: per-teacher lora_adapter_path / guidance_scale, a pure-distillation fast path, and the v1 trainer.

Checklist Before Starting

• Duplicate check: teacher and distillation return only [trainer, worker, cfg] feat: OPD end-to-end draft — teacher-in-ref + distill verification (#293) #304 and [trainer, rollout] Support pickscore for DiffusionOPD #355.
  • [trainer, worker, cfg] feat: OPD end-to-end draft — teacher-in-ref + distill verification (#293) #304 (mine, draft): this PR combined the 1st and 3rd of the split discussed there, which implements verl-style teacher registry and routing + standalone teacher-pool.
  • [trainer, rollout] Support pickscore for DiffusionOPD #355: routes by the task_id of its multi-task dataset to frozen LoRA teachers on the actor backbone. This PR is the config, registry, routing, and placement layer in verl's shape for separate-weight teachers. The two compose: [trainer, rollout] Support pickscore for DiffusionOPD #355's teacher_name column can serve directly as teacher_key.
• PR title format: [worker, trainer, cfg, tests, doc] feat: ...

API and Usage Example

Default-off; a single-teacher config needs no change. Two colocated teachers routed by data_source:

distillation:
  enabled: true
  teacher_key: data_source
  teacher_models:
    ocr:
      key: ocr
      model_path: /ckpt/sd35m-ocr-teacher
    aesthetic:
      key: aesthetic
      model_path: /ckpt/sd35m-aesthetic-teacher
actor_rollout_ref:
  actor:
    diffusion_loss:
      loss_mode: distill_kl

The same teachers on their own GPUs (the actor keeps trainer.n_gpus_per_node):

distillation:
  n_gpus_per_node: 2
  nnodes: 1
  teacher_models:
    ocr: {key: ocr, model_path: /ckpt/sd35m-ocr-teacher, world_size: 1}
    aesthetic: {key: aesthetic, model_path: /ckpt/sd35m-aesthetic-teacher, world_size: 1}

Full reference, including the teacher_model-is-popped pitfall inherited from verl: docs/algo/diffusion_opd.md.

Test

CPU (python -m pytest … -q, verl 0.10.0.dev @ c4b389a, vllm-omni 0.27.0rc1):

tests/trainer/diffusion/test_distillation_config_on_cpu.py
tests/trainer/diffusion/test_distillation_trainer_wiring_on_cpu.py
tests/trainer/diffusion/test_diffusion_teacher_manager_on_cpu.py   (new)
tests/workers/test_distillation_teacher_config_on_cpu.py
tests/special_sanity/test_config_docs.py
→ 34 passed; tests/trainer/diffusion + tests/workers -k on_cpu → 290 passed, 1 skipped
pre-commit run --files <all changed> → all hooks pass

GPU (4× RTX PRO 6000, vllm-omni 0.27.0rc1, tiny random SD3 checkpoints unless noted; tests/special_e2e/run_diffusion_teacher_smoke.sh). With rollout.seed=0 the rollout engine is close to, but not bit-exact across runs: four same-code runs of the same smoke span 2e-5 relative in distill_kl_loss, and their all_latents differ by up to 1.7e-2. That floor is the yardstick for every scalar below; the placement question is settled on tensors instead.

All smokes, the unknown-key negative, and the real-checkpoint run were repeated after rebasing onto current main (verl pin c4b389a, vllm-omni 0.27.0rc1): step-1 distill_kl_loss 0.12816553 (single-teacher) and 0.11859834 (standalone) are bit-identical to the pre-rebase values, mopd 0.11859706 sits at the rollout floor; the negative still fails at startup with verl's message; the real SD3.5-M colocated 3-step run gives 3.18e-3 → 2.05e-3 → 2.06e-3 vs. 3.21e-3 → 2.06e-3 → 2.05e-3 pre-rebase (within the floor), timing_s/teacher 2.01 s vs. 2.07 s.

tests/gpu_smoke/run_gpu_smoke_diffusion_e2e.sh gains the mopd (test 7) and standalone (test 8) cases.

Reward Curves

Setup: the DiffusionOPD paper's Stage-2 scenario — SD3.5-M LoRA student on mixed two-task data (flow-GRPO OCR prompts + PickScore-SFW prompts, data_source column as the routing key), distilling from the paper's two task teachers (quanhaol/DiffusionOPD OCR / Aes LoRAs merged into full SD3.5-M checkpoints). Loss is distill_kl only; rewards (GenRM-OCR for ocr rows, PickScore for pickscore rows — the same scorer the reference implementation uses for the Aes task) are monitored, never in the loss. 300 steps, val every 5 steps on a fixed 64-prompt set (32/task), colocated placement (standalone equivalence is settled bit-exactly above).

Four runs, one sentence each:

• Teacher bands (val-only): OCR-Teacher scores ocr 0.719 / pickscore 0.736; Aes-Teacher scores ocr 0.405 / pickscore 0.826 — orthogonal specialties.
• Routed two teachers (this PR's teacher_models + default teacher_key=data_source): OCR reward climbs 0.20 → the 0.719 teacher band by step ~25 and oscillates around it; PickScore climbs 0.755 → 0.819 (holding 0.818–0.822 over the last 60 steps, peak 0.826 at step 280 — the Aes band) — each task converges toward its own teacher.
• Single-teacher control (OCR teacher only, the unchanged auto-key default path): OCR also reaches its band, but PickScore drops from 0.755 and stays pinned at 0.732–0.739 for all 300 steps — exactly the OCR-Teacher's own PickScore level; without routing, the off-task reward is dragged to the single teacher's level and longer training does not recover it.

On the distill_kl wobble of the routed run (left panel): each logged point is the mean over one 8-pr
ompt step whose task mix is a fresh random draw, and the two teachers' KL magnitudes differ, so the ag
gregate swings with the mix. One LoRA also cannot sit tightly on two teachers at once, so the two runs
' loss levels are not comparable. The per-task reward curves are the primary evidence.

Command (the new example script in this PR; the control run uses the same mixed data with a single teacher_model entry, i.e. the unchanged auto-key default path):

OCR_TEACHER_PATH=$MODELS/sd35-ocr-teacher AES_TEACHER_PATH=$MODELS/sd35-aes-teacher \
bash examples/diffusionopd_trainer/sd35/run_sd35_medium_mopd_distill.sh \
    trainer.total_training_steps=300 trainer.test_freq=5

Design & Code Changes

• verl_omni/workers/config/diffusion/distillation.py: DiffusionDistillationTeacherModelConfig (key, model_path, world_size), _resolve_teacher_models line for line from verl, pool-size check when nnodes > 0.
• verl_omni/workers/engine_workers.py: self.teachers: dict[str, TrainingWorker], role="teacher", infer_teacher_batch routes on teacher_key.
• verl_omni/trainer/diffusion/teacher_manager.py (new): DiffusionTeacherManager; routed sub-batches pad to world_size * infer_micro_batch_size_per_gpu. _to_diffusion_worker_tensordict (the [diffusion, trainer, tests] feat: strip responses from worker batches #430 responses-stripping projection) moves to diffusion_trainer_utils.py so the manager applies it to teacher RPCs too.
• examples/diffusionopd_trainer/sd35/run_sd35_medium_mopd_distill.sh + mixed_task_reward.py (new): the mixed OCR + PickScore recipe behind the reward curves.
• verl_omni/trainer/diffusion/ray_diffusion_trainer.py, main_diffusion.py: teacher_pool registration, per-teacher sub-pools and worker groups, the manager in the fit loop, profiler start/stop for standalone groups.
• Tests, smokes, and docs as listed above; docs/examples/config.md synced.

Checklist Before Submitting

• Read the Contribute Guide.
• Apply pre-commit checks.
• Add / Update the documentation: docs/algo/diffusion_opd.md, docs/examples/config.md.
• Add unit or end-to-end test(s) to the CI workflow: CPU tests above; ci-e2e-diffusion tests 7 and 8.

Note on AI assistance

AI assistance (Claude) was used for this change. I reviewed every changed line, ran the CPU suites and the GPU runs above myself on the hardware listed, and can defend the design end to end.