[trainer, worker, fsdp, cfg, tests] feat: add multi-role distillation runtime

[NancyFyong]

What does this PR do?

Draft — PR 2 of the distribution-matching architecture, stacked on #545.

Bind PR 1's control-plane executor to the existing Ray/FSDP lifecycle with independently optimized semantic roles, shared-base LoRA or colocated independent storage, student EMA, and composite checkpoint/resume.

Parent design: #519. Concrete data-plane architecture: #535. Base contracts: #520. This slice has no architecture-specific phase computer and does not, by itself, make DMD training runnable. Qwen and other model adapters remain separate follow-ups.

Updated and locally validated head: 1a0edbb99d1ebf0d324ff7033dc0666530a833c6 (2026-09-07). Incremental diff against PR 1: 26 files, +3,285 / −63. The base branch distillation-acceleration is at ae2c34c; its fixes were merged without rewriting published history. Merge PR 1 first, then retarget/rebase this PR onto main as needed; do not merge this slice independently of its contracts.

Checklist Before Starting

• Searched PRs referencing #535, multi-role distillation PRs, and the issue comments. No competing upstream implementation was found before publication.
• Title follows [{modules}] {type}: {description}.

Non-duplication: this is the dependent runtime layer, not a second implementation of PR 1. Existing OPD work (#293, #498, #375, #495) does not provide independent student/fake-score optimizers, shared-base role composition, or the composite checkpoint lifecycle. No OPD objective/scheduler is replaced.

Test

Current-head local verification. The merge result is tree-identical to the tested repair snapshot:

CPU_PYTEST_INI=$(mktemp)
printf '[pytest]\npython_files = *_on_cpu.py\nasyncio_mode = auto\n' > "$CPU_PYTEST_INI"
PYTHONPATH=. TORCH_COMPILE_DISABLE=1 TORCHINDUCTOR_DISABLE=1 \
  python -m pytest -c "$CPU_PYTEST_INI" -p no:cacheprovider --no-cov -q \
  tests/trainer/diffusion \
  tests/workers/test_diffusion_distillation_runtime_on_cpu.py \
  tests/workers/test_diffusion_distillation_lora_on_cpu.py
# 470 passed

bash scripts/generate_trainer_config.sh
# generated configs unchanged

pre-commit run --files $(git diff --name-only distillation-acceleration...HEAD)
# all applicable hooks passed across the complete incremental PR file list

git diff --check distillation-acceleration...HEAD
# clean

Historical GPU evidence at 80d7b92 — not rerun for the current repair:

PYTHONPATH=. python -m pytest tests/workers/test_distillation_fsdp_roles.py --no-cov -q
# 2 passed: one-rank FSDP1 and FSDP2

The GPU tests use a tiny model, real FSDP wrappers, named LoRA, and a small test checkpoint manager. They exercise pending-student-graph preservation across frozen role switches, gradient/optimizer isolation, EMA, role export, and model/optimizer/scheduler round-trip. They are not a full production Ray checkpoint or multi-rank inference test. Architecture-specific multi-rank/real-model evidence belongs to the later adapter PR, not to this slice.

Current CPU verification used Python 3.11 / PyTorch 2.13.0+cu129 with repository-matching verl and vLLM-Omni pins. Applicable pre-commit hooks (including mypy), docs/config checks and generated-config verification passed. The dependent PR 5 stack also passed 1,521 full L1 CPU tests. The new denominator collective is CPU-simulated; no new GPU, real multi-rank FSDP or pinned-runtime inference validation is claimed here. Historical GPU evidence above does not establish those properties. Upstream CI results are separate from these local checks. The current CPU workflow filters PR base branches to main/v0.*; this stacked base is therefore excluded until retargeting or an explicit maintainer-supported run. The ci-core label requests the applicable upstream checks without claiming that every workflow has run.

API and Usage Example

Schema fragment for a future registered architecture adapter; not a standalone launch recipe:

algorithm:
  trainer_type: distillation
  sample_source: offline

actor_rollout_ref:
  actor:
    strategy: fsdp2
    optim:
      lr: 0.0001
      lr_scheduler_type: constant
  model:
    lora_rank: 32

distillation:
  enabled: false  # existing OPD switch, deliberately not used
  distribution_matching:
    recipe: dmd2
    role_storage: shared_base_adapters
    fake_update_ratio: 2
    student_micro_batch_size_per_gpu: 1
    fake_score_micro_batch_size_per_gpu: 1
    fake_score_optim:
      lr: 0.00002
      lr_scheduler_type: constant
    ema_decay: 0.999
    ema_start_step: 0
    export_role: student

The architecture opts in via DistributionMatchingModelAdapter.build_distribution_matching_computer(model_config, plan). Its returned object must implement compute_phase(request, batch, runtime), state_dict(), and load_state_dict(state). compute_phase returns one graph-bearing scalar role loss plus detached numeric metrics and optional loss_normalizer; the generic runtime owns backward and optimization.

Review fix — explicit reduction contract:

• loss_normalizer=None keeps sample-mean accumulation unchanged.
• A positive finite denominator accumulates loss * loss_normalizer numerator gradients, then divides by the engine gradient-DP group's averaged count before clipping/stepping. Wan ODE in PR 5 uses its active-element count, preserving physical-batch equivalence when frame masks differ.
• Invalid denominators and mixed reduction modes for a role fail closed; zero_grad clears normalization state.
• Role-loss and /loss metrics use the same local micro-batch denominators; existing rank-level metric reduction and other count/timing policies remain unchanged.
• New CPU tests cover physical/micro-batch losses, gradients and updates, reset behavior, invalid/mixed modes, worker loss logging and a simulated DP count collective. Real multi-rank validation remains outstanding.

Design & Code Changes

main_diffusion / TaskRunner
  -> DistillationRayTrainer(BaseRayDiffusionTrainer)
     -> immutable plan + PR 1 control plane + stateful batch provider
     -> DiffusionDistillationWorkerGroup (executor facade)
        -> RayWorkerGroup / blocking ND-dispatched phase RPC
           -> DiffusionDistillationWorker
              | architecture-owned DistributionMatchingComputer
              ` DistillationRoleRuntime
                  -> one DistillationRoleGroupEngine per physical group
                     -> existing DiffusersFSDPEngine / FSDP1 or FSDP2
                     -> disjoint role optimizers and LR schedulers

Physical storage and logical roles are separate:

shared frozen base                 colocated independent layout
  student adapter + optimizer       student model + optimizer
  fake_score adapter + optimizer    fake_score model + optimizer
  teacher_score: adapters disabled   frozen teacher model
  student_ema adapter                student_ema model

• DiffusionDistillationWorker reuses verl's worker/dispatch/profiler APIs and EngineRegistry; the role engine extends the existing Diffusers FSDP engine rather than adding another model-loader/backend stack.
• Actor FSDP/optimizer/checkpoint subconfigs are instantiated directly, avoiding accidental interpretation of dmd2 as a PPO loss name. Profiler/tool configs are typed before DistProfiler construction.
• One phase RPC encloses local micro-batch accumulation and one successful role optimizer/scheduler step. Non-finite updates do not advance scheduler, EMA, or optimizer counters.
• Shared FSDP1 storage requires use_orig_params=true. Optimizer parameter ownership is disjoint; inactive-role gradients fail closed. Role contexts restore prior adapter and train/eval state.
• Blocking dispatch surfaces rank exceptions before lazy collection metadata RPCs; the facade also resolves a returned DataProtoFuture before extracting results. Architecture computers must still synchronize collective-dependent rollout decisions.
• Each physical model is checkpointed once per rank, not once per optimizer. Secondary role states, computer state, dataloader, driver/worker RNG and control-plane counters are included. Publication uses a temporary directory and atomic rename; canonical JSON fingerprints reject plan/config drift.
• Semantic export resolves student or student_ema and uses that adapter's own PEFT config. This is an export API, not completed CheckpointEngine/validation-replica integration.

Detailed ownership, initialization and phase sequence, checkpoint layout, failure/recovery behavior, configuration mapping, and test boundaries are in #535.

Non-goals: a Qwen/Wan implementation, standalone score transport, multi-optimizer adversarial phases, causal attention/KV caches, pipeline parallelism, or vLLM validation orchestration. No inference server handles the differentiable student rollout.

Checklist Before Submitting

• Read the Contribute Guide.
• Applied pre-commit checks to the complete incremental PR file list. Explicit --files avoids the local Git version's unsupported ls-files --deduplicate path used by --all-files.
• Updated the public architecture/configuration/checkpoint specification in [RFC] Multi-role distillation data plane, LoRA EMA, and checkpointing #535 and parent [RFC] Unified Distribution-Matching Distillation Runtime for DMD, DMD2, CausVid, and Self-Forcing #519.
• Added CPU regressions to existing L1 discovery and explicit GPU tests outside L1. GPU coverage is manual here; a model-level runner belongs to the adapter integration.
• The human submitter confirmed reviewing every changed line, personally running relevant tests, and accepting responsibility for the fixes and submission.

AI assistance: Claude and OpenAI via pi assisted implementation, review and test construction. The human submitter understands and owns the change. This remains Draft for upstream architecture review.

Why a dedicated worker/engine instead of the PPO actor/critic/ref workers

The PPO actor/critic/ref roles are three physically separate models, each wrapped by its own FSDP unit, each with its own optimizer, and each occupying its own resource pool. They interact only by passing DataProto tensors (log-probs, values, KL) — there is no shared backbone and, crucially, no shared autograd graph across roles.

The DMD-family roles are the opposite: in the recommended shared_base_adapters layout, student, fake_score, and student_ema are named LoRA adapters on one shared base, and teacher_score is that same base with adapters disabled. Within a single student phase they are tightly coupled on the same graph:

x0 = student_rollout(...)                 # gradient-bearing
with torch.no_grad():
    v_fake = fake_score(renoise(x0))       # switch to fake adapter on the SAME base module
    v_real = teacher(renoise(x0))          # disable adapters on the SAME base module
grad = dmd_gradient(v_real, v_fake, ...)   # backprops through x0 into the student

This requires switching the active role on one physical module while (a) selecting that role's optimizer/scheduler and (b) not disturbing the student's autograd graph. That is exactly what DistillationRoleGroupEngine.use_role provides — and it has no equivalent in the three-independent-model PPO worker topology:

@contextmanager
def use_role(self, role, *, grad_enabled=None):
    previous_role = self._active_role
    previous_training = self.module.training
    binding = self.role_bindings[role]
    # frozen roles stay frozen even if a caller asks for grads
    effective_grad = binding.trainable if grad_enabled is None else binding.trainable and grad_enabled
    self.activate_role(role)               # set_adapter(...) or disable adapters; route this role's optimizer
    self.module.train(effective_grad)
    grad_context = nullcontext() if effective_grad else torch.no_grad()
    try:
        with grad_context:
            yield self.module
    finally:
        self.module.train(previous_training)
        self.activate_role(previous_role or self._primary_role)   # always restore the prior adapter/optimizer

activate_role toggles the adapter (set_adapter, or disable_adapters for the frozen teacher) and re-points self.optimizer/lr_scheduler/optimizer_config to that role. Gradient ownership is enforced separately by assert_gradient_isolation, and each stepped role zeroes only its own optimizer.

Consequences that the PPO worker stack cannot express without being rebuilt:

	PPO actor/critic/ref	DMD student/fake/teacher/EMA
Physical models	3 independent models	1 shared base + named adapters
Parameter sharing	none	shared-base LoRA
Autograd across roles	independent	teacher/fake forwards must preserve the student graph
Role interaction	DataProto tensor hand-off	in-forward adapter switch on one module
Update cadence	each role once per PPO step	student ×1 : fake ×K, plus fake-only warmup
Checkpoint	one shard set per model	base saved once + per-role optimizer/scheduler/EMA/RNG

So the new worker/engine is not a parallel reimplementation: DistillationRoleGroupEngine subclasses DiffusersFSDPEngine and reuses its FSDP1/FSDP2 wrap, LoRA adapter switch/copy/EMA, and export-tensor iteration; the worker subclasses verl's Worker + DistProfilerExtension and reuses the Ray lifecycle, EngineRegistry, FSDPCheckpointManager shard format, and DataProtoFuture dispatch. Only the multi-role composition layer (role routing, dual optimizers, alternating cadence, composite checkpoint) is new, because that layer is what the actor-centric single-model/single-optimizer topology cannot represent.

Review follow-up: naming propagation

Rebased onto PR 1's naming cleanup: distillation.control_plane → distillation.controller, distillation.equations → distillation.utils, and DistillationTrainerControlPlane → DistillationTrainerController. All imports, accessors (controller/build_controller), checkpoint save/load paths, and tests were updated accordingly. Newly introduced runtime state uses descriptive snake_case names; framework overrides and Python protocol dunders are unchanged.