[trainer, cfg, tests] feat: add architecture-neutral distillation trainer control plane

[NancyFyong]

What does this PR do?

Draft — architecture review for PR 1 of the distribution-matching distillation stack.

Add one architecture-neutral control plane, immutable execution plans, recipe registries, pure tensor utilities, and an OPD-separated configuration seam. Related design: #519; detailed architecture and acceptance criteria: #520.

This slice deliberately does not train a model. The production constructor accepts the existing diffusion trainer interface, but model-worker initialization stops at the explicit PR 2 boundary. PR 2 adds the data plane; model adapters are subsequent work.

Reviewed head: d7a6b99752dadcb42d7538d8e790438a7a9c7a6a. Diff: 19 files, +3,251 / −4. The companion data-plane Draft PR #546 is stacked on this head, not included here.

Checklist Before Starting

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

Non-duplication: existing OPD work (#293, #498, #375, #495) concerns teacher scoring/scheduling around policy-gradient trajectories. It does not implement the student/fake-score optimizer topology or this phase-control contract. Existing fork-only review PRs are development staging, not another upstream submission.

Test

At the reviewed head, rerun before publication:

PYTHONPATH=. TORCH_COMPILE_DISABLE=1 TORCHINDUCTOR_DISABLE=1 \
  python -m pytest tests/trainer/diffusion --no-cov -q
# 402 passed

bash scripts/generate_trainer_config.sh
# generated configs unchanged

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

git diff --check origin/main...HEAD
# clean

Tests cover deep immutability/pickling, recipe and role-layout validation, strict phase counts, normal/warmup schedules, skipped and failed phases, terminal recovery semantics, hook order, production constructor/config routing, OPD separation, and fp32 equation boundaries. DMD finite differences hold the stop-gradient target fixed.

CPU execution used Python 3.12 and the configured development environment. No GPU/model-execution claim is made. Core-module import tests check direct dependencies; they do not claim the repository's pre-existing package initialization imports no runtime libraries transitively. Upstream CI has not yet been established by these local results.

API and Usage Example

A complete in-process control-plane example; synthetic-model is a plan identifier and is not loaded:

from verl_omni.trainer.diffusion.distillation.controller import (
    DistillationTrainerController,
    FakeBatchProvider,
    FakeDistillationHooks,
    FakePhaseExecutor,
)
from verl_omni.trainer.diffusion.distillation.recipes import build_plan

plan = build_plan(
    "dmd2",
    config={"model_path": "synthetic-model", "fake_update_ratio": 2},
    capabilities=frozenset({"distribution_matching"}),
)
hooks = FakeDistillationHooks()
controller = DistillationTrainerController(
    plan, FakePhaseExecutor(), FakeBatchProvider(num_batches=6), hooks
)
controller.run(2)
assert controller.counters.global_step == 2
assert controller.counters.optimizer_steps == {"student": 2, "fake_score": 4}
assert [call["global_step"] for call in hooks.calls] == [1, 2]

Configuration routing is algorithm.trainer_type=distillation, with new settings under distillation.distribution_matching. The existing OPD flag distillation.enabled stays false. Selecting this route alone is not a runnable training recipe in PR 1.

Design & Code Changes

composed config + declared capabilities
                  |
          recipe/config factory
                  |
       validated DistillationPlan
                  |
  DistillationTrainerController
    | UpdateSchedule -> UpdateCycle -> PhaseRequest
    | BatchProvider.next(request)
    | DistillationPhaseExecutor.execute_phase(request, batch)
    | validate PhaseResult; accumulate role counters
    | commit completed cycle; advance global_step once
    ` after_completed_step(counters, metrics, executor)

The six-file subpackage has explicit boundaries:

Module under verl_omni/trainer/diffusion/distillation/	Responsibility
contracts.py	Plans, tensor containers, role/storage/transport/export contracts, schedules and validation
utils.py	Pure fp32 x0 conversion, DMD normalization/surrogate, detached fake-score and ODE losses, CFG, Euler/re-noise transitions
recipes.py	Registered declarations and factories; no model forwards or objective computation
controller.py	Phase state machine, executor/batch/hook protocols and CPU fakes
ray_trainer.py	Production-compatible constructor/lifecycle shell; no worker allocation yet
__init__.py	Public exports

Key invariants:

• The trainer never branches on recipe or architecture names.
• New controller state uses descriptive snake_case names; Python protocol dunders and BaseConfig _mutable_fields remain unchanged because they are framework contracts.
• One normal cycle has exactly one first-position student phase with repeats=1, then K fake-score phases. Counts are integers, not booleans or silently truncated floats.
• Fake-only warmup advances role counters and completed_cycles, not global_step.
• global_step commits only after all required normal-cycle phases succeed. Each completed phase reports exactly one optimizer step for each requested role and no extras.
• On failure, driver counters/metrics roll back and the driver becomes terminal. This is not an in-memory model rollback or permission to retry partial optimizer updates.
• Checkpoint/validation/export hooks see committed counters; their implementations belong to later layers.
• Causal/bidirectional role layouts are declarations only. No causal model, GAN implementation, remote scorer, or inference server is introduced.

The full implementation architecture, state transitions, tensor boundaries and extension rules are documented in #520 rather than a private/local design document.

Checklist Before Submitting

• Read the Contribute Guide.
• Applied pre-commit checks to the complete PR file list. Explicit --files avoids the local Git version's unsupported ls-files --deduplicate path used by --all-files.
• Updated the public architecture/usage specification in [RFC] Generic distillation trainer control plane, contracts, and pure DMD math #520 and parent [RFC] Unified Distribution-Matching Distillation Runtime for DMD, DMD2, CausVid, and Self-Forcing #519.
• Added/updated CPU tests discovered by the existing *_on_cpu.py workflow; no GPU workflow is applicable to this slice.
• The human submitter confirmed reviewing every changed line and personally running relevant tests.

AI assistance: Claude and OpenAI via pi assisted implementation, review and test construction. The human submitter understands and owns the change. This PR remains Draft for upstream architecture feedback, not because it is presented as an autonomous agent submission.

Review follow-up: naming cleanup

Applied after review annotations and propagated through the whole stack (PR #8/#9). Modules and identifiers now use descriptive, repository-convention names; Python protocol dunders and framework contracts are deliberately left untouched.

Before	After	Kind
control_plane.py	controller.py	module
equations.py	utils.py	module (single-noun subpackage convention)
DistillationTrainerControlPlane	DistillationTrainerController	class
build_control_plane() / self._control_plane	build_controller() / self.controller_instance	trainer accessor/state
self._metrics	self.phase_metrics	controller state
self._failed	self.failed	controller state
_registry	registered_classes	recipe registry storage
FrozenDict._data	FrozenDict.data (wrapped in MappingProxyType)	contract state
CPU-fake _num / _batch_size / _sent / _skip_student / _fail_on	num_batches / batch_size / sent_batches / skip_student / fail_on	test doubles

Deliberately not renamed: __init__, __post_init__, __all__, __slots__, and other Python protocol dunders; BaseConfig._mutable_fields (framework API required to keep config fields mutable); and inherited framework overrides such as _save_checkpoint / _load_checkpoint.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.