[ckpt, trainer] fix: persist device and condition-model RNG across DiT resume

.agents/knowledge/testing.md

@@ -55,6 +55,7 @@ by a `pytest.param` table, so a new case is a few lines:
 | `torch.compile` support | `tests/distributed/test_torch_compile.py` |
 | Ulysses SP behaviour | `tests/parallel/ulysses/test_ulysses.py` and siblings |
 | Checkpoint callback / DCP round-trip | `tests/checkpoints/test_checkpoint_callback.py`, `test_trainer_saveload.py` |
+| DiT resume / noise-stream equivalence | `tests/checkpoints/test_dit_resume_align.py` |
 | Weight loading / broadcast / EP shard | `tests/utils/test_rank0_load_and_broadcast_weights.py`, `test_moe_ep_sharded_load_matrix.py` |
 | Grad clipping with ExtraParallel | `tests/utils/test_extra_parallel_clip_grad_norm.py` |
 | FLOPs / MFU accounting | `tests/utils/test_count_flops.py` |

tests/checkpoints/test_checkpoint_callback.py

@@ -21,6 +21,24 @@
     CheckpointCallback,
 )
 from veomni.trainer.callbacks.global_state_callback import GlobalStateCallback
+from veomni.utils.device import get_device_rng_state
+
+
+class _StubConditionModel:
+    """Stand-in for a DiT condition model that owns a noise/timestep generator."""
+
+    def __init__(self, seed: int = 2024):
+        self.generator = torch.Generator(device="cpu")
+        self.generator.manual_seed(seed)
+
+    def rng_state_dict(self):
+        return {"generator": self.generator.get_state()}
+
+    def load_rng_state_dict(self, state):
+        self.generator.set_state(state["generator"])
+
+    def draw_noise(self):
+        return torch.randn(4, generator=self.generator)
 
 
 def _make_mock_trainer(save_path="/tmp/test_ckpt", save_async=False):
@@ -69,6 +87,10 @@ def _make_mock_trainer(save_path="/tmp/test_ckpt", save_async=False):
     # Single-model job: the real manager's class default, which keeps the
     # manifest's module list empty.
     trainer.checkpoint.module_name = ""
+    # Non-DiT trainers have no condition model. Set it explicitly: an
+    # auto-created MagicMock attribute is not picklable and would slip a mock
+    # into the checkpoint payload instead of the ``None`` the loader skips.
+    trainer.condition_model = None
 
     return trainer
 
@@ -319,19 +341,18 @@ def test_the_logged_destination_is_the_one_save_writes(self, mock_helper, mock_d
         assert rebuilt == "/remote/run/global_step_10/model"
         assert rebuilt == manager.save_dir(state)
 
-    def test_save_forwards_lr_scheduler_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps):
+    def test_save_forwards_extra_state_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps):
         trainer = _make_mock_trainer()
         mock_build_ckpt.return_value = MagicMock()
         manager = ModelCheckpointManager(trainer)
 
         manager.save_dcp(TrainerState(global_step=10))
 
         saved = manager.checkpointer.save.call_args.args[1]
-        assert saved["lr_scheduler"] is trainer.lr_scheduler
+        assert set(saved["extra_state"]) == {"lr_scheduler", "condition_model_rng_state"}
         assert saved["optimizer"] is trainer.optimizer
-        assert "extra_state" not in saved
 
-    def test_load_forwards_lr_scheduler_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps):
+    def test_load_forwards_extra_state_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps):
         trainer = _make_mock_trainer()
         trainer.args.train.checkpoint.load_path = "/tmp/ckpt"
         mock_checkpointer = MagicMock()
