From 51bed0761a8c77e0eaa17c5590417599f932b7c2 Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Tue, 15 Sep 2026 11:08:51 +0800 Subject: [PATCH 1/9] [ckpt, trainer] fix: persist device and condition-model RNG across DiT resume DiT condition models draw noise and timesteps from the accelerator default RNG and a per-run ``torch.Generator``. The job-level checkpoint only restored the CPU generator, so a resumed run replayed the noise stream from the start instead of continuing it. Persist the device RNG state and any condition-model ``rng_state_dict`` in the global state, and restore both on load, with backward compatibility for checkpoints written before the fields existed. Adds unit tests in test_checkpoint_callback and an end-to-end resume-equivalence test that runs the real DiTTrainer twice against a dummy dataset with dropout active (p=0.5). --- .agents/knowledge/testing.md | 1 + tests/checkpoints/test_checkpoint_callback.py | 116 +++++ tests/checkpoints/test_dit_resume_align.py | 414 ++++++++++++++++++ .../modeling_qwen_image_condition.py | 13 + .../wan_condition/modeling_wan_condition.py | 13 + .../callbacks/global_state_callback.py | 41 +- veomni/trainer/dit_trainer.py | 4 + veomni/utils/device.py | 38 ++ 8 files changed, 639 insertions(+), 1 deletion(-) create mode 100644 tests/checkpoints/test_dit_resume_align.py 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 55eee8107c..0e14c5da38 100644 --- a/tests/checkpoints/test_checkpoint_callback.py +++ b/tests/checkpoints/test_checkpoint_callback.py @@ -17,6 +17,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): @@ -60,6 +78,10 @@ def _make_mock_trainer(save_path="/tmp/test_ckpt", save_async=False): trainer.start_epoch = 0 trainer.start_step = 0 trainer.checkpoint = MagicMock() + # 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 @@ -362,6 +384,100 @@ 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_omits_condition_model_rng_without_condition_model(self, mock_dist): + trainer = _make_mock_trainer() + cb = GlobalStateCallback(trainer) + + global_state = cb.state_dict(TrainerState(global_step=10)) + + assert global_state["condition_model_rng_state"] is None + + def test_state_dict_persists_condition_model_rng(self, mock_dist): + trainer = _make_mock_trainer() + trainer.condition_model = _StubConditionModel() + cb = GlobalStateCallback(trainer) + + global_state = cb.state_dict(TrainerState(global_step=10)) + + assert set(global_state["condition_model_rng_state"]) == {"generator"} + + def test_load_resumes_noise_and_device_rng_streams(self, mock_dist, tmp_path): + """A resumed run must continue the noise stream, not replay its start.""" + 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 = {} + condition_model = _StubConditionModel() + trainer.condition_model = condition_model + + # Three steps' worth of noise, then checkpoint mid-run. + for _ in range(3): + condition_model.draw_noise() + cb = GlobalStateCallback(trainer) + payload = cb.state_dict(TrainerState(global_step=3)) + + # The uninterrupted run would have drawn exactly this next. + reference = torch.Generator(device="cpu") + reference.set_state(condition_model.generator.get_state()) + expected_next_noise = torch.randn(4, generator=reference) + torch.save(payload, tmp_path / "trainer_state_rank_0.pt") + + # Drift both streams the way the rest of a continuing run would. + condition_model.draw_noise() + torch.randn(8) + + assert GlobalStateCallback(trainer).load_global_state() is not None + + assert torch.equal(condition_model.draw_noise(), expected_next_noise) + 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_load_warns_when_condition_model_cannot_restore_rng(self, mock_dist, tmp_path): + """State for an absent condition model is a warning, not a crash.""" + 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(), + "condition_model_rng_state": {"generator": _StubConditionModel().rng_state_dict()["generator"]}, + }, + tmp_path / "trainer_state_rank_0.pt", + ) + + GlobalStateCallback(trainer).load_global_state() + + assert trainer.state.global_step == 7 def test_save_waits_for_pending_dcp(self, mock_dist, tmp_path): mock_dist.is_initialized.return_value = False diff --git a/tests/checkpoints/test_dit_resume_align.py b/tests/checkpoints/test_dit_resume_align.py new file mode 100644 index 0000000000..3fdf0940bc --- /dev/null +++ b/tests/checkpoints/test_dit_resume_align.py @@ -0,0 +1,414 @@ +# 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, +SeedVR2). 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) + + # -- job-level checkpoint protocol ------------------------------------ + 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"]) + + # -- per-step sampling ------------------------------------------------- + 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 / SeedVR2 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 DiTTrainer, VeOmniDiTArguments + + 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 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_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) + # ``GlobalStateCallback`` reaches the condition model through the + # reused ``BaseTrainer`` instance, exactly as the real DiT trainer + # arranges it. + self.base.condition_model = self.condition_model + + def _freeze_model_module(self) -> None: + self.base.lora = False + + def _build_model_assets(self) -> None: + self.base.model_assets = [self.base.model.config] + + 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() + + +# --------------------------------------------------------------------------- +# pytest side +# --------------------------------------------------------------------------- + + +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 + + apply_ops_config(OpsImplementationConfig()) + 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: + # /tests/checkpoints/test_dit_resume_align.py -> + 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}", + "--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/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/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index 2de731b666..2f0238fc34 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -17,6 +17,12 @@ Nothing here belongs to a model: where the dataloader is, the rng, the metric meters. Model weights, optimizer, HF/LoRA export, and the tokenizer/config sidecars are scheduled by :mod:`~veomni.trainer.callbacks.checkpoint_callback`. + +"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, as this callback once did, resumes a different noise stream than an +uninterrupted run. """ import os @@ -26,7 +32,7 @@ import torch.distributed as dist 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 .base import Callback, TrainerState @@ -65,6 +71,15 @@ def __init__(self, trainer: "BaseTrainer"): def rank(self) -> int: return self.trainer.args.train.global_rank + def _condition_model(self) -> Any: + """Return the diffusion condition model, if this job has one. + + DiT trainers reuse a ``BaseTrainer`` instance they never subclass, so the + condition model — which owns the noise/timestep generators — is attached + as ``condition_model`` on that instance. Other trainers never set it. + """ + return getattr(self.trainer, "condition_model", None) + def on_train_begin(self, state: TrainerState, **kwargs) -> None: self.load_global_state() @@ -88,12 +103,23 @@ def state_dict(self, state: TrainerState) -> Dict[str, Any]: channel_loss_callback = getattr(self.trainer, "channel_loss_callback", None) channel_loss_state = channel_loss_callback.state_dict() if channel_loss_callback is not None else {} + # Diffusion condition models sample noise and timesteps from the device + # default RNG or from a per-run ``torch.Generator``; neither is model + # state, so both belong in the job cursor alongside the CPU generator. + # Without them a resumed run replays the start of the run's noise stream + # instead of continuing it. ``None`` means "not applicable / unsupported" + # and the loader skips it. + condition_model = self._condition_model() + rng_state_dict = getattr(condition_model, "rng_state_dict", None) + return { "global_step": state.global_step, "train_dataloader": train_dataloader_state, "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(), + "condition_model_rng_state": None if rng_state_dict is None else rng_state_dict(), } def save_global_state(self, state: TrainerState) -> None: @@ -145,6 +171,19 @@ def load_global_state(self) -> Optional[Dict[str, Any]]: self.trainer.environ_meter.load_state_dict(global_state["environ_meter"]) torch.set_rng_state(global_state["torch_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")) + condition_model_rng_state = global_state.get("condition_model_rng_state") + if condition_model_rng_state is not None: + 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) 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 c813eada5b..13e9f7991d 100644 --- a/veomni/trainer/dit_trainer.py +++ b/veomni/trainer/dit_trainer.py @@ -321,6 +321,10 @@ def _build_condition_model( self.condition_model = model_class._from_config(condition_cfg) self.condition_model.to(get_device_type()) logger.info_rank0("Condition model loaded.") + # The condition model owns the noise/timestep generators, which are part + # of the job cursor. ``GlobalStateCallback`` runs on the reused + # ``BaseTrainer`` instance, so expose the model there as well. + self.base.condition_model = self.condition_model def _freeze_model_module(self): self.condition_model.requires_grad_(False) 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" From a4e9a3c33d72836f1703df987f473f256a2d6df6 Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Tue, 15 Sep 2026 11:25:27 +0800 Subject: [PATCH 2/9] [ckpt, trainer] fix: warn when a condition model's generator is not persisted MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A condition model that owns a ``torch.Generator`` but exposes no ``rng_state_dict`` would silently replay its noise/timestep stream from the construction-time seed on every resume. Warn on save so the miss is visible instead of silent. Only fires for models holding their own ``generator`` attribute — device-default-RNG models (LTX2, MiniMax-H3) are already covered by ``device_rng_state`` and stay quiet. --- tests/checkpoints/test_checkpoint_callback.py | 16 ++++++++++++++++ tests/checkpoints/test_dit_resume_align.py | 5 ++--- .../trainer/callbacks/global_state_callback.py | 13 +++++++++++++ 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/tests/checkpoints/test_checkpoint_callback.py b/tests/checkpoints/test_checkpoint_callback.py index 0e14c5da38..3486290f89 100644 --- a/tests/checkpoints/test_checkpoint_callback.py +++ b/tests/checkpoints/test_checkpoint_callback.py @@ -403,6 +403,22 @@ def test_state_dict_persists_condition_model_rng(self, mock_dist): assert set(global_state["condition_model_rng_state"]) == {"generator"} + def test_state_dict_warns_when_condition_model_misses_rng_dict(self, mock_dist): + """A generator-owning condition model with no rng_state_dict is worth a warning.""" + trainer = _make_mock_trainer() + + class _GeneratorOnly: + generator = torch.Generator() + + trainer.condition_model = _GeneratorOnly() + cb = GlobalStateCallback(trainer) + + with patch("veomni.trainer.callbacks.global_state_callback.logger.warning_rank0") as mock_warn: + global_state = cb.state_dict(TrainerState(global_step=10)) + + mock_warn.assert_called_once() + assert global_state["condition_model_rng_state"] is None + def test_load_resumes_noise_and_device_rng_streams(self, mock_dist, tmp_path): """A resumed run must continue the noise stream, not replay its start.""" mock_dist.is_initialized.return_value = False diff --git a/tests/checkpoints/test_dit_resume_align.py b/tests/checkpoints/test_dit_resume_align.py index 3fdf0940bc..47d4402752 100644 --- a/tests/checkpoints/test_dit_resume_align.py +++ b/tests/checkpoints/test_dit_resume_align.py @@ -16,8 +16,7 @@ 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, -SeedVR2). Neither is model state, so the job-level checkpoint has to carry both +(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. @@ -127,7 +126,7 @@ def process_condition(self, hidden_states=None, encoder_hidden_states=None, late 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 / SeedVR2 rely on. + # 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. diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index 2f0238fc34..0217d7366a 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -111,6 +111,19 @@ def state_dict(self, state: TrainerState) -> Dict[str, Any]: # and the loader skips it. condition_model = self._condition_model() rng_state_dict = getattr(condition_model, "rng_state_dict", None) + if ( + condition_model is not None + and rng_state_dict is None + and getattr(condition_model, "generator", None) is not None + ): + # A condition model holding its own ``torch.Generator`` but not exposing + # ``rng_state_dict`` would silently replay its stream from the + # construction-time seed on every resume. The device default stream, the + # other source, is already covered by ``device_rng_state``. + 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 { "global_step": state.global_step, From f614bde16b099f321e3ca019c80d973d16660d70 Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Tue, 15 Sep 2026 16:30:01 +0800 Subject: [PATCH 3/9] [ckpt, trainer] test: drop decorative divider comments in resume-alignment test CodeRabbit flagged the section-divider comments (job-level checkpoint protocol, per-step sampling, pytest side) as non-conforming to the repo-wide no-decorative- comment convention. Remove them; behavior and coverage unchanged. --- tests/checkpoints/test_dit_resume_align.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/checkpoints/test_dit_resume_align.py b/tests/checkpoints/test_dit_resume_align.py index 47d4402752..5a44978727 100644 --- a/tests/checkpoints/test_dit_resume_align.py +++ b/tests/checkpoints/test_dit_resume_align.py @@ -105,14 +105,12 @@ def __init__(self, seed: int, dp_rank: int = 0) -> None: self.generator.manual_seed(seed + dp_rank) self.dropout = torch.nn.Dropout(DROPOUT_P) - # -- job-level checkpoint protocol ------------------------------------ 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"]) - # -- per-step sampling ------------------------------------------------- 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 @@ -245,11 +243,6 @@ def main(): _build_trainer_class()(args).train() -# --------------------------------------------------------------------------- -# pytest side -# --------------------------------------------------------------------------- - - def _output_dir(run: str) -> str: return os.path.abspath(f"./test_dit_resume_align_{run}") From 48d8b8994b3f5decbb7e3adb19aa704973c5587d Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Wed, 16 Sep 2026 17:01:54 +0800 Subject: [PATCH 4/9] [ckpt, trainer] refactor: move condition-model rng state into DCP extra_state Condition-model RNG is model-bound state, not job-cursor state, so persist it in the DCP extra_state alongside the lr scheduler instead of the GlobalStateCallback cursor. torch/device RNG stay in the cursor. --- tests/checkpoints/test_checkpoint_callback.py | 139 +++++++++--------- veomni/models/checkpoint_manager.py | 33 ++++- .../callbacks/global_state_callback.py | 42 ------ veomni/trainer/dit_trainer.py | 6 +- 4 files changed, 107 insertions(+), 113 deletions(-) diff --git a/tests/checkpoints/test_checkpoint_callback.py b/tests/checkpoints/test_checkpoint_callback.py index 3486290f89..9d78d2bc31 100644 --- a/tests/checkpoints/test_checkpoint_callback.py +++ b/tests/checkpoints/test_checkpoint_callback.py @@ -324,7 +324,9 @@ def test_the_logged_destination_is_the_one_save_writes(self, mock_helper, mock_d call = manager.checkpointer.save.call_args assert f"{call.args[0]}/{_GLOBAL_STEP_PREFIX}{call.kwargs['global_steps']}" == "/remote/run/global_step_10" - def test_extra_state_is_only_the_scheduler(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps): + def test_extra_state_is_only_the_scheduler_and_condition_model_rng( + self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps + ): trainer = _make_mock_trainer() mock_build_ckpt.return_value = MagicMock() manager = ModelCheckpointManager(trainer) @@ -332,8 +334,73 @@ def test_extra_state_is_only_the_scheduler(self, mock_helper, mock_dist, mock_bu manager.save_dcp(TrainerState(global_step=10)) extra_state = manager.checkpointer.save.call_args.args[1]["extra_state"] - assert set(extra_state) == {"lr_scheduler"} + assert set(extra_state) == {"lr_scheduler", "condition_model_rng_state"} assert extra_state["lr_scheduler"] == {"lr": 1e-4} + assert extra_state["condition_model_rng_state"] is None + + 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_extra_state_warns_when_condition_model_misses_rng_dict( + self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps + ): + trainer = _make_mock_trainer() + + class _GeneratorOnly: + generator = torch.Generator() + + trainer.condition_model = _GeneratorOnly() + mock_build_ckpt.return_value = MagicMock() + manager = ModelCheckpointManager(trainer) + + with patch("veomni.models.checkpoint_manager.logger.warning_rank0") as mock_warn: + extra_state = manager._extra_state(TrainerState(global_step=10)) + + mock_warn.assert_called_once() + assert extra_state["condition_model_rng_state"] is None + + 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_load_extra_state_warns_when_condition_model_cannot_restore_rng( + self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps + ): + trainer = _make_mock_trainer() + mock_build_ckpt.return_value = MagicMock() + manager = ModelCheckpointManager(trainer) + + with patch("veomni.models.checkpoint_manager.logger.warning_rank0") as mock_warn: + manager._load_extra_state( + { + "lr_scheduler": {"lr": 1e-4}, + "condition_model_rng_state": {"generator": torch.Generator().get_state()}, + } + ) + + mock_warn.assert_called_once() def test_legacy_extra_state_restores_the_job_cursor(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps): """Checkpoints written by CheckpointerCallback still resume the cursor.""" @@ -386,69 +453,31 @@ def test_state_dict_includes_channel_loss_callback_state(self, mock_dist): assert "torch_rng_state" in global_state assert "device_rng_state" in global_state - def test_state_dict_omits_condition_model_rng_without_condition_model(self, mock_dist): + 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 global_state["condition_model_rng_state"] is None + assert "condition_model_rng_state" not in global_state - def test_state_dict_persists_condition_model_rng(self, mock_dist): - trainer = _make_mock_trainer() - trainer.condition_model = _StubConditionModel() - cb = GlobalStateCallback(trainer) - - global_state = cb.state_dict(TrainerState(global_step=10)) - - assert set(global_state["condition_model_rng_state"]) == {"generator"} - - def test_state_dict_warns_when_condition_model_misses_rng_dict(self, mock_dist): - """A generator-owning condition model with no rng_state_dict is worth a warning.""" - trainer = _make_mock_trainer() - - class _GeneratorOnly: - generator = torch.Generator() - - trainer.condition_model = _GeneratorOnly() - cb = GlobalStateCallback(trainer) - - with patch("veomni.trainer.callbacks.global_state_callback.logger.warning_rank0") as mock_warn: - global_state = cb.state_dict(TrainerState(global_step=10)) - - mock_warn.assert_called_once() - assert global_state["condition_model_rng_state"] is None - - def test_load_resumes_noise_and_device_rng_streams(self, mock_dist, tmp_path): - """A resumed run must continue the noise stream, not replay its start.""" + 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 = {} - condition_model = _StubConditionModel() - trainer.condition_model = condition_model - - # Three steps' worth of noise, then checkpoint mid-run. - for _ in range(3): - condition_model.draw_noise() cb = GlobalStateCallback(trainer) payload = cb.state_dict(TrainerState(global_step=3)) - - # The uninterrupted run would have drawn exactly this next. - reference = torch.Generator(device="cpu") - reference.set_state(condition_model.generator.get_state()) - expected_next_noise = torch.randn(4, generator=reference) torch.save(payload, tmp_path / "trainer_state_rank_0.pt") # Drift both streams the way the rest of a continuing run would. - condition_model.draw_noise() torch.randn(8) assert GlobalStateCallback(trainer).load_global_state() is not None - assert torch.equal(condition_model.draw_noise(), expected_next_noise) assert torch.equal(get_device_rng_state(), payload["device_rng_state"]) assert torch.equal(torch.get_rng_state(), payload["torch_rng_state"]) @@ -473,28 +502,6 @@ def test_load_tolerates_checkpoint_without_rng_fields(self, mock_dist, tmp_path) assert trainer.state.global_step == 7 - def test_load_warns_when_condition_model_cannot_restore_rng(self, mock_dist, tmp_path): - """State for an absent condition model is a warning, not a crash.""" - 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(), - "condition_model_rng_state": {"generator": _StubConditionModel().rng_state_dict()["generator"]}, - }, - tmp_path / "trainer_state_rank_0.pt", - ) - - GlobalStateCallback(trainer).load_global_state() - - assert trainer.state.global_step == 7 - def test_save_waits_for_pending_dcp(self, mock_dist, tmp_path): mock_dist.is_initialized.return_value = False trainer = _make_mock_trainer(save_path=str(tmp_path)) diff --git a/veomni/models/checkpoint_manager.py b/veomni/models/checkpoint_manager.py index 06675f4dfb..ff3d637104 100644 --- a/veomni/models/checkpoint_manager.py +++ b/veomni/models/checkpoint_manager.py @@ -103,9 +103,27 @@ def load_dir(self) -> Optional[str]: return os.path.join(load_path, self.checkpoint_subfolder) if self.checkpoint_subfolder else load_path def _extra_state(self, state: "TrainerState") -> Dict[str, Any]: - """Model-bound state to store beside the weights.""" + """Model-bound state to store beside the weights. + + The condition model's noise/timestep generator is model-bound state, not + job-cursor state, so it is persisted here alongside the lr scheduler. + """ lr_scheduler = self.trainer.lr_scheduler - return {"lr_scheduler": None if lr_scheduler is None else lr_scheduler.state_dict()} + condition_model = getattr(self.trainer, "condition_model", None) + rng_state_dict = getattr(condition_model, "rng_state_dict", None) + if ( + condition_model is not None + and rng_state_dict is None + and getattr(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 { + "lr_scheduler": None if lr_scheduler is None else lr_scheduler.state_dict(), + "condition_model_rng_state": None if rng_state_dict is None else rng_state_dict(), + } def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: lr_state = extra_state.get("lr_scheduler") @@ -113,6 +131,17 @@ def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: if lr_state is not None and lr_scheduler is not None: lr_scheduler.load_state_dict(lr_state) + condition_model_rng_state = extra_state.get("condition_model_rng_state") + if condition_model_rng_state is not None: + loader = getattr(getattr(self.trainer, "condition_model", None), "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) + # Pre-split DCP extra_state also held the job cursor. New writes do not; # GlobalStateCallback owns that file. Restore the old blob so a mid-job # resume from a CheckpointerCallback checkpoint does not silently restart diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index 0217d7366a..0ad16b75df 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -71,15 +71,6 @@ def __init__(self, trainer: "BaseTrainer"): def rank(self) -> int: return self.trainer.args.train.global_rank - def _condition_model(self) -> Any: - """Return the diffusion condition model, if this job has one. - - DiT trainers reuse a ``BaseTrainer`` instance they never subclass, so the - condition model — which owns the noise/timestep generators — is attached - as ``condition_model`` on that instance. Other trainers never set it. - """ - return getattr(self.trainer, "condition_model", None) - def on_train_begin(self, state: TrainerState, **kwargs) -> None: self.load_global_state() @@ -103,28 +94,6 @@ def state_dict(self, state: TrainerState) -> Dict[str, Any]: channel_loss_callback = getattr(self.trainer, "channel_loss_callback", None) channel_loss_state = channel_loss_callback.state_dict() if channel_loss_callback is not None else {} - # Diffusion condition models sample noise and timesteps from the device - # default RNG or from a per-run ``torch.Generator``; neither is model - # state, so both belong in the job cursor alongside the CPU generator. - # Without them a resumed run replays the start of the run's noise stream - # instead of continuing it. ``None`` means "not applicable / unsupported" - # and the loader skips it. - condition_model = self._condition_model() - rng_state_dict = getattr(condition_model, "rng_state_dict", None) - if ( - condition_model is not None - and rng_state_dict is None - and getattr(condition_model, "generator", None) is not None - ): - # A condition model holding its own ``torch.Generator`` but not exposing - # ``rng_state_dict`` would silently replay its stream from the - # construction-time seed on every resume. The device default stream, the - # other source, is already covered by ``device_rng_state``. - 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 { "global_step": state.global_step, "train_dataloader": train_dataloader_state, @@ -132,7 +101,6 @@ def state_dict(self, state: TrainerState) -> Dict[str, Any]: "channel_loss_callback": channel_loss_state, "torch_rng_state": torch.get_rng_state(), "device_rng_state": get_device_rng_state(), - "condition_model_rng_state": None if rng_state_dict is None else rng_state_dict(), } def save_global_state(self, state: TrainerState) -> None: @@ -187,16 +155,6 @@ def load_global_state(self) -> Optional[Dict[str, Any]]: # ``.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")) - condition_model_rng_state = global_state.get("condition_model_rng_state") - if condition_model_rng_state is not None: - 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) 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 13e9f7991d..1d3530c2c8 100644 --- a/veomni/trainer/dit_trainer.py +++ b/veomni/trainer/dit_trainer.py @@ -321,9 +321,9 @@ def _build_condition_model( self.condition_model = model_class._from_config(condition_cfg) self.condition_model.to(get_device_type()) logger.info_rank0("Condition model loaded.") - # The condition model owns the noise/timestep generators, which are part - # of the job cursor. ``GlobalStateCallback`` runs on the reused - # ``BaseTrainer`` instance, so expose the model there as well. + # The condition model owns the noise/timestep generator, persisted in the + # DCP extra_state. ``ModelCheckpointManager`` reaches it through the + # reused ``BaseTrainer`` instance, so expose it there as well. self.base.condition_model = self.condition_model def _freeze_model_module(self): From 94efa2a164b5f83ce0ea05f17f9e25a94eebbdf2 Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Thu, 17 Sep 2026 12:51:40 +0800 Subject: [PATCH 5/9] [ckpt, trainer] refactor: store lr_scheduler and condition-model RNG in model extra_state Restore the model-bound extra_state blob (replacing the replicated lr_scheduler.pt sidecar) so the DiT condition model's per-rank noise/timestep generator persists beside the lr scheduler, not in the job cursor. --- tests/checkpoints/test_checkpoint_callback.py | 51 +++++++++---- veomni/checkpoint/dcp_checkpointer.py | 73 ++++++------------- veomni/models/checkpoint_manager.py | 61 +++++++++++++--- .../callbacks/global_state_callback.py | 36 --------- 4 files changed, 111 insertions(+), 110 deletions(-) diff --git a/tests/checkpoints/test_checkpoint_callback.py b/tests/checkpoints/test_checkpoint_callback.py index 5c04478685..cfc0e5933b 100644 --- a/tests/checkpoints/test_checkpoint_callback.py +++ b/tests/checkpoints/test_checkpoint_callback.py @@ -341,7 +341,7 @@ 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) @@ -349,11 +349,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 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() @@ -363,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 ): @@ -466,22 +494,13 @@ def test_state_dict_includes_channel_loss_callback_state(self, mock_dist): assert "torch_rng_state" in global_state assert "device_rng_state" in global_state - def test_state_dict_omits_condition_model_rng_without_condition_model(self, mock_dist): - trainer = _make_mock_trainer() - cb = GlobalStateCallback(trainer) - - global_state = cb.state_dict(TrainerState(global_step=10)) - - assert global_state["condition_model_rng_state"] is None - - def test_state_dict_persists_condition_model_rng(self, mock_dist): + def test_state_dict_has_no_condition_model_rng(self, mock_dist): trainer = _make_mock_trainer() - trainer.condition_model = _StubConditionModel() cb = GlobalStateCallback(trainer) global_state = cb.state_dict(TrainerState(global_step=10)) - assert set(global_state["condition_model_rng_state"]) == {"generator"} + 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.""" diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 8a98cea60d..f351362036 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): @@ -887,7 +886,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( @@ -1027,7 +1026,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 +1047,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 +1233,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: - return - - 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)) + 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) + extra_state_path = os.path.join(extra_state_dir, _EXTRA_STATE_FORMAT.format(dist.get_rank())) + torch.save(state[_EXTRA_STATE_KEY], extra_state_path) - # 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)." + extra_state_path = os.path.join( + checkpoint_dir, _EXTRA_STATE_DIRNAME, _EXTRA_STATE_FORMAT.format(dist.get_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 896a96f1bc..9c40ffe95d 100644 --- a/veomni/models/checkpoint_manager.py +++ b/veomni/models/checkpoint_manager.py @@ -26,7 +26,7 @@ On-disk layout: ``docs/usage/checkpoint.md``. """ -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Dict, Optional import torch.distributed as dist @@ -136,44 +136,87 @@ def load_dir(self) -> Optional[str]: """ return self.config.load_path + def _extra_state(self, state: "TrainerState") -> Dict[str, Any]: + """Model-bound state to store beside the weights. + + The condition model's noise/timestep generator and the lr scheduler are + model-bound state, persisted here for the checkpointer to write per rank. + """ + lr_scheduler = self.trainer.lr_scheduler + condition_model = getattr(self.trainer, "condition_model", None) + rng_state_dict = getattr(condition_model, "rng_state_dict", None) + if ( + condition_model is not None + and rng_state_dict is None + and getattr(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 { + "lr_scheduler": None if lr_scheduler is None else lr_scheduler.state_dict(), + "condition_model_rng_state": None if rng_state_dict is None else rng_state_dict(), + } + + def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: + lr_state = extra_state.get("lr_scheduler") + lr_scheduler = self.trainer.lr_scheduler + if lr_state is not None and lr_scheduler is not None: + lr_scheduler.load_state_dict(lr_state) + + condition_model_rng_state = extra_state.get("condition_model_rng_state") + if condition_model_rng_state is not None: + loader = getattr(getattr(self.trainer, "condition_model", None), "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 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 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.trainer.model, + "optimizer": self.trainer.optimizer, + "extra_state": {}, + } self.checkpointer.load( load_dir, - { - "model": self.trainer.model, - "optimizer": self.trainer.optimizer, - "lr_scheduler": self.trainer.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.trainer.model, "optimizer": self.trainer.optimizer, - "lr_scheduler": self.trainer.lr_scheduler, + "extra_state": extra_state, }, global_steps=state.global_step, module=self.module_name, diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index d06969af20..872a0e81fe 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -90,15 +90,6 @@ def __init__(self, trainer: "BaseTrainer"): def rank(self) -> int: return self.trainer.args.train.global_rank - def _condition_model(self) -> Any: - """Return the diffusion condition model, if this job has one. - - DiT trainers reuse a ``BaseTrainer`` instance they never subclass, so the - condition model — which owns the noise/timestep generators — is attached - as ``condition_model`` on that instance. Other trainers never set it. - """ - return getattr(self.trainer, "condition_model", None) - def on_train_begin(self, state: TrainerState, **kwargs) -> None: self.load_global_state() @@ -141,22 +132,6 @@ def state_dict(self, state: TrainerState) -> Dict[str, Any]: channel_loss_callback = getattr(self.trainer, "channel_loss_callback", None) channel_loss_state = channel_loss_callback.state_dict() if channel_loss_callback is not None else {} - condition_model = self._condition_model() - rng_state_dict = getattr(condition_model, "rng_state_dict", None) - if ( - condition_model is not None - and rng_state_dict is None - and getattr(condition_model, "generator", None) is not None - ): - # A condition model holding its own ``torch.Generator`` but not exposing - # ``rng_state_dict`` would silently replay its stream from the - # construction-time seed on every resume. The device default stream, the - # other source, is already covered by ``device_rng_state``. - 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 { "global_step": state.global_step, "train_dataloader": train_dataloader_state, @@ -164,7 +139,6 @@ def state_dict(self, state: TrainerState) -> Dict[str, Any]: "channel_loss_callback": channel_loss_state, "torch_rng_state": torch.get_rng_state(), "device_rng_state": get_device_rng_state(), - "condition_model_rng_state": None if rng_state_dict is None else rng_state_dict(), } def module_names(self) -> List[str]: @@ -327,16 +301,6 @@ def load_global_state(self) -> Optional[Dict[str, Any]]: # ``.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")) - condition_model_rng_state = global_state.get("condition_model_rng_state") - if condition_model_rng_state is not None: - 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) if self.trainer.start_step == 0 and self.trainer.train_dataloader is not None: iter(self.trainer.train_dataloader) From 4eb350b29a6dadfb9c5b6cd8a76664d762764626 Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Fri, 18 Sep 2026 16:24:16 +0800 Subject: [PATCH 6/9] [ckpt, trainer] refactor: delegate model extra state to the trainer ModelCheckpointManager no longer hard-codes the DiT condition model; it calls trainer.extra_state()/load_extra_state() so each trainer contributes its own model-bound state (the condition model's noise/timestep generator for DiT). --- tests/checkpoints/test_checkpoint_callback.py | 47 +++++-------------- veomni/models/checkpoint_manager.py | 31 +++--------- veomni/trainer/base.py | 13 +++++ veomni/trainer/dit_trainer.py | 34 +++++++++++++- 4 files changed, 64 insertions(+), 61 deletions(-) diff --git a/tests/checkpoints/test_checkpoint_callback.py b/tests/checkpoints/test_checkpoint_callback.py index cfc0e5933b..ad29134002 100644 --- a/tests/checkpoints/test_checkpoint_callback.py +++ b/tests/checkpoints/test_checkpoint_callback.py @@ -24,23 +24,6 @@ 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): """Build a minimal mock trainer for CheckpointCallback / manager tests.""" checkpoint_cfg = SimpleNamespace( @@ -91,6 +74,9 @@ def _make_mock_trainer(save_path="/tmp/test_ckpt", save_async=False): # 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 + # The manager merges ``trainer.extra_state()`` into the model blob; the mock + # contributes nothing by default. + trainer.extra_state.return_value = {} return trainer @@ -349,7 +335,7 @@ def test_save_forwards_extra_state_like_optimizer(self, mock_helper, mock_dist, manager.save_dcp(TrainerState(global_step=10)) saved = manager.checkpointer.save.call_args.args[1] - assert set(saved["extra_state"]) == {"lr_scheduler", "condition_model_rng_state"} + assert set(saved["extra_state"]) == {"lr_scheduler"} assert saved["optimizer"] is trainer.optimizer def test_load_forwards_extra_state_like_optimizer(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps): @@ -367,34 +353,27 @@ def test_load_forwards_extra_state_like_optimizer(self, mock_helper, mock_dist, 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): + def test_extra_state_merges_trainer_contribution(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps): trainer = _make_mock_trainer() - trainer.condition_model = _StubConditionModel() + trainer.extra_state.return_value = {"condition_model_rng_state": {"generator": "state"}} 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"} + assert extra_state["condition_model_rng_state"] == {"generator": "state"} + assert extra_state["lr_scheduler"] == {"lr": 1e-4} - def test_load_extra_state_restores_condition_model_rng(self, mock_helper, mock_dist, mock_build_ckpt, mock_get_ps): + def test_load_extra_state_delegates_to_trainer(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}) + extra = {"lr_scheduler": {"lr": 1e-4}, "condition_model_rng_state": {"generator": "state"}} + manager._load_extra_state(extra) - assert torch.equal(condition_model.draw_noise(), expected_next_noise) + trainer.load_extra_state.assert_called_once_with(extra) + trainer.lr_scheduler.load_state_dict.assert_called_once_with({"lr": 1e-4}) 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 diff --git a/veomni/models/checkpoint_manager.py b/veomni/models/checkpoint_manager.py index 9c40ffe95d..5dd98e4bfe 100644 --- a/veomni/models/checkpoint_manager.py +++ b/veomni/models/checkpoint_manager.py @@ -139,25 +139,15 @@ def load_dir(self) -> Optional[str]: def _extra_state(self, state: "TrainerState") -> Dict[str, Any]: """Model-bound state to store beside the weights. - The condition model's noise/timestep generator and the lr scheduler are - model-bound state, persisted here for the checkpointer to write per rank. + The lr scheduler plus whatever the trainer contributes via + ``extra_state()``. """ lr_scheduler = self.trainer.lr_scheduler - condition_model = getattr(self.trainer, "condition_model", None) - rng_state_dict = getattr(condition_model, "rng_state_dict", None) - if ( - condition_model is not None - and rng_state_dict is None - and getattr(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 { + extra_state = { "lr_scheduler": None if lr_scheduler is None else lr_scheduler.state_dict(), - "condition_model_rng_state": None if rng_state_dict is None else rng_state_dict(), } + extra_state.update(self.trainer.extra_state()) + return extra_state def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: lr_state = extra_state.get("lr_scheduler") @@ -165,16 +155,7 @@ def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: if lr_state is not None and lr_scheduler is not None: lr_scheduler.load_state_dict(lr_state) - condition_model_rng_state = extra_state.get("condition_model_rng_state") - if condition_model_rng_state is not None: - loader = getattr(getattr(self.trainer, "condition_model", None), "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) + self.trainer.load_extra_state(extra_state) def wait_for_pending_save(self) -> None: """Block until the in-flight async save is on disk, if there is one.""" diff --git a/veomni/trainer/base.py b/veomni/trainer/base.py index b0822cfbba..b71e870b98 100644 --- a/veomni/trainer/base.py +++ b/veomni/trainer/base.py @@ -684,6 +684,19 @@ def save_dcp(self, state: TrainerState) -> None: """Write this job'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. + + Trainers with model-specific 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 job's weights in whichever format the model was trained in.""" self.checkpoint.save_hf_or_lora(state, stage=stage) diff --git a/veomni/trainer/dit_trainer.py b/veomni/trainer/dit_trainer.py index 1d3530c2c8..5351d99535 100644 --- a/veomni/trainer/dit_trainer.py +++ b/veomni/trainer/dit_trainer.py @@ -322,9 +322,39 @@ def _build_condition_model( self.condition_model.to(get_device_type()) logger.info_rank0("Condition model loaded.") # The condition model owns the noise/timestep generator, persisted in the - # DCP extra_state. ``ModelCheckpointManager`` reaches it through the - # reused ``BaseTrainer`` instance, so expose it there as well. + # DCP extra_state. ``ModelCheckpointManager`` holds the reused + # ``BaseTrainer`` instance, so expose the model and its extra-state hooks + # through it. self.base.condition_model = self.condition_model + self.base.extra_state = self.extra_state + self.base.load_extra_state = self.load_extra_state + + 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 _freeze_model_module(self): self.condition_model.requires_grad_(False) From e27c398a87d8b363b6a1c37489e3272234397461 Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Sun, 20 Sep 2026 11:27:55 +0800 Subject: [PATCH 7/9] [ckpt] refactor: adapt tests and single-process rank to per-rank extra_state Keep the lr scheduler and condition-model RNG in the per-rank extra_state/ layout rather than the upstream lr_scheduler.pt sidecar: gate the extra_state rank on dist.is_initialized() so single-process saves read/write rank 0, and rewrite the checkpoint tests that still asserted the replicated sidecar. --- tests/checkpoints/test_dcp_checkpointer.py | 200 +++++++-------------- veomni/checkpoint/dcp_checkpointer.py | 22 +-- 2 files changed, 78 insertions(+), 144 deletions(-) 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/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index f351362036..941a4bdfdd 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -610,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 @@ -793,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. @@ -977,13 +978,14 @@ def load( load training state from distributed checkpoint Mirrors :meth:`save`: weights from ``model//ckpt``, optimizer from - ``model//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//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`. @@ -1243,7 +1245,8 @@ def _save_extra_state(cls, checkpoint_dir: str, state: Dict[str, Any]) -> None: return extra_state_dir = os.path.join(checkpoint_dir, _EXTRA_STATE_DIRNAME) os.makedirs(extra_state_dir, exist_ok=True) - extra_state_path = os.path.join(extra_state_dir, _EXTRA_STATE_FORMAT.format(dist.get_rank())) + 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) @classmethod @@ -1252,9 +1255,8 @@ def _load_extra_state(cls, checkpoint_dir: str, state: Dict[str, Any]) -> None: if _EXTRA_STATE_KEY not in state: logger.warning_rank0("extra_state not found in state, skipping extra_state load") return - extra_state_path = os.path.join( - checkpoint_dir, _EXTRA_STATE_DIRNAME, _EXTRA_STATE_FORMAT.format(dist.get_rank()) - ) + 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) From b7cf4f0480a4029c71e4317948cea7cc8cea0e6b Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Sun, 20 Sep 2026 13:02:28 +0800 Subject: [PATCH 8/9] [ckpt] fix: install the resume-align stand-in condition model on the runtime The condition model is built by DiTModelRuntime, not DiTTrainer, so the test's trainer-level overrides were dead code and the real from_pretrained() ran with a null condition_model_path. Swap the stand-in in via _build_model_runtime(). --- tests/checkpoints/test_dit_resume_align.py | 42 +++++++++++++--------- 1 file changed, 26 insertions(+), 16 deletions(-) diff --git a/tests/checkpoints/test_dit_resume_align.py b/tests/checkpoints/test_dit_resume_align.py index 5a44978727..f163dc734f 100644 --- a/tests/checkpoints/test_dit_resume_align.py +++ b/tests/checkpoints/test_dit_resume_align.py @@ -151,7 +151,7 @@ def process_condition(self, hidden_states=None, encoder_hidden_states=None, late 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 DiTTrainer, VeOmniDiTArguments + from veomni.trainer.dit_trainer import DiTModelRuntime, DiTTrainer, VeOmniDiTArguments, VeOmniModelRuntime class SignatureCallback(Callback): def __init__(self, trainer) -> None: @@ -172,6 +172,29 @@ def on_train_end(self, state: TrainerState, **kwargs) -> None: 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" @@ -184,21 +207,8 @@ def __init__(self, args: "VeOmniDiTArguments"): # consumer on every forward pass. self._model_dropout = torch.nn.Dropout(DROPOUT_P) - 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) - # ``GlobalStateCallback`` reaches the condition model through the - # reused ``BaseTrainer`` instance, exactly as the real DiT trainer - # arranges it. - self.base.condition_model = self.condition_model - - def _freeze_model_module(self) -> None: - self.base.lora = False - - def _build_model_assets(self) -> None: - self.base.model_assets = [self.base.model.config] + 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): From 6d60499ea973385550b75c07a0eea7856bb714a5 Mon Sep 17 00:00:00 2001 From: Feng0w0 Date: Sun, 20 Sep 2026 15:31:15 +0800 Subject: [PATCH 9/9] [ckpt] fix: pin eager RoPE in the DiT resume-align test Wan's ``rope_apply`` has a non-standard signature, so its device patch explicitly disables the liger RoPE backend and raises if it is selected. The framework default for ``rotary_pos_emb_implementation`` is ``liger_kernel``, and Wan's production YAMLs pin ``eager`` to avoid it, but the resume-align test builds the toy Wan model from a config path and never went through a YAML. On the GPU CI that made both the in-process weight materialization and the trainer subprocess fail with ValueError: rotary_pos_emb_implementation='liger_kernel' is explicitly disabled for Wan Pin ``eager`` in the test, matching the Wan configs. --- tests/checkpoints/test_dit_resume_align.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/checkpoints/test_dit_resume_align.py b/tests/checkpoints/test_dit_resume_align.py index f163dc734f..2a5fc2eed1 100644 --- a/tests/checkpoints/test_dit_resume_align.py +++ b/tests/checkpoints/test_dit_resume_align.py @@ -275,7 +275,11 @@ def _materialize_toy_dit_weights(target_dir: str) -> None: from veomni.models.auto import build_foundation_model from veomni.ops import apply_ops_config - apply_ops_config(OpsImplementationConfig()) + # 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, @@ -335,6 +339,9 @@ def _run_trainer(run: str, train_path: str, weights_dir: str, load_path: str | N "--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