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

tests/special_sanity/check_device_api_usage.py

@@ -25,6 +25,7 @@
 # directory or file path must contain keyword ".cuda" or "cuda"
 CUDA_KEYWORD_CHECK_WHITELIST = [
     "verl_omni/workers/engine/fsdp/diffusers_impl.py",  # appear in default device_name
+    "verl_omni/workers/engine/fsdp/distillation_impl.py",  # device=[...] registry declaration
     "verl_omni/trainer/diffusion/ray_diffusion_trainer.py",  # appear in default device_name
     "verl_omni/workers/engine/fsdp/omni_impl.py",  # device=[...] registry declaration
     "verl_omni/workers/engine/veomni/diffusion_impl.py",  # device=[...] registry declaration

tests/trainer/diffusion/test_distillation_checkpoint_on_cpu.py

@@ -0,0 +1,205 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""CPU tests for atomic multi-role distillation checkpoint orchestration."""
+
+import os
+import random
+from dataclasses import replace
+
+import numpy as np
+import pytest
+import torch
+from omegaconf import OmegaConf
+
+from verl_omni.trainer.diffusion.distillation.contracts import PhaseRequest
+from verl_omni.trainer.diffusion.distillation.controller import (
+    DistillationTrainerController,
+    FakeBatchProvider,
+    FakeDistillationHooks,
+    FakePhaseExecutor,
+)
+from verl_omni.trainer.diffusion.distillation.ray_trainer import DistillationBatchProvider, DistillationRayTrainer
+from verl_omni.trainer.diffusion.distillation.recipes import build_plan
+
+
+class CheckpointExecutor(FakePhaseExecutor):
+    def __init__(self, *, fail_save=False):
+        super().__init__()
+        self.role_state = {
+            "student": 1.0,
+            "fake_score": 2.0,
+            "student_ema": 1.5,
+            "student_optimizer": 3,
+            "fake_optimizer": 4,
+            "student_scheduler": 5,
+            "fake_scheduler": 6,
+        }
+        self.fail_save = fail_save
+        self.loaded_path = None
+
+    def save_checkpoint(self, local_path, global_step):
+        os.makedirs(local_path, exist_ok=True)
+        torch.save({"global_step": global_step, "role_state": self.role_state}, os.path.join(local_path, "roles.pt"))
+        if self.fail_save:
+            raise RuntimeError("injected save failure")
+
+    def load_checkpoint(self, local_path):
+        state = torch.load(os.path.join(local_path, "roles.pt"), weights_only=False)
+        self.role_state = state["role_state"]
+        self.loaded_path = local_path
+
+
+class StatefulLoader:
+    def __init__(self):
+        self.position = 0
+
+    def state_dict(self):
+        return {"position": self.position}
+
+    def load_state_dict(self, state):
+        self.position = state["position"]
+
+
+class TestDistillationBatchProvider:
+    def test_reuse_student_returns_cached_batch_without_advancing(self):
+        batches = [
+            {"values": torch.tensor([[1.0]]), "responses": torch.tensor([[9.0]])},
+            {"values": torch.tensor([[2.0]]), "responses": torch.tensor([[8.0]])},
+        ]
+        provider = DistillationBatchProvider(batches)
+        student = PhaseRequest("student", 0, 0, "fresh", ("student",), True)
+        reused = PhaseRequest("fake_score", 0, 0, "reuse_student", ("fake_score",), False)
+        fresh = PhaseRequest("fake_score", 0, 1, "fresh", ("fake_score",), False)
+        student_batch = provider.next(student)
+        reused_batch = provider.next(reused)
+        fresh_batch = provider.next(fresh)
+        torch.testing.assert_close(student_batch["values"], reused_batch["values"])
+        torch.testing.assert_close(fresh_batch["values"], torch.tensor([[2.0]]))
+        assert "responses" not in student_batch
+
+    def test_reuse_before_student_fails(self):
+        provider = DistillationBatchProvider([{"values": torch.tensor([[1.0]])}])
+        request = PhaseRequest("fake_score", 0, 0, "reuse_student", ("fake_score",), False)
+        with pytest.raises(RuntimeError, match="before a student batch"):
+            provider.next(request)
+
+
+class TestDistillationCheckpoint:
+    @staticmethod
+    def make_trainer(tmp_path, *, fail_save=False):
+        plan = build_plan(
+            "dmd2",
+            {"model_path": "/m", "fake_update_ratio": 1},
+            frozenset({"distribution_matching"}),
+        )
+        executor = CheckpointExecutor(fail_save=fail_save)
+        hooks = FakeDistillationHooks()
+        controller = DistillationTrainerController(
+            plan=plan,
+            executor=executor,
+            batch_provider=FakeBatchProvider(num_batches=10),
+            hooks=hooks,
+        )
+        controller.run_cycle()
+
+        trainer = DistillationRayTrainer(
+            plan=plan, executor=executor, batch_provider=FakeBatchProvider(10), hooks=hooks
+        )
+        trainer.controller_instance = controller
+        trainer._production = True
+        trainer.global_steps = controller.counters.global_step
+        trainer.train_dataloader = StatefulLoader()
+        trainer.train_dataloader.position = 7
+        trainer.config = OmegaConf.create(
+            {
+                "trainer": {
+                    "default_local_dir": str(tmp_path),
+                    "default_hdfs_dir": None,
+                    "resume_mode": "auto",
+                    "resume_from_path": None,
+                }
+            }
+        )
+        return trainer, executor
+
+    def test_round_trip_restores_roles_counters_dataloader_and_rng(self, tmp_path):
+        trainer, executor = self.make_trainer(tmp_path)
+        random.seed(11)
+        np.random.seed(12)
+        torch.manual_seed(13)
+        trainer._save_checkpoint()
+        expected_random = random.random()
+        expected_numpy = float(np.random.random())
+        expected_torch = float(torch.rand(()))
+
+        executor.role_state = {"corrupt": True}
+        trainer.controller.counters.global_step = 99
+        trainer.controller.counters.optimizer_steps = {"student": 99}
+        trainer.train_dataloader.position = 99
+        random.seed(101)
+        np.random.seed(102)
+        torch.manual_seed(103)
+
+        restored_step = trainer._load_checkpoint()
+        assert restored_step == 1
+        assert trainer.controller.counters.global_step == 1
+        assert trainer.controller.counters.optimizer_steps == {"student": 1, "fake_score": 1}
+        assert trainer.controller.counters.completed_cycles == 1
+        assert trainer.train_dataloader.position == 7
+        assert executor.role_state["student_optimizer"] == 3
+        assert executor.role_state["fake_scheduler"] == 6
+        assert random.random() == expected_random
+        assert float(np.random.random()) == expected_numpy
+        assert float(torch.rand(())) == expected_torch
+
+        checkpoint = tmp_path / "global_step_1"
+        assert (checkpoint / "manifest.json").is_file()
+        assert (checkpoint / "trainer_state.pt").is_file()
+        assert (checkpoint / "data.pt").is_file()
+        assert (checkpoint / "rng.pt").is_file()
+        assert executor.loaded_path == str(checkpoint / "workers")
+
+    def test_failed_save_never_publishes_a_checkpoint(self, tmp_path):
+        trainer, _ = self.make_trainer(tmp_path, fail_save=True)
+        with pytest.raises(RuntimeError, match="injected save failure"):
+            trainer._save_checkpoint()
+        assert not (tmp_path / "global_step_1").exists()
+        assert not list(tmp_path.glob(".global_step_1_*"))
+        assert not (tmp_path / "latest_checkpointed_iteration.txt").exists()
+
+    def test_equivalent_plan_mappings_have_identical_fingerprints(self, tmp_path):
+        trainer, _ = self.make_trainer(tmp_path)
+        fingerprint = trainer.checkpoint_fingerprint()
+        trainer.plan = replace(trainer.plan, objective=dict(reversed(list(trainer.plan.objective.items()))))
+        assert trainer.checkpoint_fingerprint() == fingerprint
+
+    def test_changed_plan_is_rejected_before_worker_restore(self, tmp_path):
+        trainer, executor = self.make_trainer(tmp_path)
+        trainer._save_checkpoint()
+        trainer.plan = build_plan(
+            "dmd2",
+            {"model_path": "/m", "fake_update_ratio": 2},
+            frozenset({"distribution_matching"}),
+        )
+        with pytest.raises(ValueError, match="does not match the active run"):
+            trainer._load_checkpoint()
+        assert executor.loaded_path is None
+
+    def test_incomplete_checkpoint_is_rejected(self, tmp_path):
+        trainer, _ = self.make_trainer(tmp_path)
+        incomplete = tmp_path / "global_step_1"
+        incomplete.mkdir()
+        (tmp_path / "latest_checkpointed_iteration.txt").write_text("1")
+        with pytest.raises(FileNotFoundError, match="Incomplete distillation checkpoint"):
+            trainer._load_checkpoint()

tests/trainer/diffusion/test_distillation_config_on_cpu.py

@@ -36,6 +36,12 @@ def test_defaults_select_dmd2_without_enabling_opd(self):
         assert config.distribution_matching.recipe == "dmd2"
         assert config.distribution_matching.profile is None
         assert config.distribution_matching.fake_update_ratio is None
+        assert config.distribution_matching.role_storage == "shared_base_adapters"
+        assert config.distribution_matching.student_micro_batch_size_per_gpu == 1
+        assert config.distribution_matching.fake_score_micro_batch_size_per_gpu == 1
+        assert config.distribution_matching.ema_decay == pytest.approx(0.999)
+        assert config.distribution_matching.ema_start_step == 0
+        assert config.distribution_matching.fake_score_optim.lr == pytest.approx(2e-5)
 
     @pytest.mark.parametrize(
         "kwargs,error",
@@ -47,6 +53,12 @@ def test_defaults_select_dmd2_without_enabling_opd(self):
             ({"rollout_strategy": "typo"}, "Invalid rollout_strategy"),
             ({"data_mode": "typo"}, "Invalid data_mode"),
             ({"export_role": "teacher_score"}, "Invalid export_role"),
+            ({"role_storage": "remote"}, "Invalid role_storage"),
+            ({"student_micro_batch_size_per_gpu": 0}, "greater than 0"),
+            ({"fake_score_micro_batch_size_per_gpu": 0}, "greater than 0"),
+            ({"ema_decay": -0.1}, "ema_decay"),
+            ({"ema_decay": 1.1}, "ema_decay"),
+            ({"ema_start_step": -1}, "non-negative"),
         ],
     )
     def test_invalid_values_fail_closed(self, kwargs, error):
@@ -194,6 +206,7 @@ def test_cli_distribution_matching_overrides_do_not_enable_opd(self):
         cfg = self._compose(
             [
                 "algorithm.trainer_type=distillation",
+                "algorithm.sample_source=offline",
                 "distillation.distribution_matching.recipe=dmd2",
                 "distillation.distribution_matching.fake_update_ratio=2",
                 "distillation.distribution_matching.rollout_strategy=consistency_renoise",
@@ -203,13 +216,15 @@ def test_cli_distribution_matching_overrides_do_not_enable_opd(self):
         assert config.enabled is False
         assert config.distribution_matching.fake_update_ratio == 2
         assert config.distribution_matching.rollout_strategy == "consistency_renoise"
+        assert config.distribution_matching.fake_score_optim.lr == pytest.approx(2e-5)
 
     def test_composed_config_builds_validated_plan(self):
         from verl_omni.trainer.diffusion.distillation.recipes import build_plan_from_config
 
         cfg = self._compose(
             [
                 "algorithm.trainer_type=distillation",
+                "algorithm.sample_source=offline",
                 "actor_rollout_ref.model.path=/m",
                 "distillation.distribution_matching.fake_update_ratio=2",
             ]
@@ -225,6 +240,7 @@ def test_null_overrides_use_each_recipe_default(self):
         cfg = self._compose(
             [
                 "algorithm.trainer_type=distillation",
+                "algorithm.sample_source=offline",
                 "actor_rollout_ref.model.path=/m",
                 "distillation.distribution_matching.recipe=dmd",
             ]

tests/trainer/diffusion/test_distillation_contracts_on_cpu.py

@@ -280,6 +280,16 @@ def test_causal_recipes_use_separate_causal_and_bidirectional_groups(self):
             assert role_groups["student"] == role_groups["student_ema"] == "causal_base"
             assert role_groups["teacher_score"] == role_groups["fake_score"] == "bidirectional_base"
 
+    def test_colocated_independent_materializes_one_group_per_role(self):
+        plan = build_plan(
+            "dmd2",
+            {"model_path": "/m", "role_storage": "colocated_independent"},
+            ALL_CAPS,
+        )
+        assert {group.storage for group in plan.role_layout.groups} == {"independent_module"}
+        assert len(plan.role_layout.groups) == len(plan.role_layout.bindings) == 4
+        assert all(binding.group == f"{binding.role}_model" for binding in plan.role_layout.bindings)
+
     @pytest.mark.parametrize(
         "name,config,error",
         [
@@ -291,6 +301,7 @@ def test_causal_recipes_use_separate_causal_and_bidirectional_groups(self):
             ("dmd2", {"fake_update_ratio": 1.5, "model_path": "/m"}, "integer"),
             ("dmd2", {"fake_update_ratio": True, "model_path": "/m"}, "integer"),
             ("dmd2", {"fake_warmup_cycles": 1.5, "model_path": "/m"}, "integer"),
+            ("dmd2", {"role_storage": "remote", "model_path": "/m"}, "role_storage"),
         ],
     )
     def test_invalid_recipe_values_fail_closed(self, name, config, error):
@@ -368,6 +379,20 @@ def test_warmup_requires_fake_only_phases(self):
                 warmup_cycles=1,
             )
 
+    def test_warmup_cannot_reuse_a_missing_student_batch(self):
+        with pytest.raises(ValueError, match="cannot reuse a student batch"):
+            UpdateSchedule(
+                phases=(self.student_phase(), self.fake_phase()),
+                warmup_phases=(
+                    UpdatePhaseSpec(
+                        kind="fake_score",
+                        batch_policy="reuse_student",
+                        trainable_roles=("fake_score",),
+                    ),
+                ),
+                warmup_cycles=1,
+            )
+
     def test_warmup_cycles_require_warmup_phases(self):
         with pytest.raises(ValueError, match="requires at least one warmup phase"):
             UpdateSchedule(phases=(self.student_phase(), self.fake_phase()), warmup_cycles=1)

tests/trainer/diffusion/test_distillation_controller_on_cpu.py

@@ -217,9 +217,33 @@ def test_reset_clears_healthy_driver_state(self):
         assert controller.counters.optimizer_steps == {}
         assert controller.metrics == {}
 
-    def test_failed_driver_cannot_be_reset(self):
+    def test_state_dict_round_trip_restores_completed_counters(self):
+        controller, _, _ = make_controller(make_plan(fake_repeats=2))
+        controller.run(3)
+        state = controller.state_dict()
+
+        restored, _, _ = make_controller(make_plan(fake_repeats=2))
+        restored.load_state_dict(state)
+        assert restored.counters.global_step == 3
+        assert restored.counters.completed_cycles == 3
+        assert restored.counters.optimizer_steps == {"student": 3, "fake_score": 6}
+
+    def test_invalid_checkpoint_state_is_rejected(self):
+        controller, _, _ = make_controller(make_plan())
+        with pytest.raises(ValueError, match="exactly"):
+            controller.load_state_dict({"global_step": 1})
+        with pytest.raises(ValueError, match="non-negative integer"):
+            controller.load_state_dict({"global_step": -1, "optimizer_steps": {}, "completed_cycles": 0})
+        with pytest.raises(ValueError, match="unknown optimizer roles"):
+            controller.load_state_dict({"global_step": 0, "optimizer_steps": {"unknown": 1}, "completed_cycles": 1})
+        with pytest.raises(ValueError, match="must equal global_step"):
+            controller.load_state_dict({"global_step": 2, "optimizer_steps": {"student": 1}, "completed_cycles": 2})
+
+    def test_failed_driver_cannot_be_checkpointed_or_reset(self):
         controller, _, _ = make_controller(make_plan(), executor=FakePhaseExecutor(fail_on="student"))
         with pytest.raises(RuntimeError):
             controller.run_cycle()
+        with pytest.raises(RuntimeError, match="Cannot checkpoint"):
+            controller.state_dict()
         with pytest.raises(RuntimeError, match="cannot be reset"):
             controller.reset()