diff --git a/.agents/knowledge/testing.md b/.agents/knowledge/testing.md
index 24b42e4501..8a818492d6 100644
--- a/.agents/knowledge/testing.md
+++ b/.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` |
diff --git a/tests/checkpoints/test_checkpoint_callback.py b/tests/checkpoints/test_checkpoint_callback.py
index 3f75408adf..b2941f9ddd 100644
--- a/tests/checkpoints/test_checkpoint_callback.py
+++ b/tests/checkpoints/test_checkpoint_callback.py
@@ -81,6 +81,8 @@ def _make_mock_runtime(save_path="/tmp/test_ckpt", save_async=False):
         parallel_state=SimpleNamespace(global_rank=trainer.args.train.global_rank),
         args=trainer.args.model,
         train_args=trainer.args.train,
+        extra_state=MagicMock(return_value={}),
+        load_extra_state=MagicMock(),
     )
 
 
@@ -340,7 +342,7 @@ class Named(ModelCheckpointManager):
         named = Named(runtime)
         assert named.assets_dir() == f"{runtime.train_args.checkpoint.model_assets_dir}/vision_encoder"
 
-    def test_save_forwards_lr_scheduler_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt):
+    def test_save_forwards_extra_state_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt):
         runtime = _make_mock_runtime()
         mock_build_ckpt.return_value = MagicMock()
         manager = ModelCheckpointManager(runtime)
@@ -348,11 +350,10 @@ def test_save_forwards_lr_scheduler_like_optimizer(self, mock_helper, mock_dist,
         manager.save_dcp(TrainerState(global_step=10))
 
         saved = manager.checkpointer.save.call_args.args[1]
-        assert saved["lr_scheduler"] is runtime.lr_scheduler
+        assert set(saved["extra_state"]) == {"lr_scheduler"}
         assert saved["optimizer"] is runtime.optimizer
-        assert "extra_state" not in saved
 
-    def test_load_forwards_lr_scheduler_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt):
+    def test_load_forwards_extra_state_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt):
         runtime = _make_mock_runtime()
         runtime.train_args.checkpoint.load_path = "/tmp/ckpt"
         mock_checkpointer = MagicMock()
@@ -362,7 +363,7 @@ 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 runtime.lr_scheduler
+        assert loaded["extra_state"] == {}
         assert loaded["optimizer"] is runtime.optimizer
         assert mock_checkpointer.load.call_args.kwargs["parallel_state"] is runtime.parallel_state
 
diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py
index 61735ab681..eebc2b3ded 100644
--- a/tests/checkpoints/test_dcp_checkpointer.py
+++ b/tests/checkpoints/test_dcp_checkpointer.py
@@ -1,7 +1,7 @@
 """Unit tests for distributed checkpoint and resume behavior.
 
 Covers: OptimizerState (no placeholder synthesis), key normalization,
-lr_scheduler sidecar persistence, allow_partial_load planner, skip-HF resume, and
+per-rank extra_state persistence, allow_partial_load planner, skip-HF resume, and
 trainer step-counting correctness. Tests marked ``xfail`` document known
 in-tree bugs — they become regression guards once the fix lands.
 """
@@ -285,14 +285,14 @@ def test_weights_and_optimizer_are_separable_on_disk(self, mock_gps, tmp_path):
 
         DistributedCheckpointer.save(
             path=str(tmp_path),
-            state={"model": model, "optimizer": optimizer, "lr_scheduler": scheduler},
+            state={"model": model, "optimizer": optimizer, "extra_state": {"lr_scheduler": scheduler.state_dict()}},
             global_steps=7,
         )
 
         model_root = tmp_path / "global_step_7" / "model"
         assert (model_root / "ckpt" / ".metadata").is_file()
         assert (model_root / "optimizer" / ".metadata").is_file()
-        assert (model_root / "lr_scheduler.pt").is_file()
+        assert (model_root / "extra_state" / "extra_state_rank_0.pt").is_file()
         # The pre-split marker at the step root is gone; only the manifest, which
         # GlobalStateCallback writes, marks a step complete now.
         assert not (tmp_path / "global_step_7" / ".metadata").exists()
@@ -312,7 +312,7 @@ def test_resume_restores_weights_optimizer_and_scheduler(self, mock_gps, tmp_pat
 
         DistributedCheckpointer.save(
             path=str(tmp_path),
-            state={"model": model, "optimizer": optimizer, "lr_scheduler": scheduler},
+            state={"model": model, "optimizer": optimizer, "extra_state": {"lr_scheduler": scheduler.state_dict()}},
             global_steps=7,
         )
 
@@ -320,10 +320,9 @@ def test_resume_restores_weights_optimizer_and_scheduler(self, mock_gps, tmp_pat
         resumed, resumed_optimizer, resumed_scheduler = self._build(seed=1)
         assert not torch.equal(resumed.weight, model.weight)
 
-        DistributedCheckpointer.load(
-            path=str(tmp_path / "global_step_7"),
-            state={"model": resumed, "optimizer": resumed_optimizer, "lr_scheduler": resumed_scheduler},
-        )
+        state = {"model": resumed, "optimizer": resumed_optimizer, "extra_state": {}}
+        DistributedCheckpointer.load(path=str(tmp_path / "global_step_7"), state=state)
+        resumed_scheduler.load_state_dict(state["extra_state"]["lr_scheduler"])
 
         torch.testing.assert_close(resumed.weight, model.weight, atol=0.0, rtol=0.0)
         expected_exp_avg = optimizer.state_dict()["state"][0]["exp_avg"]
@@ -341,7 +340,11 @@ def test_modules_of_one_job_do_not_share_a_directory(self, mock_gps, tmp_path):
             model, optimizer, scheduler = self._build(seed=index)
             DistributedCheckpointer.save(
                 path=str(tmp_path),
-                state={"model": model, "optimizer": optimizer, "lr_scheduler": scheduler},
+                state={
+                    "model": model,
+                    "optimizer": optimizer,
+                    "extra_state": {"lr_scheduler": scheduler.state_dict()},
+                },
                 global_steps=7,
                 module=module,
             )
@@ -351,6 +354,7 @@ def test_modules_of_one_job_do_not_share_a_directory(self, mock_gps, tmp_path):
         for module in ("vision", "llm"):
             assert (model_root / module / "ckpt" / ".metadata").is_file()
             assert (model_root / module / "optimizer" / ".metadata").is_file()
+            assert (model_root / module / "extra_state" / "extra_state_rank_0.pt").is_file()
 
     def test_resume_reads_a_pre_split_checkpoint(self, mock_gps, tmp_path):
         """A checkpoint written before the split keeps both in one directory.
@@ -358,7 +362,6 @@ def test_resume_reads_a_pre_split_checkpoint(self, mock_gps, tmp_path):
         Delete this along with ``veomni/checkpoint/legacy_v0_1_12.py``.
         """
         from veomni.checkpoint.dcp_checkpointer import (
-            _LR_SCHEDULER_FILENAME,
             DistributedCheckpointer,
             ModelState,
             OptimizerState,
@@ -377,16 +380,18 @@ def test_resume_reads_a_pre_split_checkpoint(self, mock_gps, tmp_path):
             },
             checkpoint_id=step,
         )
-        torch.save(scheduler.state_dict(), step / _LR_SCHEDULER_FILENAME)
 
-        resumed, resumed_optimizer, resumed_scheduler = self._build(seed=1)
+        resumed, resumed_optimizer, _ = self._build(seed=1)
         DistributedCheckpointer.load(
             path=str(step),
-            state={"model": resumed, "optimizer": resumed_optimizer, "lr_scheduler": resumed_scheduler},
+            state={"model": resumed, "optimizer": resumed_optimizer, "extra_state": {}},
         )
 
         torch.testing.assert_close(resumed.weight, model.weight, atol=0.0, rtol=0.0)
-        assert resumed_scheduler.state_dict()["last_epoch"] == scheduler.state_dict()["last_epoch"]
+        expected_exp_avg = optimizer.state_dict()["state"][0]["exp_avg"]
+        torch.testing.assert_close(
+            resumed_optimizer.state_dict()["state"][0]["exp_avg"], expected_exp_avg, atol=0.0, rtol=0.0
+        )
 
 
 class TestAllowPartialLoad:
@@ -408,11 +413,11 @@ def test_load_passes_partial_planner_to_dcp(self, mock_dcp, mock_gps):
         model._fqn2spec_info = None
         optimizer = MagicMock()
 
-        state = {"model": model, "optimizer": optimizer, "lr_scheduler": MagicMock()}
+        state = {"model": model, "optimizer": optimizer, "extra_state": {}}
 
         mock_dcp.load = MagicMock()
 
-        with patch.object(DistributedCheckpointer, "_load_lr_scheduler"):
+        with patch.object(DistributedCheckpointer, "_load_extra_state"):
             with patch.object(DistributedCheckpointer, "_create_storage_reader") as mock_reader:
                 mock_reader.return_value = MagicMock()
                 DistributedCheckpointer.load(path="/fake", state=state)
@@ -1387,119 +1392,44 @@ def test_peft_lora_base_model_key(self):
 
 
 @patch("veomni.checkpoint.dcp_checkpointer.dist")
-class TestLrSchedulerSaveLoad:
+class TestExtraStateSaveLoad:
     def test_roundtrip(self, mock_dist, tmp_path):
         mock_dist.is_initialized.return_value = False
-        mock_dist.get_rank.return_value = 0
-        from veomni.checkpoint.dcp_checkpointer import (
-            _LR_SCHEDULER_FILENAME,
-            DistributedCheckpointer,
-        )
-
-        saved_scheduler = MagicMock()
-        saved_scheduler.state_dict.return_value = {"last_epoch": 10, "base_lrs": [1e-4]}
-        DistributedCheckpointer._save_lr_scheduler(str(tmp_path), {"lr_scheduler": saved_scheduler})
+        from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
 
-        assert (tmp_path / _LR_SCHEDULER_FILENAME).is_file()
-        assert {p.name for p in tmp_path.iterdir()} == {_LR_SCHEDULER_FILENAME}
+        payload = {"lr_scheduler": {"last_epoch": 10}, "condition_model_rng_state": {"generator": "state"}}
+        DistributedCheckpointer._save_extra_state(str(tmp_path), {"extra_state": payload})
 
-        loaded_scheduler = MagicMock()
-        DistributedCheckpointer._load_lr_scheduler(str(tmp_path), {"lr_scheduler": loaded_scheduler})
+        assert (tmp_path / "extra_state" / "extra_state_rank_0.pt").is_file()
+        assert {p.name for p in tmp_path.iterdir()} == {"extra_state"}
 
-        loaded_scheduler.load_state_dict.assert_called_once_with({"last_epoch": 10, "base_lrs": [1e-4]})
+        state = {"extra_state": {}}
+        DistributedCheckpointer._load_extra_state(str(tmp_path), state)
+        assert state["extra_state"] == payload
 
-    def test_only_rank_zero_writes(self, mock_dist, tmp_path):
+    def test_each_rank_writes_its_own_file(self, mock_dist, tmp_path):
         mock_dist.is_initialized.return_value = True
         mock_dist.get_rank.return_value = 3
-        from veomni.checkpoint.dcp_checkpointer import (
-            _LR_SCHEDULER_FILENAME,
-            DistributedCheckpointer,
-        )
-
-        saved_scheduler = MagicMock()
-        saved_scheduler.state_dict.return_value = {"last_epoch": 10}
-        with patch(
-            "veomni.checkpoint.dcp_checkpointer.any_rank_failed", side_effect=lambda failed, group=None: failed
-        ):
-            DistributedCheckpointer._save_lr_scheduler(str(tmp_path), {"lr_scheduler": saved_scheduler})
-
-        saved_scheduler.state_dict.assert_not_called()
-        assert not (tmp_path / _LR_SCHEDULER_FILENAME).exists()
-
-    def test_missing_lr_scheduler_key_save(self, mock_dist, tmp_path):
-        mock_dist.is_initialized.return_value = False
-        from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
-
-        state = {"model": MagicMock()}
-        DistributedCheckpointer._save_lr_scheduler(str(tmp_path), state)
-
-    def test_missing_lr_scheduler_key_load(self, mock_dist, tmp_path):
-        mock_dist.is_initialized.return_value = False
-        from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
-
-        state = {"model": MagicMock()}
-        DistributedCheckpointer._load_lr_scheduler(str(tmp_path), state)
-
-    def test_none_scheduler_is_a_no_op(self, mock_dist, tmp_path):
-        mock_dist.is_initialized.return_value = False
         from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
 
-        DistributedCheckpointer._save_lr_scheduler(str(tmp_path), {"lr_scheduler": None})
-        DistributedCheckpointer._load_lr_scheduler(str(tmp_path), {"lr_scheduler": None})
-
-    def test_missing_sidecar_raises_when_scheduler_is_expected(self, mock_dist, tmp_path):
-        mock_dist.is_initialized.return_value = False
-        from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
+        DistributedCheckpointer._save_extra_state(str(tmp_path), {"extra_state": {"rng": "state"}})
 
-        loaded = MagicMock()
-        with pytest.raises(FileNotFoundError, match="lr_scheduler sidecar"):
-            DistributedCheckpointer._load_lr_scheduler(str(tmp_path), {"lr_scheduler": loaded})
-        loaded.load_state_dict.assert_not_called()
+        assert (tmp_path / "extra_state" / "extra_state_rank_3.pt").is_file()
 
-    def test_load_falls_back_to_extra_state_scheduler(self, mock_dist, tmp_path):
+    def test_missing_extra_state_key_skips(self, mock_dist, tmp_path):
         mock_dist.is_initialized.return_value = False
-        mock_dist.get_rank.return_value = 0
         from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
-        from veomni.checkpoint.legacy_v0_1_12 import extra_state_path
-
-        extra = tmp_path / "extra_state"
-        extra.mkdir()
-        torch.save({"lr_scheduler": {"last_epoch": 4}}, extra_state_path(str(tmp_path), 0))
 
-        loaded = MagicMock()
-        DistributedCheckpointer._load_lr_scheduler(str(tmp_path), {"lr_scheduler": loaded})
-        loaded.load_state_dict.assert_called_once_with({"last_epoch": 4})
+        DistributedCheckpointer._save_extra_state(str(tmp_path), {"model": MagicMock()})
+        DistributedCheckpointer._load_extra_state(str(tmp_path), {"model": MagicMock()})
 
-    def test_sidecar_wins_over_extra_state(self, mock_dist, tmp_path):
+    def test_load_missing_file_keeps_state_empty(self, mock_dist, tmp_path):
         mock_dist.is_initialized.return_value = False
-        from veomni.checkpoint.dcp_checkpointer import (
-            _LR_SCHEDULER_FILENAME,
-            DistributedCheckpointer,
-        )
-        from veomni.checkpoint.legacy_v0_1_12 import extra_state_path
-
-        torch.save({"last_epoch": 1}, tmp_path / _LR_SCHEDULER_FILENAME)
-        extra = tmp_path / "extra_state"
-        extra.mkdir()
-        torch.save({"lr_scheduler": {"last_epoch": 99}}, extra_state_path(str(tmp_path), 0))
-
-        loaded = MagicMock()
-        DistributedCheckpointer._load_lr_scheduler(str(tmp_path), {"lr_scheduler": loaded})
-        loaded.load_state_dict.assert_called_once_with({"last_epoch": 1})
-
-    def test_rank_nonzero_reads_rank0_extra_state_for_scheduler(self, mock_dist, tmp_path):
-        mock_dist.is_initialized.return_value = True
-        mock_dist.get_rank.return_value = 3
         from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
-        from veomni.checkpoint.legacy_v0_1_12 import extra_state_path
-
-        extra = tmp_path / "extra_state"
-        extra.mkdir()
-        torch.save({"lr_scheduler": {"last_epoch": 8}}, extra_state_path(str(tmp_path), 0))
 
-        loaded = MagicMock()
-        DistributedCheckpointer._load_lr_scheduler(str(tmp_path), {"lr_scheduler": loaded})
-        loaded.load_state_dict.assert_called_once_with({"last_epoch": 8})
+        state = {"extra_state": {}}
+        DistributedCheckpointer._load_extra_state(str(tmp_path), state)
+        assert state["extra_state"] == {}
 
 
 class TestPromoteStagedCheckpoint:
@@ -1582,7 +1512,7 @@ def test_nested_sidecar_is_copied_before_metadata(self, tmp_path):
         """A nested file is copied with the shards, before any ``.metadata`` is."""
         import shutil as _shutil
 
-        from veomni.checkpoint.dcp_checkpointer import _LR_SCHEDULER_FILENAME, _promote_staged_checkpoint
+        from veomni.checkpoint.dcp_checkpointer import _promote_staged_checkpoint
 
         stage_path = tmp_path / "stage"
         final_path = tmp_path / "final"
@@ -1590,7 +1520,7 @@ def test_nested_sidecar_is_copied_before_metadata(self, tmp_path):
         nested.mkdir(parents=True)
         (stage_path / "__0_0.distcp").write_text("weights")
         (stage_path / ".metadata").write_text("meta")
-        (stage_path / _LR_SCHEDULER_FILENAME).write_text("scheduler-v2")
+        (stage_path / "extra_state.pt").write_text("extra-state")
         (nested / "sidecar.pt").write_text("nested")
 
         order = []
@@ -1606,12 +1536,12 @@ def spy(src, dst):
                 _promote_staged_checkpoint(str(stage_path), str(final_path))
 
         nested_rel = os.path.join("nested", "sidecar.pt")
-        assert _LR_SCHEDULER_FILENAME in order
+        assert "extra_state.pt" in order
         assert nested_rel in order
         assert order[-1] == ".metadata"
-        assert order.index(_LR_SCHEDULER_FILENAME) < order.index(".metadata")
+        assert order.index("extra_state.pt") < order.index(".metadata")
         assert order.index(nested_rel) < order.index(".metadata")
-        assert (final_path / _LR_SCHEDULER_FILENAME).read_text() == "scheduler-v2"
+        assert (final_path / "extra_state.pt").read_text() == "extra-state"
         assert (final_path / nested_rel).read_text() == "nested"
 
     def test_stale_metadata_is_gone_before_any_data_is_copied(self, staged):
@@ -2081,16 +2011,16 @@ def test_local_staging_failure_raises_its_own_error(self, tmp_path):
                 with pytest.raises(OSError, match="No space left"):
                     _prepare_stage_dir(str(tmp_path), "/remote/ckpt")
 
-    def test_failed_staged_overwrite_does_not_change_previous_scheduler(self, tmp_path):
-        """A failed staged save must not mutate the live checkpoint's scheduler.
+    def test_failed_staged_overwrite_does_not_change_previous_extra_state(self, tmp_path):
+        """A failed staged save must not mutate the live checkpoint's extra state.
 
         ``.metadata`` still advertises the previous checkpoint until promotion.
-        Writing the new ``lr_scheduler.pt`` into the destination first would let
-        a resume load the old model and optimizer with the new scheduler; a
+        Writing the new ``extra_state`` into the destination first would let a
+        resume load the old model and optimizer with the new scheduler/rng; a
         failed ``dcp.save`` would leave the same mix. The sidecar has to live
         under ``stage_path`` until promotion copies it with the shards.
         """
-        from veomni.checkpoint.dcp_checkpointer import _LR_SCHEDULER_FILENAME, DistributedCheckpointer
+        from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
 
         final = tmp_path / "ckpt"
         model_root = final / "global_step_10" / "model"
@@ -2098,13 +2028,12 @@ def test_failed_staged_overwrite_does_not_change_previous_scheduler(self, tmp_pa
         weights.mkdir(parents=True)
         (weights / ".metadata").write_text("previous")
         (weights / "__0_0.distcp").write_text("old-weights")
-        sidecar = model_root / _LR_SCHEDULER_FILENAME
-        previous = {"last_epoch": 10, "base_lrs": [1e-4]}
+        extra_state_dir = model_root / "extra_state"
+        extra_state_dir.mkdir(parents=True)
+        sidecar = extra_state_dir / "extra_state_rank_0.pt"
+        previous = {"lr_scheduler": {"last_epoch": 10, "base_lrs": [1e-4]}}
         torch.save(previous, sidecar)
 
-        new_scheduler = MagicMock()
-        new_scheduler.state_dict.return_value = {"last_epoch": 99, "base_lrs": [1e-3]}
-
         with (
             patch("veomni.checkpoint.dcp_checkpointer.dist.is_initialized", return_value=False),
             patch("veomni.checkpoint.dcp_checkpointer.dist.get_rank", return_value=0),
@@ -2117,7 +2046,10 @@ def test_failed_staged_overwrite_does_not_change_previous_scheduler(self, tmp_pa
             with pytest.raises(OSError, match="dcp write failed"):
                 DistributedCheckpointer.save(
                     path=str(final),
-                    state={"model": MagicMock(), "lr_scheduler": new_scheduler},
+                    state={
+                        "model": MagicMock(),
+                        "extra_state": {"lr_scheduler": {"last_epoch": 99, "base_lrs": [1e-3]}},
+                    },
                     save_async=False,
                     global_steps=10,
                     stage_dir=str(tmp_path / "stage"),
@@ -2149,7 +2081,7 @@ def test_a_failed_staged_save_over_a_legacy_step_leaves_it_resumable(self, tmp_p
             patch("veomni.checkpoint.dcp_checkpointer.any_rank_failed", side_effect=lambda failed, group=None: failed),
             patch.object(DistributedCheckpointer, "execute_save", side_effect=OSError("dcp write failed")),
             patch.object(DistributedCheckpointer, "_create_storage_writer"),
-            patch.object(DistributedCheckpointer, "_save_lr_scheduler"),
+            patch.object(DistributedCheckpointer, "_save_extra_state"),
             patch("veomni.checkpoint.dcp_checkpointer.ModelState"),
         ):
             with pytest.raises(OSError, match="dcp write failed"):
@@ -2251,7 +2183,7 @@ def test_unset_stage_dir_writes_straight_to_the_destination(self, tmp_path):
         with (
             patch.object(DistributedCheckpointer, "execute_save") as execute_save,
             patch.object(DistributedCheckpointer, "_create_storage_writer") as create_writer,
-            patch.object(DistributedCheckpointer, "_save_lr_scheduler") as save_sched,
+            patch.object(DistributedCheckpointer, "_save_extra_state") as save_sched,
             patch("veomni.checkpoint.dcp_checkpointer.ModelState"),
             patch("veomni.checkpoint.dcp_checkpointer.OptimizerState"),
             patch("veomni.checkpoint.dcp_checkpointer._prepare_stage_dir") as prepare,
@@ -2298,7 +2230,7 @@ def test_rewriting_a_step_leaves_the_manifest_to_its_owner(self, tmp_path):
         with (
             patch.object(DistributedCheckpointer, "execute_save"),
             patch.object(DistributedCheckpointer, "_create_storage_writer"),
-            patch.object(DistributedCheckpointer, "_save_lr_scheduler"),
+            patch.object(DistributedCheckpointer, "_save_extra_state"),
             patch("veomni.checkpoint.dcp_checkpointer.ModelState"),
         ):
             DistributedCheckpointer.save(
@@ -2328,7 +2260,7 @@ def test_rewriting_a_legacy_step_drops_its_marker_too(self, tmp_path):
         with (
             patch.object(DistributedCheckpointer, "execute_save"),
             patch.object(DistributedCheckpointer, "_create_storage_writer"),
-            patch.object(DistributedCheckpointer, "_save_lr_scheduler"),
+            patch.object(DistributedCheckpointer, "_save_extra_state"),
             patch("veomni.checkpoint.dcp_checkpointer.ModelState"),
         ):
             DistributedCheckpointer.save(
@@ -2363,7 +2295,7 @@ def test_rewriting_a_step_drops_the_modules_own_markers(self, tmp_path):
                 side_effect=lambda **kw: seen.append([(d / layout.DCP_MARKER_FILENAME).exists() for d in markers]),
             ),
             patch.object(DistributedCheckpointer, "_create_storage_writer"),
-            patch.object(DistributedCheckpointer, "_save_lr_scheduler"),
+            patch.object(DistributedCheckpointer, "_save_extra_state"),
             patch("veomni.checkpoint.dcp_checkpointer.ModelState"),
             patch("veomni.checkpoint.dcp_checkpointer.OptimizerState"),
         ):
@@ -2392,7 +2324,7 @@ def test_rewriting_a_step_leaves_another_modules_markers_alone(self, tmp_path):
         with (
             patch.object(DistributedCheckpointer, "execute_save"),
             patch.object(DistributedCheckpointer, "_create_storage_writer"),
-            patch.object(DistributedCheckpointer, "_save_lr_scheduler"),
+            patch.object(DistributedCheckpointer, "_save_extra_state"),
             patch("veomni.checkpoint.dcp_checkpointer.ModelState"),
         ):
             DistributedCheckpointer.save(
@@ -2438,7 +2370,7 @@ def test_save_without_optimizer_writes_only_the_weights(self, tmp_path):
         with (
             patch.object(DistributedCheckpointer, "execute_save") as execute_save,
             patch.object(DistributedCheckpointer, "_create_storage_writer"),
-            patch.object(DistributedCheckpointer, "_save_lr_scheduler"),
+            patch.object(DistributedCheckpointer, "_save_extra_state"),
             patch("veomni.checkpoint.dcp_checkpointer.ModelState"),
         ):
             DistributedCheckpointer.save(
@@ -2468,7 +2400,7 @@ def test_staging_runs_on_the_stage_group(self, tmp_path):
         with (
             patch.object(DistributedCheckpointer, "execute_save") as execute_save,
             patch.object(DistributedCheckpointer, "_create_storage_writer"),
-            patch.object(DistributedCheckpointer, "_save_lr_scheduler"),
+            patch.object(DistributedCheckpointer, "_save_extra_state"),
             patch.object(DistributedCheckpointer, "_get_stage_process_group", return_value=group) as get_group,
             patch("veomni.checkpoint.dcp_checkpointer.ModelState"),
             patch(
diff --git a/tests/checkpoints/test_dit_resume_align.py b/tests/checkpoints/test_dit_resume_align.py
new file mode 100644
index 0000000000..2a5fc2eed1
--- /dev/null
+++ b/tests/checkpoints/test_dit_resume_align.py
@@ -0,0 +1,423 @@
+# Copyright 2025 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.
+
+"""DiT resume equivalence: the noise/timestep stream must continue, not restart.
+
+A diffusion trainer draws its noise and timestep ids inside
+``condition_model.process_condition``, either from a per-run ``torch.Generator``
+(Wan, Qwen-Image) or from the accelerator's default RNG (LTX2, MiniMax-H3). Neither is model state, so the job-level checkpoint has to carry both
+across a resume. When it does not, the resumed process — whose condition model
+was just constructed and re-seeded — replays the stream from the *start* of the
+run at the first resumed step.
+
+This runs the real ``DiTTrainer`` twice against the same dummy dataset:
+
+* **run A** trains ``TOTAL_STEPS`` steps uninterrupted and records a signature of
+  the noise/timestep drawn at every step;
+* **run B** loads run A's mid-run checkpoint, so it starts past the checkpoint
+  and its condition model is built fresh in a new process, never having seen the
+  steps it resumes past.
+
+Run B's signatures must equal run A's on the resumed steps. The stand-in
+condition model is what makes this observable: the real Wan one needs gated
+VAE/text-encoder weights, but the property under test is the RNG contract, so
+the stand-in reproduces it — a generator seeded at construction, plus per-step
+draws from both that generator and the device default RNG.
+
+Dropout is on by default (``VEOMNI_ALIGN_DROPOUT_P``) in both the condition model
+and the forward path. That is deliberate, because it is the case most likely to
+look like a hole in the fix:
+
+* Dropout adds device-RNG consumers, but it does **not** break resume on its own.
+  The checkpoint restores the whole device generator *state*, so as long as the
+  two runs consume the stream in the same order, the dropout masks come back
+  identical. Verified with ``p=0.5`` in both places.
+* What does break it is an *asymmetric* consumer: anything that draws from the
+  device stream after the restore point without a counterpart at the same
+  logical position in the uninterrupted run (a random buffer init, a random
+  sampler or augmentation, an extra warmup forward). Measured with such an
+  injected draw: the generator-backed ``timestep_id`` still matched (that stream
+  is restored exactly), while the dropout-affected ``noise_sum``/``device_sum``
+  diverged.
+
+So the invariant this guards is consumption-order symmetry from the restore
+point onward — not the absence of dropout.
+
+Note that ``tests/train_scripts/train_dit_test.py`` sets the condition model to
+``None`` and skips ``process_condition`` entirely, which is why the resume path
+had no coverage before.
+"""
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+from typing import Any, Dict
+
+import pytest
+import torch
+
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from tools.data_generators import DummyDataset  # noqa: E402
+from tools.launch_utils import find_free_port  # noqa: E402
+
+
+TOTAL_STEPS = 6
+SAVE_STEPS = 3
+# 1-based global_step numbering, so these are the steps after the checkpoint.
+RESUMED_STEPS = list(range(SAVE_STEPS + 1, TOTAL_STEPS + 1))
+SIGNATURE_FILE = "step_signatures.json"
+
+# Mirrors ``WanTransformer3DConditionModel``: same seed derivation, same
+# generator, and the same two randomness sources per step.
+STANDIN_SEED = 2024
+TIMESTEP_CHOICES = 1000
+
+# Dropout is the interesting second-order case: it is never disabled on the
+# frozen condition model (only the model itself is put in ``train()`` and the
+# condition model is merely ``requires_grad_(False)``), and every active dropout
+# consumes the *device* RNG on top of the per-step ``process_condition`` draws.
+# Restoring the device RNG state is therefore not enough on its own — the two
+# runs must also consume the stream in the same order. Keeping dropout on here
+# makes the test guard that, not just the "one draw per step" case.
+DROPOUT_P = float(os.environ.get("VEOMNI_ALIGN_DROPOUT_P", "0.5"))
+
+
+class _AlignConditionModel:
+    """Stands in for a diffusion condition model's noise/timestep sampling."""
+
+    def __init__(self, seed: int, dp_rank: int = 0) -> None:
+        self.generator = torch.Generator(device=torch.device("cpu"))
+        self.generator.manual_seed(seed + dp_rank)
+        self.dropout = torch.nn.Dropout(DROPOUT_P)
+
+    def rng_state_dict(self) -> Dict[str, torch.Tensor]:
+        return {"generator": self.generator.get_state()}
+
+    def load_rng_state_dict(self, state: Dict[str, torch.Tensor]) -> None:
+        self.generator.set_state(state["generator"])
+
+    def process_condition(self, hidden_states=None, encoder_hidden_states=None, latents=None, **kwargs):
+        # The DiT model consumes one tensor per sample (it zips these lists in
+        # its forward), so the outputs keep the same list layout as the real
+        # condition models.
+        reference = hidden_states if hidden_states is not None else latents
+        context = encoder_hidden_states if encoder_hidden_states is not None else [None] * len(reference)
+
+        noises, timesteps, device_sums = [], [], []
+        for sample in reference:
+            noise = torch.randn(sample.shape, dtype=sample.dtype, generator=self.generator).to(sample.device)
+            timestep_id = int(torch.randint(0, TIMESTEP_CHOICES, (1,), generator=self.generator).item())
+            # A second draw from the device default RNG, so the accelerator
+            # stream is part of the signature too. This is the branch that
+            # LTX2 / MiniMax-H3 rely on.
+            device_draw = torch.randn(4, device=sample.device)
+            # Active dropout (the condition model is never ``eval()``-ed) draws
+            # another slice out of the same device stream, before backward.
+            noise = self.dropout(noise)
+            noises.append(noise)
+            timesteps.append(torch.full((sample.shape[0],), timestep_id / TIMESTEP_CHOICES, device=sample.device))
+            device_sums.append(float(device_draw.float().sum()))
+
+        return {
+            "latents": latents,
+            "hidden_states": noises,
+            "training_target": noises,
+            "timestep": timesteps,
+            "encoder_hidden_states": list(context),
+            "_signature": {
+                "noise_sum": float(sum(noise.float().sum() for noise in noises)),
+                "noise_first": float(noises[0].flatten()[0].float()),
+                "timestep_id": int(round(timesteps[0][0].item() * TIMESTEP_CHOICES)),
+                "device_sum": float(sum(device_sums)),
+            },
+        }
+
+
+def _build_trainer_class():
+    """Import the trainer lazily so pytest can import this module standalone."""
+    from veomni.trainer.callbacks import Callback, TrainerState
+    from veomni.trainer.dit_trainer import DiTModelRuntime, DiTTrainer, VeOmniDiTArguments, VeOmniModelRuntime
+
+    class SignatureCallback(Callback):
+        def __init__(self, trainer) -> None:
+            super().__init__(trainer)
+            self.signatures: Dict[str, Dict[str, float]] = {}
+
+        def on_step_end(self, state: TrainerState, **kwargs) -> None:
+            signature = getattr(self.trainer, "_last_signature", None)
+            if signature is not None:
+                self.signatures[str(state.global_step)] = signature
+                self.trainer._last_signature = None
+
+        def on_train_end(self, state: TrainerState, **kwargs) -> None:
+            if self.trainer.args.train.global_rank != 0:
+                return
+            output_dir = self.trainer.args.train.checkpoint.output_dir
+            os.makedirs(output_dir, exist_ok=True)
+            with open(os.path.join(output_dir, SIGNATURE_FILE), "w") as handle:
+                json.dump(self.signatures, handle, indent=2, sort_keys=True)
+
+    class AlignDiTModelRuntime(DiTModelRuntime):
+        """A DiT runtime whose condition model is the RNG stand-in.
+
+        The condition model is built by the runtime, not the trainer, so
+        swapping in the stand-in means overriding here. Its per-rank seed
+        derivation and the ``rng_state_dict``/``load_rng_state_dict`` pair are
+        what the resume path under test actually persists.
+        """
+
+        def _build_condition_model(self, condition_model_type: str) -> None:
+            from veomni.distributed.parallel_state import get_parallel_state
+
+            # Same per-rank seed derivation as the real condition models.
+            self.condition_model = _AlignConditionModel(seed=STANDIN_SEED, dp_rank=get_parallel_state().dp_rank)
+
+        def _freeze_model_module(self) -> None:
+            """Freeze the DiT only.
+
+            The real runtime freezes the condition model alongside it, but the
+            stand-in is deliberately not an ``nn.Module`` and owns no parameters.
+            """
+            VeOmniModelRuntime._freeze_model_module(self)
+
+    class ResumeAlignDiTTrainer(DiTTrainer):
+        def __init__(self, args: "VeOmniDiTArguments"):
+            args.train.training_task = "offline_training"
+            super().__init__(args)
+            # Callbacks are handed the reused ``BaseTrainer`` instance, so the
+            # per-step signature has to live there to be observable.
+            self.base._last_signature = None
+            self._signature_callback = SignatureCallback(self.base)
+            # Stands in for dropout inside the DiT transformer: another device-RNG
+            # consumer on every forward pass.
+            self._model_dropout = torch.nn.Dropout(DROPOUT_P)
+
+        def _build_model_runtime(self) -> DiTModelRuntime:
+            return AlignDiTModelRuntime(self.base.args.model, "base", train=self.base.args.train)
+
+        def _build_data_transform(self) -> None:
+            def process_dummy_example(example: dict, **kwargs):
+                return [{key: torch.tensor(value) for key, value in example.items()}]
+
+            self.base.data_transform = process_dummy_example
+
+        def forward_backward_step(self, micro_batch: Dict[str, Any]) -> tuple:
+            from veomni.distributed.parallel_state import use_parallel_state
+
+            micro_batch = self.preforward(micro_batch)
+            with torch.no_grad():
+                micro_batch = self.condition_model.process_condition(**micro_batch)
+            self.base._last_signature = micro_batch.pop("_signature")
+
+            # Model-side dropout, i.e. a device-RNG consumer in the forward path.
+            micro_batch["hidden_states"] = [self._model_dropout(h) for h in micro_batch["hidden_states"]]
+
+            with use_parallel_state("base"), self.base.model_fwd_context:
+                outputs = self.base.model(**micro_batch)
+            loss, loss_dict = self.postforward(outputs, micro_batch)
+            with use_parallel_state("base"), self.base.model_bwd_context:
+                loss.backward()
+            return loss, loss_dict
+
+        def on_step_end(self, loss=None, loss_dict=None, grad_norm=None, aux_metrics=None):
+            self._signature_callback.on_step_end(self.base.state)
+            super().on_step_end(loss=loss, loss_dict=loss_dict, grad_norm=grad_norm, aux_metrics=aux_metrics)
+
+        def on_train_end(self):
+            super().on_train_end()
+            self._signature_callback.on_train_end(self.base.state)
+
+    return ResumeAlignDiTTrainer
+
+
+def main():
+    from veomni.arguments import parse_args
+    from veomni.trainer.dit_trainer import VeOmniDiTArguments
+
+    args: VeOmniDiTArguments = parse_args(VeOmniDiTArguments)
+    _build_trainer_class()(args).train()
+
+
+def _output_dir(run: str) -> str:
+    return os.path.abspath(f"./test_dit_resume_align_{run}")
+
+
+def _checkpoint_dir() -> str:
+    return os.path.join(_output_dir("full"), "checkpoints", f"global_step_{SAVE_STEPS}")
+
+
+def _materialize_toy_dit_weights(target_dir: str) -> None:
+    """Write HF-format weights for the toy Wan transformer.
+
+    The FSDP2 path materializes a meta-init model through ``init_weights()``,
+    which under transformers v5 reaches ``tie_weights`` and trips over the
+    diffusers-based model's ``_tied_weights_keys`` (v5 renamed it to
+    ``all_tied_weights_keys``). Supplying a ``model_path`` sidesteps that
+    entirely: the weights are loaded instead of randomly initialized. Building
+    the same model directly (no meta device, no FSDP) does not hit the issue.
+    """
+    from veomni.arguments.arguments_types import OpsImplementationConfig
+    from veomni.models.auto import build_foundation_model
+    from veomni.ops import apply_ops_config
+
+    # Wan's ``rope_apply`` has a non-standard signature, so its device patch
+    # explicitly disables the liger RoPE backend. The framework default for
+    # ``rotary_pos_emb_implementation`` is ``liger_kernel``, which would raise
+    # here; every Wan YAML pins ``eager`` for the same reason.
+    apply_ops_config(OpsImplementationConfig(rotary_pos_emb_implementation="eager"))
+    model = build_foundation_model(
+        config_path="tests/toy_config/wan_t2v_toy/config.json",
+        weights_path=None,
+        torch_dtype="float32",
+        init_device="cpu",
+    )
+    model.save_pretrained(target_dir)
+
+
+def _repo_root() -> str:
+    # <repo>/tests/checkpoints/test_dit_resume_align.py -> <repo>
+    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+def _child_env() -> Dict[str, str]:
+    """Environment for the trainer subprocess, pinned to this checkout.
+
+    ``torch.distributed.run`` puts the *script's* directory (not the working
+    directory) on the child's ``sys.path[0]``, so an unrelated ``veomni``
+    installed in the environment — an editable install pointing at some other
+    checkout, say — would be imported instead of the one under test. Putting the
+    repo root first on ``PYTHONPATH`` makes the test validate the code it lives
+    in, whatever the ambient environment says.
+    """
+    env = dict(os.environ)
+    existing = env.get("PYTHONPATH")
+    env["PYTHONPATH"] = _repo_root() + (os.pathsep + existing if existing else "")
+    return env
+
+
+def _run_trainer(run: str, train_path: str, weights_dir: str, load_path: str | None) -> None:
+    output_dir = _output_dir(run)
+    shutil.rmtree(output_dir, ignore_errors=True)
+
+    cmd = [
+        sys.executable,
+        "-m",
+        "torch.distributed.run",
+        "--nnodes=1",
+        "--nproc_per_node=2",
+        f"--master-port={find_free_port()}",
+        os.path.abspath(__file__),
+        "--model.config_path=tests/toy_config/wan_t2v_toy/config.json",
+        f"--model.model_path={weights_dir}",
+        f"--data.train_path={train_path}",
+        "--data.max_seq_len=64",
+        "--train.training_task=offline_training",
+        # FSDP needs world_size > 1 (``fsdp_enabled`` is ``fsdp_size > 1``), so
+        # keep the global batch at the dp size.
+        "--train.global_batch_size=2",
+        "--train.micro_batch_size=1",
+        f"--train.max_steps={TOTAL_STEPS}",
+        "--train.bsz_warmup_ratio=0",
+        "--train.enable_full_determinism=True",
+        "--train.checkpoint.manager=dcp",
+        f"--train.checkpoint.save_steps={SAVE_STEPS}",
+        "--train.checkpoint.save_epochs=0",
+        "--train.checkpoint.save_hf_weights=False",
+        f"--train.checkpoint.output_dir={output_dir}",
+        # Wan's RoPE signature is non-standard, so the liger backend is
+        # explicitly disabled for it; pin the eager backend the Wan YAMLs use.
+        "--model.ops_implementation.rotary_pos_emb_implementation=eager",
+        "--model.accelerator.fsdp_config.fsdp_mode=fsdp2",
+        "--model.accelerator.init_device=meta",
+        # Gradient checkpointing is orthogonal to the RNG question, and its
+        # non-reentrant path forwards ``early_stop`` into the Wan block, which
+        # does not accept it.
+        "--model.accelerator.gradient_checkpointing.enable=False",
+    ]
+    if load_path is not None:
+        cmd.append(f"--train.checkpoint.load_path={load_path}")
+
+    result = subprocess.run(cmd, capture_output=True, text=True, env=_child_env())
+    os.makedirs(output_dir, exist_ok=True)
+    with open(os.path.join(output_dir, "trainer.log"), "w") as handle:
+        handle.write(result.stdout)
+        handle.write("\n----- stderr -----\n")
+        handle.write(result.stderr)
+    if result.returncode != 0:
+        raise AssertionError(
+            f"trainer run {run!r} failed (exit {result.returncode})\n"
+            f"--- stdout tail ---\n{result.stdout[-8000:]}\n"
+            f"--- stderr tail ---\n{result.stderr[-8000:]}"
+        )
+
+
+def _read_signatures(run: str) -> Dict[str, Dict[str, float]]:
+    with open(os.path.join(_output_dir(run), SIGNATURE_FILE)) as handle:
+        return json.load(handle)
+
+
+def test_dit_resume_continues_noise_and_timestep_stream():
+    """A resumed DiT run draws the same noise/timestep as the uninterrupted one."""
+    from veomni.utils.device import get_device_type
+
+    if get_device_type() == "cpu":
+        pytest.skip("needs an accelerator: the trainer cannot run on CPU")
+
+    dummy = DummyDataset(seq_len=64, dataset_type="wan_t2v")
+    weights_dir = os.path.abspath("./test_dit_resume_align_weights")
+    try:
+        shutil.rmtree(weights_dir, ignore_errors=True)
+        _materialize_toy_dit_weights(weights_dir)
+        _run_trainer("full", dummy.save_path, weights_dir, load_path=None)
+        assert os.path.isdir(dummy.save_path), f"dataset vanished after run A: {dummy.save_path}"
+        _run_trainer("resumed", dummy.save_path, weights_dir, load_path=_checkpoint_dir())
+
+        full = _read_signatures("full")
+        resumed = _read_signatures("resumed")
+
+        # Guard the test's own premise: the resumed run must not have replayed
+        # the steps before the checkpoint, or nothing about a resume is exercised.
+        pre_checkpoint = {str(step) for step in range(1, SAVE_STEPS + 1)}
+        assert pre_checkpoint.isdisjoint(resumed), (
+            f"resumed run replayed pre-checkpoint steps, so it did not resume: {sorted(resumed, key=int)}"
+        )
+        assert sorted(int(step) for step in resumed) == RESUMED_STEPS, (
+            f"resumed run covered steps {sorted(resumed, key=int)}, expected {RESUMED_STEPS}"
+        )
+        assert all(str(step) in full for step in RESUMED_STEPS), (
+            f"uninterrupted run did not record steps {RESUMED_STEPS}: {sorted(full, key=int)}"
+        )
+
+        for step in RESUMED_STEPS:
+            assert full[str(step)] == resumed[str(step)], (
+                f"step {step}: the resumed run diverged from the uninterrupted run\n"
+                f"  uninterrupted = {full[str(step)]}\n"
+                f"  resumed       = {resumed[str(step)]}"
+            )
+    finally:
+        if os.environ.get("VEOMNI_KEEP_ALIGN_LOGS"):
+            print(f"kept: {_output_dir('full')}, {_output_dir('resumed')}, {weights_dir}")
+        else:
+            shutil.rmtree(_output_dir("full"), ignore_errors=True)
+            shutil.rmtree(_output_dir("resumed"), ignore_errors=True)
+            shutil.rmtree(weights_dir, ignore_errors=True)
+        dummy.clean_cache()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py
index 8a98cea60d..941a4bdfdd 100644
--- a/veomni/checkpoint/dcp_checkpointer.py
+++ b/veomni/checkpoint/dcp_checkpointer.py
@@ -58,14 +58,13 @@
     step_dir,
     weights_dir,
 )
-from .layout import (
-    LR_SCHEDULER_FILENAME as _LR_SCHEDULER_FILENAME,
-)
 
 
 logger = logging.get_logger(__name__)
 
-_LR_SCHEDULER_KEY = "lr_scheduler"
+_EXTRA_STATE_KEY = "extra_state"
+_EXTRA_STATE_DIRNAME = "extra_state"
+_EXTRA_STATE_FORMAT = "extra_state_rank_{}.pt"
 
 
 class _ModelStrictLoadPlanner(DefaultLoadPlanner):
@@ -611,7 +610,7 @@ def _promote_staged_checkpoint(
     ``step_root`` is the step directory, given when the destination may hold a
     pre-split checkpoint whose marker sits there rather than inside
     ``final_path``. Nested files are copied in the data phase, before any
-    ``.metadata`` is. ``lr_scheduler.pt`` is one of those files.
+    ``.metadata`` is. The per-rank ``extra_state/`` files are among those files.
     """
     metadata_name = DCP_MARKER_FILENAME
     is_node_leader = _local_rank() == 0
@@ -794,7 +793,8 @@ def save(
 
         Writes three things under ``model/`` (see ``veomni.checkpoint.layout``):
         ``ckpt/`` for the weights, ``optimizer/`` for the optimizer state, and a
-        replicated ``lr_scheduler.pt``. Weights and optimizer are separate DCP
+        per-rank ``extra_state/`` for model-bound extra state (lr scheduler and
+        condition-model RNG). Weights and optimizer are separate DCP
         directories so the weights can be shipped or converted on their own; a
         single directory interleaves both into the same ``.distcp`` files.
 
@@ -887,7 +887,7 @@ def save(
         # and it is written only after every module's save has returned — but
         # writing the small replicated file first still means a save that dies
         # part-way leaves less behind.
-        cls._save_lr_scheduler(checkpoint_dir=write_root, state=state)
+        cls._save_extra_state(checkpoint_dir=write_root, state=state)
 
         try:
             cls.execute_save(
@@ -978,13 +978,14 @@ def load(
         load training state from distributed checkpoint
 
         Mirrors :meth:`save`: weights from ``model/<module>/ckpt``, optimizer from
-        ``model/<module>/optimizer``, scheduler from the sidecar beside them. A
-        checkpoint written before the split has no ``model/`` at all and keeps
-        both in one directory; that shape is detected and read as-is.
+        ``model/<module>/optimizer``, model-bound extra state from the per-rank
+        ``extra_state/`` beside them. A checkpoint written before the split has
+        no ``model/`` at all and keeps both in one directory; that shape is
+        detected and read as-is.
 
         args:
             path: step directory to load from
-            state: state to load, "model" is required; "optimizer" and "lr_scheduler" are optional
+            state: state to load, "model" is required; "optimizer" and "extra_state" are optional
             process_group: process group for loading checkpoint
             module: name of the model to load, for a job that trains several.
                 Empty for a single-model job. See :meth:`save`.
@@ -1027,7 +1028,7 @@ def _optimizer_state() -> "OptimizerState":
                 process_group=process_group,
                 planner=_ModelStrictLoadPlanner(strict_model=not trainable_only),
             )
-            cls._load_lr_scheduler(checkpoint_dir=fused_dir, state=state)
+            cls._load_extra_state(checkpoint_dir=fused_dir, state=state)
             logger.info_rank0(f"Loaded pre-split checkpoint from {fused_dir}")
             return state
 
@@ -1048,7 +1049,7 @@ def _optimizer_state() -> "OptimizerState":
                 planner=_ModelStrictLoadPlanner(strict_model=False),
             )
 
-        cls._load_lr_scheduler(checkpoint_dir=model_root, state=state)
+        cls._load_extra_state(checkpoint_dir=model_root, state=state)
 
         logger.info_rank0(f"Loaded checkpoint from {model_root}")
 
@@ -1234,54 +1235,30 @@ def _create_storage_writer(cls, checkpoint_dir: str) -> FileSystemWriter:
         )
 
     @classmethod
-    def _save_lr_scheduler(cls, checkpoint_dir: str, state: Dict[str, Any]) -> None:
-        """Pickle ``lr_scheduler.state_dict`` into a single ``lr_scheduler.pt``.
+    def _save_extra_state(cls, checkpoint_dir: str, state: Dict[str, Any]) -> None:
+        """Pickle this rank's ``extra_state`` dict into ``extra_state/``.
 
-        The scheduler is replicated across ranks, so only rank 0 writes. Every
-        rank still joins the reduction afterwards: a failed write must not let
-        peers enter the DCP collective alone.
+        The condition-model RNG is rank-local, so every rank writes its own file.
         """
-        error: Optional[Exception] = None
-        is_writer = (not dist.is_initialized()) or dist.get_rank() == 0
-        if is_writer:
-            try:
-                if _LR_SCHEDULER_KEY not in state:
-                    logger.warning_rank0("lr_scheduler not found in state, skipping lr_scheduler save")
-                else:
-                    lr_scheduler = state[_LR_SCHEDULER_KEY]
-                    if lr_scheduler is not None:
-                        torch.save(lr_scheduler.state_dict(), os.path.join(checkpoint_dir, _LR_SCHEDULER_FILENAME))
-            except Exception as e:  # noqa: BLE001 - raised once every rank has agreed
-                error = e
-        if any_rank_failed(error is not None):
-            raise error or RuntimeError("another rank could not save lr_scheduler")
-
-    @classmethod
-    def _load_lr_scheduler(cls, checkpoint_dir: str, state: Dict[str, Any]) -> None:
-        """Load ``lr_scheduler.pt`` into ``lr_scheduler``. Every rank reads the same file."""
-        if _LR_SCHEDULER_KEY not in state:
-            logger.warning_rank0("lr_scheduler not found in state, skipping lr_scheduler load")
-            return
-        lr_scheduler = state[_LR_SCHEDULER_KEY]
-        if lr_scheduler is None:
+        if _EXTRA_STATE_KEY not in state:
+            logger.warning_rank0("extra_state not found in state, skipping extra_state save")
             return
+        extra_state_dir = os.path.join(checkpoint_dir, _EXTRA_STATE_DIRNAME)
+        os.makedirs(extra_state_dir, exist_ok=True)
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        extra_state_path = os.path.join(extra_state_dir, _EXTRA_STATE_FORMAT.format(rank))
+        torch.save(state[_EXTRA_STATE_KEY], extra_state_path)
 
-        lr_scheduler_path = os.path.join(checkpoint_dir, _LR_SCHEDULER_FILENAME)
-        if os.path.exists(lr_scheduler_path):
-            lr_scheduler.load_state_dict(torch.load(lr_scheduler_path, weights_only=False))
-            return
-
-        # Delete this import (and veomni/checkpoint/legacy_v0_1_12.py) to drop 0.1.12 extra_state resume.
-        from .legacy_v0_1_12 import apply_legacy_lr_scheduler
-
-        if apply_legacy_lr_scheduler(checkpoint_dir, lr_scheduler):
+    @classmethod
+    def _load_extra_state(cls, checkpoint_dir: str, state: Dict[str, Any]) -> None:
+        """Load this rank's ``extra_state`` dict from ``extra_state/``."""
+        if _EXTRA_STATE_KEY not in state:
+            logger.warning_rank0("extra_state not found in state, skipping extra_state load")
             return
-
-        raise FileNotFoundError(
-            f"lr_scheduler sidecar not found at {lr_scheduler_path}. "
-            "This layout writes lr_scheduler.pt next to the DCP shards "
-            "(see docs/usage/checkpoint.md)."
-        )
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        extra_state_path = os.path.join(checkpoint_dir, _EXTRA_STATE_DIRNAME, _EXTRA_STATE_FORMAT.format(rank))
+        if os.path.exists(extra_state_path):
+            state[_EXTRA_STATE_KEY] = torch.load(extra_state_path, weights_only=False)
 
 
 def get_dtype_size(dtype: torch.dtype) -> int:
diff --git a/veomni/models/checkpoint_manager.py b/veomni/models/checkpoint_manager.py
index 4edbb6c7e7..bc2d2f8bb3 100644
--- a/veomni/models/checkpoint_manager.py
+++ b/veomni/models/checkpoint_manager.py
@@ -29,7 +29,7 @@
   :class:`~veomni.trainer.callbacks.global_state_callback.GlobalStateCallback`.
 """
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import torch.distributed as dist
 
@@ -160,40 +160,64 @@ def wait_for_pending_save(self) -> None:
         """Block until the in-flight async save is on disk, if there is one."""
         self.checkpointer.wait_for_pending_save()
 
+    def _extra_state(self, state: "TrainerState") -> Dict[str, Any]:
+        """Model-bound state to store beside the weights.
+
+        The lr scheduler plus whatever the runtime contributes via
+        ``extra_state()`` (e.g. the DiT condition model's generator).
+        """
+        lr_scheduler = self.runtime.lr_scheduler
+        extra_state = {
+            "lr_scheduler": None if lr_scheduler is None else lr_scheduler.state_dict(),
+        }
+        extra_state.update(self.runtime.extra_state())
+        return extra_state
+
+    def _load_extra_state(self, extra_state: Dict[str, Any]) -> None:
+        lr_state = extra_state.get("lr_scheduler")
+        lr_scheduler = self.runtime.lr_scheduler
+        if lr_state is not None and lr_scheduler is not None:
+            lr_scheduler.load_state_dict(lr_state)
+
+        self.runtime.load_extra_state(extra_state)
+
     def load(self) -> None:
-        """Restore model, optimizer and lr_scheduler from ``load_path``."""
+        """Restore model, optimizer and model-bound extra state from ``load_path``."""
         load_dir = self.load_dir()
         if load_dir is None:
             return
 
         self.wait_for_pending_save()
+        state: Dict[str, Any] = {
+            "model": self.runtime.model,
+            "optimizer": self.runtime.optimizer,
+            "extra_state": {},
+        }
         self.checkpointer.load(
             load_dir,
-            {
-                "model": self.runtime.model,
-                "optimizer": self.runtime.optimizer,
-                "lr_scheduler": self.runtime.lr_scheduler,
-            },
+            state,
             module=self.module_name,
             trainable_only=self.trainable_only,
             parallel_state=self.parallel_state,
         )
+        self._load_extra_state(state["extra_state"])
         dist.barrier()
         logger.info_rank0(f"Load distributed checkpoint from {load_dir} successfully!")
 
     def save_dcp(self, state: "TrainerState") -> None:
-        """Write model, optimizer and lr_scheduler for ``state.global_step``.
+        """Write model, optimizer and model-bound extra state for ``state.global_step``.
 
         Only model-bound state goes in here. Job-level state — where the
         dataloader is, the rng — has its own writer.
         """
+        extra_state = self._extra_state(state)
         helper.empty_cache()
         self.checkpointer.save(
             self.config.save_path,
             {
                 "model": self.runtime.model,
                 "optimizer": self.runtime.optimizer,
-                "lr_scheduler": self.runtime.lr_scheduler,
+                "extra_state": extra_state,
             },
             global_steps=state.global_step,
             module=self.module_name,
diff --git a/veomni/models/diffusers/qwen_image/qwen_image_condition/modeling_qwen_image_condition.py b/veomni/models/diffusers/qwen_image/qwen_image_condition/modeling_qwen_image_condition.py
index 783b1e5fe1..d30fa0ee01 100644
--- a/veomni/models/diffusers/qwen_image/qwen_image_condition/modeling_qwen_image_condition.py
+++ b/veomni/models/diffusers/qwen_image/qwen_image_condition/modeling_qwen_image_condition.py
@@ -236,6 +236,19 @@ def get_condition(self, inputs, images, **kwargs) -> dict[str, Any]:
             "img_shapes": img_shapes_list,
         }
 
+    def rng_state_dict(self) -> dict[str, torch.Tensor]:
+        """Snapshot the noise/timestep generator for the job-level checkpoint.
+
+        ``process_condition`` draws both the noise and the timestep ids from
+        ``self.generator``, so a resumed run that does not restore it replays the
+        stream from the construction-time seed instead of continuing it.
+        """
+        return {"generator": self.generator.get_state()}
+
+    def load_rng_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Restore a snapshot produced by :meth:`rng_state_dict`."""
+        self.generator.set_state(state["generator"])
+
     def process_condition(
         self,
         latents=None,
diff --git a/veomni/models/diffusers/wan_t2v/wan_condition/modeling_wan_condition.py b/veomni/models/diffusers/wan_t2v/wan_condition/modeling_wan_condition.py
index a2dadf426d..8f81da15d3 100644
--- a/veomni/models/diffusers/wan_t2v/wan_condition/modeling_wan_condition.py
+++ b/veomni/models/diffusers/wan_t2v/wan_condition/modeling_wan_condition.py
@@ -132,6 +132,19 @@ def get_condition(self, inputs, videos, **kwargs) -> dict[str, Any]:
 
         return {"latents": latents_list, "context": context_list}
 
+    def rng_state_dict(self) -> dict[str, torch.Tensor]:
+        """Snapshot the noise/timestep generator for the job-level checkpoint.
+
+        ``process_condition`` draws both the noise and the timestep ids from
+        ``self.generator``, so a resumed run that does not restore it replays the
+        stream from the construction-time seed instead of continuing it.
+        """
+        return {"generator": self.generator.get_state()}
+
+    def load_rng_state_dict(self, state: dict[str, torch.Tensor]) -> None:
+        """Restore a snapshot produced by :meth:`rng_state_dict`."""
+        self.generator.set_state(state["generator"])
+
     def process_condition(self, latents: list[torch.Tensor], context: list[torch.Tensor]) -> dict[str, Any]:
         if not self._timesteps_ready:
             self.scheduler.set_timesteps(self.config.num_train_timesteps, device=latents[0].device)
diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py
index 1410896df0..4fd3239867 100644
--- a/veomni/models/model_runtime.py
+++ b/veomni/models/model_runtime.py
@@ -568,6 +568,19 @@ def save_dcp(self, state: "TrainerState") -> None:
         """Write this model's resumable checkpoint for ``state.global_step``."""
         self.checkpoint.save_dcp(state)
 
+    def extra_state(self) -> Dict[str, Any]:
+        """Model-bound state to persist beside the weights.
+
+        Models with extra state (e.g. the DiT condition model's noise/timestep
+        generator) contribute it here. The checkpoint manager merges the result
+        into the model's extra_state blob. Default: nothing.
+        """
+        return {}
+
+    def load_extra_state(self, extra_state: Dict[str, Any]) -> None:
+        """Restore what :meth:`extra_state` produced. Default: nothing."""
+        return
+
     def save_hf_or_lora(self, state: "TrainerState", stage: str = "step_end") -> None:
         """Export this model in whichever format it was trained in.
 
diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py
index d38a5f41d3..60588a06a0 100644
--- a/veomni/trainer/callbacks/global_state_callback.py
+++ b/veomni/trainer/callbacks/global_state_callback.py
@@ -21,6 +21,11 @@
 multisource-sampler state is large, and resuming weights onto a different
 dataset means keeping ``extra_state/`` while discarding ``loader/``.
 
+"the rng" spans three streams: the CPU default generator, the accelerator
+default generator (diffusion condition models sample noise and timesteps from
+it), and any per-run ``torch.Generator`` a condition model owns. Restoring only
+the first resumes a different noise stream than an uninterrupted run.
+
 This callback also writes ``checkpoint_manifest.json``, which records that the
 files above are down. It says nothing about the model state: that is covered by
 DCP's own ``.metadata``, one per directory under ``model/``, and a step counts
@@ -39,7 +44,7 @@
 
 from ...checkpoint import layout
 from ...utils import helper
-from ...utils.device import get_device_type
+from ...utils.device import get_device_rng_state, get_device_type, set_device_rng_state
 from ...utils.dist_utils import raise_if_any_rank_failed
 from .base import Callback, TrainerState
 
@@ -133,6 +138,7 @@ def state_dict(self, state: TrainerState) -> Dict[str, Any]:
             "environ_meter": self.trainer.environ_meter.state_dict(),
             "channel_loss_callback": channel_loss_state,
             "torch_rng_state": torch.get_rng_state(),
+            "device_rng_state": get_device_rng_state(),
         }
 
     def save_global_state(self, state: TrainerState) -> None:
@@ -282,6 +288,9 @@ def load_global_state(self) -> Optional[Dict[str, Any]]:
         rng_state = global_state.get("torch_rng_state")
         if rng_state is not None:
             torch.set_rng_state(rng_state)
+        # ``.get`` keeps checkpoints written before these fields existed loadable;
+        # the device namespace may also not support RNG state at all.
+        set_device_rng_state(global_state.get("device_rng_state"))
         if self.trainer.start_step == 0 and self.trainer.train_dataloader is not None:
             iter(self.trainer.train_dataloader)
 
diff --git a/veomni/trainer/dit_trainer.py b/veomni/trainer/dit_trainer.py
index 57d07c630b..e5dec2bfdf 100644
--- a/veomni/trainer/dit_trainer.py
+++ b/veomni/trainer/dit_trainer.py
@@ -240,6 +240,33 @@ def _freeze_model_module(self):
         if self.train_args.training_task != "offline_embedding":
             super()._freeze_model_module()
 
+    def extra_state(self) -> dict[str, Any]:
+        """Model-bound state: the condition model's noise/timestep generator."""
+        rng_state_dict = getattr(self.condition_model, "rng_state_dict", None)
+        if (
+            self.condition_model is not None
+            and rng_state_dict is None
+            and getattr(self.condition_model, "generator", None) is not None
+        ):
+            logger.warning_rank0(
+                "Condition model owns a ``generator`` but exposes no ``rng_state_dict``; "
+                "its noise/timestep stream will not be restored across a resume."
+            )
+        return {} if rng_state_dict is None else {"condition_model_rng_state": rng_state_dict()}
+
+    def load_extra_state(self, extra_state: dict[str, Any]) -> None:
+        condition_model_rng_state = extra_state.get("condition_model_rng_state")
+        if condition_model_rng_state is None:
+            return
+        loader = getattr(self.condition_model, "load_rng_state_dict", None)
+        if loader is None:
+            logger.warning_rank0(
+                "Checkpoint carries condition-model RNG state but the model cannot restore it; "
+                "the resumed run may replay its initial noise stream."
+            )
+        else:
+            loader(condition_model_rng_state)
+
     def _build_parallelized_model(self) -> None:
         """``offline_embedding`` builds no DiT, so there is nothing to wrap."""
         if self.train_args.training_task != "offline_embedding":
diff --git a/veomni/utils/device.py b/veomni/utils/device.py
index d009feb36a..fe7f56b7ca 100644
--- a/veomni/utils/device.py
+++ b/veomni/utils/device.py
@@ -115,6 +115,44 @@ def set_device(device: torch.types.Device) -> None:
     get_torch_device().set_device(device)
 
 
+def get_device_rng_state() -> Any:
+    """Snapshot the accelerator's default RNG state, or ``None`` if unsupported.
+
+    ``torch.get_rng_state()`` only covers the CPU generator, while diffusion
+    condition models draw noise and timesteps from the device default RNG. A
+    checkpoint that restores only the CPU state therefore resumes a different
+    device random stream than an uninterrupted run.
+
+    The accelerator namespaces differ (``torch.cuda`` / ``torch.npu`` /
+    ``torch.mlu``), and a namespace may not expose the accessor at all, so this
+    returns ``None`` rather than failing the checkpoint.
+    """
+    if get_device_type() == "cpu":
+        return torch.get_rng_state()
+
+    getter = getattr(get_torch_device(), "get_rng_state", None)
+    if getter is None:
+        logger.warning(f"Device RNG state is not available on {get_device_type()}; checkpoint will not restore it.")
+        return None
+    return getter()
+
+
+def set_device_rng_state(state: Any) -> None:
+    """Restore a snapshot produced by :func:`get_device_rng_state`."""
+    if state is None:
+        return
+
+    if get_device_type() == "cpu":
+        torch.set_rng_state(state)
+        return
+
+    setter = getattr(get_torch_device(), "set_rng_state", None)
+    if setter is None:
+        logger.warning(f"Device RNG state cannot be restored on {get_device_type()}; ignoring the snapshot.")
+        return
+    setter(state)
+
+
 def is_nccl_backend(backend: str | None = None) -> bool:
     """Check if the distributed communication backend is NCCL."""
     return (backend or get_dist_comm_backend()) == "nccl"