@@ -341,11 +362,40 @@ def test_load_forwards_lr_scheduler_like_optimizer(self, mock_helper, mock_dist,
         manager.load()
 
         loaded = mock_checkpointer.load.call_args.args[1]
-        assert loaded["lr_scheduler"] is trainer.lr_scheduler
+        assert loaded["extra_state"] == {}
         assert loaded["optimizer"] is trainer.optimizer
         assert trainer.state.global_step == 0
         assert mock_checkpointer.load.call_args.kwargs["parallel_state"] is mock_get_ps.return_value
 
+    def test_extra_state_persists_condition_model_rng(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps):
+        trainer = _make_mock_trainer()
+        trainer.condition_model = _StubConditionModel()
+        mock_build_ckpt.return_value = MagicMock()
+        manager = ModelCheckpointManager(trainer)
+
+        extra_state = manager._extra_state(TrainerState(global_step=10))
+
+        assert set(extra_state["condition_model_rng_state"]) == {"generator"}
+
+    def test_load_extra_state_restores_condition_model_rng(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps):
+        trainer = _make_mock_trainer()
+        condition_model = _StubConditionModel()
+        trainer.condition_model = condition_model
+        mock_build_ckpt.return_value = MagicMock()
+        manager = ModelCheckpointManager(trainer)
+
+        for _ in range(3):
+            condition_model.draw_noise()
+        rng_state = condition_model.rng_state_dict()
+        reference = torch.Generator(device="cpu")
+        reference.set_state(condition_model.generator.get_state())
+        expected_next_noise = torch.randn(4, generator=reference)
+
+        condition_model.draw_noise()
+        manager._load_extra_state({"lr_scheduler": {"lr": 1e-4}, "condition_model_rng_state": rng_state})
+
+        assert torch.equal(condition_model.draw_noise(), expected_next_noise)
+
     def test_save_lora_writes_the_adapter_to_its_own_export_dir(
         self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps, tmp_path
     ):
@@ -442,6 +492,56 @@ def test_state_dict_includes_channel_loss_callback_state(self, mock_dist):
         assert "train_dataloader" in global_state
         assert "environ_meter" in global_state
         assert "torch_rng_state" in global_state
+        assert "device_rng_state" in global_state
+
+    def test_state_dict_has_no_condition_model_rng(self, mock_dist):
+        trainer = _make_mock_trainer()
+        cb = GlobalStateCallback(trainer)
+
+        global_state = cb.state_dict(TrainerState(global_step=10))
+
+        assert "condition_model_rng_state" not in global_state
+
+    def test_load_resumes_device_and_cpu_rng_streams(self, mock_dist, tmp_path):
+        """A resumed run must continue the device and CPU rng streams, not replay."""
+        mock_dist.is_initialized.return_value = False
+        trainer = _make_mock_trainer()
+        trainer.args.train.checkpoint.load_path = str(tmp_path)
+        trainer.train_dataloader = None
+        trainer.data_iterator = None
+        trainer.environ_meter.state_dict.return_value = {}
+        cb = GlobalStateCallback(trainer)
+        payload = cb.state_dict(TrainerState(global_step=3))
+        torch.save(payload, tmp_path / "trainer_state_rank_0.pt")
+
+        # Drift both streams the way the rest of a continuing run would.
+        torch.randn(8)
+
+        assert GlobalStateCallback(trainer).load_global_state() is not None
+
+        assert torch.equal(get_device_rng_state(), payload["device_rng_state"])
+        assert torch.equal(torch.get_rng_state(), payload["torch_rng_state"])
+
+    def test_load_tolerates_checkpoint_without_rng_fields(self, mock_dist, tmp_path):
+        """Checkpoints written before the rng fields existed must still load."""
+        mock_dist.is_initialized.return_value = False
+        trainer = _make_mock_trainer()
+        trainer.args.train.checkpoint.load_path = str(tmp_path)
+        trainer.train_dataloader = None
+        torch.save(
+            {
+                "global_step": 7,
+                "train_dataloader": None,
+                "environ_meter": {},
+                "channel_loss_callback": {},
+                "torch_rng_state": torch.get_rng_state(),
+            },
+            tmp_path / "trainer_state_rank_0.pt",
+        )
+
+        GlobalStateCallback(trainer).load_global_state()
+
+        assert trainer.state.global_step == 7
 
     def test_save_does_not_wait_for_a_pending_dcp(self, mock_dist, tmp_path):
         """The cursor does not depend on the shards, and the manifest beside it