From 8773090325de7a86475669adb4354441a9c86de9 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 29 Oct 2026 13:34:41 +0800 Subject: [PATCH 001/120] [ops, dist] fix: compute NPU expert weight gradients in FP32 --- docs/design/op_selection.md | 35 ++++ tests/ci/_npu_moe_fp32_cold_worker.py | 181 ++++++++++++++++++ tests/ci/test_npu_moe_fp32_cold.py | 52 +++++ tests/ci/test_npu_moe_precision_consensus.py | 161 ++++++++++++++++ tests/ops/moe_experts/test_npu_fp32_wgrad.py | 138 +++++++++++++ veomni/arguments/arguments_types.py | 13 ++ veomni/distributed/npu_moe_precision.py | 116 +++++++++++ veomni/distributed/torch_parallelize.py | 83 +++++++- veomni/models/model_runtime.py | 2 + .../omni_module/omni_module_runtime.py | 3 + .../patched_modeling_qwen4_exp_gpu.diff | 78 ++++---- .../patched_modeling_qwen4_exp_gpu.py | 2 + .../patched_modeling_qwen4_exp_npu.diff | 78 ++++---- .../patched_modeling_qwen4_exp_npu.py | 2 + .../qwen4_exp_gpu_patch_gen_config.py | 2 + .../ops/kernels/moe_experts/standard/npu.py | 43 +++-- veomni/trainer/base.py | 7 + 17 files changed, 901 insertions(+), 95 deletions(-) create mode 100644 tests/ci/_npu_moe_fp32_cold_worker.py create mode 100644 tests/ci/test_npu_moe_fp32_cold.py create mode 100644 tests/ci/test_npu_moe_precision_consensus.py create mode 100644 tests/ops/moe_experts/test_npu_fp32_wgrad.py create mode 100644 veomni/distributed/npu_moe_precision.py diff --git a/docs/design/op_selection.md b/docs/design/op_selection.md index 01ea9e4f9c..4ca6df151c 100644 --- a/docs/design/op_selection.md +++ b/docs/design/op_selection.md @@ -591,6 +591,41 @@ path instead. VeOmni's product-based Ascend images install and verify `triton-ascend`; other environments can install a release compatible with their CANN and `torch_npu` stack to enable the fused activation. +#### FP32 expert weight gradients on NPU + +Qwen4-Exp SFT can opt into FP32 expert weight gradients with: + +```yaml +model: + accelerator: + fsdp_config: + low_precision_reduce_scatter_comm: false + mixed_precision: + enable: true + param_dtype: bfloat16 + reduce_dtype: float32 + npu_expert_fp32_wgrad: true +``` + +The default is `false`. With the option enabled, expert master weights and their +FSDP unsharded representation stay FP32; the fused NPU MoE casts weights to +BF16 inside the custom autograd function for the forward and input gradient, +while its grouped weight-gradient GEMM and gradient reduction use FP32. This +avoids rounding each local expert gradient to BF16 before reduction; it does +not make the rest of the model's BF16 computation equivalent to FP32. + +The current integration requires Qwen4-Exp SiLU `fused_npu`, FSDP2, TP=PP=1, +FP32 reduction, native communication, and `torch.compile` disabled. The +common NPU trainer initialization checks the precision config and MoE backend +across all job ranks before mesh creation; direct runtimes, composed Omni wraps, +and custom model parallelization without that check reject the opt-in. DPO reference-model isolation and resume +configuration checks are not wired by this option yet. Keep the setting fixed +across checkpoint save and resume when comparing training trajectories. + +FP32 expert all-gathers and temporary FP32 GEMM operands can increase memory, +communication and compute time. Measure them on the target model and topology +before adopting this experimental option in a production configuration. + ### Key files - Config: `veomni/arguments/arguments_types.py` — `OpsImplementationConfig` diff --git a/tests/ci/_npu_moe_fp32_cold_worker.py b/tests/ci/_npu_moe_fp32_cold_worker.py new file mode 100644 index 0000000000..4ffb0ea5cb --- /dev/null +++ b/tests/ci/_npu_moe_fp32_cold_worker.py @@ -0,0 +1,181 @@ +# 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. + +"""Two-rank native FSDP2/EP expert-gradient and fresh-process DCP probe.""" + +import os +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch import nn +from torch.distributed._tensor import DTensor, Shard +from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict + +from veomni.arguments.arguments_types import MixedPrecisionConfig +from veomni.distributed.npu_moe_precision import validate_expert_precision_consensus +from veomni.distributed.parallel_plan import ParallelPlan +from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state +from veomni.distributed.torch_parallelize import build_parallelize_model +from veomni.models.transformers.qwen4_exp.generated.patched_modeling_qwen4_exp_npu import Qwen4ExpTextExperts +from veomni.ops.config import set_ops_config + + +class ToyDecoder(nn.Module): + def __init__(self): + super().__init__() + config = SimpleNamespace(num_experts=4, hidden_size=256, moe_intermediate_size=32, hidden_act="silu") + self.experts = Qwen4ExpTextExperts(config) + + def forward(self, x, indices, weights): + return self.experts(x, indices, weights) + + +class ToyQwen4(nn.Module): + _no_split_modules = ["ToyDecoder"] + + def __init__(self): + super().__init__() + self.config = SimpleNamespace(model_type="qwen4_exp") + self.layer = ToyDecoder() + + def forward(self, x, indices, weights): + return self.layer(x, indices, weights) + + def init_weights(self): + with torch.no_grad(): + for name, param in self.named_parameters(): + param.fill_(0.01 if "gate" in name else 0.02) + + def get_parallel_plan(self): + plan = ParallelPlan( + extra_parallel_plan={"ep": {"layer.experts.gate_up_proj": Shard(0), "layer.experts.down_proj": Shard(0)}} + ) + plan.extra_parallel_fsdp_no_shard_module = {"ep": {"layer.experts"}} + return plan + + +def _local(tensor): + return tensor.to_local() if isinstance(tensor, DTensor) else tensor + + +def _snapshot(model, optimizer): + result = {} + for name, param in model.named_parameters(): + state = optimizer.state[param] + result[name] = { + "param": _local(param.detach()).cpu().clone(), + "m": _local(state["exp_avg"]).detach().cpu().clone(), + "v": _local(state["exp_avg_sq"]).detach().cpu().clone(), + "step": _local(state["step"]).detach().cpu().clone(), + } + assert all(result[name][part].dtype == torch.float32 for part in ("param", "m", "v")) + return result + + +def _step(model, optimizer, rank, ep, step): + generator = torch.Generator().manual_seed(9900 + rank * 10 + step) + x = torch.randn(8, 256, generator=generator).to(device="npu", dtype=torch.bfloat16) + # Each EP rank owns a zero-token expert while experts 0 and 2 receive tokens. + indices = torch.tensor([[0, 2]] * 8, device="npu", dtype=torch.int64) + weights = torch.full((8, 2), 0.5, device="npu", dtype=torch.float32) + optimizer.zero_grad(set_to_none=True) + loss = model(x, indices, weights).float().square().mean() + loss.backward() + grads = {} + for name, param in model.named_parameters(): + assert param.grad is not None and param.grad.dtype == torch.float32 + grads[name] = _local(param.grad).detach().cpu().clone() + assert torch.isfinite(grads[name]).all() + assert grads[name].shape[0] == 2 + assert torch.count_nonzero(grads[name][0]).item() > 0 + assert torch.count_nonzero(grads[name][1]).item() == 0 + assert torch.any(grads[name][0] != grads[name][0].bfloat16().float()).item() + optimizer.step() + return { + "input": x.cpu(), + "indices": indices.cpu(), + "weights": weights.cpu(), + "loss": loss.detach().cpu(), + "grads": grads, + "state": _snapshot(model, optimizer), + } + + +def run(rank, phase, output_dir, ep, endpoint=None): + import torch_npu # noqa: F401 + + torch.npu.set_device(rank) + kwargs = {"backend": "hccl", "rank": rank, "world_size": 2, "timeout": timedelta(seconds=180)} + if endpoint is not None: + kwargs["store"] = endpoint.connect() + dist.init_process_group(**kwargs) + try: + mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=True) + fsdp = SimpleNamespace(mixed_precision=mixed, fsdp_mode="fsdp2", low_precision_reduce_scatter_comm=False) + accelerator = SimpleNamespace( + fsdp_config=fsdp, tp_size=1, pp_size=1, torch_compile=SimpleNamespace(enable=False) + ) + validate_expert_precision_consensus(accelerator, "fused_npu") + _init_parallel_state( + dp_size=2, + dp_shard_size=2, + device_type="npu", + name="npu-fp32-experts", + extra_parallel_sizes=(ep,), + extra_parallel_names=("ep",), + extra_parallel_placement_innermost=(False,), + ) + set_ops_config(SimpleNamespace(moe_implementation="fused_npu")) + with torch.device("meta"): + model = ToyQwen4() + model = build_parallelize_model( + model, + mixed_precision=mixed, + enable_gradient_checkpointing=False, + init_device="meta", + enable_forward_prefetch=False, + ) + assert model.layer.experts._veomni_fp32_expert_wgrad + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.0) + ckpt = str(Path(output_dir) / f"ep{ep}-dcp") + if phase == "fresh": + _step(model, optimizer, rank, ep, 1) + model_state, optim_state = get_state_dict(model, optimizer) + dcp.save({"model": model_state, "optimizer": optim_state}, checkpoint_id=ckpt) + else: + model_state, optim_state = get_state_dict(model, optimizer) + dcp.load({"model": model_state, "optimizer": optim_state}, checkpoint_id=ckpt) + set_state_dict(model, optimizer, model_state_dict=model_state, optim_state_dict=optim_state) + before = _snapshot(model, optimizer) + after = _step(model, optimizer, rank, ep, 2) + for name in before: + assert not torch.equal(before[name]["param"], after["state"][name]["param"]) + torch.save({"before": before, "after": after}, Path(output_dir) / f"ep{ep}-{phase}-rank{rank}.pt") + dist.barrier() + finally: + clear_parallel_state() + dist.destroy_process_group() + + +if __name__ == "__main__": + run( + int(os.environ["RANK"]), + os.environ["FP32_EXPERT_PHASE"], + os.environ["FP32_EXPERT_OUTPUT"], + int(os.environ["FP32_EXPERT_EP"]), + ) diff --git a/tests/ci/test_npu_moe_fp32_cold.py b/tests/ci/test_npu_moe_fp32_cold.py new file mode 100644 index 0000000000..08a4813de5 --- /dev/null +++ b/tests/ci/test_npu_moe_fp32_cold.py @@ -0,0 +1,52 @@ +# 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. + +"""FSDP2 + expert-parallel optimizer/DCP cold-recovery regression on Ascend.""" + +import os + +import pytest +import torch + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from tests.ci._npu_moe_fp32_cold_worker import run +from veomni.utils.device import IS_NPU_AVAILABLE + + +def _assert_equal(left, right): + if isinstance(left, dict): + assert left.keys() == right.keys() + for key in left: + _assert_equal(left[key], right[key]) + elif isinstance(left, torch.Tensor): + assert torch.equal(left, right) + else: + assert left == right + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2)) +@pytest.mark.parametrize("ep", [1, 2]) +def test_native_fsdp_expert_fp32_dcp_cold_resume(tmp_path, ep): + if not IS_NPU_AVAILABLE or torch.npu.device_count() < 2: + if os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1": + pytest.fail("VEOMNI_REQUIRE_NPU_TESTS=1 requires two native NPU devices") + pytest.skip("requires two native Ascend devices") + from tests.tools.launch_utils import spawn_with_rendezvous + + for phase in ("fresh", "resume"): + spawn_with_rendezvous(run, args=(phase, str(tmp_path), ep), nprocs=2) + for rank in range(2): + fresh = torch.load(tmp_path / f"ep{ep}-fresh-rank{rank}.pt", map_location="cpu", weights_only=True) + resumed = torch.load(tmp_path / f"ep{ep}-resume-rank{rank}.pt", map_location="cpu", weights_only=True) + _assert_equal(fresh, resumed) diff --git a/tests/ci/test_npu_moe_precision_consensus.py b/tests/ci/test_npu_moe_precision_consensus.py new file mode 100644 index 0000000000..b89dc12c62 --- /dev/null +++ b/tests/ci/test_npu_moe_precision_consensus.py @@ -0,0 +1,161 @@ +# 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. + +"""CPU/Gloo coverage for the job-wide NPU expert precision handshake.""" + +import copy +import json +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch.distributed as dist +import torch.multiprocessing as mp + +from veomni.arguments.arguments_types import MixedPrecisionConfig +from veomni.distributed.npu_moe_precision import ( + require_expert_precision_consensus, + validate_expert_precision_consensus, + validate_qwen4_experts, +) + + +def _accelerator(enabled, *, reduce_dtype="float32", low_comm=False): + mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=enabled, reduce_dtype=reduce_dtype) + fsdp = SimpleNamespace(mixed_precision=mixed, fsdp_mode="fsdp2", low_precision_reduce_scatter_comm=low_comm) + return SimpleNamespace(fsdp_config=fsdp, tp_size=1, pp_size=1, torch_compile=SimpleNamespace(enable=False)) + + +def _worker(rank, directory, case): + dist.init_process_group( + "gloo", init_method=f"file://{directory}/rendezvous", rank=rank, world_size=2, timeout=timedelta(seconds=20) + ) + try: + enabled = case != "disabled" and (case != "mismatch_flag" or rank == 0) + dtype = "bfloat16" if case in ("disabled", "mismatch_dtype") and rank == 1 else "float32" + accelerator = _accelerator(enabled, reduce_dtype=dtype, low_comm=case == "low_comm") + if case == "mismatch_output_dtype" and rank == 1: + accelerator.fsdp_config.mixed_precision.output_dtype = "float32" + try: + backend = "eager" if case == "mismatch_backend" and rank == 1 else "fused_npu" + validate_expert_precision_consensus(accelerator, backend) + status = "ok" + if enabled: + mixed = accelerator.fsdp_config.mixed_precision + require_expert_precision_consensus(mixed, "fsdp2", False, 1, 1, False) + altered = copy.deepcopy(mixed) + altered.enable = False + with pytest.raises(ValueError, match="unchanged job-wide"): + require_expert_precision_consensus(altered, "fsdp2", False, 1, 1, False) + altered = copy.deepcopy(mixed) + altered.output_dtype = "float32" + with pytest.raises(ValueError, match="unchanged job-wide"): + require_expert_precision_consensus(altered, "fsdp2", False, 1, 1, False) + if case == "mismatch_module": + import torch.nn as nn + + class Qwen4ExpTextExperts(nn.Module): + def __init__(self): + super().__init__() + self.veomni_moe = SimpleNamespace(impl="eager" if rank == 1 else "fused_npu") + self.use_swiglu_mlp = True + + model = nn.Module() + model.config = SimpleNamespace(model_type="qwen4_exp") + model.experts = Qwen4ExpTextExperts() + validate_qwen4_experts(model) + except ValueError as exc: + status = str(exc) + (Path(directory) / f"rank{rank}.json").write_text(json.dumps({"status": status})) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + ("case", "expected"), + [ + ("disabled", "ok"), + ("enabled", "ok"), + ("mismatch_flag", "identical precision configuration"), + ("mismatch_dtype", "identical precision configuration"), + ("mismatch_output_dtype", "identical precision configuration"), + ("mismatch_backend", "identical precision configuration"), + ("mismatch_module", "identical Qwen4-Exp expert modules"), + ("low_comm", "native communication"), + ], +) +def test_all_ranks_agree_before_mesh(tmp_path, case, expected): + mp.spawn(_worker, args=(str(tmp_path), case), nprocs=2, join=True) + for rank in range(2): + result = json.loads((tmp_path / f"rank{rank}.json").read_text()) + assert expected in result["status"] + + +def test_caller_without_consensus_cannot_opt_in(): + mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=True) + with pytest.raises(ValueError, match="job-wide precision consensus"): + require_expert_precision_consensus(mixed, "fsdp2", False, 1, 1, False) + + +@pytest.mark.parametrize("value", ["false", 0, 1, None]) +def test_expert_precision_option_requires_real_boolean(value): + with pytest.raises(ValueError, match="must be a boolean"): + MixedPrecisionConfig(npu_expert_fp32_wgrad=value) + + +def test_direct_fsdp2_cannot_bypass_consensus(monkeypatch): + from veomni.distributed import torch_parallelize + + state = SimpleNamespace(dp_mode="fsdp2", tp_size=1, pp_size=1, fsdp_enabled=True) + monkeypatch.setattr(torch_parallelize, "get_parallel_state", lambda: state) + with pytest.raises(ValueError, match="job-wide precision consensus"): + torch_parallelize.parallelize_model_fsdp2( + object(), mixed_precision=MixedPrecisionConfig(npu_expert_fp32_wgrad=True) + ) + + +def test_qwen4_vlm_decoder_layers_are_fsdp_targets(): + import torch + from transformers import Qwen4ExpConfig + + from tests.models.compare import eager_ops_config, ops_config_scope + from veomni.distributed.torch_parallelize import _fsdp_target_classes, _is_fsdp_wrap_target + from veomni.models.transformers.qwen4_exp.generated.patched_modeling_qwen4_exp_npu import ( + Qwen4ExpForConditionalGeneration, + ) + + config = Qwen4ExpConfig.from_pretrained("tests/toy_config/qwen4_exp_toy/config.json") + config._attn_implementation = "eager" + with ops_config_scope(eager_ops_config()), torch.device("meta"): + model = Qwen4ExpForConditionalGeneration(config) + targets = _fsdp_target_classes(model, [], True) + blocks = { + name for name, module in model.named_modules() if _is_fsdp_wrap_target(name, type(module).__name__, targets) + } + experts = [name for name, module in model.named_modules() if type(module).__name__ == "Qwen4ExpTextExperts"] + assert experts and all(any(name.startswith(block + ".") for block in blocks) for name in experts) + # Also cover a root that has not inherited the child model's no-split list. + assert "Qwen4ExpTextDecoderLayer" in _fsdp_target_classes(SimpleNamespace(_no_split_modules=None), [], True) + assert not _fsdp_target_classes(SimpleNamespace(_no_split_modules=None), [], False) + + +@pytest.mark.parametrize("for_inference", [False, True]) +def test_seed_omni_module_does_not_silently_ignore_opt_in(for_inference): + from veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime import ModuleRuntime + + accelerator = _accelerator(True) + args = SimpleNamespace(accelerator=accelerator) + with pytest.raises(ValueError, match="unsupported by SeedOmni"): + ModuleRuntime(args, "test", module_config=None, global_accelerator=accelerator, for_inference=for_inference) diff --git a/tests/ops/moe_experts/test_npu_fp32_wgrad.py b/tests/ops/moe_experts/test_npu_fp32_wgrad.py new file mode 100644 index 0000000000..a68874cc5a --- /dev/null +++ b/tests/ops/moe_experts/test_npu_fp32_wgrad.py @@ -0,0 +1,138 @@ +# 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. + +"""NPU expert GEMM backward precision, including the BF16 default path.""" + +import os +import sys +import types + +import pytest +import torch +import torch.utils.checkpoint + +from veomni.ops.kernels.moe_experts.standard.npu import GmmFunction +from veomni.utils.device import IS_NPU_AVAILABLE + + +def _grouped(xs, weights, *, group_list, group_type, **kwargs): + x, weight = xs[0], weights[0] + assert x.dtype == weight.dtype + out, offset = [], 0 + for expert, count in enumerate(group_list): + count = int(count) + if group_type == 0: + out.append(x[offset : offset + count] @ weight[expert]) + else: + assert group_type == 2 + out.append(x[:, offset : offset + count] @ weight[offset : offset + count]) + offset += count + return [torch.cat(out) if group_type == 0 else torch.stack(out)] + + +def _reference(x, dy, counts): + start, grads = 0, [] + for count in counts: + grads.append(x[start : start + count].float().T @ dy[start : start + count].float()) + start += count + return torch.stack(grads) + + +@pytest.fixture +def fake_npu(monkeypatch): + monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_grouped_matmul=_grouped)) + + +def _cpu_inputs(counts=(6, 7)): + torch.manual_seed(714) + return ( + torch.randn(sum(counts), 11, dtype=torch.bfloat16), + torch.randn(len(counts), 11, 9), + torch.randn(sum(counts), 9, dtype=torch.bfloat16), + ) + + +def test_fp32_wgrad_keeps_bf16_forward_and_input_gradient(fake_npu): + x, master, dy = _cpu_inputs() + counts = [6, 7] + inp, param = x.clone().requires_grad_(True), master.clone().requires_grad_(True) + output = GmmFunction.apply(inp, param, counts, True) + output.backward(dy) + assert torch.equal(output.detach(), _grouped([x], [master.bfloat16()], group_list=counts, group_type=0)[0]) + expected_dx = _grouped([dy], [master.bfloat16().transpose(1, 2)], group_list=counts, group_type=0)[0] + assert torch.equal(inp.grad, expected_dx) + assert param.grad.dtype == torch.float32 + assert torch.equal(param.grad, _reference(x, dy, counts)) + assert not torch.equal(param.grad, param.grad.bfloat16().float()) + + +def test_bf16_default_and_cast_before_autograd_preserve_legacy_rounding(fake_npu): + x, master, dy = _cpu_inputs() + counts = [6, 7] + param = master.clone().requires_grad_(True) + GmmFunction.apply(x, param.to(torch.bfloat16), counts).backward(dy) + assert torch.equal(param.grad, param.grad.bfloat16().float()) + assert not torch.equal(param.grad, _reference(x, dy, counts)) + + +@pytest.mark.parametrize("reentrant", [True, False]) +def test_checkpoint_recomputation_keeps_fp32_weight_gradient(fake_npu, reentrant): + x, master, dy = _cpu_inputs() + counts = [6, 7] + inp, param = x.clone().requires_grad_(True), master.clone().requires_grad_(True) + out = torch.utils.checkpoint.checkpoint( + lambda a, b: GmmFunction.apply(a, b, counts, True), inp, param, use_reentrant=reentrant + ) + out.backward(dy) + assert torch.equal(param.grad, _reference(x, dy, counts)) + + +@pytest.mark.parametrize("counts", [[4, 0, 4], [0, 0, 0]]) +def test_empty_experts_keep_zero_fp32_gradient(fake_npu, counts): + x, master, dy = _cpu_inputs(counts) + param = master.requires_grad_(True) + GmmFunction.apply(x, param, counts, True).backward(dy) + assert param.grad.dtype == torch.float32 + assert torch.equal(param.grad, _reference(x, dy, counts)) + + +def test_bad_compute_or_master_dtype_is_rejected(fake_npu): + x, master, _ = _cpu_inputs() + with pytest.raises(ValueError, match="FP32 expert wgrad"): + GmmFunction.apply(x.float(), master, [6, 7], True) + with pytest.raises(ValueError, match="FP32 expert wgrad"): + GmmFunction.apply(x, master.bfloat16(), [6, 7], True) + + +def test_native_npu_fp32_wgrad_matches_cpu_reference(): + if not IS_NPU_AVAILABLE: + if os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1": + pytest.fail("VEOMNI_REQUIRE_NPU_TESTS=1 requires a native NPU") + pytest.skip("native NPU is unavailable") + import torch_npu # noqa: F401 + + torch.npu.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + torch.manual_seed(20261029) + counts = [13, 0, 17] + x = torch.randn(30, 256, dtype=torch.bfloat16, device="npu", requires_grad=True) + weight = torch.randn(3, 256, 32, dtype=torch.float32, device="npu", requires_grad=True) + dy = torch.randn(30, 32, dtype=torch.bfloat16, device="npu") + groups = torch.tensor(counts, dtype=torch.int64, device="npu") + output = GmmFunction.apply(x, weight, groups, True) + output.backward(dy) + torch.npu.synchronize() + expected = _reference(x.detach().cpu(), dy.cpu(), counts) + assert weight.grad.dtype == torch.float32 + torch.testing.assert_close(weight.grad.cpu(), expected, atol=1e-5, rtol=1e-5) + assert torch.count_nonzero(weight.grad[1]).item() == 0 diff --git a/veomni/arguments/arguments_types.py b/veomni/arguments/arguments_types.py index f689579aca..a9ae87fb58 100644 --- a/veomni/arguments/arguments_types.py +++ b/veomni/arguments/arguments_types.py @@ -567,8 +567,21 @@ class MixedPrecisionConfig: default=True, metadata={"help": "Enable mixed precision cast forward inputs (FSDP2)."}, ) + npu_expert_fp32_wgrad: bool = field( + default=False, + metadata={ + "help": ( + "Compute NPU fused MoE expert weight gradients in FP32 while keeping BF16 forward and " + "input gradients. Currently supports Qwen4-Exp with FSDP2, BF16 parameters, FP32 " + "reduction, and native communication; increases expert gather and temporary memory." + ) + }, + ) def __post_init__(self): + if not isinstance(self.npu_expert_fp32_wgrad, bool): + raise ValueError("npu_expert_fp32_wgrad must be a boolean (true or false).") + def _check_dtype(dtype: str): if dtype is not None and dtype not in ["bfloat16", "float32", "float16"]: raise ValueError(f"Invalid dtype {dtype} for mixed precision training.") diff --git a/veomni/distributed/npu_moe_precision.py b/veomni/distributed/npu_moe_precision.py new file mode 100644 index 0000000000..334c387006 --- /dev/null +++ b/veomni/distributed/npu_moe_precision.py @@ -0,0 +1,116 @@ +# 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. + +"""Job-wide agreement for the optional NPU expert weight-gradient precision.""" + +import torch.distributed as dist + + +def expert_precision_fingerprint(mixed, fsdp_mode, low_precision_comm, tp_size, pp_size, compile_enabled, moe_impl): + return ( + mixed.npu_expert_fp32_wgrad, + mixed.enable, + mixed.param_dtype, + mixed.reduce_dtype, + mixed.output_dtype, + mixed.cast_forward_inputs, + fsdp_mode, + low_precision_comm, + tp_size, + pp_size, + compile_enabled, + moe_impl, + ) + + +def validate_expert_precision_consensus(accelerator, moe_implementation): + """Agree before mesh creation; every NPU training rank must call this even when disabled.""" + fsdp = accelerator.fsdp_config + mixed = fsdp.mixed_precision + local = expert_precision_fingerprint( + mixed, + fsdp.fsdp_mode, + fsdp.low_precision_reduce_scatter_comm, + accelerator.tp_size, + accelerator.pp_size, + accelerator.torch_compile.enable, + moe_implementation or "fused_npu", + ) + peers = [None] * dist.get_world_size() + dist.all_gather_object(peers, local) + if not any(peer[0] for peer in peers): + return + if any(peer != local for peer in peers): + raise ValueError("NPU expert FP32 wgrad requires identical precision configuration on all ranks") + if ( + not mixed.enable + or mixed.param_dtype != "bfloat16" + or mixed.reduce_dtype != "float32" + or fsdp.fsdp_mode != "fsdp2" + or fsdp.low_precision_reduce_scatter_comm + or accelerator.tp_size != 1 + or accelerator.pp_size != 1 + or accelerator.torch_compile.enable + or local[-1] != "fused_npu" + ): + raise ValueError( + "NPU expert FP32 wgrad requires FSDP2, BF16 compute, FP32 reduction, native communication, " + "TP=PP=1, fused_npu MoE, and torch.compile disabled" + ) + mixed._npu_expert_fp32_wgrad_fingerprint = local + + +def require_expert_precision_consensus(mixed, fsdp_mode, low_precision_comm, tp_size, pp_size, compile_enabled): + """Refuse unsupported callers or a config mutated after the job-wide check.""" + if not mixed.npu_expert_fp32_wgrad: + return + original = getattr(mixed, "_npu_expert_fp32_wgrad_fingerprint", None) + current = expert_precision_fingerprint( + mixed, fsdp_mode, low_precision_comm, tp_size, pp_size, compile_enabled, "fused_npu" + ) + if original != current: + raise ValueError("NPU expert FP32 wgrad requires the unchanged job-wide precision consensus") + + +def validate_qwen4_experts(model): + """Check the constructed expert modules on every rank before any FSDP wrapping.""" + experts = [ + (name, module) for name, module in model.named_modules() if type(module).__name__ == "Qwen4ExpTextExperts" + ] + report = ( + getattr(getattr(model, "config", None), "model_type", None), + tuple( + ( + name, + getattr(getattr(module, "veomni_moe", None), "impl", None), + getattr(module, "use_swiglu_mlp", False), + tuple( + (param_name, str(param.dtype), tuple(param.shape)) + for param_name, param in module.named_parameters() + ), + ) + for name, module in experts + ), + ) + peers = [None] * dist.get_world_size() + dist.all_gather_object(peers, report) + if any(peer != report for peer in peers): + raise ValueError("NPU expert FP32 wgrad requires identical Qwen4-Exp expert modules on all ranks") + if report[0] != "qwen4_exp" or not experts: + raise ValueError("NPU expert FP32 wgrad currently requires Qwen4-Exp expert modules") + if any(impl != "fused_npu" or not swiglu for _, impl, swiglu, _ in report[1]): + raise ValueError("NPU expert FP32 wgrad requires Qwen4-Exp SiLU fused_npu experts") + if any(dtype != "torch.float32" for _, _, _, params in report[1] for _, dtype, _ in params): + raise ValueError("NPU expert FP32 wgrad requires FP32 expert master parameters") + return {id(module) for _, module in experts} diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py index 422196c605..591145f31a 100644 --- a/veomni/distributed/torch_parallelize.py +++ b/veomni/distributed/torch_parallelize.py @@ -415,6 +415,14 @@ def _is_fsdp_wrap_target(fqn: str, class_name: str, targets: Collection[str]) -> return False +def _fsdp_target_classes(model, basic_modules, expert_fp32_wgrad): + classes = set(getattr(model, "_no_split_modules", None) or []) | set(basic_modules or []) + if expert_fp32_wgrad: + # The VLM root may not expose the text model's own no-split modules. + classes.add("Qwen4ExpTextDecoderLayer") + return classes + + def _configure_fsdp_gradient_reduction( module: FSDPModule, *, @@ -465,6 +473,20 @@ def parallelize_model_fsdp2( We will use this model for illustration of Expert Parallel + Embed Parallel below. """ parallel_state = get_parallel_state() + compile_config = compile_config or CompileConfig() + if mixed_precision.npu_expert_fp32_wgrad: + from .npu_moe_precision import require_expert_precision_consensus + + require_expert_precision_consensus( + mixed_precision, + parallel_state.dp_mode, + low_precision_reduce_scatter_comm, + parallel_state.tp_size, + parallel_state.pp_size, + compile_config.enable, + ) + if not parallel_state.fsdp_enabled or parallel_state.dp_mode != "fsdp2": + raise ValueError("FP32 expert weight gradients require FSDP2") use_low_precision_transport = validate_low_precision_reduce_scatter_comm( low_precision_reduce_scatter_comm, mixed_precision @@ -475,8 +497,18 @@ def parallelize_model_fsdp2( elif low_precision_reduce_scatter_comm: logger.info_rank0("Parameter dtype matches reduce dtype; using the native PyTorch collective.") - model_no_split_modules = getattr(model, "_no_split_modules", None) or [] - target_classes = set(model_no_split_modules) | set(basic_modules or []) + expert_fp32_wgrad = mixed_precision.npu_expert_fp32_wgrad + if expert_fp32_wgrad: + if not IS_NPU_AVAILABLE or get_device_type() != "npu": + raise ValueError("FP32 expert weight gradients require an NPU") + from .npu_moe_precision import validate_qwen4_experts + + expert_ids = validate_qwen4_experts(model) + else: + expert_ids = set() + wrapped_expert_ids = set() + + target_classes = _fsdp_target_classes(model, basic_modules, expert_fp32_wgrad) # Make a list of tuples that contains target classes' name and module # Note that all target classes should include all ExtraParallel modules. @@ -493,8 +525,6 @@ def parallelize_model_fsdp2( logger.info_rank0(f"target classes to shard: {target_classes}") model._veomni_compile_enabled = False - compile_config = compile_config or CompileConfig() - # Step 1: Apply ExtraParallel # e.g. Apply expert parallelism (slice expert tensors [128,H,I] -> [16,H,I]) # Apply embed parallelism (slice embed tensors [64,H] -> [16,H]) @@ -653,6 +683,16 @@ def parallelize_model_fsdp2( cast_forward_inputs=mixed_precision.cast_forward_inputs, ) fsdp_kwargs["mp_policy"] = mp_policy + + if expert_fp32_wgrad: + expert_mp_policy = MixedPrecisionPolicy( + param_dtype=torch.float32, reduce_dtype=torch.float32, cast_forward_inputs=False + ) + + def configure_fp32_expert(expert): + expert._veomni_fp32_expert_wgrad = True + wrapped_expert_ids.add(id(expert)) + # prepare offload_policy kwargs enable_fsdp_cpu_offload = kwargs.pop("enable_fsdp_offload", False) offload_pin_memory = kwargs.pop("fsdp_offload_pin_memory", True) @@ -791,7 +831,13 @@ def parallelize_model_fsdp2( # decoder/root FSDP group through the normal bottom-up walk. continue # shard para module (e.g. expert/decoder.moe, embed_tokens/decoder.embed_tokens) - fully_shard(_para_mod, **extra_parallel_fsdp_kwargs[para]) + para_kwargs = dict(extra_parallel_fsdp_kwargs[para]) + if expert_fp32_wgrad and id(_para_mod) in expert_ids: + if para != "ep" or id(_para_mod) in wrapped_expert_ids: + raise RuntimeError("Unexpected or duplicate expert FSDP wrap") + configure_fp32_expert(_para_mod) + para_kwargs["mp_policy"] = expert_mp_policy + fully_shard(_para_mod, **para_kwargs) # average para (e.g. ep) grads across para (e.g. ep) ranks # NOTE: in torch 2.8 and later we should use # experts_mod.set_gradient_divide_factor(parallel_state.ep_size) @@ -812,6 +858,14 @@ def parallelize_model_fsdp2( ) layer_mod._fsdp_modules.append(_para_mod) + if expert_fp32_wgrad and not parallel_state.extra_parallel_enabled("ep"): + for expert in layer_mod.modules(): + if id(expert) not in expert_ids or id(expert) in wrapped_expert_ids: + continue + configure_fp32_expert(expert) + fully_shard(expert, **{**fsdp_kwargs, "mp_policy": expert_mp_policy}) + layer_mod._fsdp_modules.append(expert) + # shard module that needs to ignore mixed precision control if mp_ignored_classes: for sub_mod in layer_mod.modules(): @@ -834,6 +888,9 @@ def parallelize_model_fsdp2( layer_mod._fsdp_modules.append(layer_mod) logger.info_rank0(f"{layer_fqn=}, {layer_mod._fsdp_modules=}") + if expert_fp32_wgrad and wrapped_expert_ids != expert_ids: + raise RuntimeError("Not every Qwen4-Exp expert was wrapped for FP32 weight gradients") + # Shard root WITHOUT explicit `reshard_after_forward` so FSDP2's # auto-no-reshard for the root kicks in (`_fsdp_state.py:_lazy_init`). # This keeps the root's params (lm_head + embeddings + final norm) @@ -1032,13 +1089,25 @@ def build_parallelize_model( enabled gradient checkpointing and before FSDP2 shards the blocks. """ parallel_state = get_parallel_state() + compile_config = compile_config or CompileConfig() + if mixed_precision.npu_expert_fp32_wgrad: + from .npu_moe_precision import require_expert_precision_consensus + + require_expert_precision_consensus( + mixed_precision, + parallel_state.dp_mode, + low_precision_reduce_scatter_comm, + parallel_state.tp_size, + parallel_state.pp_size, + compile_config.enable, + ) + if not parallel_state.fsdp_enabled or parallel_state.dp_mode != "fsdp2": + raise ValueError("FP32 expert weight gradients require FSDP2") if low_precision_reduce_scatter_comm is not False: # Only literal False bypasses validation; false-like non-booleans must still raise. validate_low_precision_reduce_scatter_comm( low_precision_reduce_scatter_comm, mixed_precision, fsdp_mode=parallel_state.dp_mode ) - compile_config = compile_config or CompileConfig() - if not parallel_state.fsdp_enabled: if kwargs.get("init_device") not in ["cuda", "npu", "mlu"]: raise ValueError("Only FSDP training supports `init_device=meta`.") diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index 450dbbdf45..afd32d7aaf 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -419,6 +419,8 @@ def _build_parallelized_model(self) -> None: # Customized parallelize model. customized_parallelize_model_function = getattr(self.model, "build_parallelize_model", None) if callable(customized_parallelize_model_function): + if args.accelerator.fsdp_config.mixed_precision.npu_expert_fp32_wgrad: + raise ValueError("NPU expert FP32 wgrad is unsupported with custom model parallelization") parallelized_model = customized_parallelize_model_function(weights_path=self.weights_path, args=self.args) if parallelized_model is not None: self.model = parallelized_model diff --git a/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py b/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py index 24bb07e061..f2d11c30e0 100644 --- a/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py +++ b/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py @@ -138,6 +138,9 @@ def __init__( self.lr_scheduler = None self._global_accelerator = global_accelerator + if args.accelerator.fsdp_config.mixed_precision.npu_expert_fp32_wgrad: + raise ValueError("NPU expert FP32 wgrad is unsupported by SeedOmni module runtime") + if for_inference: if self.is_eager: self._init_eager_inference() diff --git a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.diff b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.diff index b938996f2a..51e4e7594f 100644 --- a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.diff +++ b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.diff @@ -1356,7 +1356,7 @@ def forward( self, -@@ -874,25 +1432,29 @@ +@@ -874,25 +1432,31 @@ top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: @@ -1391,6 +1391,7 @@ + self.num_experts, + ) + unused = self.gate_up_proj.new_empty(0) ++ expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} + return self.veomni_moe( + hidden_states, + top_k_weights, @@ -1400,12 +1401,13 @@ + self.down_proj, + self.gate_up_proj, + num_experts=self.num_experts, ++ **expert_options, + ) + # --- Patch.1 --- class Qwen4ExpTextTopKRouter(nn.Module): -@@ -916,6 +1478,12 @@ +@@ -916,6 +1480,12 @@ return router_logits, router_scores, router_indices @@ -1418,7 +1420,7 @@ class Qwen4ExpTextSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() -@@ -924,11 +1492,33 @@ +@@ -924,11 +1494,33 @@ self.shared_expert = Qwen4ExpTextMLP(config, intermediate_size=config.shared_expert_intermediate_size) self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) @@ -1454,7 +1456,7 @@ expert_output = self.experts(hidden_states_reshaped, selected_experts, routing_weights) shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states_reshaped)) * shared_expert_output -@@ -970,9 +1560,13 @@ +@@ -970,9 +1562,13 @@ _MASK64 = (1 << 64) - 1 @@ -1468,7 +1470,7 @@ _PRIME_1 = 10007 -@@ -1015,8 +1609,14 @@ +@@ -1015,8 +1611,14 @@ return prime @@ -1484,7 +1486,7 @@ super().__init__() self.layer_idx = layer_idx self.ngram_size = config.ngram_size -@@ -1028,7 +1628,6 @@ +@@ -1028,7 +1630,6 @@ self.ngram_vocab_size_base = config.ngram_vocab_size_base head_dim_per_ngram = embedding_dim // self.ngram_heads self.seed = config.seed @@ -1492,7 +1494,7 @@ self.eos_token_id = config.eos_token_id[0] if isinstance(config.eos_token_id, list) else config.eos_token_id self.head_vocab_sizes = [] -@@ -1046,9 +1645,29 @@ +@@ -1046,9 +1647,29 @@ ) self.ngram_heads_vocab_sizes = nn.Buffer(torch.tensor(self.head_vocab_sizes, dtype=torch.long)) self.ngram_heads_offsets = nn.Buffer(torch.tensor(self.head_offsets, dtype=torch.long)) @@ -1525,7 +1527,7 @@ def _shift_right_ignore_eos(self, token_ids: torch.Tensor, shift: int) -> torch.Tensor: if shift == 0: -@@ -1066,19 +1685,183 @@ +@@ -1066,19 +1687,183 @@ valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0) return torch.where(valid, shifted, token_ids.new_full((), self.eos_token_id)) @@ -1716,7 +1718,7 @@ if ( not past_key_values.has_previous_state(self.layer_idx, state_idx=2) and input_ids.shape[1] < self.context_len -@@ -1090,28 +1873,48 @@ +@@ -1090,28 +1875,48 @@ input_ids_to_cache, self.layer_idx, state_idx=2, conv_kernel_size=self.context_len ) @@ -1773,7 +1775,7 @@ class Qwen4ExpTextPLELayer(nn.Module): -@@ -1147,24 +1950,91 @@ +@@ -1147,24 +1952,91 @@ bias=False, ) @@ -1881,7 +1883,7 @@ def forward( self, -@@ -1172,8 +2042,25 @@ +@@ -1172,8 +2044,25 @@ input_ids: torch.Tensor, past_key_values: Cache | None, conv_mask: torch.Tensor | None = None, @@ -1908,7 +1910,7 @@ key_normed = self.norm_key(self.key_proj(embeddings)).unflatten(-1, (self.hc_count, self.hidden_size)) value = self.value_proj(embeddings) query_normed = self.norm_query(hidden_states).unflatten(-1, (self.hc_count, self.hidden_size)) -@@ -1185,8 +2072,21 @@ +@@ -1185,8 +2074,21 @@ if conv_mask is not None: gated_value = apply_mask_to_padding_states(gated_value, conv_mask) gated_value_normed = apply_mask_to_padding_states(gated_value_normed, conv_mask) @@ -1931,7 +1933,7 @@ class Qwen4ExpTextDecoderLayer(GradientCheckpointingLayer): -@@ -1204,6 +2104,15 @@ +@@ -1204,6 +2106,15 @@ self.attn_hyper_connection = Qwen4ExpTextGatedResidual(config) self.mlp_hyper_connection = Qwen4ExpTextGatedResidual(config) @@ -1947,7 +1949,7 @@ def forward( self, hidden_states: torch.Tensor, -@@ -1212,17 +2121,42 @@ +@@ -1212,17 +2123,42 @@ conv_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, ple_input_ids: torch.LongTensor | None = None, @@ -1993,7 +1995,7 @@ ) else: hidden_states, _ = self.self_attn( -@@ -1230,6 +2164,9 @@ +@@ -1230,6 +2166,9 @@ position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, @@ -2003,7 +2005,7 @@ **kwargs, ) -@@ -1308,6 +2245,12 @@ +@@ -1308,6 +2247,12 @@ rope_deltas: torch.LongTensor | None = None router_logits: tuple[torch.FloatTensor] | None = None @@ -2016,7 +2018,7 @@ @auto_docstring -@@ -1352,20 +2295,73 @@ +@@ -1352,20 +2297,73 @@ """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -2092,7 +2094,7 @@ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1) -@@ -1381,31 +2377,45 @@ +@@ -1381,31 +2379,45 @@ else: text_position_ids = None @@ -2144,7 +2146,7 @@ # CODEPATH: @ArthurZucker fix flagging for no reason here if self.config.ple_layer_ids and conv_mask is not None: eos_token_id = self.config.eos_token_id -@@ -1416,16 +2426,40 @@ +@@ -1416,16 +2428,40 @@ position_embeddings = self.rotary_emb(hidden_states, position_ids) hidden_states = hidden_states.repeat(1, 1, self.config.hc_count) @@ -2187,7 +2189,7 @@ hidden_states = self.hyper_connection_mixer(hidden_states) -@@ -1435,12 +2469,6 @@ +@@ -1435,12 +2471,6 @@ ) def reverse_embedding(self, inputs_embeds: torch.Tensor) -> torch.Tensor: @@ -2200,7 +2202,7 @@ with torch.no_grad(): input_ids = ( (inputs_embeds[:, :, None, :] == self.embed_tokens.weight[None, None, :, :]).all(dim=3).nonzero()[:, 2] -@@ -1448,11 +2476,15 @@ +@@ -1448,11 +2478,15 @@ try: input_ids = input_ids.view(inputs_embeds.shape[:2]) except RuntimeError: @@ -2220,7 +2222,7 @@ return input_ids -@@ -1718,21 +2750,14 @@ +@@ -1718,21 +2752,14 @@ return x @@ -2247,7 +2249,7 @@ def __init__(self, config: Qwen4ExpVisionConfig) -> None: super().__init__() self.dim = config.hidden_size -@@ -1745,13 +2770,16 @@ +@@ -1745,13 +2772,16 @@ self.config = config self.attention_dropout = 0.0 self.is_causal = False @@ -2265,7 +2267,7 @@ **kwargs, ) -> torch.Tensor: seq_length = hidden_states.shape[0] -@@ -1759,20 +2787,29 @@ +@@ -1759,20 +2789,29 @@ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) ) cos, sin = position_embeddings @@ -2304,7 +2306,7 @@ self, query_states, key_states, -@@ -1787,28 +2824,6 @@ +@@ -1787,28 +2826,6 @@ is_causal=False, **kwargs, ) @@ -2333,7 +2335,7 @@ attn_output = attn_output.reshape(seq_length, -1).contiguous() attn_output = self.proj(attn_output) -@@ -1845,6 +2860,12 @@ +@@ -1845,6 +2862,12 @@ return hidden_states @@ -2346,7 +2348,7 @@ @auto_docstring class Qwen4ExpVisionModel(Qwen4ExpPreTrainedModel): config: Qwen4ExpVisionConfig -@@ -1884,30 +2905,127 @@ +@@ -1884,30 +2907,127 @@ self.post_init() @@ -2496,7 +2498,7 @@ @merge_with_config_defaults @capture_outputs -@@ -1922,33 +3040,145 @@ +@@ -1922,33 +3042,145 @@ Returns: `torch.Tensor`: hidden_states. """ @@ -2657,7 +2659,7 @@ position_embeddings=position_embeddings, **kwargs, ) -@@ -1959,6 +3189,39 @@ +@@ -1959,6 +3191,39 @@ last_hidden_state=hidden_states, pooler_output=merged_hidden_states, ) @@ -2697,7 +2699,7 @@ @auto_docstring -@@ -1969,12 +3232,17 @@ +@@ -1969,12 +3234,17 @@ _no_split_modules = ["Qwen4ExpDecoderLayer", "Qwen4ExpVisionBlock"] def __init__(self, config): @@ -2720,7 +2722,7 @@ self.post_init() def get_vision_position_ids( -@@ -2260,53 +3528,226 @@ +@@ -2260,53 +3530,226 @@ mm_token_type_ids: torch.IntTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen4ExpModelOutputWithPast: @@ -2978,7 +2980,7 @@ outputs = self.language_model( input_ids=None, position_ids=position_ids, -@@ -2316,11 +3757,7 @@ +@@ -2316,11 +3759,7 @@ ple_input_ids=ple_input_ids, **kwargs, ) @@ -2991,7 +2993,7 @@ @auto_docstring -@@ -2337,6 +3774,27 @@ +@@ -2337,6 +3776,27 @@ aux_loss: torch.FloatTensor | None = None @@ -3019,7 +3021,7 @@ @auto_docstring class Qwen4ExpForConditionalGeneration(Qwen4ExpPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} -@@ -2346,11 +3804,23 @@ +@@ -2346,11 +3806,23 @@ _fsdp_plan = {"lm_head": "keep_full_weight"} @@ -3044,7 +3046,7 @@ self.post_init() @auto_docstring -@@ -2396,47 +3866,10 @@ +@@ -2396,47 +3868,10 @@ mm_token_type_ids: torch.IntTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -3096,7 +3098,7 @@ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, -@@ -2448,33 +3881,64 @@ +@@ -2448,33 +3883,64 @@ attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, @@ -3176,7 +3178,7 @@ loss=loss, aux_loss=aux_loss, logits=logits, -@@ -2483,6 +3947,7 @@ +@@ -2483,6 +3949,7 @@ attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, router_logits=outputs.router_logits, @@ -3184,7 +3186,7 @@ ) def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): -@@ -2696,6 +4161,32 @@ +@@ -2696,6 +4163,32 @@ } return causal_mask_mapping diff --git a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py index d2f3b6f00f..5d1eaa9364 100644 --- a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py +++ b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py @@ -1444,6 +1444,7 @@ def forward( self.num_experts, ) unused = self.gate_up_proj.new_empty(0) + expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} return self.veomni_moe( hidden_states, top_k_weights, @@ -1453,6 +1454,7 @@ def forward( self.down_proj, self.gate_up_proj, num_experts=self.num_experts, + **expert_options, ) # --- Patch.1 --- diff --git a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.diff b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.diff index 274e129f8c..439b1d923d 100644 --- a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.diff +++ b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.diff @@ -1356,7 +1356,7 @@ def forward( self, -@@ -874,25 +1432,29 @@ +@@ -874,25 +1432,31 @@ top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: @@ -1391,6 +1391,7 @@ + self.num_experts, + ) + unused = self.gate_up_proj.new_empty(0) ++ expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} + return self.veomni_moe( + hidden_states, + top_k_weights, @@ -1400,12 +1401,13 @@ + self.down_proj, + self.gate_up_proj, + num_experts=self.num_experts, ++ **expert_options, + ) + # --- Patch.1 --- class Qwen4ExpTextTopKRouter(nn.Module): -@@ -916,6 +1478,12 @@ +@@ -916,6 +1480,12 @@ return router_logits, router_scores, router_indices @@ -1418,7 +1420,7 @@ class Qwen4ExpTextSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() -@@ -924,11 +1492,33 @@ +@@ -924,11 +1494,33 @@ self.shared_expert = Qwen4ExpTextMLP(config, intermediate_size=config.shared_expert_intermediate_size) self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) @@ -1454,7 +1456,7 @@ expert_output = self.experts(hidden_states_reshaped, selected_experts, routing_weights) shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states_reshaped)) * shared_expert_output -@@ -970,9 +1560,13 @@ +@@ -970,9 +1562,13 @@ _MASK64 = (1 << 64) - 1 @@ -1468,7 +1470,7 @@ _PRIME_1 = 10007 -@@ -1015,8 +1609,14 @@ +@@ -1015,8 +1611,14 @@ return prime @@ -1484,7 +1486,7 @@ super().__init__() self.layer_idx = layer_idx self.ngram_size = config.ngram_size -@@ -1028,7 +1628,6 @@ +@@ -1028,7 +1630,6 @@ self.ngram_vocab_size_base = config.ngram_vocab_size_base head_dim_per_ngram = embedding_dim // self.ngram_heads self.seed = config.seed @@ -1492,7 +1494,7 @@ self.eos_token_id = config.eos_token_id[0] if isinstance(config.eos_token_id, list) else config.eos_token_id self.head_vocab_sizes = [] -@@ -1046,9 +1645,29 @@ +@@ -1046,9 +1647,29 @@ ) self.ngram_heads_vocab_sizes = nn.Buffer(torch.tensor(self.head_vocab_sizes, dtype=torch.long)) self.ngram_heads_offsets = nn.Buffer(torch.tensor(self.head_offsets, dtype=torch.long)) @@ -1525,7 +1527,7 @@ def _shift_right_ignore_eos(self, token_ids: torch.Tensor, shift: int) -> torch.Tensor: if shift == 0: -@@ -1066,19 +1685,183 @@ +@@ -1066,19 +1687,183 @@ valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0) return torch.where(valid, shifted, token_ids.new_full((), self.eos_token_id)) @@ -1716,7 +1718,7 @@ if ( not past_key_values.has_previous_state(self.layer_idx, state_idx=2) and input_ids.shape[1] < self.context_len -@@ -1090,28 +1873,48 @@ +@@ -1090,28 +1875,48 @@ input_ids_to_cache, self.layer_idx, state_idx=2, conv_kernel_size=self.context_len ) @@ -1773,7 +1775,7 @@ class Qwen4ExpTextPLELayer(nn.Module): -@@ -1147,24 +1950,91 @@ +@@ -1147,24 +1952,91 @@ bias=False, ) @@ -1881,7 +1883,7 @@ def forward( self, -@@ -1172,8 +2042,25 @@ +@@ -1172,8 +2044,25 @@ input_ids: torch.Tensor, past_key_values: Cache | None, conv_mask: torch.Tensor | None = None, @@ -1908,7 +1910,7 @@ key_normed = self.norm_key(self.key_proj(embeddings)).unflatten(-1, (self.hc_count, self.hidden_size)) value = self.value_proj(embeddings) query_normed = self.norm_query(hidden_states).unflatten(-1, (self.hc_count, self.hidden_size)) -@@ -1185,8 +2072,21 @@ +@@ -1185,8 +2074,21 @@ if conv_mask is not None: gated_value = apply_mask_to_padding_states(gated_value, conv_mask) gated_value_normed = apply_mask_to_padding_states(gated_value_normed, conv_mask) @@ -1931,7 +1933,7 @@ class Qwen4ExpTextDecoderLayer(GradientCheckpointingLayer): -@@ -1204,6 +2104,15 @@ +@@ -1204,6 +2106,15 @@ self.attn_hyper_connection = Qwen4ExpTextGatedResidual(config) self.mlp_hyper_connection = Qwen4ExpTextGatedResidual(config) @@ -1947,7 +1949,7 @@ def forward( self, hidden_states: torch.Tensor, -@@ -1212,17 +2121,42 @@ +@@ -1212,17 +2123,42 @@ conv_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, ple_input_ids: torch.LongTensor | None = None, @@ -1993,7 +1995,7 @@ ) else: hidden_states, _ = self.self_attn( -@@ -1230,6 +2164,9 @@ +@@ -1230,6 +2166,9 @@ position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, @@ -2003,7 +2005,7 @@ **kwargs, ) -@@ -1308,6 +2245,12 @@ +@@ -1308,6 +2247,12 @@ rope_deltas: torch.LongTensor | None = None router_logits: tuple[torch.FloatTensor] | None = None @@ -2016,7 +2018,7 @@ @auto_docstring -@@ -1352,20 +2295,73 @@ +@@ -1352,20 +2297,73 @@ """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -2092,7 +2094,7 @@ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1) -@@ -1381,31 +2377,45 @@ +@@ -1381,31 +2379,45 @@ else: text_position_ids = None @@ -2144,7 +2146,7 @@ # CODEPATH: @ArthurZucker fix flagging for no reason here if self.config.ple_layer_ids and conv_mask is not None: eos_token_id = self.config.eos_token_id -@@ -1416,16 +2426,40 @@ +@@ -1416,16 +2428,40 @@ position_embeddings = self.rotary_emb(hidden_states, position_ids) hidden_states = hidden_states.repeat(1, 1, self.config.hc_count) @@ -2187,7 +2189,7 @@ hidden_states = self.hyper_connection_mixer(hidden_states) -@@ -1435,12 +2469,6 @@ +@@ -1435,12 +2471,6 @@ ) def reverse_embedding(self, inputs_embeds: torch.Tensor) -> torch.Tensor: @@ -2200,7 +2202,7 @@ with torch.no_grad(): input_ids = ( (inputs_embeds[:, :, None, :] == self.embed_tokens.weight[None, None, :, :]).all(dim=3).nonzero()[:, 2] -@@ -1448,11 +2476,15 @@ +@@ -1448,11 +2478,15 @@ try: input_ids = input_ids.view(inputs_embeds.shape[:2]) except RuntimeError: @@ -2220,7 +2222,7 @@ return input_ids -@@ -1718,21 +2750,14 @@ +@@ -1718,21 +2752,14 @@ return x @@ -2247,7 +2249,7 @@ def __init__(self, config: Qwen4ExpVisionConfig) -> None: super().__init__() self.dim = config.hidden_size -@@ -1745,13 +2770,16 @@ +@@ -1745,13 +2772,16 @@ self.config = config self.attention_dropout = 0.0 self.is_causal = False @@ -2265,7 +2267,7 @@ **kwargs, ) -> torch.Tensor: seq_length = hidden_states.shape[0] -@@ -1759,20 +2787,29 @@ +@@ -1759,20 +2789,29 @@ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) ) cos, sin = position_embeddings @@ -2304,7 +2306,7 @@ self, query_states, key_states, -@@ -1787,28 +2824,6 @@ +@@ -1787,28 +2826,6 @@ is_causal=False, **kwargs, ) @@ -2333,7 +2335,7 @@ attn_output = attn_output.reshape(seq_length, -1).contiguous() attn_output = self.proj(attn_output) -@@ -1845,6 +2860,12 @@ +@@ -1845,6 +2862,12 @@ return hidden_states @@ -2346,7 +2348,7 @@ @auto_docstring class Qwen4ExpVisionModel(Qwen4ExpPreTrainedModel): config: Qwen4ExpVisionConfig -@@ -1884,30 +2905,127 @@ +@@ -1884,30 +2907,127 @@ self.post_init() @@ -2496,7 +2498,7 @@ @merge_with_config_defaults @capture_outputs -@@ -1922,33 +3040,145 @@ +@@ -1922,33 +3042,145 @@ Returns: `torch.Tensor`: hidden_states. """ @@ -2657,7 +2659,7 @@ position_embeddings=position_embeddings, **kwargs, ) -@@ -1959,6 +3189,39 @@ +@@ -1959,6 +3191,39 @@ last_hidden_state=hidden_states, pooler_output=merged_hidden_states, ) @@ -2697,7 +2699,7 @@ @auto_docstring -@@ -1969,12 +3232,17 @@ +@@ -1969,12 +3234,17 @@ _no_split_modules = ["Qwen4ExpDecoderLayer", "Qwen4ExpVisionBlock"] def __init__(self, config): @@ -2720,7 +2722,7 @@ self.post_init() def get_vision_position_ids( -@@ -2260,53 +3528,226 @@ +@@ -2260,53 +3530,226 @@ mm_token_type_ids: torch.IntTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen4ExpModelOutputWithPast: @@ -2978,7 +2980,7 @@ outputs = self.language_model( input_ids=None, position_ids=position_ids, -@@ -2316,11 +3757,7 @@ +@@ -2316,11 +3759,7 @@ ple_input_ids=ple_input_ids, **kwargs, ) @@ -2991,7 +2993,7 @@ @auto_docstring -@@ -2337,6 +3774,27 @@ +@@ -2337,6 +3776,27 @@ aux_loss: torch.FloatTensor | None = None @@ -3019,7 +3021,7 @@ @auto_docstring class Qwen4ExpForConditionalGeneration(Qwen4ExpPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} -@@ -2346,11 +3804,23 @@ +@@ -2346,11 +3806,23 @@ _fsdp_plan = {"lm_head": "keep_full_weight"} @@ -3044,7 +3046,7 @@ self.post_init() @auto_docstring -@@ -2396,47 +3866,10 @@ +@@ -2396,47 +3868,10 @@ mm_token_type_ids: torch.IntTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -3096,7 +3098,7 @@ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, -@@ -2448,33 +3881,64 @@ +@@ -2448,33 +3883,64 @@ attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, @@ -3176,7 +3178,7 @@ loss=loss, aux_loss=aux_loss, logits=logits, -@@ -2483,6 +3947,7 @@ +@@ -2483,6 +3949,7 @@ attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, router_logits=outputs.router_logits, @@ -3184,7 +3186,7 @@ ) def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): -@@ -2696,6 +4161,32 @@ +@@ -2696,6 +4163,32 @@ } return causal_mask_mapping diff --git a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py index 17c0a8db72..e15dd4e602 100644 --- a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py +++ b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py @@ -1444,6 +1444,7 @@ def forward( self.num_experts, ) unused = self.gate_up_proj.new_empty(0) + expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} return self.veomni_moe( hidden_states, top_k_weights, @@ -1453,6 +1454,7 @@ def forward( self.down_proj, self.gate_up_proj, num_experts=self.num_experts, + **expert_options, ) # --- Patch.1 --- diff --git a/veomni/models/transformers/qwen4_exp/qwen4_exp_gpu_patch_gen_config.py b/veomni/models/transformers/qwen4_exp/qwen4_exp_gpu_patch_gen_config.py index 1e121ec4bc..c0bb89532e 100644 --- a/veomni/models/transformers/qwen4_exp/qwen4_exp_gpu_patch_gen_config.py +++ b/veomni/models/transformers/qwen4_exp/qwen4_exp_gpu_patch_gen_config.py @@ -982,6 +982,7 @@ def forward( self.num_experts, ) unused = self.gate_up_proj.new_empty(0) + expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} return self.veomni_moe( hidden_states, top_k_weights, @@ -991,6 +992,7 @@ def forward( self.down_proj, self.gate_up_proj, num_experts=self.num_experts, + **expert_options, ) # --- Patch.1 --- diff --git a/veomni/ops/kernels/moe_experts/standard/npu.py b/veomni/ops/kernels/moe_experts/standard/npu.py index 40d93ffd3d..32a4f761e0 100644 --- a/veomni/ops/kernels/moe_experts/standard/npu.py +++ b/veomni/ops/kernels/moe_experts/standard/npu.py @@ -29,15 +29,22 @@ class GmmFunction(torch.autograd.Function): """NPU grouped matmul with dgrad and wgrad.""" @staticmethod - def forward(ctx, x, weight, group_list): + def forward(ctx, x, weight, group_list, fp32_wgrad=False): """Grouped matmul ``x @ weight`` per expert.""" import torch_npu - ctx.save_for_backward(x, weight) + # Cast inside the autograd Function so its weight input stays FP32. + # Casting before the Function would make autograd round its gradient + # to BF16, even if the grouped GEMM itself returned FP32. + ctx.fp32_wgrad = fp32_wgrad + if fp32_wgrad and (weight.dtype != torch.float32 or x.dtype != torch.bfloat16): + raise ValueError("FP32 expert wgrad requires FP32 expert weight and BF16 activations") + compute_weight = weight.to(x.dtype) if ctx.fp32_wgrad else weight + ctx.save_for_backward(x, compute_weight) ctx.group_list = group_list fwd_output = torch_npu.npu_grouped_matmul( - [x], [weight], bias=None, group_list=group_list, split_item=2, group_type=0, group_list_type=1 + [x], [compute_weight], bias=None, group_list=group_list, split_item=2, group_type=0, group_list_type=1 )[0] return fwd_output @@ -55,8 +62,8 @@ def backward(ctx, grad_output): )[0] grad_weight = torch_npu.npu_grouped_matmul( - [input_tensor.T], - [grad_output], + [input_tensor.float().contiguous().T if ctx.fp32_wgrad else input_tensor.T], + [grad_output.float() if ctx.fp32_wgrad else grad_output], bias=None, group_list=group_list, split_item=3, @@ -64,12 +71,15 @@ def backward(ctx, grad_output): group_list_type=1, )[0] - return grad_input, grad_weight, None + if ctx.fp32_wgrad and grad_weight.dtype != torch.float32: + raise RuntimeError("FP32 expert weight-gradient GEMM did not return FP32") + return (grad_input, grad_weight, None, None)[: len(ctx.needs_input_grad)] -def npu_group_gemm(x, weight, group_list): + +def npu_group_gemm(x, weight, group_list, fp32_wgrad=False): """Apply ``GmmFunction`` to ``(x, weight, group_list)``.""" - output = GmmFunction.apply(x, weight, group_list) + output = GmmFunction.apply(x, weight, group_list, fp32_wgrad) return output @@ -136,6 +146,7 @@ def _npu_fused_moe_forward( fc2_weight: torch.Tensor, fc1_1_2_weight: torch.Tensor | None = None, swiglu_limit: float | None = None, + fp32_wgrad: bool = False, ) -> torch.Tensor: """NPU single-device fused MoE forward pass (non-EP). @@ -151,9 +162,9 @@ def _npu_fused_moe_forward( tokens_per_expert = torch.histc(selected_experts, bins=num_experts, min=0, max=num_experts) fc1_weight = _fc1_weight(fc1_1_weight, fc1_2_weight, fc1_1_2_weight).transpose(1, 2) - intermediate_hidden_states = npu_group_gemm(permuted_hidden_states, fc1_weight, tokens_per_expert) + intermediate_hidden_states = npu_group_gemm(permuted_hidden_states, fc1_weight, tokens_per_expert, fp32_wgrad) intermediate_activations = _swiglu(intermediate_hidden_states, swiglu_limit) - output = npu_group_gemm(intermediate_activations, fc2_weight.transpose(1, 2), tokens_per_expert) + output = npu_group_gemm(intermediate_activations, fc2_weight.transpose(1, 2), tokens_per_expert, fp32_wgrad) # NPU unpermute applies routing after fc2. Equivalent for a biasless down # projection; bf16 rounding can differ from the pre-fc2 fused order. hidden_states = torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=routing_weights) @@ -171,6 +182,7 @@ def npu_ep_fused_moe_forward( fc1_1_2_weight: torch.Tensor | None = None, ep_group: dist.ProcessGroup | None = None, swiglu_limit: float | None = None, + fp32_wgrad: bool = False, ) -> torch.Tensor: """NPU expert-parallel fused MoE forward pass. @@ -192,10 +204,12 @@ def npu_ep_fused_moe_forward( ) fc1_weight = _fc1_weight(fc1_1_weight, fc1_2_weight, fc1_1_2_weight).transpose(1, 2) - intermediate_hidden_states = npu_group_gemm(hidden_states, fc1_weight, num_global_sum_tokens_per_local_expert) + intermediate_hidden_states = npu_group_gemm( + hidden_states, fc1_weight, num_global_sum_tokens_per_local_expert, fp32_wgrad + ) intermediate_activations = _swiglu(intermediate_hidden_states, swiglu_limit) hidden_states = npu_group_gemm( - intermediate_activations, fc2_weight.transpose(1, 2), num_global_sum_tokens_per_local_expert + intermediate_activations, fc2_weight.transpose(1, 2), num_global_sum_tokens_per_local_expert, fp32_wgrad ) hidden_states = alltoall_combine( @@ -330,6 +344,7 @@ def npu_fused_moe_forward( fc1_1_2_weight: torch.Tensor | None = None, swiglu_limit: float | None = None, assume_distinct_experts: bool = False, # intentionally unused; see the del below + fp32_wgrad: bool = False, ): """NPU fused MoE entry. Branches on EP.""" # Router scores stay float32. Cast before the fused scale so the NPU GEMM @@ -352,6 +367,7 @@ def npu_fused_moe_forward( fc1_1_2_weight, ep_group=get_parallel_state().ep_group, swiglu_limit=swiglu_limit, + fp32_wgrad=fp32_wgrad, ) else: final_hidden_states = _npu_fused_moe_forward( @@ -364,6 +380,7 @@ def npu_fused_moe_forward( fc2_weight, fc1_1_2_weight, swiglu_limit=swiglu_limit, + fp32_wgrad=fp32_wgrad, ) return final_hidden_states @@ -380,6 +397,7 @@ def wrapper( num_experts: int, swiglu_limit: float | None = None, assume_distinct_experts: bool = False, + fp32_wgrad: bool = False, ) -> Tensor: """Call the NPU fused MoE path. Empty weights are ``None``.""" # Empty tensor = unused layout. Registry args stay tensors. @@ -394,4 +412,5 @@ def wrapper( fc1_1_2_weight if fc1_1_2_weight.numel() else None, swiglu_limit=swiglu_limit, assume_distinct_experts=assume_distinct_experts, + fp32_wgrad=fp32_wgrad, ) diff --git a/veomni/trainer/base.py b/veomni/trainer/base.py index bbe9d4f5de..808d69dae8 100644 --- a/veomni/trainer/base.py +++ b/veomni/trainer/base.py @@ -354,6 +354,13 @@ def _setup(args: VeOmniArguments) -> torch.device: logger.info(f"Process rank: {args.train.global_rank}, world size: {args.train.world_size}") + if get_device_type() == "npu": + from ..distributed.npu_moe_precision import validate_expert_precision_consensus + + # This is unconditional on the flag: otherwise differing rank configs could hang. + moe_impl = getattr(getattr(args.model, "ops_implementation", None), "moe_implementation", None) + validate_expert_precision_consensus(args.model.accelerator, moe_impl) + # Register ParallelState before seed/determinism env vars. Mesh creation # must not run under NCCL_DETERMINISTIC=1 on some GPU platforms (L20). init_parallel_state_from_config(args.model.accelerator, name="base") From 3d94278d92737698263ad32d8524a6b677f139ac Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:09:40 +0800 Subject: [PATCH 002/120] [ops, dist] simplify NPU expert precision validation --- docs/design/op_selection.md | 10 +- tests/ci/_npu_moe_fp32_cold_worker.py | 6 - tests/ci/test_npu_moe_precision.py | 166 +++++++++++++++++++ tests/ci/test_npu_moe_precision_consensus.py | 161 ------------------ tests/ops/moe_experts/test_npu_fp32_wgrad.py | 6 + veomni/distributed/npu_moe_precision.py | 101 +++-------- veomni/distributed/torch_parallelize.py | 8 +- veomni/trainer/base.py | 7 - 8 files changed, 204 insertions(+), 261 deletions(-) create mode 100644 tests/ci/test_npu_moe_precision.py delete mode 100644 tests/ci/test_npu_moe_precision_consensus.py diff --git a/docs/design/op_selection.md b/docs/design/op_selection.md index 4ca6df151c..6e7b88d96d 100644 --- a/docs/design/op_selection.md +++ b/docs/design/op_selection.md @@ -615,10 +615,12 @@ avoids rounding each local expert gradient to BF16 before reduction; it does not make the rest of the model's BF16 computation equivalent to FP32. The current integration requires Qwen4-Exp SiLU `fused_npu`, FSDP2, TP=PP=1, -FP32 reduction, native communication, and `torch.compile` disabled. The -common NPU trainer initialization checks the precision config and MoE backend -across all job ranks before mesh creation; direct runtimes, composed Omni wraps, -and custom model parallelization without that check reject the opt-in. DPO reference-model isolation and resume +FP32 reduction, native communication, and `torch.compile` disabled. Each rank +validates its local precision settings and constructed expert modules before +FSDP wrapping. As with the rest of the trainer configuration, all ranks must +use the same settings. Composed Omni wraps and custom model parallelization +reject the opt-in. MoE LoRA expert wrappers also reject the opt-in because +their separate kernel is not covered. DPO reference-model isolation and resume configuration checks are not wired by this option yet. Keep the setting fixed across checkpoint save and resume when comparing training trajectories. diff --git a/tests/ci/_npu_moe_fp32_cold_worker.py b/tests/ci/_npu_moe_fp32_cold_worker.py index 4ffb0ea5cb..4897a575ee 100644 --- a/tests/ci/_npu_moe_fp32_cold_worker.py +++ b/tests/ci/_npu_moe_fp32_cold_worker.py @@ -27,7 +27,6 @@ from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict from veomni.arguments.arguments_types import MixedPrecisionConfig -from veomni.distributed.npu_moe_precision import validate_expert_precision_consensus from veomni.distributed.parallel_plan import ParallelPlan from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state from veomni.distributed.torch_parallelize import build_parallelize_model @@ -126,11 +125,6 @@ def run(rank, phase, output_dir, ep, endpoint=None): dist.init_process_group(**kwargs) try: mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=True) - fsdp = SimpleNamespace(mixed_precision=mixed, fsdp_mode="fsdp2", low_precision_reduce_scatter_comm=False) - accelerator = SimpleNamespace( - fsdp_config=fsdp, tp_size=1, pp_size=1, torch_compile=SimpleNamespace(enable=False) - ) - validate_expert_precision_consensus(accelerator, "fused_npu") _init_parallel_state( dp_size=2, dp_shard_size=2, diff --git a/tests/ci/test_npu_moe_precision.py b/tests/ci/test_npu_moe_precision.py new file mode 100644 index 0000000000..0bc4cd7102 --- /dev/null +++ b/tests/ci/test_npu_moe_precision.py @@ -0,0 +1,166 @@ +# 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. + +"""Local validation for optional NPU expert gradient precision.""" + +from types import SimpleNamespace + +import pytest +import torch + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.arguments.arguments_types import MixedPrecisionConfig +from veomni.distributed.npu_moe_precision import validate_expert_precision_config, validate_qwen4_experts + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +def _accelerator(enabled): + mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=enabled) + fsdp = SimpleNamespace(mixed_precision=mixed, fsdp_mode="fsdp2", low_precision_reduce_scatter_comm=False) + return SimpleNamespace(fsdp_config=fsdp, tp_size=1, pp_size=1, torch_compile=SimpleNamespace(enable=False)) + + +def _validate_config(mixed, *, mode="fsdp2", low_comm=False, tp=1, pp=1, compile_enabled=False): + validate_expert_precision_config(mixed, mode, low_comm, tp, pp, compile_enabled) + + +def test_valid_or_disabled_configuration_needs_no_process_group(): + _validate_config(MixedPrecisionConfig(npu_expert_fp32_wgrad=True)) + _validate_config(MixedPrecisionConfig(npu_expert_fp32_wgrad=False), mode="ddp", tp=2) + + +@pytest.mark.parametrize( + "mixed_changes,parallel_changes", + [ + ({"enable": False}, {}), + ({"param_dtype": "float32"}, {}), + ({"reduce_dtype": "bfloat16"}, {}), + ({}, {"mode": "ddp"}), + ({}, {"low_comm": True}), + ({}, {"tp": 2}), + ({}, {"pp": 2}), + ({}, {"compile_enabled": True}), + ], +) +def test_unsupported_configuration_is_rejected(mixed_changes, parallel_changes): + mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=True, **mixed_changes) + with pytest.raises(ValueError, match="FP32 wgrad requires"): + _validate_config(mixed, **parallel_changes) + + +@pytest.mark.parametrize("value", ["false", 0, 1, None]) +def test_expert_precision_option_requires_real_boolean(value): + with pytest.raises(ValueError, match="must be a boolean"): + MixedPrecisionConfig(npu_expert_fp32_wgrad=value) + + +def test_direct_fsdp2_rejects_unsupported_precision(monkeypatch): + from veomni.distributed import torch_parallelize + + state = SimpleNamespace(dp_mode="fsdp2", tp_size=1, pp_size=1, fsdp_enabled=True) + monkeypatch.setattr(torch_parallelize, "get_parallel_state", lambda: state) + mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=True, reduce_dtype="bfloat16") + with pytest.raises(ValueError, match="FP32 wgrad requires"): + torch_parallelize.parallelize_model_fsdp2(object(), mixed_precision=mixed) + + +class Qwen4ExpTextExperts(torch.nn.Module): + def __init__(self, impl="fused_npu", swiglu=True, dtype=torch.float32): + super().__init__() + self.veomni_moe = SimpleNamespace(impl=impl) + self.use_swiglu_mlp = swiglu + self.weight = torch.nn.Parameter(torch.ones(2, 2, dtype=dtype)) + + +def _model_with_experts(*, model_type="qwen4_exp", impl="fused_npu", swiglu=True, dtype=torch.float32): + model = torch.nn.Module() + model.config = SimpleNamespace(model_type=model_type) + model.experts = Qwen4ExpTextExperts(impl, swiglu, dtype) + return model + + +def test_local_qwen4_experts_are_accepted(): + model = _model_with_experts() + assert validate_qwen4_experts(model) == {id(model.experts)} + + +def test_qwen4_model_without_experts_is_rejected(): + model = torch.nn.Module() + model.config = SimpleNamespace(model_type="qwen4_exp") + with pytest.raises(ValueError, match="requires Qwen4-Exp expert modules"): + validate_qwen4_experts(model) + + +@pytest.mark.parametrize("wrapper_name", ["LoraSharedExperts", "LoraIndependentExperts"]) +def test_qwen4_model_with_partial_moe_lora_is_rejected(wrapper_name): + from veomni.lora import moe_layers + + wrapper_type = getattr(moe_layers, wrapper_name) + wrapper = wrapper_type.__new__(wrapper_type) + torch.nn.Module.__init__(wrapper) + model = _model_with_experts() + model.lora_experts = wrapper + with pytest.raises(ValueError, match="does not support MoE LoRA expert wrappers"): + validate_qwen4_experts(model) + + +@pytest.mark.parametrize( + "kwargs,match", + [ + ({"model_type": "other"}, "requires Qwen4-Exp expert modules"), + ({"impl": "eager"}, "requires Qwen4-Exp SiLU fused_npu"), + ({"swiglu": False}, "requires Qwen4-Exp SiLU fused_npu"), + ({"dtype": torch.bfloat16}, "requires FP32 expert master parameters"), + ], +) +def test_invalid_local_qwen4_experts_are_rejected(kwargs, match): + with pytest.raises(ValueError, match=match): + validate_qwen4_experts(_model_with_experts(**kwargs)) + + +def test_qwen4_vlm_decoder_layers_are_fsdp_targets(): + import torch + from transformers import Qwen4ExpConfig + + from tests.models.compare import eager_ops_config, ops_config_scope + from veomni.distributed.torch_parallelize import _fsdp_target_classes, _is_fsdp_wrap_target + from veomni.models.transformers.qwen4_exp.generated.patched_modeling_qwen4_exp_npu import ( + Qwen4ExpForConditionalGeneration, + ) + + config = Qwen4ExpConfig.from_pretrained("tests/toy_config/qwen4_exp_toy/config.json") + config._attn_implementation = "eager" + with ops_config_scope(eager_ops_config()), torch.device("meta"): + model = Qwen4ExpForConditionalGeneration(config) + targets = _fsdp_target_classes(model, [], True) + blocks = { + name for name, module in model.named_modules() if _is_fsdp_wrap_target(name, type(module).__name__, targets) + } + experts = [name for name, module in model.named_modules() if type(module).__name__ == "Qwen4ExpTextExperts"] + assert experts and all(any(name.startswith(block + ".") for block in blocks) for name in experts) + # Also cover a root that has not inherited the child model's no-split list. + assert "Qwen4ExpTextDecoderLayer" in _fsdp_target_classes(SimpleNamespace(_no_split_modules=None), [], True) + assert not _fsdp_target_classes(SimpleNamespace(_no_split_modules=None), [], False) + + +@pytest.mark.parametrize("for_inference", [False, True]) +def test_seed_omni_module_does_not_silently_ignore_opt_in(for_inference): + from veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime import ModuleRuntime + + accelerator = _accelerator(True) + args = SimpleNamespace(accelerator=accelerator) + with pytest.raises(ValueError, match="unsupported by SeedOmni"): + ModuleRuntime(args, "test", module_config=None, global_accelerator=accelerator, for_inference=for_inference) diff --git a/tests/ci/test_npu_moe_precision_consensus.py b/tests/ci/test_npu_moe_precision_consensus.py deleted file mode 100644 index b89dc12c62..0000000000 --- a/tests/ci/test_npu_moe_precision_consensus.py +++ /dev/null @@ -1,161 +0,0 @@ -# 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. - -"""CPU/Gloo coverage for the job-wide NPU expert precision handshake.""" - -import copy -import json -from datetime import timedelta -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch.distributed as dist -import torch.multiprocessing as mp - -from veomni.arguments.arguments_types import MixedPrecisionConfig -from veomni.distributed.npu_moe_precision import ( - require_expert_precision_consensus, - validate_expert_precision_consensus, - validate_qwen4_experts, -) - - -def _accelerator(enabled, *, reduce_dtype="float32", low_comm=False): - mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=enabled, reduce_dtype=reduce_dtype) - fsdp = SimpleNamespace(mixed_precision=mixed, fsdp_mode="fsdp2", low_precision_reduce_scatter_comm=low_comm) - return SimpleNamespace(fsdp_config=fsdp, tp_size=1, pp_size=1, torch_compile=SimpleNamespace(enable=False)) - - -def _worker(rank, directory, case): - dist.init_process_group( - "gloo", init_method=f"file://{directory}/rendezvous", rank=rank, world_size=2, timeout=timedelta(seconds=20) - ) - try: - enabled = case != "disabled" and (case != "mismatch_flag" or rank == 0) - dtype = "bfloat16" if case in ("disabled", "mismatch_dtype") and rank == 1 else "float32" - accelerator = _accelerator(enabled, reduce_dtype=dtype, low_comm=case == "low_comm") - if case == "mismatch_output_dtype" and rank == 1: - accelerator.fsdp_config.mixed_precision.output_dtype = "float32" - try: - backend = "eager" if case == "mismatch_backend" and rank == 1 else "fused_npu" - validate_expert_precision_consensus(accelerator, backend) - status = "ok" - if enabled: - mixed = accelerator.fsdp_config.mixed_precision - require_expert_precision_consensus(mixed, "fsdp2", False, 1, 1, False) - altered = copy.deepcopy(mixed) - altered.enable = False - with pytest.raises(ValueError, match="unchanged job-wide"): - require_expert_precision_consensus(altered, "fsdp2", False, 1, 1, False) - altered = copy.deepcopy(mixed) - altered.output_dtype = "float32" - with pytest.raises(ValueError, match="unchanged job-wide"): - require_expert_precision_consensus(altered, "fsdp2", False, 1, 1, False) - if case == "mismatch_module": - import torch.nn as nn - - class Qwen4ExpTextExperts(nn.Module): - def __init__(self): - super().__init__() - self.veomni_moe = SimpleNamespace(impl="eager" if rank == 1 else "fused_npu") - self.use_swiglu_mlp = True - - model = nn.Module() - model.config = SimpleNamespace(model_type="qwen4_exp") - model.experts = Qwen4ExpTextExperts() - validate_qwen4_experts(model) - except ValueError as exc: - status = str(exc) - (Path(directory) / f"rank{rank}.json").write_text(json.dumps({"status": status})) - finally: - dist.destroy_process_group() - - -@pytest.mark.parametrize( - ("case", "expected"), - [ - ("disabled", "ok"), - ("enabled", "ok"), - ("mismatch_flag", "identical precision configuration"), - ("mismatch_dtype", "identical precision configuration"), - ("mismatch_output_dtype", "identical precision configuration"), - ("mismatch_backend", "identical precision configuration"), - ("mismatch_module", "identical Qwen4-Exp expert modules"), - ("low_comm", "native communication"), - ], -) -def test_all_ranks_agree_before_mesh(tmp_path, case, expected): - mp.spawn(_worker, args=(str(tmp_path), case), nprocs=2, join=True) - for rank in range(2): - result = json.loads((tmp_path / f"rank{rank}.json").read_text()) - assert expected in result["status"] - - -def test_caller_without_consensus_cannot_opt_in(): - mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=True) - with pytest.raises(ValueError, match="job-wide precision consensus"): - require_expert_precision_consensus(mixed, "fsdp2", False, 1, 1, False) - - -@pytest.mark.parametrize("value", ["false", 0, 1, None]) -def test_expert_precision_option_requires_real_boolean(value): - with pytest.raises(ValueError, match="must be a boolean"): - MixedPrecisionConfig(npu_expert_fp32_wgrad=value) - - -def test_direct_fsdp2_cannot_bypass_consensus(monkeypatch): - from veomni.distributed import torch_parallelize - - state = SimpleNamespace(dp_mode="fsdp2", tp_size=1, pp_size=1, fsdp_enabled=True) - monkeypatch.setattr(torch_parallelize, "get_parallel_state", lambda: state) - with pytest.raises(ValueError, match="job-wide precision consensus"): - torch_parallelize.parallelize_model_fsdp2( - object(), mixed_precision=MixedPrecisionConfig(npu_expert_fp32_wgrad=True) - ) - - -def test_qwen4_vlm_decoder_layers_are_fsdp_targets(): - import torch - from transformers import Qwen4ExpConfig - - from tests.models.compare import eager_ops_config, ops_config_scope - from veomni.distributed.torch_parallelize import _fsdp_target_classes, _is_fsdp_wrap_target - from veomni.models.transformers.qwen4_exp.generated.patched_modeling_qwen4_exp_npu import ( - Qwen4ExpForConditionalGeneration, - ) - - config = Qwen4ExpConfig.from_pretrained("tests/toy_config/qwen4_exp_toy/config.json") - config._attn_implementation = "eager" - with ops_config_scope(eager_ops_config()), torch.device("meta"): - model = Qwen4ExpForConditionalGeneration(config) - targets = _fsdp_target_classes(model, [], True) - blocks = { - name for name, module in model.named_modules() if _is_fsdp_wrap_target(name, type(module).__name__, targets) - } - experts = [name for name, module in model.named_modules() if type(module).__name__ == "Qwen4ExpTextExperts"] - assert experts and all(any(name.startswith(block + ".") for block in blocks) for name in experts) - # Also cover a root that has not inherited the child model's no-split list. - assert "Qwen4ExpTextDecoderLayer" in _fsdp_target_classes(SimpleNamespace(_no_split_modules=None), [], True) - assert not _fsdp_target_classes(SimpleNamespace(_no_split_modules=None), [], False) - - -@pytest.mark.parametrize("for_inference", [False, True]) -def test_seed_omni_module_does_not_silently_ignore_opt_in(for_inference): - from veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime import ModuleRuntime - - accelerator = _accelerator(True) - args = SimpleNamespace(accelerator=accelerator) - with pytest.raises(ValueError, match="unsupported by SeedOmni"): - ModuleRuntime(args, "test", module_config=None, global_accelerator=accelerator, for_inference=for_inference) diff --git a/tests/ops/moe_experts/test_npu_fp32_wgrad.py b/tests/ops/moe_experts/test_npu_fp32_wgrad.py index a68874cc5a..3c53799d6b 100644 --- a/tests/ops/moe_experts/test_npu_fp32_wgrad.py +++ b/tests/ops/moe_experts/test_npu_fp32_wgrad.py @@ -22,6 +22,7 @@ import torch import torch.utils.checkpoint +from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.ops.kernels.moe_experts.standard.npu import GmmFunction from veomni.utils.device import IS_NPU_AVAILABLE @@ -63,6 +64,7 @@ def _cpu_inputs(counts=(6, 7)): ) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) def test_fp32_wgrad_keeps_bf16_forward_and_input_gradient(fake_npu): x, master, dy = _cpu_inputs() counts = [6, 7] @@ -77,6 +79,7 @@ def test_fp32_wgrad_keeps_bf16_forward_and_input_gradient(fake_npu): assert not torch.equal(param.grad, param.grad.bfloat16().float()) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) def test_bf16_default_and_cast_before_autograd_preserve_legacy_rounding(fake_npu): x, master, dy = _cpu_inputs() counts = [6, 7] @@ -86,6 +89,7 @@ def test_bf16_default_and_cast_before_autograd_preserve_legacy_rounding(fake_npu assert not torch.equal(param.grad, _reference(x, dy, counts)) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) @pytest.mark.parametrize("reentrant", [True, False]) def test_checkpoint_recomputation_keeps_fp32_weight_gradient(fake_npu, reentrant): x, master, dy = _cpu_inputs() @@ -98,6 +102,7 @@ def test_checkpoint_recomputation_keeps_fp32_weight_gradient(fake_npu, reentrant assert torch.equal(param.grad, _reference(x, dy, counts)) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) @pytest.mark.parametrize("counts", [[4, 0, 4], [0, 0, 0]]) def test_empty_experts_keep_zero_fp32_gradient(fake_npu, counts): x, master, dy = _cpu_inputs(counts) @@ -107,6 +112,7 @@ def test_empty_experts_keep_zero_fp32_gradient(fake_npu, counts): assert torch.equal(param.grad, _reference(x, dy, counts)) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) def test_bad_compute_or_master_dtype_is_rejected(fake_npu): x, master, _ = _cpu_inputs() with pytest.raises(ValueError, match="FP32 expert wgrad"): diff --git a/veomni/distributed/npu_moe_precision.py b/veomni/distributed/npu_moe_precision.py index 334c387006..ddce0fac9e 100644 --- a/veomni/distributed/npu_moe_precision.py +++ b/veomni/distributed/npu_moe_precision.py @@ -12,105 +12,48 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Job-wide agreement for the optional NPU expert weight-gradient precision.""" +"""Local validation for the optional NPU expert weight-gradient precision.""" -import torch.distributed as dist +import torch -def expert_precision_fingerprint(mixed, fsdp_mode, low_precision_comm, tp_size, pp_size, compile_enabled, moe_impl): - return ( - mixed.npu_expert_fp32_wgrad, - mixed.enable, - mixed.param_dtype, - mixed.reduce_dtype, - mixed.output_dtype, - mixed.cast_forward_inputs, - fsdp_mode, - low_precision_comm, - tp_size, - pp_size, - compile_enabled, - moe_impl, - ) - - -def validate_expert_precision_consensus(accelerator, moe_implementation): - """Agree before mesh creation; every NPU training rank must call this even when disabled.""" - fsdp = accelerator.fsdp_config - mixed = fsdp.mixed_precision - local = expert_precision_fingerprint( - mixed, - fsdp.fsdp_mode, - fsdp.low_precision_reduce_scatter_comm, - accelerator.tp_size, - accelerator.pp_size, - accelerator.torch_compile.enable, - moe_implementation or "fused_npu", - ) - peers = [None] * dist.get_world_size() - dist.all_gather_object(peers, local) - if not any(peer[0] for peer in peers): +def validate_expert_precision_config(mixed, fsdp_mode, low_precision_comm, tp_size, pp_size, compile_enabled): + """Reject locally unsupported precision and parallelization combinations.""" + if not mixed.npu_expert_fp32_wgrad: return - if any(peer != local for peer in peers): - raise ValueError("NPU expert FP32 wgrad requires identical precision configuration on all ranks") if ( not mixed.enable or mixed.param_dtype != "bfloat16" or mixed.reduce_dtype != "float32" - or fsdp.fsdp_mode != "fsdp2" - or fsdp.low_precision_reduce_scatter_comm - or accelerator.tp_size != 1 - or accelerator.pp_size != 1 - or accelerator.torch_compile.enable - or local[-1] != "fused_npu" + or fsdp_mode != "fsdp2" + or low_precision_comm + or tp_size != 1 + or pp_size != 1 + or compile_enabled ): raise ValueError( "NPU expert FP32 wgrad requires FSDP2, BF16 compute, FP32 reduction, native communication, " "TP=PP=1, fused_npu MoE, and torch.compile disabled" ) - mixed._npu_expert_fp32_wgrad_fingerprint = local - - -def require_expert_precision_consensus(mixed, fsdp_mode, low_precision_comm, tp_size, pp_size, compile_enabled): - """Refuse unsupported callers or a config mutated after the job-wide check.""" - if not mixed.npu_expert_fp32_wgrad: - return - original = getattr(mixed, "_npu_expert_fp32_wgrad_fingerprint", None) - current = expert_precision_fingerprint( - mixed, fsdp_mode, low_precision_comm, tp_size, pp_size, compile_enabled, "fused_npu" - ) - if original != current: - raise ValueError("NPU expert FP32 wgrad requires the unchanged job-wide precision consensus") def validate_qwen4_experts(model): - """Check the constructed expert modules on every rank before any FSDP wrapping.""" + """Check the local constructed expert modules before any FSDP wrapping.""" + from ..lora.moe_layers import has_lora_moe_experts + + if has_lora_moe_experts(model): + raise ValueError("NPU expert FP32 wgrad does not support MoE LoRA expert wrappers") experts = [ (name, module) for name, module in model.named_modules() if type(module).__name__ == "Qwen4ExpTextExperts" ] - report = ( - getattr(getattr(model, "config", None), "model_type", None), - tuple( - ( - name, - getattr(getattr(module, "veomni_moe", None), "impl", None), - getattr(module, "use_swiglu_mlp", False), - tuple( - (param_name, str(param.dtype), tuple(param.shape)) - for param_name, param in module.named_parameters() - ), - ) - for name, module in experts - ), - ) - peers = [None] * dist.get_world_size() - dist.all_gather_object(peers, report) - if any(peer != report for peer in peers): - raise ValueError("NPU expert FP32 wgrad requires identical Qwen4-Exp expert modules on all ranks") - if report[0] != "qwen4_exp" or not experts: + if getattr(getattr(model, "config", None), "model_type", None) != "qwen4_exp" or not experts: raise ValueError("NPU expert FP32 wgrad currently requires Qwen4-Exp expert modules") - if any(impl != "fused_npu" or not swiglu for _, impl, swiglu, _ in report[1]): + if any( + getattr(getattr(module, "veomni_moe", None), "impl", None) != "fused_npu" + or not getattr(module, "use_swiglu_mlp", False) + for _, module in experts + ): raise ValueError("NPU expert FP32 wgrad requires Qwen4-Exp SiLU fused_npu experts") - if any(dtype != "torch.float32" for _, _, _, params in report[1] for _, dtype, _ in params): + if any(param.dtype != torch.float32 for _, module in experts for param in module.parameters()): raise ValueError("NPU expert FP32 wgrad requires FP32 expert master parameters") return {id(module) for _, module in experts} diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py index 591145f31a..be3444444c 100644 --- a/veomni/distributed/torch_parallelize.py +++ b/veomni/distributed/torch_parallelize.py @@ -475,9 +475,9 @@ def parallelize_model_fsdp2( parallel_state = get_parallel_state() compile_config = compile_config or CompileConfig() if mixed_precision.npu_expert_fp32_wgrad: - from .npu_moe_precision import require_expert_precision_consensus + from .npu_moe_precision import validate_expert_precision_config - require_expert_precision_consensus( + validate_expert_precision_config( mixed_precision, parallel_state.dp_mode, low_precision_reduce_scatter_comm, @@ -1091,9 +1091,9 @@ def build_parallelize_model( parallel_state = get_parallel_state() compile_config = compile_config or CompileConfig() if mixed_precision.npu_expert_fp32_wgrad: - from .npu_moe_precision import require_expert_precision_consensus + from .npu_moe_precision import validate_expert_precision_config - require_expert_precision_consensus( + validate_expert_precision_config( mixed_precision, parallel_state.dp_mode, low_precision_reduce_scatter_comm, diff --git a/veomni/trainer/base.py b/veomni/trainer/base.py index 808d69dae8..bbe9d4f5de 100644 --- a/veomni/trainer/base.py +++ b/veomni/trainer/base.py @@ -354,13 +354,6 @@ def _setup(args: VeOmniArguments) -> torch.device: logger.info(f"Process rank: {args.train.global_rank}, world size: {args.train.world_size}") - if get_device_type() == "npu": - from ..distributed.npu_moe_precision import validate_expert_precision_consensus - - # This is unconditional on the flag: otherwise differing rank configs could hang. - moe_impl = getattr(getattr(args.model, "ops_implementation", None), "moe_implementation", None) - validate_expert_precision_consensus(args.model.accelerator, moe_impl) - # Register ParallelState before seed/determinism env vars. Mesh creation # must not run under NCCL_DETERMINISTIC=1 on some GPU platforms (L20). init_parallel_state_from_config(args.model.accelerator, name="base") From 0cb2364fb57417a72aac594f22cfc0017f80e59f Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:17:11 +0800 Subject: [PATCH 003/120] [dist, optim] feat: selective expert CPU offload with resident modules --- docs/design/selective_expert_cpu_offload.md | 62 ++++++++++ docs/index.md | 1 + scripts/ci/accelerator/suites/gpu.json | 10 ++ scripts/ci/accelerator/suites/npu.json | 10 ++ .../runtime/test_omni_model_runtime.py | 2 + .../test_extra_parallel_clip_grad_norm.py | 108 +++++++++++++++++- tests/utils/test_fsdp2_clip_grad_norm.py | 40 ++++++- veomni/arguments/arguments_types.py | 15 +++ veomni/distributed/fsdp2/clip_grad_norm.py | 40 +++++-- veomni/distributed/torch_parallelize.py | 51 ++++++++- veomni/models/model_runtime.py | 2 + .../omni_model/omni_model_runtime.py | 2 + veomni/optim/muon.py | 2 + veomni/optim/optimizer.py | 5 + 14 files changed, 336 insertions(+), 14 deletions(-) create mode 100644 docs/design/selective_expert_cpu_offload.md diff --git a/docs/design/selective_expert_cpu_offload.md b/docs/design/selective_expert_cpu_offload.md new file mode 100644 index 0000000000..a8da6f686e --- /dev/null +++ b/docs/design/selective_expert_cpu_offload.md @@ -0,0 +1,62 @@ +# Selective expert CPU offload + +FSDP2 can keep selected expert modules' compute-weight shards and accumulated +gradients on CPU while the rest of the model stays on the accelerator. This +uses native PyTorch `CPUOffloadPolicy`; it is separate from optimizer-state +offload and does not change expert ownership, gradient division or precision. + +```yaml +model: + accelerator: + fsdp_config: + fsdp_mode: fsdp2 + offload: false + offload_pin_memory: false + cpu_offload_extra_parallel_names: [ep] + cpu_offload_exclude_module_patterns: + - "*.layers.5.mlp.experts" + - "*.layers.11.mlp.experts" +``` + +The exclusion list keeps complete matching expert modules resident on the +accelerator. Patterns match complete module FQNs using the parallel plan's `*` +wildcard rules; every exclusion must match a selected expert module. An empty +list offloads all selected experts. Keeping 8 or 16 layers resident changes this +list, not the implementation or checkpoint layout. + +## Placement and update lifecycle + +[`VeOmniModelRuntime._build_parallelized_model()`](../../veomni/models/model_runtime.py#L401) +forwards these settings to +[`parallelize_model_fsdp2()`](../../veomni/distributed/torch_parallelize.py#L432). +The latter applies CPU offload only to non-excluded +EP FSDP groups, loads the model, and moves those groups to CPU before optimizer +construction. FSDP stages their weights for forward/backward and accumulates +their reduced gradients on CPU. Non-expert groups and persistent ExtraParallel +parameters, such as PLE tables, retain their existing placement. + +Global gradient clipping computes the norm with the existing ownership groups. +CPU gradient shards are scaled locally with the same global coefficient as +resident accelerator gradients. Full model, optimizer-state and checkpoint +precision policies are unchanged; FP32 masters and moment offload must be +configured through the chosen optimizer separately. + +The model's `_cpu_offload_resident_module_fqns` lists the matched resident +modules. Offloaded parameters carry `_veomni_optimizer_compute_device` so a +CPU-capable optimizer can distinguish storage placement from the compute device. +Standard `DistributedMuon` does not support this placement and rejects it; +specialized CPU-capable optimizer recipes are independent of this feature. + +## Supported configuration + +- FSDP2, EP enabled, and expert FSDP mesh size 1. +- Global FSDP offload disabled; explicit `ep` selection only. +- Eager execution; `torch.compile` is rejected for this placement. +- `offload_pin_memory` applies to the selected CPU groups as well. + +CPU RAM holds both the compute-weight shadow and accumulated gradients in +addition to any optimizer state. Initial loading may still materialize the +compute shards on the accelerator before they move to CPU. Residency does not +change packing, Ulysses, CP, routing, parameter names or checkpoint ownership. +The caller remains responsible for choosing a host-memory budget and validating +the complete training/checkpoint configuration. diff --git a/docs/index.md b/docs/index.md index 1a8b7ed590..b275c726ad 100644 --- a/docs/index.md +++ b/docs/index.md @@ -122,6 +122,7 @@ design/op_selection.md design/fused_moe_ops.md design/deepgemm_fp8_moe.md design/local_parallel_state.md +design/selective_expert_cpu_offload.md design/patchgen.md design/qwen4_exp_ple_2d_parallelism.md design/verl_topk_distill_integration.md diff --git a/scripts/ci/accelerator/suites/gpu.json b/scripts/ci/accelerator/suites/gpu.json index 422591c7bb..62dbb190c1 100644 --- a/scripts/ci/accelerator/suites/gpu.json +++ b/scripts/ci/accelerator/suites/gpu.json @@ -643,5 +643,15 @@ "-x", "tests/models/test_omni_fsdp.py" ] + }, + { + "name": "gpu-selective-expert-cpu-clip-unit", + "suite": "unit", + "args": [ + "-v", + "-s", + "-x", + "tests/utils/test_fsdp2_clip_grad_norm.py" + ] } ] diff --git a/scripts/ci/accelerator/suites/npu.json b/scripts/ci/accelerator/suites/npu.json index e8644c5baf..eca306f7ab 100644 --- a/scripts/ci/accelerator/suites/npu.json +++ b/scripts/ci/accelerator/suites/npu.json @@ -437,5 +437,15 @@ "-x", "tests/distributed/test_fsdp_equivalence.py" ] + }, + { + "name": "npu-selective-expert-cpu-clip-unit", + "suite": "unit", + "args": [ + "-v", + "-s", + "-x", + "tests/utils/test_fsdp2_clip_grad_norm.py" + ] } ] diff --git a/tests/seed_omni/runtime/test_omni_model_runtime.py b/tests/seed_omni/runtime/test_omni_model_runtime.py index 2c0e2953e9..b5f7f8449e 100644 --- a/tests/seed_omni/runtime/test_omni_model_runtime.py +++ b/tests/seed_omni/runtime/test_omni_model_runtime.py @@ -442,6 +442,8 @@ def test_composed_wrap_leaves_cache_reader_modules_on_meta(monkeypatch): forward_prefetch=False, offload=False, offload_pin_memory=False, + cpu_offload_extra_parallel_names=[], + cpu_offload_exclude_module_patterns=[], max_load_broadcast_size=20.0, ), torch_compile=SimpleNamespace( diff --git a/tests/utils/test_extra_parallel_clip_grad_norm.py b/tests/utils/test_extra_parallel_clip_grad_norm.py index dac6bc2fd0..08d540e8c1 100644 --- a/tests/utils/test_extra_parallel_clip_grad_norm.py +++ b/tests/utils/test_extra_parallel_clip_grad_norm.py @@ -39,6 +39,7 @@ _CI_LAUNCH_4 = LaunchSpec(nproc=4) _CI_LAUNCH_8 = LaunchSpec(nproc=8) +_CI_LAUNCH_2 = LaunchSpec(nproc=2) # from veomni.optim.optimizer import build_optimizer @@ -129,6 +130,100 @@ def forward(self) -> torch.Tensor: return self.experts.sum() +class SelectiveOffloadModel(ToyMoeAndEmbedModel): + """Two expert groups let one remain resident while the other is offloaded.""" + + def __init__(self): + super().__init__() + self.decoder2 = ToyMoeAndEmbedDecoderLayer() + + def forward(self, x): + return super().forward(x) + self.decoder2() + + def init_weights(self): + super().init_weights() + for parameter in self.decoder2.parameters(): + parameter.data.fill_(1.0) + + def get_parallel_plan(self): + plan = super().get_parallel_plan() + plan.extra_parallel_plan["ep"]["decoder2.moe.experts"] = Shard(0) + plan.extra_parallel_plan["emb"]["decoder2.embed_tokens.weight"] = Shard(0) + plan.extra_parallel_fsdp_no_shard_module["ep"].add("decoder2.moe") + plan.extra_parallel_fsdp_no_shard_module["emb"].add("decoder2.embed_tokens") + return plan + + +def _run_selective_expert_offload(): + world = dist.get_world_size() + ps = _init_parallel_state( + dp_size=world, + dp_shard_size=world, + extra_parallel_names=("ep", "emb"), + extra_parallel_sizes=(world, 1), + extra_parallel_placement_innermost=(False, False), + dp_mode="fsdp2", + name="selective_expert_offload", + ) + device = torch.device(get_device_type(), get_device_id()) + config = ModelArguments(config_path="test").accelerator.fsdp_config + + def build(selected, resident): + return build_parallelize_model( + SelectiveOffloadModel(), + init_device="meta", + weights_path=None, + basic_modules=[], + mixed_precision=config.mixed_precision, + enable_gradient_checkpointing=False, + enable_fsdp_offload=False, + fsdp_offload_pin_memory=False, + cpu_offload_extra_parallel_names=["ep"] if selected else [], + cpu_offload_exclude_module_patterns=resident, + broadcast_model_weights_from_rank0=False, + ) + + def local(value): + return value.to_local() if isinstance(value, DTensor) else value + + with use_parallel_state(ps): + for resident in ([], ["decoder.moe"]): + reference, candidate = build(False, []), build(True, resident) + assert candidate._cpu_offload_resident_module_fqns == tuple(resident) + for name, parameter in candidate.named_parameters(): + offloaded = ( + name in ("decoder.moe.experts", "decoder2.moe.experts") + and name.removesuffix(".experts") not in resident + ) + assert local(parameter).device.type == ("cpu" if offloaded else device.type) + # Register the real EP/non-EP ownership buckets used by clipping. + optimizers = [ + build_optimizer(model, lr=0.125, weight_decay=0.0, fused=False, optimizer_type="adamw") + for model in (reference, candidate) + ] + for _step in range(2): + for scale in (1.0, 2.0, 0.5): + x = torch.ones(16, device=device) + left, right = reference(x) * scale, candidate(x) * scale + torch.testing.assert_close(left, right, rtol=0, atol=0) + left.backward() + right.backward() + get_torch_device().synchronize() + for left, right in zip(reference.parameters(), candidate.parameters(), strict=True): + torch.testing.assert_close(local(left.grad).cpu(), local(right.grad).cpu(), rtol=0, atol=0) + norms = [veomni_clip_grad_norm(model, 0.5) for model in (reference, candidate)] + torch.testing.assert_close(norms[0], norms[1], rtol=1e-6, atol=1e-6) + for optimizer in optimizers: + optimizer.step() + optimizer.zero_grad() + for left, right in zip(reference.parameters(), candidate.parameters(), strict=True): + torch.testing.assert_close(local(left).cpu(), local(right).cpu(), rtol=1e-6, atol=1e-7) + del reference, candidate, optimizers + gc.collect() + with pytest.raises(ValueError, match="did not match selected CPU-offload modules"): + build(True, ["missing.experts"]) + + # (label, ep_size, emb_size, dp_replicate_size, cpu_offload) _FSDP2_4GPU_SCENARIOS: tuple[tuple[str, int, int, int | None, bool], ...] = ( ("fsdp2_no_extra_no_offload", 1, 1, None, False), @@ -335,7 +430,7 @@ def check_model_param_grad_one_by_one(target_model, expected_grad, ep_expected_g def main(): parser = argparse.ArgumentParser() - parser.add_argument("--suite", choices=["fsdp2_4gpu", "hsdp_8gpu"], required=True) + parser.add_argument("--suite", choices=["fsdp2_4gpu", "hsdp_8gpu", "selective_offload"], required=True) parser.add_argument("--reverse-scenarios", action="store_true") cli_args = parser.parse_args() @@ -344,7 +439,9 @@ def main(): dist.init_process_group(backend=get_dist_comm_backend()) world_size = dist.get_world_size() - if cli_args.suite == "fsdp2_4gpu": + if cli_args.suite == "selective_offload": + scenarios = [] + elif cli_args.suite == "fsdp2_4gpu": scenarios = list(_FSDP2_4GPU_SCENARIOS) else: scenarios = [ @@ -356,6 +453,8 @@ def main(): scenarios = list(reversed(scenarios)) try: + if cli_args.suite == "selective_offload": + _run_selective_expert_offload() for label, ep_size, emb_size, dp_replicate_size, cpu_offload in scenarios: _run_clip_scenario( label=label, @@ -395,6 +494,11 @@ def test_clip_grad_norm_fsdp2_4gpu(): _run_clip_grad_norm_suite("fsdp2_4gpu", nproc=_CI_LAUNCH_4.nproc) +@pytest.mark.launch_plan(provider=_CI_LAUNCH_2) +def test_selective_expert_cpu_offload_preserves_accumulation_and_update(): + _run_clip_grad_norm_suite("selective_offload", nproc=_CI_LAUNCH_2.nproc) + + # --------------------------------------------------------------------------- # HSDP regression: dp = (dp_replicate=2, dp_shard=4) on 8 GPUs. # diff --git a/tests/utils/test_fsdp2_clip_grad_norm.py b/tests/utils/test_fsdp2_clip_grad_norm.py index 01b50283e3..c3e366340c 100644 --- a/tests/utils/test_fsdp2_clip_grad_norm.py +++ b/tests/utils/test_fsdp2_clip_grad_norm.py @@ -3,7 +3,15 @@ import pytest import torch -from veomni.distributed.fsdp2.clip_grad_norm import _fsdp_grad_norm_reduce_groups, _local_pth_sum +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.distributed.fsdp2.clip_grad_norm import ( + _clip_grads_with_norm, + _fsdp_grad_norm_reduce_groups, + _local_pth_sum, +) + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=1)) def _parallel_state( @@ -111,3 +119,33 @@ def test_local_pth_sum_matches_previous_fp32_conversion(dtype, p): assert actual.dtype == torch.float32 torch.testing.assert_close(actual.cpu(), expected.cpu()) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_cpu_shards_use_one_external_global_clip_coefficient(dtype): + # The supplied norm includes other ownership buckets/ranks. Recomputing a + # local-only norm would produce a different update. + params = [torch.nn.Parameter(torch.ones(2, dtype=dtype)), torch.nn.Parameter(torch.ones(1, dtype=dtype))] + params[0].grad = torch.tensor([3.0, 4.0], dtype=dtype) + params[1].grad = torch.tensor([12.0], dtype=dtype) + norm = torch.tensor(26.0) + expected = [p.grad.clone().mul_(torch.clamp(13.0 / (norm + 1e-6), max=1.0)) for p in params] + _clip_grads_with_norm(params, 13.0, norm) + for parameter, reference in zip(params, expected, strict=True): + assert torch.equal(parameter.grad, reference) + + +@pytest.mark.parametrize( + "settings", + [ + {"cpu_offload_extra_parallel_names": ["ep"], "offload": True}, + {"cpu_offload_extra_parallel_names": ["ep"], "fsdp_mode": "ddp"}, + {"cpu_offload_extra_parallel_names": ["ple"]}, + {"cpu_offload_exclude_module_patterns": ["layers.0.experts"]}, + ], +) +def test_selective_expert_offload_rejects_incompatible_configuration(settings): + from veomni.arguments.arguments_types import FSDPConfig + + with pytest.raises(ValueError): + FSDPConfig(**settings) diff --git a/veomni/arguments/arguments_types.py b/veomni/arguments/arguments_types.py index f689579aca..d206422a47 100644 --- a/veomni/arguments/arguments_types.py +++ b/veomni/arguments/arguments_types.py @@ -666,6 +666,14 @@ class FSDPConfig: ) }, ) + cpu_offload_extra_parallel_names: List[str] = field( + default_factory=list, + metadata={"help": "CPU-offload selected ExtraParallel parameter/gradient groups (currently ep only)."}, + ) + cpu_offload_exclude_module_patterns: List[str] = field( + default_factory=list, + metadata={"help": "Keep matching selected expert modules resident on the accelerator; complete FQN patterns."}, + ) low_precision_reduce_scatter_comm: bool = field( default=False, metadata={ @@ -689,6 +697,13 @@ class FSDPConfig: mixed_precision: MixedPrecisionConfig = field(default_factory=MixedPrecisionConfig) def __post_init__(self): + if self.cpu_offload_exclude_module_patterns and not self.cpu_offload_extra_parallel_names: + raise ValueError("Resident module exclusions require selective ExtraParallel CPU offload") + if self.cpu_offload_extra_parallel_names: + if self.fsdp_mode != "fsdp2" or self.offload: + raise ValueError("Selective ExtraParallel CPU offload requires FSDP2 and global offload=False") + if self.cpu_offload_extra_parallel_names != ["ep"]: + raise ValueError("Only the ep ExtraParallel axis supports selective CPU offload") if self.fsdp_mode not in ("ddp", "fsdp2", "eager"): raise ValueError( f"Unsupported fsdp_mode={self.fsdp_mode!r}. FSDP1 has been removed; " diff --git a/veomni/distributed/fsdp2/clip_grad_norm.py b/veomni/distributed/fsdp2/clip_grad_norm.py index e9eaf12f6e..42ded03f2d 100644 --- a/veomni/distributed/fsdp2/clip_grad_norm.py +++ b/veomni/distributed/fsdp2/clip_grad_norm.py @@ -14,6 +14,32 @@ logger = get_logger(__name__) +def _clip_grads_with_norm(params, max_norm, total_norm, foreach=None): + """Clip CPU shards locally; communicate only the global norm statistic. + + Offloaded DTensors keep the accelerator mesh for ownership. Passing these + through a DTensor foreach kernel can accidentally dispatch CPU payloads + onto that mesh. Use their local gradients for the pointwise scaling. + """ + resident = [] + cpu_grads = [] + for parameter in params: + grad = parameter.grad + if grad is None: + continue + local = grad.to_local() if isinstance(grad, DTensor) else grad + if local.device.type == "cpu": + cpu_grads.append(local) + else: + resident.append(parameter) + if cpu_grads: + coefficient = torch.clamp(max_norm / (total_norm.detach().to("cpu") + 1e-6), max=1.0) + for grad in cpu_grads: + grad.mul_(coefficient) + if resident: + torch.nn.utils.clip_grads_with_norm_(resident, max_norm, total_norm, foreach=foreach) + + def clip_grad_norm( model, max_norm: float, norm_type: float = 2.0, error_if_nonfinite: bool = False, foreach: bool | None = None ) -> torch.Tensor: @@ -77,7 +103,7 @@ def _cpu_offload_fsdp2_clip_grad_norm( total_norm = _finalize_total_norm(total_norm_or_pth_sum, norm_type) _raise_if_nonfinite(total_norm, norm_type, error_if_nonfinite) - torch.nn.utils.clip_grads_with_norm_(params, max_norm, total_norm, foreach=foreach) + _clip_grads_with_norm(params, max_norm, total_norm, foreach=foreach) return total_norm @@ -249,14 +275,10 @@ def extra_parallel_fsdp2_clip_grad_norm( # Apply the same clip coefficient to all groups for para in ps.extra_parallel_names: - torch.nn.utils.clip_grads_with_norm_(extra_parallel_params[para], max_norm, total_norm, foreach=foreach) - torch.nn.utils.clip_grads_with_norm_( - extra_parallel_replicated_params[para], max_norm, total_norm, foreach=foreach - ) - torch.nn.utils.clip_grads_with_norm_( - persistent_extra_parallel_params[para], max_norm, total_norm, foreach=foreach - ) - torch.nn.utils.clip_grads_with_norm_(non_extra_parallel_params, max_norm, total_norm, foreach=foreach) + _clip_grads_with_norm(extra_parallel_params[para], max_norm, total_norm, foreach=foreach) + _clip_grads_with_norm(extra_parallel_replicated_params[para], max_norm, total_norm, foreach=foreach) + _clip_grads_with_norm(persistent_extra_parallel_params[para], max_norm, total_norm, foreach=foreach) + _clip_grads_with_norm(non_extra_parallel_params, max_norm, total_norm, foreach=foreach) return total_norm diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py index 422196c605..44f01778fa 100644 --- a/veomni/distributed/torch_parallelize.py +++ b/veomni/distributed/torch_parallelize.py @@ -656,6 +656,21 @@ def parallelize_model_fsdp2( # prepare offload_policy kwargs enable_fsdp_cpu_offload = kwargs.pop("enable_fsdp_offload", False) offload_pin_memory = kwargs.pop("fsdp_offload_pin_memory", True) + cpu_offload_axes = kwargs.pop("cpu_offload_extra_parallel_names", []) + resident_patterns = kwargs.pop("cpu_offload_exclude_module_patterns", []) + if resident_patterns and not cpu_offload_axes: + raise ValueError("Resident module exclusions require selective ExtraParallel CPU offload") + matched_resident_patterns = set() + resident_module_fqns = [] + module_fqns = {id(module): name for name, module in model.named_modules()} if resident_patterns else {} + if cpu_offload_axes: + if cpu_offload_axes != ["ep"] or enable_fsdp_cpu_offload: + raise ValueError("Selective CPU offload supports ep only, without global CPU offload") + if not parallel_state.extra_parallel_enabled("ep"): + raise ValueError("Selective expert CPU offload requires expert parallelism") + if compile_config.enable: + raise ValueError("Selective expert CPU offload is not qualified with torch.compile") + cpu_offloaded_modules = [] model._fsdp_cpu_offload_enabled = enable_fsdp_cpu_offload if enable_fsdp_cpu_offload: if persistent_extra_parallel_params: @@ -705,6 +720,10 @@ def parallelize_model_fsdp2( para_fsdp_mesh = para_mesh[para_mesh_dim_names[:-1]] para_fsdp_kwargs = dict(fsdp_kwargs) para_fsdp_kwargs["mesh"] = para_fsdp_mesh + if para in cpu_offload_axes: + if para_fsdp_mesh.size() != 1: + raise ValueError("Selective expert CPU offload currently requires ep_fsdp_size=1") + para_fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(pin_memory=offload_pin_memory) shard_dim_for_para = 1 # Muon zero-comm needs whole experts per rank; otherwise keep the # default hidden-dim sharding. A para this model's plan does not @@ -791,7 +810,20 @@ def parallelize_model_fsdp2( # decoder/root FSDP group through the normal bottom-up walk. continue # shard para module (e.g. expert/decoder.moe, embed_tokens/decoder.embed_tokens) - fully_shard(_para_mod, **extra_parallel_fsdp_kwargs[para]) + para_kwargs = extra_parallel_fsdp_kwargs[para] + resident_matches = ( + [pattern for pattern in resident_patterns if check_fqn_match(pattern, module_fqns[id(_para_mod)])] + if para in cpu_offload_axes + else [] + ) + if resident_matches: + para_kwargs = dict(para_kwargs) + para_kwargs.pop("offload_policy") + matched_resident_patterns.update(resident_matches) + resident_module_fqns.append(module_fqns[id(_para_mod)]) + fully_shard(_para_mod, **para_kwargs) + if para in cpu_offload_axes and not resident_matches: + cpu_offloaded_modules.append(_para_mod) # average para (e.g. ep) grads across para (e.g. ep) ranks # NOTE: in torch 2.8 and later we should use # experts_mod.set_gradient_divide_factor(parallel_state.ep_size) @@ -834,6 +866,12 @@ def parallelize_model_fsdp2( layer_mod._fsdp_modules.append(layer_mod) logger.info_rank0(f"{layer_fqn=}, {layer_mod._fsdp_modules=}") + unmatched_resident_patterns = set(resident_patterns) - matched_resident_patterns + if unmatched_resident_patterns: + raise ValueError( + f"Resident patterns did not match selected CPU-offload modules: {sorted(unmatched_resident_patterns)}" + ) + # Shard root WITHOUT explicit `reshard_after_forward` so FSDP2's # auto-no-reshard for the root kicks in (`_fsdp_state.py:_lazy_init`). # This keeps the root's params (lm_head + embeddings + final norm) @@ -907,7 +945,16 @@ def parallelize_model_fsdp2( module_adapter_path=kwargs.pop("module_adapter_path", None), ) - if materialize_device == "cpu": + # FSDP stages the selected expert shards and their reduced gradients on CPU. + # Move them before optimizer construction so its parameter references agree. + for module in cpu_offloaded_modules: + module.cpu() + for parameter in module.parameters(): + parameter._veomni_optimizer_compute_device = get_device_type() + model._cpu_offloaded_extra_parallel_names = tuple(cpu_offload_axes) + model._cpu_offload_resident_module_fqns = tuple(sorted(resident_module_fqns)) + + if materialize_device == "cpu" or cpu_offloaded_modules: _move_buffers_to_device(model, get_device_type()) # Register grad norm clipping method for FSDP2 diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index 450dbbdf45..973cb8d029 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -477,6 +477,8 @@ def _build_parallelized_model(self) -> None: enable_forward_prefetch=args.accelerator.fsdp_config.forward_prefetch, enable_fsdp_offload=args.accelerator.fsdp_config.offload, fsdp_offload_pin_memory=args.accelerator.fsdp_config.offload_pin_memory, + cpu_offload_extra_parallel_names=args.accelerator.fsdp_config.cpu_offload_extra_parallel_names, + cpu_offload_exclude_module_patterns=args.accelerator.fsdp_config.cpu_offload_exclude_module_patterns, broadcast_model_weights_from_rank0=args.broadcast_model_weights_from_rank0, ep_sharded_stream_load=args.ep_sharded_stream_load, max_load_broadcast_size=args.accelerator.fsdp_config.max_load_broadcast_size, diff --git a/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py b/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py index 4371d50e2f..76874e967e 100644 --- a/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py +++ b/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py @@ -317,6 +317,8 @@ def _parallelize_composed_model(self, *, for_inference: bool = False) -> None: enable_forward_prefetch=acc.fsdp_config.forward_prefetch, enable_fsdp_offload=acc.fsdp_config.offload, fsdp_offload_pin_memory=acc.fsdp_config.offload_pin_memory, + cpu_offload_extra_parallel_names=acc.fsdp_config.cpu_offload_extra_parallel_names, + cpu_offload_exclude_module_patterns=acc.fsdp_config.cpu_offload_exclude_module_patterns, broadcast_model_weights_from_rank0=args.broadcast_model_weights_from_rank0, ep_sharded_stream_load=args.ep_sharded_stream_load, max_load_broadcast_size=acc.fsdp_config.max_load_broadcast_size, diff --git a/veomni/optim/muon.py b/veomni/optim/muon.py index dc9bb54e71..621e809dbd 100644 --- a/veomni/optim/muon.py +++ b/veomni/optim/muon.py @@ -879,6 +879,8 @@ def _validate_group(group: Dict[str, Any]) -> None: if head_blocks < 1: raise ValueError(f"head_blocks must be >= 1 but is: {head_blocks}") for p in group["params"]: + if getattr(p, "_veomni_optimizer_compute_device", None) is not None: + raise ValueError("DistributedMuon does not support selectively CPU-offloaded parameters") if not _is_muon_eligible_ndim(p): raise ValueError( "DistributedMuon supports only 2D and 3D parameters; " diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index c8af1818ad..64d23bbc11 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -450,6 +450,11 @@ def build_optimizer( param groups. ``muon_kwargs`` overrides the config key by key, for callers that build Muon without an ``OptimizerConfig`` at all. """ + if optimizer_type == "muon" and any( + parameter.requires_grad and getattr(parameter, "_veomni_optimizer_compute_device", None) is not None + for parameter in model.parameters() + ): + raise ValueError("DistributedMuon does not support selectively CPU-offloaded parameters") if optimizer_type == "muon": if optimizer_config is not None: muon_kwargs = {**_collect_muon_kwargs(optimizer_config), **(muon_kwargs or {})} From ebbb5133193250ad75c3311b76bd47857e11373e Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:29:02 +0800 Subject: [PATCH 004/120] [ckpt, tests] fix: preserve CPU expert checkpoint storage --- tests/checkpoints/test_dcp_checkpointer.py | 82 +++++++++++++++++++ .../test_extra_parallel_clip_grad_norm.py | 75 +++++++++++++++-- veomni/checkpoint/dcp_checkpointer.py | 39 ++++++++- 3 files changed, 185 insertions(+), 11 deletions(-) diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py index da4a2db256..c58d2469b6 100644 --- a/tests/checkpoints/test_dcp_checkpointer.py +++ b/tests/checkpoints/test_dcp_checkpointer.py @@ -39,6 +39,88 @@ _CI_ACCELERATOR_2 = LaunchSpec(nproc=2, device="accelerator") +def _offloaded_extra_parallel_checkpoint_worker(rank, world_size, directory, phase, fused, endpoint): + from torch.distributed.device_mesh import DeviceMesh + from torch.distributed.tensor import DTensor, Shard + + from veomni.checkpoint.dcp_checkpointer import drop_extra_parallel_dim, restore_extra_parallel_dim + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=world_size) + try: + singleton_groups = [dist.new_group([i]) for i in range(world_size)] + # An accelerator mesh can describe CPU-resident offloaded shards. Gloo + # supplies the real checkpoint collectives without requiring a GPU. + local_mesh = DeviceMesh.from_group(singleton_groups[rank], "cuda", mesh_dim_names=("ep_fsdp",)) + full_mesh = DeviceMesh.from_group( + [singleton_groups[rank], dist.group.WORLD], + "cuda", + mesh=torch.arange(world_size).reshape(1, world_size), + mesh_dim_names=("ep_fsdp", "ep"), + ) + spec = DTensor.from_local(torch.empty(2, 4, device="meta"), local_mesh, [Shard(1)])._spec + weight = nn.Parameter((torch.arange(8, dtype=torch.float32).reshape(2, 4) + rank * 8) / 100) + optimizer = torch.optim.AdamW([weight], lr=0.01, fused=fused) + + def update(step): + weight.grad = torch.full_like(weight, rank + step + 0.25) + optimizer.step() + optimizer.zero_grad() + + def snapshot(): + return {"weight": weight.detach().clone(), **{k: v.clone() for k, v in optimizer.state[weight].items()}} + + def checkpoint_view(value): + origin = DTensor(value.detach(), spec, requires_grad=False) + result = restore_extra_parallel_dim(origin, full_mesh, local_mesh) + assert result.to_local().device.type == "cpu" + assert result.to_local().data_ptr() == value.data_ptr() + assert result.dtype == torch.float32 + return result + + update(1) + state = optimizer.state[weight] + if phase == "resume": + with torch.no_grad(): + weight.add_(100) + for value in state.values(): + value.add_(100) + tensors = { + "weight": checkpoint_view(weight), + "m": checkpoint_view(state["exp_avg"]), + "v": checkpoint_view(state["exp_avg_sq"]), + "step": state["step"], + } + directory = Path(directory) + if phase == "save": + dcp.save(tensors, checkpoint_id=str(directory / "dcp")) + before = snapshot() + update(2) + torch.save({"before": before, "after": snapshot()}, directory / f"rank-{rank}.pt") + else: + dcp.load(tensors, checkpoint_id=str(directory / "dcp")) + for key in ("weight", "m", "v"): + restored = drop_extra_parallel_dim(tensors[key], local_mesh) + assert restored.to_local().device.type == "cpu" + assert restored.dtype == torch.float32 + expected = torch.load(directory / f"rank-{rank}.pt", weights_only=True) + for key, value in snapshot().items(): + torch.testing.assert_close(value, expected["before"][key], rtol=0, atol=0) + update(2) + for key, value in snapshot().items(): + torch.testing.assert_close(value, expected["after"][key], rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) +@pytest.mark.parametrize("fused", [False, True]) +def test_offloaded_extra_parallel_checkpoint_preserves_cpu_storage_and_update(tmp_path, fused): + for phase in ("save", "resume"): + spawn_with_rendezvous( + _offloaded_extra_parallel_checkpoint_worker, args=(2, str(tmp_path), phase, fused), nprocs=2 + ) + + def _fsdp2_multi_optimizer_worker(rank: int, world_size: int, tmp_path: Path, endpoint): """Worker for the FSDP2 MultiOptimizer DCP round-trip test.""" from veomni.checkpoint.dcp_checkpointer import ModelState, OptimizerState diff --git a/tests/utils/test_extra_parallel_clip_grad_norm.py b/tests/utils/test_extra_parallel_clip_grad_norm.py index 08d540e8c1..7cdf579742 100644 --- a/tests/utils/test_extra_parallel_clip_grad_norm.py +++ b/tests/utils/test_extra_parallel_clip_grad_norm.py @@ -7,6 +7,7 @@ import os import subprocess from dataclasses import dataclass, field +from pathlib import Path from unittest import mock import pytest @@ -154,7 +155,12 @@ def get_parallel_plan(self): return plan -def _run_selective_expert_offload(): +def _run_selective_expert_offload(fused, checkpoint_dir): + import torch.distributed.checkpoint as dcp + from torch.utils._pytree import tree_map + + from veomni.checkpoint.dcp_checkpointer import ModelState, OptimizerState + world = dist.get_world_size() ps = _init_parallel_state( dp_size=world, @@ -186,6 +192,12 @@ def build(selected, resident): def local(value): return value.to_local() if isinstance(value, DTensor) else value + def optimizer_snapshot(optimizer): + return tree_map( + lambda tensor: local(tensor).detach().cpu().clone() if torch.is_tensor(tensor) else tensor, + optimizer.state_dict(), + ) + with use_parallel_state(ps): for resident in ([], ["decoder.moe"]): reference, candidate = build(False, []), build(True, resident) @@ -198,7 +210,7 @@ def local(value): assert local(parameter).device.type == ("cpu" if offloaded else device.type) # Register the real EP/non-EP ownership buckets used by clipping. optimizers = [ - build_optimizer(model, lr=0.125, weight_decay=0.0, fused=False, optimizer_type="adamw") + build_optimizer(model, lr=0.125, weight_decay=0.0, fused=fused, optimizer_type="adamw") for model in (reference, candidate) ] for _step in range(2): @@ -218,7 +230,45 @@ def local(value): optimizer.zero_grad() for left, right in zip(reference.parameters(), candidate.parameters(), strict=True): torch.testing.assert_close(local(left).cpu(), local(right).cpu(), rtol=1e-6, atol=1e-7) - del reference, candidate, optimizers + checkpoint = str(Path(checkpoint_dir) / f"resident-{len(resident)}") + saved = {"model": ModelState(candidate), "optimizer": OptimizerState(candidate, optimizers[1])} + model_state, optimizer_state = saved["model"].state_dict(), saved["optimizer"].state_dict() + for name, parameter in candidate.named_parameters(): + if local(parameter).device.type == "cpu": + for value in ( + model_state[name], + optimizer_state[f"state.{name}.exp_avg"], + optimizer_state[f"state.{name}.exp_avg_sq"], + ): + assert local(value).device.type == "cpu" + assert value.dtype == torch.float32 + dcp.save(saved, checkpoint_id=checkpoint) + resumed = build(True, resident) + resumed_optimizer = build_optimizer( + resumed, lr=0.125, weight_decay=0.0, fused=fused, optimizer_type="adamw" + ) + dcp.load( + {"model": ModelState(resumed), "optimizer": OptimizerState(resumed, resumed_optimizer, load=True)}, + checkpoint_id=checkpoint, + ) + + for left, right in zip(candidate.parameters(), resumed.parameters(), strict=True): + assert local(left).device == local(right).device + torch.testing.assert_close(local(left).cpu(), local(right).cpu(), rtol=0, atol=0) + torch.testing.assert_close( + optimizer_snapshot(optimizers[1]), optimizer_snapshot(resumed_optimizer), rtol=0, atol=0 + ) + for model, optimizer in ((candidate, optimizers[1]), (resumed, resumed_optimizer)): + model(torch.ones(16, device=device)).backward() + veomni_clip_grad_norm(model, 0.5) + optimizer.step() + optimizer.zero_grad() + for left, right in zip(candidate.parameters(), resumed.parameters(), strict=True): + torch.testing.assert_close(local(left).cpu(), local(right).cpu(), rtol=0, atol=0) + torch.testing.assert_close( + optimizer_snapshot(optimizers[1]), optimizer_snapshot(resumed_optimizer), rtol=0, atol=0 + ) + del reference, candidate, resumed, resumed_optimizer, optimizers gc.collect() with pytest.raises(ValueError, match="did not match selected CPU-offload modules"): build(True, ["missing.experts"]) @@ -432,7 +482,11 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("--suite", choices=["fsdp2_4gpu", "hsdp_8gpu", "selective_offload"], required=True) parser.add_argument("--reverse-scenarios", action="store_true") + parser.add_argument("--fused", action="store_true") + parser.add_argument("--checkpoint-dir") cli_args = parser.parse_args() + if cli_args.suite == "selective_offload" and cli_args.checkpoint_dir is None: + parser.error("--checkpoint-dir is required for the selective_offload save/resume test") local_rank = int(os.environ.get("LOCAL_RANK", "0")) get_torch_device().set_device(f"{get_device_type()}:{local_rank}") @@ -454,7 +508,7 @@ def main(): try: if cli_args.suite == "selective_offload": - _run_selective_expert_offload() + _run_selective_expert_offload(cli_args.fused, cli_args.checkpoint_dir) for label, ep_size, emb_size, dp_replicate_size, cpu_offload in scenarios: _run_clip_scenario( label=label, @@ -473,7 +527,9 @@ def main(): clear_parallel_state() -def _run_clip_grad_norm_suite(suite: str, nproc: int, reverse_scenarios: bool = False) -> None: +def _run_clip_grad_norm_suite( + suite: str, nproc: int, reverse_scenarios: bool = False, fused: bool = False, checkpoint_dir=None +) -> None: command = [ "torchrun", "--nnodes=1", @@ -484,6 +540,10 @@ def _run_clip_grad_norm_suite(suite: str, nproc: int, reverse_scenarios: bool = ] if reverse_scenarios: command.append("--reverse-scenarios") + if fused: + command.append("--fused") + if checkpoint_dir is not None: + command.extend(["--checkpoint-dir", str(checkpoint_dir)]) result = subprocess.run(command, check=True) assert result.returncode == 0 @@ -495,8 +555,9 @@ def test_clip_grad_norm_fsdp2_4gpu(): @pytest.mark.launch_plan(provider=_CI_LAUNCH_2) -def test_selective_expert_cpu_offload_preserves_accumulation_and_update(): - _run_clip_grad_norm_suite("selective_offload", nproc=_CI_LAUNCH_2.nproc) +@pytest.mark.parametrize("fused", [False, True]) +def test_selective_expert_cpu_offload_preserves_accumulation_and_update(fused, tmp_path): + _run_clip_grad_norm_suite("selective_offload", nproc=_CI_LAUNCH_2.nproc, fused=fused, checkpoint_dir=tmp_path) # --------------------------------------------------------------------------- diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index c51cca5e79..91d85a63f2 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -374,6 +374,35 @@ def get_state_dict_with_extra_parallel_dim_preprocess(self, state_dict, action): ) +def _from_local_preserve_device(local_tensor, device_mesh, placements, *, shape=None, stride=None): + """Wrap optimizer CPU state without moving it onto its accelerator mesh. + + DTensor.from_local implicitly copies CPU input to device_mesh.device_type. + A metadata-only construction computes the same global shape/stride/spec; + attaching the actual local shard then preserves CPU storage and the original + mesh identity. DTensor's spec constructor is intentionally isolated here. + No collectives or accelerator allocations are needed (run_check=False). + """ + if local_tensor.device.type == "cpu" and device_mesh.device_type != "cpu": + metadata = DTensor.from_local( + local_tensor.to("meta"), + device_mesh=device_mesh, + placements=placements, + run_check=False, + shape=shape, + stride=stride, + ) + return DTensor(local_tensor, metadata._spec, requires_grad=local_tensor.requires_grad) + return DTensor.from_local( + local_tensor, + device_mesh=device_mesh, + placements=placements, + run_check=False, + shape=shape, + stride=stride, + ) + + def drop_extra_parallel_dim(loaded_tensor: torch.Tensor, device_mesh: DeviceMesh, fsdp_shard_dim: int = 1): """ Drop ExtraParallel dims after loading from DCP so that ExtraParallel-FSDP would not be confused. @@ -395,11 +424,11 @@ def drop_extra_parallel_dim(loaded_tensor: torch.Tensor, device_mesh: DeviceMesh if num_placements == 1: tensor_to_put = loaded_tensor.to_local() elif num_placements == 2: - tensor_to_put = DTensor.from_local( + tensor_to_put = _from_local_preserve_device( loaded_tensor._local_tensor, device_mesh=device_mesh, placements=[Shard(fsdp_shard_dim)] ) elif num_placements == 3: - tensor_to_put = DTensor.from_local( + tensor_to_put = _from_local_preserve_device( loaded_tensor._local_tensor, device_mesh=device_mesh, placements=[Replicate(), Shard(fsdp_shard_dim)] ) else: @@ -457,10 +486,12 @@ def restore_extra_parallel_dim( placements = [Replicate(), Shard(fsdp_shard_dim), Shard(ep_shard_dim)] else: placements = [Shard(fsdp_shard_dim), Shard(ep_shard_dim)] - dtensor = DTensor.from_local(orgin_tensor._local_tensor, device_mesh=fsdp_mesh, placements=placements) + dtensor = _from_local_preserve_device(orgin_tensor._local_tensor, device_mesh=fsdp_mesh, placements=placements) elif torch.is_tensor(orgin_tensor): # If there is no FSDP but only ExtraParallel - dtensor = DTensor.from_local(orgin_tensor, device_mesh=extra_parallel_fsdp_mesh, placements=[Shard(0)]) + dtensor = _from_local_preserve_device( + orgin_tensor, device_mesh=extra_parallel_fsdp_mesh, placements=[Shard(0)] + ) else: raise RuntimeError(f"origin_tensor - {orgin_tensor} is not a tensor!") From 204465c144d5411484f4b5abb7e7f16d2b7f277a Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:02:12 +0800 Subject: [PATCH 005/120] [ckpt] fix: match full optimizer parameter namespaces --- .../test_extra_parallel_optimizer_keys.py | 63 +++++++++++++++++++ veomni/checkpoint/dcp_checkpointer.py | 17 +++-- 2 files changed, 71 insertions(+), 9 deletions(-) create mode 100644 tests/checkpoints/test_extra_parallel_optimizer_keys.py diff --git a/tests/checkpoints/test_extra_parallel_optimizer_keys.py b/tests/checkpoints/test_extra_parallel_optimizer_keys.py new file mode 100644 index 0000000000..6dc9763265 --- /dev/null +++ b/tests/checkpoints/test_extra_parallel_optimizer_keys.py @@ -0,0 +1,63 @@ +"""Optimizer checkpoint keys must preserve their full parameter namespace.""" + +import pytest +import torch +from torch.distributed.tensor import Replicate + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.checkpoint import dcp_checkpointer + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +@pytest.mark.parametrize("action", ["drop", "restore"]) +@pytest.mark.parametrize( + "key, expected, scalar", + [ + ("state.embed_tokens.weight.exp_avg", "embed_tokens.weight", False), + ("state.decoder.embed_tokens.weight.exp_avg", "decoder.embed_tokens.weight", False), + ("state.decoder2.embed_tokens.weight.exp_avg_sq", "decoder2.embed_tokens.weight", False), + ("state.decoder.embed_tokens.weight.step", None, True), + ("state.decoder.embed_tokens.weight_extra.exp_avg", None, False), + ("state.unrelated.decoder.embed_tokens.weight.exp_avg", None, False), + ("param_groups.decoder.embed_tokens.weight.lr", None, False), + ("decoder.embed_tokens.weight.exp_avg", None, False), + ], +) +def test_optimizer_spec_matching_uses_the_full_namespace(monkeypatch, action, key, expected, scalar): + selected = [] + + class Spec: + def __init__(self, name): + self.name = name + + @property + def placement(self): + selected.append(self.name) + return Replicate() + + specs = { + name: Spec(name) + for name in ("embed_tokens.weight", "decoder.embed_tokens.weight", "decoder2.embed_tokens.weight") + } + value = torch.ones(()) if scalar else torch.ones(2, 3) + state = {key: value, "param_groups.decoder.embed_tokens.weight.amsgrad": False} + monkeypatch.setattr(dcp_checkpointer, "_validate_extra_parallel_meshes", lambda _: None) + result = dcp_checkpointer._apply_extra_parallel_dim(state, specs, None, action, key_match="substring") + assert selected == ([] if expected is None else [expected]) + assert result is state and result[key] is value + assert result["param_groups.decoder.embed_tokens.weight.amsgrad"] is False + + +@pytest.mark.parametrize("action", ["drop", "restore"]) +def test_optimizer_spec_matching_still_rejects_ambiguous_specs(monkeypatch, action): + monkeypatch.setattr(dcp_checkpointer, "_validate_extra_parallel_meshes", lambda _: None) + with pytest.raises(AssertionError, match="Ambiguous ExtraParallel spec match"): + dcp_checkpointer._apply_extra_parallel_dim( + {"state.embedding.weight.exp_avg": torch.ones(2, 3)}, + dict.fromkeys(("embedding", "embedding.weight")), + None, + action, + key_match="substring", + ) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 91d85a63f2..482fcee98c 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -138,8 +138,7 @@ def _apply_extra_parallel_dim( * ``"exact"`` (model): the state-dict key IS the FQN, e.g. ``"model.layers.0.mlp.experts.gate_proj"``. - * ``"substring"`` (optimizer): the state-dict key contains the FQN - with extra prefix/suffix, e.g. + * ``"substring"`` (optimizer): the key starts with ``state..``, e.g. ``"state.model.layers.0.mlp.experts.gate_proj.exp_avg"``. Non-tensor values and 0-D tensors are skipped unconditionally — they @@ -156,12 +155,18 @@ def _apply_extra_parallel_dim( extra_parallel_keys = list(extra_parallel_fqn2spec_info.keys()) if key_match == "substring" else None for name in sorted(state_dict.keys()): + tensor = state_dict[name] + if not torch.is_tensor(tensor) or tensor.ndim == 0: + continue + if key_match == "exact": if name not in extra_parallel_fqn2spec_info: continue spec_info = extra_parallel_fqn2spec_info[name] else: # "substring" - matches = [k for k in extra_parallel_keys if k in name] + # Anchor the optimizer namespace: a nested parameter may end with + # another parameter's FQN. Param-group fields are not tensor state. + matches = [k for k in extra_parallel_keys if name.startswith(f"state.{k}.")] if not matches: continue assert len(matches) == 1, f"Ambiguous ExtraParallel spec match for state key '{name}': {matches}" @@ -177,12 +182,6 @@ def _apply_extra_parallel_dim( if getattr(spec_info, "persistent_fsdp_shard_dim", None) is not None: continue - tensor = state_dict[name] - if not torch.is_tensor(tensor): - continue - if tensor.ndim == 0: - continue - assert spec_info.para_fsdp_mesh is not None, f"ExtraParallel spec {name} must have an ExtraParallel FSDP mesh" # Drop the innermost ExtraParallel (e.g. ``ep``) dim and keep the FSDP From 0e0e57c8e11a62a47ffad49f9ca54c446f5eb2b7 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 30 Oct 2026 16:02:19 +0800 Subject: [PATCH 006/120] [test] use complete mixed precision fixtures --- tests/checkpoints/test_dcp_checkpointer.py | 9 +++++---- .../seed_omni/runtime/test_module_runtime.py | 20 +++++++++++++------ .../runtime/test_module_weight_load.py | 4 ++-- .../runtime/test_omni_model_runtime.py | 3 ++- .../utils/test_moe_ep_sharded_load_matrix.py | 3 ++- 5 files changed, 25 insertions(+), 14 deletions(-) diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py index da4a2db256..8a2fc62e43 100644 --- a/tests/checkpoints/test_dcp_checkpointer.py +++ b/tests/checkpoints/test_dcp_checkpointer.py @@ -21,6 +21,7 @@ from scripts.ci.accelerator.launch_plan import LaunchSpec from tests.tools.launch_utils import spawn_with_rendezvous +from veomni.arguments import MixedPrecisionConfig from veomni.distributed import torch_parallelize from veomni.distributed.parallel_state import _init_parallel_state, get_parallel_state from veomni.distributed.torch_parallelize import ( @@ -497,7 +498,7 @@ def test_build_parallelize_model_forwards_should_skip_hf_weight_load(self, monke result = build_parallelize_model( model, - mixed_precision=SimpleNamespace(enable=False), + mixed_precision=MixedPrecisionConfig(enable=False), enable_gradient_checkpointing=False, should_skip_hf_weight_load=True, ) @@ -539,7 +540,7 @@ def init_weights(self): resumed = parallelize_model_fsdp2( resumed, weights_path="unused-hf-path", - mixed_precision=SimpleNamespace(enable=False), + mixed_precision=MixedPrecisionConfig(enable=False), should_skip_hf_weight_load=True, init_device="meta", ) @@ -608,7 +609,7 @@ def test_fsdp2_forwards_the_loader_flags(self, loaders, monkeypatch): parallelize_model_fsdp2( nn.Linear(2, 2), weights_path="hf-path", - mixed_precision=SimpleNamespace(enable=False), + mixed_precision=MixedPrecisionConfig(enable=False), init_device="meta", broadcast_model_weights_from_rank0=True, ) @@ -848,7 +849,7 @@ def test_build_parallelize_model_forwards_should_skip_hf_weight_load_to_ddp(self build_parallelize_model( MagicMock(), - mixed_precision=SimpleNamespace(enable=False), + mixed_precision=MixedPrecisionConfig(enable=False), enable_gradient_checkpointing=False, should_skip_hf_weight_load=True, init_device="meta", diff --git a/tests/seed_omni/runtime/test_module_runtime.py b/tests/seed_omni/runtime/test_module_runtime.py index 737389356c..01e84d2f3c 100644 --- a/tests/seed_omni/runtime/test_module_runtime.py +++ b/tests/seed_omni/runtime/test_module_runtime.py @@ -28,7 +28,7 @@ from torch.utils.checkpoint import noop_context_fn from scripts.ci.accelerator.launch_plan import LaunchSpec -from veomni.arguments import GradientCheckpointingConfig +from veomni.arguments import GradientCheckpointingConfig, MixedPrecisionConfig from veomni.distributed import parallel_state from veomni.distributed.parallel_state import ParallelState, use_parallel_state from veomni.models.model_runtime import VeOmniModelRuntime @@ -42,7 +42,11 @@ def _fsdp(scope: str, mode: str = "fsdp2") -> SimpleNamespace: - return SimpleNamespace(fsdp_config=SimpleNamespace(fsdp_scope=scope, fsdp_mode=mode)) + return SimpleNamespace( + fsdp_config=SimpleNamespace( + fsdp_scope=scope, fsdp_mode=mode, mixed_precision=MixedPrecisionConfig(enable=False) + ) + ) def _unbuilt(model: nn.Module | None = None, **args_fields) -> ModuleRuntime: @@ -141,7 +145,7 @@ def fake_build_foundation_model(**kwargs): accelerator=SimpleNamespace( init_device="meta", fsdp_config=SimpleNamespace( - fsdp_mode="fsdp2", fsdp_scope="module", mixed_precision=SimpleNamespace(enable=False) + fsdp_mode="fsdp2", fsdp_scope="module", mixed_precision=MixedPrecisionConfig(enable=False) ), ), ) @@ -277,7 +281,7 @@ def _offline_cache_runtime(training_task): accelerator=SimpleNamespace( init_device="cpu", fsdp_config=SimpleNamespace( - fsdp_mode="fsdp2", fsdp_scope="module", mixed_precision=SimpleNamespace(enable=False) + fsdp_mode="fsdp2", fsdp_scope="module", mixed_precision=MixedPrecisionConfig(enable=False) ), ), ) @@ -592,7 +596,7 @@ def test_distributed_inference_wraps_lora_before_loading_weights(monkeypatch): monkeypatch.setattr(ModuleRuntime, "_build_model", lambda self: setattr(self, "model", nn.Linear(2, 2))) monkeypatch.setattr(ModuleRuntime, "_scoped", lambda self: nullcontext()) accelerator = _fsdp("module") - accelerator.fsdp_config.mixed_precision = SimpleNamespace(enable=True) + accelerator.fsdp_config.mixed_precision = MixedPrecisionConfig(enable=True) ModuleRuntime( SimpleNamespace(accelerator=accelerator), @@ -672,7 +676,11 @@ def test_an_eager_inference_module_forwards_without_a_parallel_state(monkeypatch enter; its scope is a no-op rather than a registry lookup that would fail.""" model = _RecordAmbientState() monkeypatch.setattr(ModuleRuntime, "_init_eager_inference", lambda self: setattr(self, "model", model)) - args = SimpleNamespace(accelerator=SimpleNamespace(fsdp_config=SimpleNamespace(fsdp_mode="eager"))) + args = SimpleNamespace( + accelerator=SimpleNamespace( + fsdp_config=SimpleNamespace(fsdp_mode="eager", mixed_precision=MixedPrecisionConfig(enable=False)) + ) + ) runtime = ModuleRuntime( args, "vision_encoder", module_config=SimpleNamespace(), global_accelerator=_fsdp("module"), for_inference=True diff --git a/tests/seed_omni/runtime/test_module_weight_load.py b/tests/seed_omni/runtime/test_module_weight_load.py index a469fa05a3..3912898fff 100644 --- a/tests/seed_omni/runtime/test_module_weight_load.py +++ b/tests/seed_omni/runtime/test_module_weight_load.py @@ -14,7 +14,7 @@ import torch.nn as nn from scripts.ci.accelerator.launch_plan import LaunchSpec -from veomni.arguments import GradientCheckpointingConfig +from veomni.arguments import GradientCheckpointingConfig, MixedPrecisionConfig from veomni.distributed.torch_compile import CompileConfig from veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime import ModuleRuntime @@ -94,7 +94,7 @@ def test_parallelize_forwards_module_skip_decision(monkeypatch: pytest.MonkeyPat fsdp_mode="fsdp2", fsdp_scope="module", reshard_after_forward=True, - mixed_precision=SimpleNamespace(enable=False), + mixed_precision=MixedPrecisionConfig(enable=False), forward_prefetch=False, offload=False, offload_pin_memory=False, diff --git a/tests/seed_omni/runtime/test_omni_model_runtime.py b/tests/seed_omni/runtime/test_omni_model_runtime.py index 7b61916772..de8f655628 100644 --- a/tests/seed_omni/runtime/test_omni_model_runtime.py +++ b/tests/seed_omni/runtime/test_omni_model_runtime.py @@ -16,6 +16,7 @@ import torch from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.arguments import MixedPrecisionConfig from veomni.models.seed_omni.accelerated.omni_model.omni_model_runtime import ( MultiLRScheduler, OmniModelRuntime, @@ -442,7 +443,7 @@ def test_composed_wrap_leaves_cache_reader_modules_on_meta(monkeypatch): fsdp_config=SimpleNamespace( fsdp_scope="model", reshard_after_forward=True, - mixed_precision=SimpleNamespace(enable=False), + mixed_precision=MixedPrecisionConfig(enable=False), forward_prefetch=False, offload=False, offload_pin_memory=False, diff --git a/tests/utils/test_moe_ep_sharded_load_matrix.py b/tests/utils/test_moe_ep_sharded_load_matrix.py index 9a946aa5db..71c0a3c749 100644 --- a/tests/utils/test_moe_ep_sharded_load_matrix.py +++ b/tests/utils/test_moe_ep_sharded_load_matrix.py @@ -58,6 +58,7 @@ from torch.distributed.tensor import DTensor, Shard from scripts.ci.accelerator.launch_plan import LaunchSpec, runner_device_count +from veomni.arguments import MixedPrecisionConfig from veomni.distributed.parallel_plan import ParallelPlan from veomni.distributed.parallel_state import _init_parallel_state, get_parallel_state from veomni.distributed.torch_parallelize import build_parallelize_model @@ -240,7 +241,7 @@ def _build_fresh(weights_path, *, loader: str, model_cls=FakeMoeModel): kwargs = dict( weights_path=weights_path, init_device="meta", - mixed_precision=SimpleNamespace(enable=False), + mixed_precision=MixedPrecisionConfig(enable=False), enable_gradient_checkpointing=False, enable_fsdp_offload=False, basic_modules=[], From 01d29bfc5f76c0d7baa4077ee290a15f03d56e7f Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 31 Oct 2026 03:25:21 +0800 Subject: [PATCH 007/120] [dist] fix: use CPU for default Gloo process groups --- scripts/ci/accelerator/suites/gpu.json | 5 +++ scripts/ci/accelerator/suites/npu.json | 5 +++ tests/checkpoints/test_dcp_checkpointer.py | 5 ++- tests/utils/test_dist_utils_group.py | 40 ++++++++++++++++++++++ veomni/utils/dist_utils.py | 2 +- 5 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 tests/utils/test_dist_utils_group.py diff --git a/scripts/ci/accelerator/suites/gpu.json b/scripts/ci/accelerator/suites/gpu.json index 422591c7bb..1107def6d6 100644 --- a/scripts/ci/accelerator/suites/gpu.json +++ b/scripts/ci/accelerator/suites/gpu.json @@ -1,4 +1,9 @@ [ + { + "name": "distributed-default-gloo-group", + "suite": "unit", + "args": ["-v", "-s", "-x", "tests/utils/test_dist_utils_group.py"] + }, { "name": "01-unit", "suite": "unit", diff --git a/scripts/ci/accelerator/suites/npu.json b/scripts/ci/accelerator/suites/npu.json index e8644c5baf..9640a756c2 100644 --- a/scripts/ci/accelerator/suites/npu.json +++ b/scripts/ci/accelerator/suites/npu.json @@ -1,4 +1,9 @@ [ + { + "name": "distributed-default-gloo-group", + "suite": "unit", + "args": ["-v", "-s", "-x", "tests/utils/test_dist_utils_group.py"] + }, { "name": "npu-unit-01", "suite": "unit", diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py index da4a2db256..3817f647eb 100644 --- a/tests/checkpoints/test_dcp_checkpointer.py +++ b/tests/checkpoints/test_dcp_checkpointer.py @@ -1471,7 +1471,10 @@ def _no_real_collectives(self): means; cases about *another* rank failing patch `any_rank_failed` directly. """ - with patch("veomni.utils.dist_utils.get_device_type", return_value="cpu"): + with ( + patch("veomni.utils.dist_utils.get_device_type", return_value="cpu"), + patch("veomni.utils.dist_utils.dist.get_backend", return_value="gloo"), + ): with patch("veomni.utils.dist_utils.dist.all_reduce", side_effect=lambda t, op=None, group=None: None): yield diff --git a/tests/utils/test_dist_utils_group.py b/tests/utils/test_dist_utils_group.py new file mode 100644 index 0000000000..f2f87fda81 --- /dev/null +++ b/tests/utils/test_dist_utils_group.py @@ -0,0 +1,40 @@ +"""Ensure error propagation follows the backend of the actual process group.""" + +from unittest.mock import patch + +import pytest +import torch.distributed as dist + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from tests.tools.launch_utils import spawn_with_rendezvous +from veomni.utils.dist_utils import any_rank_failed, raise_if_any_rank_failed + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) + + +def _default_gloo_worker(rank, world_size, endpoint): + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=world_size) + try: + # A host may also expose CUDA/NPU devices while the default group is + # Gloo; selecting a device from hardware availability would be wrong. + with patch("veomni.utils.dist_utils.get_device_type", return_value="cuda"): + assert any_rank_failed(rank == 0) + assert not any_rank_failed(False) + expected = ValueError if rank == 0 else RuntimeError + message = "rank 0" if rank == 0 else "checking a distributed action" + with pytest.raises(expected, match=message): + raise_if_any_rank_failed(ValueError("rank 0") if rank == 0 else None, "checking a distributed action") + group = dist.new_group(ranks=list(range(world_size)), backend="gloo") + assert any_rank_failed(rank == 1, group=group) + finally: + dist.destroy_process_group() + + +def test_default_gloo_group_uses_cpu_even_when_accelerator_is_available(): + spawn_with_rendezvous(_default_gloo_worker, args=(2,), nprocs=2) + + +def test_any_rank_failed_without_process_group(): + assert any_rank_failed(True) + assert not any_rank_failed(False) diff --git a/veomni/utils/dist_utils.py b/veomni/utils/dist_utils.py index 356addf7db..218bfbd911 100644 --- a/veomni/utils/dist_utils.py +++ b/veomni/utils/dist_utils.py @@ -75,7 +75,7 @@ def any_rank_failed(failed: bool, group: Optional[Any] = None) -> bool: """ if not dist.is_initialized(): return failed - device = "cpu" if group is not None and dist.get_backend(group) == "gloo" else get_device_type() + device = "cpu" if dist.get_backend(group) == "gloo" else get_device_type() flag = torch.tensor([1 if failed else 0], dtype=torch.int32, device=device) dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=group) return bool(flag.item()) From b334b0fa9be8f1a00aa0914954c00d21225ce5bf Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 31 Oct 2026 03:39:37 +0800 Subject: [PATCH 008/120] [test] cover native default process groups --- tests/utils/test_dist_utils_group.py | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/tests/utils/test_dist_utils_group.py b/tests/utils/test_dist_utils_group.py index f2f87fda81..df06c5baa0 100644 --- a/tests/utils/test_dist_utils_group.py +++ b/tests/utils/test_dist_utils_group.py @@ -7,12 +7,10 @@ from scripts.ci.accelerator.launch_plan import LaunchSpec from tests.tools.launch_utils import spawn_with_rendezvous +from veomni.utils.device import get_device_type, get_dist_comm_backend, get_torch_device from veomni.utils.dist_utils import any_rank_failed, raise_if_any_rank_failed -pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) - - def _default_gloo_worker(rank, world_size, endpoint): dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=world_size) try: @@ -31,10 +29,34 @@ def _default_gloo_worker(rank, world_size, endpoint): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_default_gloo_group_uses_cpu_even_when_accelerator_is_available(): spawn_with_rendezvous(_default_gloo_worker, args=(2,), nprocs=2) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) def test_any_rank_failed_without_process_group(): assert any_rank_failed(True) assert not any_rank_failed(False) + + +def _accelerator_group_worker(rank, world_size, endpoint): + device = get_torch_device() + device.set_device(rank) + backend = get_dist_comm_backend() + dist.init_process_group(backend, store=endpoint.connect(), rank=rank, world_size=world_size) + try: + assert dist.get_backend() == backend + assert any_rank_failed(rank == 0) + assert not any_rank_failed(False) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="accelerator")) +@pytest.mark.skipif( + get_device_type() == "cpu" or not get_torch_device().is_available() or get_torch_device().device_count() < 2, + reason="Requires two accelerator devices", +) +def test_default_accelerator_group_failure_propagation(): + spawn_with_rendezvous(_accelerator_group_worker, args=(2,), nprocs=2) From b1bcb872dac8fb40d9da4f38eee97df42fd4b9d6 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 31 Oct 2026 23:44:32 +0800 Subject: [PATCH 009/120] [optim] preserve sparse optimizer parameter group state --- tests/optim/test_multi_optimizer_sparse_lr.py | 183 ++++++++++++++++++ veomni/optim/optimizer.py | 65 ++++++- 2 files changed, 241 insertions(+), 7 deletions(-) create mode 100644 tests/optim/test_multi_optimizer_sparse_lr.py diff --git a/tests/optim/test_multi_optimizer_sparse_lr.py b/tests/optim/test_multi_optimizer_sparse_lr.py new file mode 100644 index 0000000000..d8fe8e96d5 --- /dev/null +++ b/tests/optim/test_multi_optimizer_sparse_lr.py @@ -0,0 +1,183 @@ +"""An unused optimizer leaf still has a scheduler-controlled learning rate.""" + +from types import SimpleNamespace + +import pytest +import torch + +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer +from veomni.optim.optimizer import MultiOptimizer + + +def _build(): + model = torch.nn.Linear(2, 2) + weight = torch.optim.AdamW([model.weight], lr=0.01) + bias = torch.optim.AdamW([model.bias], lr=0.01) + optimizer = MultiOptimizer(model, {"weight": weight, "bias": bias}, ["weight", "bias"]) + return model, weight, bias, optimizer + + +def test_sparse_leaf_preserves_learning_rate_without_creating_state(): + model, weight, bias, optimizer = _build() + model.weight.grad = torch.ones_like(model.weight) + weight.step() + weight.zero_grad() + bias.param_groups[0]["lr"] = 0.0025 + result = optimizer._sparse_state_dict() + assert result["param_groups.bias.lr"] == 0.0025 + assert not bias.state + + +def test_sparse_leaf_uses_optimizer_state_dict_post_hook(): + _, _, bias, optimizer = _build() + + def change_serialized_lr(_optimizer, state): + state["param_groups"][0]["lr"] = 0.037 + return state + + bias.register_state_dict_post_hook(change_serialized_lr) + result = optimizer._sparse_state_dict() + assert result["param_groups.bias.lr"] == 0.037 + assert bias.param_groups[0]["lr"] == 0.01 + assert not bias.state + + +def test_sparse_leaf_preserves_tensor_learning_rate(): + model = torch.nn.Linear(2, 2) + learning_rate = torch.tensor(0.0025) + leaf = torch.optim.AdamW([model.bias], lr=learning_rate) + optimizer = MultiOptimizer(model, {"bias": leaf}, ["bias"]) + result = optimizer._sparse_state_dict() + torch.testing.assert_close(result["param_groups.bias.lr"], learning_rate, rtol=0, atol=0) + assert not leaf.state + + +def test_sparse_leaf_serializes_state_created_by_pre_hook(): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([model.bias], lr=0.01) + optimizer = MultiOptimizer(model, {"bias": leaf}, ["bias"]) + + def add_state(opt): + opt.state[model.bias]["custom"] = torch.ones_like(model.bias) + + leaf.register_state_dict_pre_hook(add_state) + result = optimizer._sparse_state_dict() + torch.testing.assert_close(result["state.bias.custom"], torch.ones_like(model.bias), rtol=0, atol=0) + + +def test_sparse_leaf_respects_post_hook_parameter_reordering(): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + def swap_serialized_groups(_optimizer, state): + first, second = state["param_groups"] + first["params"], second["params"] = second["params"], first["params"] + return state + + leaf.register_state_dict_post_hook(swap_serialized_groups) + result = optimizer._sparse_state_dict() + assert result["param_groups.weight.lr"] == 0.02 + assert result["param_groups.bias.lr"] == 0.01 + assert not leaf.state + + +def test_sparse_leaf_respects_live_group_reordering_in_post_hook(): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + def reverse_live_groups(opt, state): + opt.param_groups.reverse() + return state + + leaf.register_state_dict_post_hook(reverse_live_groups) + result = optimizer._sparse_state_dict() + assert result["param_groups.weight.lr"] == 0.01 + assert result["param_groups.bias.lr"] == 0.02 + assert len(leaf._optimizer_state_dict_pre_hooks) == 0 + assert len(leaf._optimizer_state_dict_post_hooks) == 1 + + +def test_sparse_leaf_rejects_custom_state_dict(): + class CustomAdamW(torch.optim.AdamW): + def state_dict(self): + return super().state_dict() + + model = torch.nn.Linear(2, 2) + leaf = CustomAdamW(model.parameters()) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + with pytest.raises(RuntimeError, match="custom state_dict"): + optimizer._sparse_state_dict() + + +def test_sparse_leaf_rejects_instance_patched_state_dict_before_mutation(): + from torch.distributed.checkpoint.state_dict import _patch_optimizer_state_dict + + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW(model.parameters()) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + _patch_optimizer_state_dict(model, optimizers=(leaf,)) + before = [param.detach().clone() for param in model.parameters()] + with pytest.raises(RuntimeError, match="custom state_dict"): + optimizer._sparse_state_dict() + assert not leaf.state + for actual, expected in zip(model.parameters(), before, strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_sparse_leaf_reentrant_hook_fails_without_leaking_temporary_hooks(): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW(model.parameters()) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + entered = False + + def reenter(opt): + nonlocal entered + if not entered: + entered = True + try: + opt.state_dict() + finally: + entered = False + + leaf.register_state_dict_pre_hook(reenter) + with pytest.raises(RuntimeError, match="Cannot reenter"): + optimizer._sparse_state_dict() + assert not leaf.state + assert len(leaf._optimizer_state_dict_pre_hooks) == 1 + assert not leaf._optimizer_state_dict_post_hooks + + +def test_sparse_leaf_learning_rate_survives_dcp_roundtrip(tmp_path): + model, weight, bias, optimizer = _build() + model.weight.grad = torch.ones_like(model.weight) + weight.step() + weight.zero_grad() + bias.param_groups[0]["lr"] = 0.0025 + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + assert not bias.state + restored_model, _, restored_bias, restored = _build() + DistributedCheckpointer.load( + str(tmp_path), {"model": restored_model, "optimizer": restored}, parallel_state=parallel + ) + assert restored_bias.param_groups[0]["lr"] == 0.0025 + + +def test_sparse_leaf_post_hook_live_reorder_survives_dcp_roundtrip(tmp_path): + def build(): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) + return model, leaf, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + model, leaf, optimizer = build() + leaf.register_state_dict_post_hook(lambda opt, state: (opt.param_groups.reverse(), state)[1]) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + restored_model, restored_leaf, restored = build() + DistributedCheckpointer.load( + str(tmp_path), {"model": restored_model, "optimizer": restored}, parallel_state=parallel + ) + assert restored_leaf.param_groups[0]["lr"] == 0.01 + assert restored_leaf.param_groups[1]["lr"] == 0.02 diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index c8af1818ad..dee3c57933 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -22,6 +22,8 @@ from torch.distributed._tensor import DTensor from torch.distributed.checkpoint.state_dict import ( StateDictOptions, + _flatten_optim_state_dict, + _get_fqns, get_optimizer_state_dict, set_optimizer_state_dict, ) @@ -255,6 +257,53 @@ def step(self, closure=None): p.data.addcdiv_(exp_avg, centered_variance, value=-step_size) +def _empty_optimizer_state_dict(model, opt): + if getattr(opt.state_dict, "__func__", None) is not Optimizer.state_dict: + raise RuntimeError("Cannot safely serialize an empty optimizer with a custom state_dict") + packed_groups = [] + id_to_name = {} + captured = False + + def capture_params(leaf): + nonlocal captured + if captured: + raise RuntimeError("Cannot reenter optimizer state_dict while saving an empty optimizer") + captured = True + packed_groups.extend([list(group["params"]) for group in leaf.param_groups]) + + def capture_ids(_leaf, packed): + names = {} + for runtime_name, param in model.named_parameters(): + fqns = _get_fqns(model, runtime_name) + if len(fqns) != 1: + raise ValueError(f"Expected one checkpoint name for parameter {runtime_name}") + names[id(param)] = next(iter(fqns)) + for params, group in zip(packed_groups, packed["param_groups"], strict=True): + for param, saved_id in zip(params, group["params"], strict=True): + param_name = names[id(param)] + if saved_id in id_to_name and id_to_name[saved_id] != param_name: + raise RuntimeError("Ambiguous parameter mapping in optimizer state_dict") + id_to_name[saved_id] = param_name + + # Capture after user pre-hooks and before user post-hooks can change either + # the serialized or live groups. Concurrent serialization must be excluded. + pre_hook = opt.register_state_dict_pre_hook(capture_params) + try: + post_hook = opt.register_state_dict_post_hook(capture_ids, prepend=True) + try: + raw = opt.state_dict() + finally: + post_hook.remove() + finally: + pre_hook.remove() + state = {id_to_name[key]: value for key, value in raw["state"].items()} + groups = [ + {**saved_group, "params": [id_to_name[key] for key in saved_group["params"]]} + for saved_group in raw["param_groups"] + ] + return _flatten_optim_state_dict({"state": state, "param_groups": groups}) + + class MultiOptimizer(Optimizer, Stateful): """ A container that handles multiple optimizers (for extra_parallel and non-extra_parallel parameters when extra_parallel+fsdp2 is enabled) @@ -309,18 +358,20 @@ def _sparse_state_dict( ) -> Dict[str, Any]: """Return a flattened optimizer state dict without materializing state. - Unlike ``state_dict()``, this skips sub-optimizers whose ``.state`` is - still empty. PyTorch's ``get_optimizer_state_dict`` would otherwise - call ``_init_optim_state`` and synthesize default state for every - parameter owned by that sub-optimizer (e.g. an unused MoE expert on - this rank), bloating the checkpoint. + Unlike ``state_dict()``, this serializes an empty leaf's parameter + groups without calling PyTorch's synthetic optimizer initialization. + The groups must still be saved because their LR may have been updated + by a scheduler even when no parameter has received a gradient. """ merged: Dict[str, Any] = {} for name in self.key_names: opt = self.optimizers_dict.get(name) if not opt.state: - continue - sd = get_optimizer_state_dict(self.model, opt, options=StateDictOptions(flatten_optimizer_state_dict=True)) + sd = _empty_optimizer_state_dict(self.model, opt) + else: + sd = get_optimizer_state_dict( + self.model, opt, options=StateDictOptions(flatten_optimizer_state_dict=True) + ) overlap = set(merged.keys()) & set(sd.keys()) if overlap: raise KeyError( From e9f451fcb2ab834bce94be7b22a75cd98680fdaf Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 31 Oct 2026 23:59:52 +0800 Subject: [PATCH 010/120] [test, optim] cover sparse groups with FSDP2 DCP --- tests/optim/test_multi_optimizer_sparse_lr.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/tests/optim/test_multi_optimizer_sparse_lr.py b/tests/optim/test_multi_optimizer_sparse_lr.py index d8fe8e96d5..669d0cf66f 100644 --- a/tests/optim/test_multi_optimizer_sparse_lr.py +++ b/tests/optim/test_multi_optimizer_sparse_lr.py @@ -5,6 +5,7 @@ import pytest import torch +from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.optim.optimizer import MultiOptimizer @@ -181,3 +182,66 @@ def build(): ) assert restored_leaf.param_groups[0]["lr"] == 0.01 assert restored_leaf.param_groups[1]["lr"] == 0.02 + + +def _fsdp2_empty_leaf_worker(rank, checkpoint, endpoint): + from datetime import timedelta + + import torch.distributed as dist + from torch.distributed.fsdp import fully_shard + + from veomni.distributed.parallel_state import _init_parallel_state, get_parallel_state + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + _init_parallel_state(dp_size=2, dp_mode="fsdp2", device_type="cpu") + parallel = get_parallel_state() + + def build(seed): + torch.manual_seed(seed) + model = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.Linear(4, 4)) + for layer in model: + fully_shard(layer, mesh=parallel.dp_shard_mesh) + fully_shard(model, mesh=parallel.dp_shard_mesh) + first = torch.optim.AdamW(model[0].parameters(), lr=0.01) + second = torch.optim.AdamW(model[1].parameters(), lr=0.01) + return model, first, second, MultiOptimizer(model, {"first": first, "second": second}, ["first", "second"]) + + model, first, second, optimizer = build(42) + scheduler = torch.optim.lr_scheduler.StepLR(second, step_size=1, gamma=0.25) + model(torch.ones(2, 4)).sum().backward() + for param in model[1].parameters(): + param.grad = None + first.step() + first.zero_grad() + second.step() + scheduler.step() + assert second.param_groups[0]["lr"] == 0.0025 + assert first.state + assert not second.state + source_shards = [param.to_local().detach().clone() for param in model.parameters()] + DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + assert not second.state + for expected, param in zip(source_shards, model.parameters(), strict=True): + torch.testing.assert_close(expected, param.to_local(), rtol=0, atol=0) + + restored_model, _, restored_second, restored = build(123) + assert any( + not torch.equal(saved, param.to_local()) + for saved, param in zip(source_shards, restored_model.parameters(), strict=True) + ) + DistributedCheckpointer.load( + checkpoint, {"model": restored_model, "optimizer": restored}, parallel_state=parallel + ) + assert restored_second.param_groups[0]["lr"] == 0.0025 + for expected, param in zip(source_shards, restored_model.parameters(), strict=True): + torch.testing.assert_close(expected, param.to_local(), rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_sparse_leaf_learning_rate_survives_two_rank_fsdp2_dcp(tmp_path): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_fsdp2_empty_leaf_worker, args=(str(tmp_path / "ckpt"),), nprocs=2) From 93c9ab57417d619051728d5e5f970f2d590c3327 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sun, 1 Nov 2026 05:44:48 +0800 Subject: [PATCH 011/120] [checkpoint] version canonical extra-parallel model layout --- docs/usage/checkpoint.md | 20 + tests/checkpoints/test_dcp_checkpointer.py | 1 + .../test_extra_parallel_layout_v2.py | 455 +++++++++++++ tests/seed_omni/model/test_hf_layout.py | 24 + veomni/checkpoint/dcp_checkpointer.py | 634 +++++++++++++++++- veomni/distributed/parallel_plan.py | 8 + veomni/models/seed_omni/utils/hf_layout.py | 64 +- veomni/utils/save_safetensor_utils.py | 101 ++- 8 files changed, 1254 insertions(+), 53 deletions(-) create mode 100644 tests/checkpoints/test_extra_parallel_layout_v2.py diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 90b021d3a2..83895d484d 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -235,6 +235,26 @@ To drop legacy resume, delete `veomni/checkpoint/legacy_v0_1_12.py` and the imports that load it (search for `legacy_v0_1_12`). An old checkpoint then fails resume instead of falling back. +## Versioned ExtraParallel layout + +The EP-outer layout is opt-in through `SpecInfo.checkpoint_layout_v2`; existing +plans remain on their original layout. A plan enabling it must retain the +original `global_shape` and must use FSDP DTensor parameters. Every +non-persistent ExtraParallel shard in a checkpoint must use the canonical +layout: mixing legacy and canonical shards is rejected. The v2 tensor namespace +identifies canonical global values independently of source geometry, so DCP can +reshard a model checkpoint across supported FSDP shard axes. Use +`DistributedCheckpointer` for save and resume: it rejects old/unversioned +layouts before copying model tensors. Migrating a legacy checkpoint requires +verified source layout information. Canonical v2 optimizer save and resume are +rejected until an owner-aware codec can preserve rank-specific optimizer +states; this opt-in therefore does not yet support resuming training. The +standard VeOmni HF export strips the namespace. + +Raw `torch.distributed.checkpoint.load(..., planner=DefaultLoadPlanner(allow_partial_load=True))` is not a +supported migration API: it can return success without reading any tensor when +the namespaces differ. Replicated values must already agree across ranks. + ## Related pages - [Checkpoint conversion](checkpoint_conversion.md) — DCP shards → HuggingFace safetensors (`scripts/merge_dcp_to_hf.py`). diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py index da4a2db256..1b6a2ba1f9 100644 --- a/tests/checkpoints/test_dcp_checkpointer.py +++ b/tests/checkpoints/test_dcp_checkpointer.py @@ -2367,6 +2367,7 @@ def test_a_rank_that_cannot_remove_the_markers_fails_the_group(self, tmp_path): with ( patch("veomni.checkpoint.dcp_checkpointer.dist.is_initialized", return_value=True), patch("veomni.checkpoint.dcp_checkpointer.dist.get_rank", return_value=1), + patch("veomni.checkpoint.dcp_checkpointer._reject_unproven_extra_parallel_optimizer"), patch("veomni.utils.dist_utils.any_rank_failed", return_value=True), patch.object(DistributedCheckpointer, "execute_save") as execute_save, ): diff --git a/tests/checkpoints/test_extra_parallel_layout_v2.py b/tests/checkpoints/test_extra_parallel_layout_v2.py new file mode 100644 index 0000000000..8eab059ca9 --- /dev/null +++ b/tests/checkpoints/test_extra_parallel_layout_v2.py @@ -0,0 +1,455 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +"""Exact global and local tensor layout for composed EP/FSDP checkpoint shards.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import Replicate, Shard, distribute_tensor + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.checkpoint.dcp_checkpointer import ( + _EXTRA_PARALLEL_LAYOUT_MARKER, + DistributedCheckpointer, + ModelState, + OptimizerState, + _apply_extra_parallel_dim, + _export_extra_parallel_layout_key, + _extra_parallel_layout_key, + _get_sharding_plan, + _ModelStrictLoadPlanner, + _needs_extra_parallel_layout_marker, + drop_extra_parallel_dim_v2, + restore_extra_parallel_dim, + restore_extra_parallel_dim_v2, +) +from veomni.distributed.parallel_plan import SpecInfo + + +def _roundtrip(rank, root, ep_shard_dim, fsdp_shard_dim): + dist.init_process_group("gloo", init_method=f"file://{root}/init", rank=rank, world_size=4) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("fsdp", "ep")) + fsdp_mesh = mesh["fsdp"] + full = torch.arange(35, dtype=torch.float32).reshape(5, 7) + ep_index = mesh.get_coordinate()[1] + local_ep = torch.chunk(full, 2, dim=ep_shard_dim)[ep_index] + local = distribute_tensor(local_ep.contiguous(), fsdp_mesh, placements=[Shard(fsdp_shard_dim)]) + with pytest.raises(RuntimeError, match="declared FSDP placement"): + restore_extra_parallel_dim_v2( + local, mesh, fsdp_mesh, ep_shard_dim, (fsdp_shard_dim + 1) % 2, global_shape=tuple(full.shape) + ) + combined = restore_extra_parallel_dim_v2( + local, mesh, fsdp_mesh, ep_shard_dim, fsdp_shard_dim, global_shape=tuple(full.shape) + ) + torch.testing.assert_close(combined.full_tensor(), full, rtol=0, atol=0) + spec = SpecInfo("ep", Shard(ep_shard_dim), "weight", mesh, fsdp_shard_dim, None, tuple(full.shape), True) + parallel = SimpleNamespace(extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh}) + converted = _apply_extra_parallel_dim( + {"weight": local}, {"weight": spec}, parallel, "restore", key_match="exact" + ) + torch.testing.assert_close(converted["weight"].full_tensor(), full, rtol=0, atol=0) + spec.checkpoint_layout_v2 = False + legacy = restore_extra_parallel_dim(local, mesh, fsdp_mesh, ep_shard_dim, fsdp_shard_dim) + actual = _apply_extra_parallel_dim({"weight": local}, {"weight": spec}, parallel, "restore", key_match="exact") + assert actual["weight"].shape == legacy.shape + assert actual["weight"].placements == legacy.placements + torch.testing.assert_close(actual["weight"].to_local(), legacy.to_local(), rtol=0, atol=0) + dcp.save({"expert": combined}, checkpoint_id=f"{root}/checkpoint") + target = restore_extra_parallel_dim_v2( + torch.zeros_like(local), mesh, fsdp_mesh, ep_shard_dim, fsdp_shard_dim, global_shape=tuple(full.shape) + ) + dcp.load({"expert": target}, checkpoint_id=f"{root}/checkpoint") + torch.testing.assert_close(target.full_tensor(), full, rtol=0, atol=0) + restored = drop_extra_parallel_dim_v2(target, fsdp_mesh, fsdp_shard_dim, extra_parallel_name="ep") + assert restored.shape == local.shape + assert restored.placements == local.placements + torch.testing.assert_close(restored.to_local(), local.to_local(), rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +@pytest.mark.parametrize("ep_shard_dim,fsdp_shard_dim", [(0, 0), (0, 1), (-1, -1), (-1, 1)]) +def test_composed_ep_fsdp_layout_roundtrip(tmp_path, ep_shard_dim, fsdp_shard_dim): + torch.multiprocessing.spawn(_roundtrip, args=(str(tmp_path), ep_shard_dim, fsdp_shard_dim), nprocs=4, join=True) + + +def _hsdp_roundtrip(rank, root): + dist.init_process_group("gloo", init_method=f"file://{root}/init", rank=rank, world_size=8) + try: + mesh = init_device_mesh("cpu", (2, 2, 2), mesh_dim_names=("replicate", "fsdp", "ep")) + fsdp_mesh = mesh["replicate", "fsdp"] + full = torch.arange(35, dtype=torch.float32).reshape(5, 7) + ep_index = mesh.get_coordinate()[-1] + local_ep = torch.chunk(full, 2, dim=0)[ep_index] + local = distribute_tensor(local_ep.contiguous(), fsdp_mesh, placements=[Replicate(), Shard(0)]) + combined = restore_extra_parallel_dim_v2(local, mesh, fsdp_mesh, 0, 0, global_shape=tuple(full.shape)) + torch.testing.assert_close(combined.full_tensor(), full, rtol=0, atol=0) + dcp.save({"expert": combined}, checkpoint_id=f"{root}/checkpoint") + target = restore_extra_parallel_dim_v2( + torch.zeros_like(local), mesh, fsdp_mesh, 0, 0, global_shape=tuple(full.shape) + ) + dcp.load({"expert": target}, checkpoint_id=f"{root}/checkpoint") + resumed = drop_extra_parallel_dim_v2(target, fsdp_mesh, 0, extra_parallel_name="ep") + torch.testing.assert_close(resumed.to_local(), local.to_local(), rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=8)) +def test_composed_ep_hsdp_layout_roundtrip(tmp_path): + torch.multiprocessing.spawn(_hsdp_roundtrip, args=(str(tmp_path),), nprocs=8, join=True) + + +def _cross_layout_roundtrip(rank, root, source_dim, target_dim): + dist.init_process_group("gloo", init_method=f"file://{root}/init", rank=rank, world_size=4) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("fsdp", "ep")) + fsdp_mesh = mesh["fsdp"] + full = torch.arange(48, dtype=torch.float32).reshape(6, 8) + local_ep = torch.chunk(full, 2, dim=0)[mesh.get_coordinate()[1]] + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + + def build(fsdp_dim, fill_zero): + model = torch.nn.Module() + value = torch.zeros_like(local_ep) if fill_zero else local_ep + shard = distribute_tensor(value.contiguous(), fsdp_mesh, placements=[Shard(fsdp_dim)]) + model.register_parameter("weight", torch.nn.Parameter(shard)) + model._fqn2spec_info = { + "weight": SpecInfo("ep", Shard(0), "weight", mesh, fsdp_dim, None, tuple(full.shape), True) + } + return model + + source = build(source_dim, False) + dcp.save({"model": ModelState(source, parallel_state=parallel)}, checkpoint_id=f"{root}/checkpoint") + target = build(target_dim, True) + dcp.load( + {"model": ModelState(target, parallel_state=parallel)}, + checkpoint_id=f"{root}/checkpoint", + planner=_ModelStrictLoadPlanner(strict_model=True), + ) + expected = distribute_tensor(local_ep.contiguous(), fsdp_mesh, placements=[Shard(target_dim)]) + torch.testing.assert_close(target.weight.to_local(), expected.to_local(), rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +@pytest.mark.parametrize("source_dim,target_dim", [(0, 1), (1, 0)]) +def test_canonical_checkpoint_can_change_fsdp_shard_axis(tmp_path, source_dim, target_dim): + torch.multiprocessing.spawn( + _cross_layout_roundtrip, args=(str(tmp_path), source_dim, target_dim), nprocs=4, join=True + ) + + +def test_spec_info_positional_shard_dimension_remains_compatible(): + spec = SpecInfo("ep", Shard(0), "weight", None, 0) + assert spec.fsdp_shard_dim == 0 + assert spec.global_shape is None + assert not _needs_extra_parallel_layout_marker({"weight": spec}) + spec.checkpoint_layout_v2 = True + spec.global_shape = (4, 4) + assert _needs_extra_parallel_layout_marker({"weight": spec}) + + +def test_negative_dimensions_are_canonicalized_for_layout_identity(): + mesh = SimpleNamespace(shape=(2, 2)) + negative = SpecInfo("ep", Shard(-1), "weight", mesh, 1, None, (4, 8), True) + positive = SpecInfo("ep", Shard(1), "weight", mesh, -1, None, (4, 8), True) + assert _needs_extra_parallel_layout_marker({"weight": negative}) + assert _extra_parallel_layout_key({"weight": negative}) == _extra_parallel_layout_key({"weight": positive}) + + +def test_v2_requires_original_shape_and_fsdp_tensor(): + spec = SpecInfo("ep", Shard(0), "weight", None, 0, checkpoint_layout_v2=True) + with pytest.raises(RuntimeError, match="requires the original global_shape"): + _extra_parallel_layout_key({"weight": spec}) + + +def test_v2_optimizer_fails_before_model_or_checkpoint_mutation(tmp_path): + model = torch.nn.Linear(2, 2, bias=False) + mesh = SimpleNamespace(shape=(2, 2)) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, (2, 2), True)} + optimizer = SimpleNamespace() + parallel = SimpleNamespace(dp_mode="fsdp2") + adapter = OptimizerState(model, optimizer, parallel_state=parallel) + with pytest.raises(RuntimeError, match="owner-aware checkpoint codec"): + adapter.state_dict() + with pytest.raises(RuntimeError, match="owner-aware checkpoint codec"): + adapter.load_state_dict({}) + with pytest.raises(RuntimeError, match="owner-aware checkpoint codec"): + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + assert not list(tmp_path.iterdir()) + + +def _rank_local_invalid_layout_worker(rank, root, failure): + from datetime import timedelta + + dist.init_process_group( + "gloo", init_method=f"file://{root}/init", rank=rank, world_size=2, timeout=timedelta(seconds=20) + ) + try: + model = torch.nn.Linear(2, 2, bias=False) + mesh = SimpleNamespace(shape=(2, 2)) + shape = None if failure == "shape" and rank == 0 else (2, 2) + v2 = not (failure == "version" and rank == 0) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, shape, v2)} + with pytest.raises(RuntimeError): + DistributedCheckpointer.save( + f"{root}/checkpoint", {"model": model}, parallel_state=SimpleNamespace(dp_mode="fsdp2") + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +@pytest.mark.parametrize("failure", ["shape", "version"]) +def test_rank_local_invalid_layout_stops_all_ranks_before_write(tmp_path, failure): + torch.multiprocessing.spawn(_rank_local_invalid_layout_worker, args=(str(tmp_path), failure), nprocs=2, join=True) + assert not (tmp_path / "checkpoint").exists() + + +def test_same_v2_wrapper_can_save_and_load(tmp_path, monkeypatch): + model = torch.nn.Linear(2, 2, bias=False) + mesh = SimpleNamespace(shape=(2, 2)) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, (2, 2), True)} + parallel = SimpleNamespace(dp_mode="fsdp2") + monkeypatch.setattr( + ModelState, "get_state_dict_with_extra_parallel_dim_preprocess", lambda self, state, action: state + ) + wrapper = ModelState(model, parallel_state=parallel) + saved = model.weight.detach().clone() + direct = wrapper.state_dict() + wrapper.load_state_dict(direct) + torch.testing.assert_close(model.weight, saved, rtol=0, atol=0) + dcp.save({"model": wrapper}, checkpoint_id=tmp_path) + with torch.no_grad(): + model.weight.zero_() + dcp.load({"model": wrapper}, checkpoint_id=tmp_path, planner=_ModelStrictLoadPlanner(strict_model=True)) + torch.testing.assert_close(model.weight, saved, rtol=0, atol=0) + + +def test_v2_tensors_cannot_be_silently_read_with_legacy_dcp_keys(tmp_path, monkeypatch): + from torch.distributed.checkpoint.api import CheckpointException + + model = torch.nn.Linear(2, 2, bias=False) + mesh = SimpleNamespace(shape=(2, 2)) + spec = SpecInfo("ep", Shard(0), "weight", mesh, 0, None, (2, 2), True) + model._fqn2spec_info = {"weight": spec} + monkeypatch.setattr( + ModelState, "get_state_dict_with_extra_parallel_dim_preprocess", lambda self, state, action: state + ) + wrapper = ModelState(model, parallel_state=SimpleNamespace(dp_mode="fsdp2")) + dcp.save({"model": wrapper}, checkpoint_id=tmp_path) + legacy = torch.zeros_like(model.weight) + with pytest.raises(CheckpointException): + dcp.load({"model": {"weight": legacy}}, checkpoint_id=tmp_path) + torch.testing.assert_close(legacy, torch.zeros_like(legacy), rtol=0, atol=0) + dcp.load( + {"model": {"weight": legacy}}, checkpoint_id=tmp_path, planner=dcp.DefaultLoadPlanner(allow_partial_load=True) + ) + torch.testing.assert_close(legacy, torch.zeros_like(legacy), rtol=0, atol=0) + marker = _extra_parallel_layout_key({"weight": spec}) + assert _export_extra_parallel_layout_key(marker) is None + assert _export_extra_parallel_layout_key(marker + "__model.foo.weight") == "model.foo.weight" + + +def test_legacy_partial_dcp_cannot_overwrite_v2_target(tmp_path, monkeypatch): + dcp.save({"model": {"weight": torch.ones(2, 2)}}, checkpoint_id=tmp_path) + model = torch.nn.Linear(2, 2, bias=False) + with torch.no_grad(): + model.weight.zero_() + mesh = SimpleNamespace(shape=(2, 2)) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, (2, 2), True)} + monkeypatch.setattr( + ModelState, "get_state_dict_with_extra_parallel_dim_preprocess", lambda self, state, action: state + ) + wrapper = ModelState(model, parallel_state=SimpleNamespace(dp_mode="fsdp2")) + dcp.load({"model": wrapper}, checkpoint_id=tmp_path, planner=dcp.DefaultLoadPlanner(allow_partial_load=True)) + torch.testing.assert_close(model.weight, torch.zeros_like(model.weight), rtol=0, atol=0) + + +def test_v2_online_hf_export_restores_fqn_and_preserves_values(monkeypatch): + from veomni.utils.save_safetensor_utils import get_model_save_state + + model = torch.nn.Linear(2, 2, bias=False) + mesh = SimpleNamespace(shape=(2, 2)) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, (2, 2), True)} + monkeypatch.setattr( + ModelState, "get_state_dict_with_extra_parallel_dim_preprocess", lambda self, state, action: state + ) + before = model.weight.detach().clone() + exported = get_model_save_state(model, None, parallel_state=SimpleNamespace(dp_mode="fsdp2")) + assert set(exported) == {"weight"} + torch.testing.assert_close(exported["weight"], before.to(torch.bfloat16), rtol=0, atol=0) + torch.testing.assert_close(model.weight, before, rtol=0, atol=0) + + +def test_opted_in_layout_rejects_unversioned_checkpoint_before_tensor_copy(tmp_path): + from torch.distributed.checkpoint.api import CheckpointException + + dcp.save({"model": {"weight": torch.ones(2)}}, checkpoint_id=tmp_path) + target = torch.zeros(2) + with pytest.raises(CheckpointException, match="Incompatible ExtraParallel checkpoint layout"): + dcp.load( + {"model": {"weight": target, _EXTRA_PARALLEL_LAYOUT_MARKER + "a" * 64: True}}, + checkpoint_id=tmp_path, + planner=_ModelStrictLoadPlanner(strict_model=False), + ) + torch.testing.assert_close(target, torch.zeros_like(target), rtol=0, atol=0) + + +def test_canonical_layout_rejects_mixed_legacy_shards(): + mesh = SimpleNamespace(shape=(2, 2)) + source_a = SpecInfo("ep", Shard(0), "a", mesh, 0, None, (4, 4), True) + source_b = SpecInfo("ep", Shard(0), "b", mesh, 0, None, (4, 4), False) + with pytest.raises(RuntimeError, match="cannot mix versioned and legacy"): + _extra_parallel_layout_key({"a": source_a, "b": source_b}) + + +def test_hf_export_rejects_duplicate_legacy_and_canonical_source(tmp_path): + mesh = SimpleNamespace(shape=(2, 2)) + spec = SpecInfo("ep", Shard(0), "weight", mesh, 0, None, (2, 2), True) + marker = _extra_parallel_layout_key({"weight": spec}) + dcp.save( + { + "model": { + "model.weight": torch.ones(2, 2), + marker + "__model.weight": torch.zeros(2, 2), + marker: "ep-outer-v1", + } + }, + checkpoint_id=tmp_path, + ) + with pytest.raises(RuntimeError, match="mixes versioned and legacy"): + _get_sharding_plan(str(tmp_path)) + + +def test_hf_export_rejects_versioned_tensor_without_marker(tmp_path): + marker = _extra_parallel_layout_key({}) + dcp.save({"model": {marker + "__model.weight": torch.ones(2, 2)}}, checkpoint_id=tmp_path) + with pytest.raises(RuntimeError, match="marker is missing"): + _get_sharding_plan(str(tmp_path)) + + +def test_hf_export_rejects_malformed_marker_value(tmp_path): + marker = _extra_parallel_layout_key({}) + dcp.save({"model": {marker: "wrong", marker + "__model.weight": torch.ones(2)}}, checkpoint_id=tmp_path) + with pytest.raises(RuntimeError, match="marker value"): + _get_sharding_plan(str(tmp_path)) + + +def test_hf_export_rejects_unknown_extra_parallel_layout(): + unknown = _EXTRA_PARALLEL_LAYOUT_MARKER + "0" * 64 + "__model.weight" + with pytest.raises(RuntimeError, match="verified migration"): + _export_extra_parallel_layout_key(unknown) + + +def test_v2_model_load_rejects_dtype_conversion(tmp_path): + marker = _extra_parallel_layout_key({}) + key = marker + "__weight" + dcp.save({"model": {marker: "ep-outer-v1", key: torch.ones(2, dtype=torch.float32)}}, checkpoint_id=tmp_path) + state = {"model": {marker: "ep-outer-v1", key: torch.zeros(2, dtype=torch.bfloat16)}} + with pytest.raises(BaseException, match="checkpoint dtype differs"): + dcp.load(state, checkpoint_id=tmp_path, planner=_ModelStrictLoadPlanner(strict_model=True)) + assert torch.equal(state["model"][key], torch.zeros(2, dtype=torch.bfloat16)) + + +def test_v2_model_load_rejects_legacy_keys_alongside_marker(tmp_path): + marker = _extra_parallel_layout_key({}) + key = marker + "__weight" + dcp.save( + {"model": {marker: "ep-outer-v1", key: torch.ones(2), "weight": torch.full((2,), 2.0)}}, checkpoint_id=tmp_path + ) + state = {"model": {marker: "ep-outer-v1", key: torch.zeros(2)}} + with pytest.raises(BaseException, match="mixes versioned and legacy"): + dcp.load(state, checkpoint_id=tmp_path, planner=_ModelStrictLoadPlanner(strict_model=True)) + assert torch.equal(state["model"][key], torch.zeros(2)) + + +def test_v2_rejects_checkpoint_wrapper_fqn_not_in_canonical_state(): + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper + + model = torch.nn.Module() + model.block = checkpoint_wrapper(torch.nn.Linear(2, 2, bias=False)) + mesh = SimpleNamespace(shape=(2, 2)) + name = "block._checkpoint_wrapped_module.weight" + model._fqn2spec_info = {name: SpecInfo("ep", Shard(0), name, mesh, 1, None, (2, 2), True)} + with pytest.raises(RuntimeError, match="not present in the canonical state dict"): + ModelState(model, parallel_state=SimpleNamespace(dp_mode="fsdp2")).state_dict() + + +def test_v2_trainable_only_allows_frozen_spec_absent(): + from veomni.checkpoint.dcp_checkpointer import _check_canonical_spec_coverage + + model = torch.nn.Linear(2, 2) + model.weight.requires_grad_(False) + mesh = SimpleNamespace(shape=(2, 2)) + model._fqn2spec_info = { + "weight": SpecInfo("ep", Shard(0), "weight", mesh, 1, None, (2, 2), True), + "bias": SpecInfo("ep", Replicate(), "bias", mesh, 1), + } + _check_canonical_spec_coverage(model, model._fqn2spec_info, {"bias"}, trainable_only=True) + with pytest.raises(RuntimeError, match="not present in the canonical state dict"): + _check_canonical_spec_coverage(model, model._fqn2spec_info, {"bias"}, trainable_only=False) + + +def test_v2_rejects_missing_runtime_parameter_spec(): + model = torch.nn.Linear(2, 2) + mesh = SimpleNamespace(shape=(2, 2)) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 1, None, (2, 2), True)} + with pytest.raises(RuntimeError, match="plan is missing parameter: bias"): + ModelState(model, parallel_state=SimpleNamespace(dp_mode="fsdp2")).state_dict() + + +def _high_level_roundtrip_worker(rank, root): + from datetime import timedelta + + dist.init_process_group( + "gloo", init_method=f"file://{root}/init-high", rank=rank, world_size=4, timeout=timedelta(seconds=45) + ) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ep_fsdp", "ep")) + fsdp_mesh = mesh["ep_fsdp"] + ep_rank = mesh.get_local_rank("ep") + full = torch.arange(64, dtype=torch.float32).reshape(8, 8) + ep_local = full.chunk(2, dim=0)[ep_rank].contiguous() + + def model_for(axis, values): + model = torch.nn.Module() + model.register_parameter("weight", torch.nn.Parameter(distribute_tensor(values, fsdp_mesh, [Shard(axis)]))) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, axis, None, (8, 8), True)} + return model + + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + source = model_for(1, ep_local) + DistributedCheckpointer.save(f"{root}/high", {"model": source}, parallel_state=parallel) + target = model_for(0, torch.zeros_like(ep_local)) + DistributedCheckpointer.load(f"{root}/high", {"model": target}, parallel_state=parallel) + expected = distribute_tensor(ep_local, fsdp_mesh, [Shard(0)]) + torch.testing.assert_close(target.weight.to_local(), expected.to_local(), rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +def test_high_level_v2_cross_layout_model_only_roundtrip(tmp_path): + torch.multiprocessing.spawn(_high_level_roundtrip_worker, args=(str(tmp_path),), nprocs=4, join=True) + + +def test_regular_nested_tensor_list_loads_with_strict_planner(tmp_path): + dcp.save({"optimizer": {"foo": [torch.ones(2), torch.full((2,), 2.0)]}}, checkpoint_id=tmp_path) + target = [torch.zeros(2), torch.zeros(2)] + dcp.load({"optimizer": {"foo": target}}, checkpoint_id=tmp_path, planner=_ModelStrictLoadPlanner(False)) + torch.testing.assert_close(target[0], torch.ones(2), rtol=0, atol=0) + torch.testing.assert_close(target[1], torch.full((2,), 2.0), rtol=0, atol=0) diff --git a/tests/seed_omni/model/test_hf_layout.py b/tests/seed_omni/model/test_hf_layout.py index a5fc449432..1f95543106 100644 --- a/tests/seed_omni/model/test_hf_layout.py +++ b/tests/seed_omni/model/test_hf_layout.py @@ -490,3 +490,27 @@ def test_an_offline_embedding_run_skips_the_export_check(tmp_path, monkeypatch): args.train.training_task = "online_training" build_omni_model_runtime_args(args) assert len(checked) == 1 + + +def test_source_export_strips_canonical_extra_parallel_namespace(tmp_path, monkeypatch): + from veomni.checkpoint.dcp_checkpointer import ( + ModelState, + _extra_parallel_layout_key, + _namespace_extra_parallel_layout_state, + ) + + root = tmp_path / "fake-hf" + original = _write_hf_checkpoint(root, sharded=False) + model = OmniModel.from_pretrained(str(root), torch_dtype=torch.bfloat16, device_map="cpu") + module = model.modules_dict[FAKE_A] + state_dict = ModelState.state_dict + marker = _extra_parallel_layout_key({}) + monkeypatch.setattr( + ModelState, "state_dict", lambda self: _namespace_extra_parallel_layout_state(state_dict(self), marker) + ) + + out = tmp_path / "export" + save_hf_source_checkpoint(model.config._hf_source, {FAKE_A: (module, None)}, str(out)) + restored, _ = _read_safetensors(root=out) + for key in original: + assert torch.equal(restored[key], original[key]) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index c51cca5e79..1f6f8deed0 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -17,6 +17,7 @@ import hashlib import os import shutil +import weakref from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor from datetime import timedelta @@ -74,6 +75,32 @@ def __init__(self, strict_model: bool): self.strict_model = strict_model def create_local_plan(self): + assert self.metadata is not None + target = {key for key, path in self.mappings.items() if path and _is_extra_parallel_layout_marker(path[-1])} + source = { + key for key in self.metadata.state_dict_metadata if _is_extra_parallel_layout_marker(key.split(".")[-1]) + } + if source != target: + raise RuntimeError("Incompatible ExtraParallel checkpoint layout; explicit layout migration is required") + if source: + if len(source) != 1: + raise RuntimeError("ExtraParallel checkpoint has multiple layout markers") + marker = next(iter(source)) + if any( + key.startswith("model.") and key != marker and not key.startswith(marker + "__") + for key in self.metadata.state_dict_metadata + ): + raise RuntimeError("ExtraParallel checkpoint mixes versioned and legacy tensor keys") + if target: + for key, value in self.state_dict.items(): + if not torch.is_tensor(value) or not key.startswith("model."): + continue + source_value = self.metadata.state_dict_metadata.get(key) + source_dtype = getattr(getattr(source_value, "properties", None), "dtype", None) + if source_dtype is not None and source_dtype != value.dtype: + raise RuntimeError( + f"ExtraParallel checkpoint dtype differs for {key}; explicit conversion required" + ) plan = super().create_local_plan() if not self.strict_model: return plan @@ -122,6 +149,311 @@ def _validate_extra_parallel_meshes(parallel_state) -> None: ) +_EXTRA_PARALLEL_LAYOUT_MARKER = "__veomni_ep_outer_layout_v1__" +_EXTRA_PARALLEL_FIRST_MESHES = {} + + +def _is_extra_parallel_layout_marker(key): + if not isinstance(key, str): + return False + suffix = key.removeprefix(_EXTRA_PARALLEL_LAYOUT_MARKER) + return ( + key.startswith(_EXTRA_PARALLEL_LAYOUT_MARKER) + and len(suffix) == 64 + and all(c in "0123456789abcdef" for c in suffix) + ) + + +def _normalized_extra_parallel_dims(spec): + shape = getattr(spec, "global_shape", None) + if shape is None: + raise RuntimeError("The v2 ExtraParallel checkpoint layout requires the original global_shape") + ndim = len(shape) + ep_dim, fsdp_dim = spec.placement.dim, spec.fsdp_shard_dim + if not (-ndim <= ep_dim < ndim and -ndim <= fsdp_dim < ndim): + raise RuntimeError("ExtraParallel shard dimension is outside the global tensor rank") + return ep_dim % ndim, fsdp_dim % ndim + + +def _needs_extra_parallel_layout_marker(specs) -> bool: + return any( + getattr(spec, "checkpoint_layout_v2", False) + and isinstance(spec.placement, Shard) + and getattr(spec, "persistent_fsdp_shard_dim", None) is None + for spec in (specs or {}).values() + ) + + +def _check_canonical_spec_coverage(model, specs, state_keys, *, trainable_only): + from torch.distributed.checkpoint.state_dict import _get_fqns + + params = dict(model.named_parameters(remove_duplicate=False)) + buffers = dict(model.named_buffers(remove_duplicate=False)) + for name, param in params.items(): + if trainable_only and not param.requires_grad: + continue + if any(canonical in state_keys for canonical in _get_fqns(model, name)) and name not in (specs or {}): + raise RuntimeError(f"Canonical ExtraParallel plan is missing parameter: {name}") + for name, spec in (specs or {}).items(): + if not isinstance(spec.placement, Shard) or getattr(spec, "persistent_fsdp_shard_dim", None) is not None: + continue + value = params.get(name, buffers.get(name)) + if value is None: + raise RuntimeError(f"ExtraParallel checkpoint spec has no runtime tensor: {name}") + if trainable_only and name in params and not value.requires_grad: + continue + if name not in state_keys: + raise RuntimeError(f"ExtraParallel checkpoint spec is not present in the canonical state dict: {name}") + + +def _reject_unproven_extra_parallel_optimizer(state, parallel_state, *, group=None, trainable_only=False): + error, signature = None, None + try: + specs = getattr(state.get("model"), "_fqn2spec_info", None) + if _needs_extra_parallel_layout_marker(specs): + _extra_parallel_layout_key(specs) + signature = _runtime_extra_parallel_layout_signature(specs) + if state.get("optimizer") is not None: + ps = parallel_state if parallel_state is not None else get_parallel_state() + else: + ps = None + if state.get("optimizer") is not None and getattr(ps, "dp_mode", None) == "fsdp2": + if _needs_extra_parallel_layout_marker(specs): + raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking canonical ExtraParallel optimizer state", group=group) + if dist.is_initialized(): + signatures = [None] * dist.get_world_size(group) + dist.all_gather_object(signatures, signature, group=group) + if any(item != signature for item in signatures): + raise RuntimeError("Ranks disagree on canonical ExtraParallel checkpoint layout") + error = None + try: + if signature is not None: + model = state["model"] + options = StateDictOptions(ignore_frozen_params=True) if trainable_only else None + canonical_state = get_model_state_dict(model, options=options) + _check_canonical_spec_coverage(model, specs, canonical_state.keys(), trainable_only=trainable_only) + tensors = dict(model.named_parameters(remove_duplicate=False)) + tensors.update(model.named_buffers(remove_duplicate=False)) + for name, spec in specs.items(): + if ( + not isinstance(spec.placement, Shard) + or getattr(spec, "persistent_fsdp_shard_dim", None) is not None + ): + continue + if name not in tensors: + raise RuntimeError(f"ExtraParallel checkpoint tensor {name} is absent from the runtime model") + mesh = spec.para_fsdp_mesh + fsdp_mesh = mesh[mesh.mesh_dim_names[:-1]] + _validate_extra_parallel_tensor_v2( + tensors[name], mesh, fsdp_mesh, spec.placement.dim, spec.fsdp_shard_dim, spec.global_shape + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking canonical ExtraParallel runtime tensors", group=group) + + +def _extra_parallel_layout_key(specs): + """Version canonical tensors independently of their source geometry.""" + for spec in (specs or {}).values(): + if not isinstance(spec.placement, Shard): + continue + if getattr(spec, "persistent_fsdp_shard_dim", None) is not None: + continue + if not getattr(spec, "checkpoint_layout_v2", False): + raise RuntimeError("Canonical ExtraParallel checkpoints cannot mix versioned and legacy sharded tensors") + _normalized_extra_parallel_dims(spec) + digest = hashlib.sha256(b"veomni-extra-parallel-canonical-v2").hexdigest() + return f"{_EXTRA_PARALLEL_LAYOUT_MARKER}{digest}" + + +def _runtime_extra_parallel_layout_signature(specs): + """Detect mutation while a single Stateful adapter is in use.""" + if _needs_extra_parallel_layout_marker(specs): + _extra_parallel_layout_key(specs) + return tuple( + ( + name, + spec.para_name, + repr(spec.placement), + spec.fsdp_shard_dim, + getattr(spec, "persistent_fsdp_shard_dim", None), + getattr(spec, "global_shape", None), + getattr(spec, "checkpoint_layout_v2", False), + tuple(spec.para_fsdp_mesh.shape) if spec.para_fsdp_mesh else None, + ) + for name, spec in sorted((specs or {}).items()) + ) + + +def _strip_extra_parallel_layout_marker(state_dict, required, marker_key): + if not required: + return state_dict + state_dict = dict(state_dict) + marker = state_dict.pop(marker_key, None) + if marker != "ep-outer-v1": + raise RuntimeError( + "Unversioned same-dimension ExtraParallel checkpoint; explicit layout migration is required" + ) + prefix = marker_key + "__" + if any(not key.startswith(prefix) for key in state_dict): + raise RuntimeError("ExtraParallel checkpoint contains an unexpected unversioned tensor key") + return {key[len(prefix) :]: value for key, value in state_dict.items()} + + +def _namespace_extra_parallel_layout_state(state_dict, marker_key): + """Give v2 tensors distinct keys so a legacy planner cannot read them as v1.""" + prefix = marker_key + "__" + result = {prefix + key: value for key, value in state_dict.items()} + result[marker_key] = "ep-outer-v1" + return result + + +def _export_extra_parallel_layout_key(key): + if not key.startswith(_EXTRA_PARALLEL_LAYOUT_MARKER): + return key + marker_length = len(_EXTRA_PARALLEL_LAYOUT_MARKER) + 64 + marker = key[:marker_length] + if not _is_extra_parallel_layout_marker(marker): + raise RuntimeError("Malformed versioned ExtraParallel checkpoint key") + if marker != _extra_parallel_layout_key({}): + raise RuntimeError("Unknown ExtraParallel checkpoint layout version; a verified migration is required") + if key == marker: + return None + if key[marker_length : marker_length + 2] != "__": + raise RuntimeError("Malformed versioned ExtraParallel checkpoint tensor key") + return key[marker_length + 2 :] + + +def _read_extra_parallel_marker(reader, marker_key): + scratch = {"model": {marker_key: ""}} + dcp.load(scratch, storage_reader=reader, no_dist=True) + if scratch["model"][marker_key] != "ep-outer-v1": + raise RuntimeError("Invalid ExtraParallel checkpoint layout marker value") + + +def _extra_parallel_first_mesh(mesh: DeviceMesh) -> DeviceMesh: + """Return a view with the ExtraParallel dim outside the FSDP dim. + + The runtime mesh is ``(para_fsdp, para)`` (or + ``(para_replicate, para_fsdp, para)`` for HSDP). When both placements + shard the same tensor dim, PyTorch composes mesh dims outer-to-inner. + Keeping the runtime order would therefore describe FSDP-outer / + ExtraParallel-inner shards, while the optimizer stores the inverse layout. + Swapping only those two mesh dims preserves rank-local storage and makes + DTensor/DCP publish the real EP-outer / FSDP-inner global offsets. + """ + key = id(mesh) + cached = _EXTRA_PARALLEL_FIRST_MESHES.get(key) + if cached is not None and cached[0]() is mesh: + return cached[1] + names = mesh.mesh_dim_names + if names is None or len(names) != mesh.ndim or mesh.ndim not in (2, 3): + raise RuntimeError("ExtraParallel checkpoint mesh must have 2 or 3 named dimensions") + order = tuple(range(mesh.ndim - 2)) + (mesh.ndim - 1, mesh.ndim - 2) + reordered = DeviceMesh.from_group( + [mesh.get_group(index) for index in order], + device_type=mesh.device_type, + mesh=mesh.mesh.permute(order), + mesh_dim_names=tuple(names[index] for index in order), + ) + + def discard(reference): + current = _EXTRA_PARALLEL_FIRST_MESHES.get(key) + if current is not None and current[0] is reference: + _EXTRA_PARALLEL_FIRST_MESHES.pop(key, None) + + _EXTRA_PARALLEL_FIRST_MESHES[key] = (weakref.ref(mesh, discard), reordered) + return reordered + + +def _from_local_preserve_device(local_tensor, device_mesh, placements, *, shape=None, stride=None): + """Wrap optimizer CPU state without moving it onto its accelerator mesh. + + DTensor.from_local implicitly copies CPU input to device_mesh.device_type. + A metadata-only construction computes the same global shape/stride/spec; + attaching the actual local shard then preserves CPU storage and the original + mesh identity. DTensor's spec constructor is intentionally isolated here. + No collectives or accelerator allocations are needed (run_check=False). + """ + if local_tensor.device.type == "cpu" and device_mesh.device_type != "cpu": + metadata = DTensor.from_local( + local_tensor.to("meta"), + device_mesh=device_mesh, + placements=placements, + run_check=False, + shape=shape, + stride=stride, + ) + return DTensor(local_tensor, metadata._spec, requires_grad=local_tensor.requires_grad) + return DTensor.from_local( + local_tensor, + device_mesh=device_mesh, + placements=placements, + run_check=False, + shape=shape, + stride=stride, + ) + + +def _chunk_extent(extent: int, num_chunks: int, index: int) -> int: + """Extent of ``index`` under ``torch.chunk`` / DTensor ``Shard`` semantics.""" + full = -(-extent // num_chunks) + start = min(full * index, extent) + return min(extent, start + full) - start + + +def _dropped_extra_parallel_shape(loaded_tensor, ep_shard_dim, fsdp_shard_dim, extra_parallel_name=None): + """Global shape/stride of one ExtraParallel slice on the FSDP sub-mesh.""" + shape = list(loaded_tensor.shape) + if ep_shard_dim == fsdp_shard_dim: + mesh = loaded_tensor.device_mesh + mesh_dim_names = mesh.mesh_dim_names + if mesh_dim_names is None or len(mesh_dim_names) != mesh.ndim: + raise RuntimeError("ExtraParallel checkpoint mesh must have named dimensions") + last_two = mesh_dim_names[-2:] + if extra_parallel_name is not None and extra_parallel_name in last_two: + ep_dim = len(mesh_dim_names) - 2 if last_two[0] == extra_parallel_name else len(mesh_dim_names) - 1 + fsdp_dim = len(mesh_dim_names) - 2 if ep_dim == len(mesh_dim_names) - 1 else len(mesh_dim_names) - 1 + elif last_two[-1] == f"{last_two[0]}_fsdp": + ep_dim, fsdp_dim = len(mesh_dim_names) - 2, len(mesh_dim_names) - 1 + elif last_two[0] == f"{last_two[-1]}_fsdp": + fsdp_dim, ep_dim = len(mesh_dim_names) - 2, len(mesh_dim_names) - 1 + else: + raise RuntimeError( + f"Cannot determine ExtraParallel/FSDP mesh order from mesh dimensions {mesh_dim_names}." + ) + ep_size = mesh.size(ep_dim) + fsdp_size = mesh.size(fsdp_dim) + coordinate = mesh.get_coordinate() + ep_rank, fsdp_rank = coordinate[ep_dim], coordinate[fsdp_dim] + extent = shape[ep_shard_dim] + if ep_dim < fsdp_dim: + # restore_extra_parallel_dim lays same-dimension checkpoints out as + # EP-outer / FSDP-inner nested chunks. Mirror that layout instead + # of flat chunking when the extent is not evenly divisible. + ep_extent = _chunk_extent(extent, ep_size, ep_rank) + expected_local = _chunk_extent(ep_extent, fsdp_size, fsdp_rank) + else: + # Legacy FSDP-outer metadata has no contiguous EP slice unless the + # extent is evenly divisible by both mesh dimensions. + if extent % (fsdp_size * ep_size): + raise RuntimeError("Unsupported uneven ExtraParallel/FSDP shard layout") + ep_extent = extent // ep_size + expected_local = ep_extent // fsdp_size + if expected_local != loaded_tensor.to_local().shape[ep_shard_dim]: + raise RuntimeError("Unsupported uneven ExtraParallel/FSDP shard layout") + shape[ep_shard_dim] = ep_extent + else: + shape[ep_shard_dim] = loaded_tensor.to_local().shape[ep_shard_dim] + stride = [1] * len(shape) + for index in range(len(shape) - 2, -1, -1): + stride[index] = stride[index + 1] * shape[index + 1] + return tuple(shape), tuple(stride) + + def _apply_extra_parallel_dim( state_dict: Dict[str, Any], extra_parallel_fqn2spec_info: Dict[str, Any], @@ -138,8 +470,7 @@ def _apply_extra_parallel_dim( * ``"exact"`` (model): the state-dict key IS the FQN, e.g. ``"model.layers.0.mlp.experts.gate_proj"``. - * ``"substring"`` (optimizer): the state-dict key contains the FQN - with extra prefix/suffix, e.g. + * ``"substring"`` (optimizer): the key starts with ``state..``, e.g. ``"state.model.layers.0.mlp.experts.gate_proj.exp_avg"``. Non-tensor values and 0-D tensors are skipped unconditionally — they @@ -156,12 +487,18 @@ def _apply_extra_parallel_dim( extra_parallel_keys = list(extra_parallel_fqn2spec_info.keys()) if key_match == "substring" else None for name in sorted(state_dict.keys()): + tensor = state_dict[name] + if not torch.is_tensor(tensor) or tensor.ndim == 0: + continue + if key_match == "exact": if name not in extra_parallel_fqn2spec_info: continue spec_info = extra_parallel_fqn2spec_info[name] else: # "substring" - matches = [k for k in extra_parallel_keys if k in name] + # Anchor the optimizer namespace: a nested parameter may end with + # another parameter's FQN. Param-group fields are not tensor state. + matches = [k for k in extra_parallel_keys if name.startswith(f"state.{k}.")] if not matches: continue assert len(matches) == 1, f"Ambiguous ExtraParallel spec match for state key '{name}': {matches}" @@ -177,12 +514,6 @@ def _apply_extra_parallel_dim( if getattr(spec_info, "persistent_fsdp_shard_dim", None) is not None: continue - tensor = state_dict[name] - if not torch.is_tensor(tensor): - continue - if tensor.ndim == 0: - continue - assert spec_info.para_fsdp_mesh is not None, f"ExtraParallel spec {name} must have an ExtraParallel FSDP mesh" # Drop the innermost ExtraParallel (e.g. ``ep``) dim and keep the FSDP @@ -197,12 +528,29 @@ def _apply_extra_parallel_dim( # one explicitly instead of a fixed ``[Shard(0), Shard(1)]``. ep_shard_dim = spec_info.placement.dim fsdp_shard_dim = spec_info.fsdp_shard_dim - if action == "drop": + versioned = getattr(spec_info, "checkpoint_layout_v2", False) + if action == "drop" and not versioned: tensor = drop_extra_parallel_dim(tensor, fsdp_submesh, fsdp_shard_dim) - else: + elif action == "restore" and not versioned: tensor = restore_extra_parallel_dim( tensor, spec_info.para_fsdp_mesh, fsdp_submesh, ep_shard_dim, fsdp_shard_dim ) + elif action == "drop": + tensor = drop_extra_parallel_dim_v2( + tensor, + fsdp_submesh, + fsdp_shard_dim, + extra_parallel_name=spec_info.para_name, + ) + else: + tensor = restore_extra_parallel_dim_v2( + tensor, + spec_info.para_fsdp_mesh, + fsdp_submesh, + ep_shard_dim, + fsdp_shard_dim, + global_shape=spec_info.global_shape, + ) state_dict[name] = tensor return state_dict @@ -228,14 +576,31 @@ def __init__(self, model, trainable_only: bool = False, parallel_state=None): # If so, we need to restore Para(e.g. EP)-dim before saving to DCP self.parallel_state = parallel_state if parallel_state is not None else get_parallel_state() self.extra_parallel_fqn2spec_info = getattr(self.model, "_fqn2spec_info", None) + self._layout_marker_key = ( + _extra_parallel_layout_key(self.extra_parallel_fqn2spec_info) + if _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info) + else None + ) + self._layout_signature = _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) self.should_extra_parallel_aware = ( self.extra_parallel_fqn2spec_info is not None and self.parallel_state.dp_mode == "fsdp2" ) @torch.no_grad() def state_dict(self): + if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( + self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) + ): + raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") options = StateDictOptions(ignore_frozen_params=True) if self.trainable_only else None model_state_dict = get_model_state_dict(model=self.model, options=options) + if self.should_extra_parallel_aware and _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info): + _check_canonical_spec_coverage( + self.model, + self.extra_parallel_fqn2spec_info, + model_state_dict.keys(), + trainable_only=self.trainable_only, + ) if self.should_extra_parallel_aware: logger.info_rank0( "Getting model state_dict from ModelState wrapper, would restore ExtraParallel dim for ExtraParallel (e.g. Experts/Embeds) module" @@ -243,7 +608,8 @@ def state_dict(self): # As fsdp+extra parallel and pure extra parallel have different placements, e.g. [Shard(0), Shard(1)] and [Shard(0)], # restoring state dict should be extra parallel aware. model_state_dict = self.get_state_dict_with_extra_parallel_dim_preprocess(model_state_dict, "restore") - + if self.should_extra_parallel_aware and _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info): + model_state_dict = _namespace_extra_parallel_layout_state(model_state_dict, self._layout_marker_key) return model_state_dict @torch.no_grad() @@ -253,7 +619,16 @@ def load_state_dict(self, state_dict): need to drop ExtraParallel-dim when loading from DCP checkpoints so that ExtraParallel-FSDP would not be confused """ - model_state_dict = state_dict + if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( + self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) + ): + raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") + model_state_dict = _strip_extra_parallel_layout_marker( + state_dict, + self.should_extra_parallel_aware + and _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info), + self._layout_marker_key, + ) if self.should_extra_parallel_aware: model_state_dict = self.get_state_dict_with_extra_parallel_dim_preprocess(model_state_dict, "drop") @@ -307,12 +682,24 @@ def __init__(self, model, optimizer, parallel_state=None, *, load: bool = False) self.optimizer = optimizer self.parallel_state = parallel_state if parallel_state is not None else get_parallel_state() self.extra_parallel_fqn2spec_info = getattr(self.model, "_fqn2spec_info", None) + self._layout_marker_key = ( + _extra_parallel_layout_key(self.extra_parallel_fqn2spec_info) + if _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info) + else None + ) + self._layout_signature = _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) self.should_extra_parallel_aware = ( self.extra_parallel_fqn2spec_info is not None and self.parallel_state.dp_mode == "fsdp2" ) self._load = load def state_dict(self): + if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( + self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) + ): + raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") + if self.should_extra_parallel_aware and self._layout_marker_key is not None: + raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") if self.should_extra_parallel_aware: logger.info_rank0( "Getting optimizer state_dict from OptimizerState wrapper, would restore ExtraParallel dim for Experts module" @@ -324,6 +711,10 @@ def state_dict(self): optim_sd_with_extra_parallel_dim = self.get_state_dict_with_extra_parallel_dim_preprocess( vanilla_optim_sd, "restore" ) + if _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info): + optim_sd_with_extra_parallel_dim = _namespace_extra_parallel_layout_state( + optim_sd_with_extra_parallel_dim, self._layout_marker_key + ) return optim_sd_with_extra_parallel_dim if getattr(self.optimizer, "_is_multi_optimizer", False): @@ -335,7 +726,18 @@ def state_dict(self): def load_state_dict(self, state_dict): """Restore optimizer state from the loaded DCP state dict.""" - optim_state_from_dcp_load = state_dict + if self.should_extra_parallel_aware and self._layout_marker_key is not None: + raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") + if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( + self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) + ): + raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") + optim_state_from_dcp_load = _strip_extra_parallel_layout_marker( + state_dict, + self.should_extra_parallel_aware + and _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info), + self._layout_marker_key, + ) if self.should_extra_parallel_aware: # we need to drop ExtraParallel dim before loading them into optimizers optim_state_without_extra_parallel_dim = self.get_state_dict_with_extra_parallel_dim_preprocess( @@ -467,6 +869,178 @@ def restore_extra_parallel_dim( return dtensor +def drop_extra_parallel_dim_v2( + loaded_tensor: torch.Tensor, + device_mesh: DeviceMesh, + fsdp_shard_dim: int = 1, + *, + extra_parallel_name: Optional[str] = None, +): + """ + Drop ExtraParallel dims after loading from DCP so that ExtraParallel-FSDP would not be confused. + + ``device_mesh`` is the FSDP sub-mesh (ExtraParallel dim already excluded): + 1D ``(ep_fsdp,)`` for plain FSDP, or 2D ``(ep_replicate, ep_fsdp)`` for HSDP. + ``fsdp_shard_dim`` is the tensor dim FSDP shards along (1 by default, 0 for + the Muon zero-comm layout). The number of placements on the loaded DTensor + reflects the full saved mesh: + + * 1 placement: pure ExtraParallel, no FSDP -> return the plain local tensor. + * 2 placements: ``(ep_fsdp, ep)`` -> keep FSDP ``Shard(fsdp_shard_dim)`` on + the 1D sub-mesh. + * 3 placements: ``(ep_replicate, ep_fsdp, ep)`` (HSDP) -> keep + ``[Replicate(), Shard(fsdp_shard_dim)]`` on the 2D sub-mesh. + """ + + if not -loaded_tensor.ndim <= fsdp_shard_dim < loaded_tensor.ndim: + raise RuntimeError("FSDP shard dimension is outside the tensor rank") + fsdp_shard_dim %= loaded_tensor.ndim + num_placements = len(loaded_tensor.placements) + if num_placements == 1: + tensor_to_put = loaded_tensor.to_local() + elif num_placements == 2: + ep_placement = loaded_tensor.placements[-1] + if not isinstance(ep_placement, Shard): + raise RuntimeError(f"ExtraParallel placement must be Shard, got {ep_placement}") + ep_shard_dim = ep_placement.dim + shape, stride = _dropped_extra_parallel_shape(loaded_tensor, ep_shard_dim, fsdp_shard_dim, extra_parallel_name) + tensor_to_put = _from_local_preserve_device( + loaded_tensor._local_tensor, + device_mesh=device_mesh, + placements=[Shard(fsdp_shard_dim)], + shape=shape, + stride=stride, + ) + elif num_placements == 3: + ep_placement = loaded_tensor.placements[-1] + if not isinstance(ep_placement, Shard): + raise RuntimeError(f"ExtraParallel placement must be Shard, got {ep_placement}") + ep_shard_dim = ep_placement.dim + shape, stride = _dropped_extra_parallel_shape(loaded_tensor, ep_shard_dim, fsdp_shard_dim, extra_parallel_name) + tensor_to_put = _from_local_preserve_device( + loaded_tensor._local_tensor, + device_mesh=device_mesh, + placements=[Replicate(), Shard(fsdp_shard_dim)], + shape=shape, + stride=stride, + ) + else: + raise RuntimeError( + "Expect ExtraParallel parameters from checkpoints to be DTensor with 1-dim (no FSDP), " + f"2-dim (ExtraParallel+FSDP) or 3-dim (ExtraParallel+HSDP), got {loaded_tensor}" + ) + + return tensor_to_put + + +def _validate_extra_parallel_tensor_v2(tensor, mesh, fsdp_mesh, ep_dim, fsdp_dim, global_shape): + if not isinstance(tensor, DTensor): + raise RuntimeError("The v2 ExtraParallel checkpoint layout requires an FSDP DTensor") + ndim = tensor.ndim + if not (-ndim <= ep_dim < ndim and -ndim <= fsdp_dim < ndim): + raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") + ep_dim, fsdp_dim = ep_dim % ndim, fsdp_dim % ndim + if tensor.device_mesh != fsdp_mesh: + raise RuntimeError("ExtraParallel source tensor does not use the declared FSDP mesh") + if tuple(tensor.placements) != (Replicate(),) * (fsdp_mesh.ndim - 1) + (Shard(fsdp_dim),): + raise RuntimeError("ExtraParallel source tensor does not use the declared FSDP placement") + if global_shape is None or len(global_shape) != ndim: + raise RuntimeError("FSDP-only ExtraParallel tensors require the original global_shape") + + expected = list(global_shape) + ep_size, ep_rank = mesh.size(-1), mesh.get_local_rank(-1) + ep_chunk = (expected[ep_dim] + ep_size - 1) // ep_size + expected[ep_dim] = max(min(expected[ep_dim] - ep_rank * ep_chunk, ep_chunk), 0) + if tuple(tensor.shape) != tuple(expected): + raise RuntimeError("ExtraParallel source tensor shape does not match the declared global shape") + local_size, local_rank = fsdp_mesh.size(-1), fsdp_mesh.get_local_rank(-1) + local_chunk = (expected[fsdp_dim] + local_size - 1) // local_size + expected[fsdp_dim] = max(min(expected[fsdp_dim] - local_rank * local_chunk, local_chunk), 0) + if tuple(tensor.to_local().shape) != tuple(expected): + raise RuntimeError("ExtraParallel local tensor shape does not match its declared FSDP placement") + + +def restore_extra_parallel_dim_v2( + orgin_tensor: torch.Tensor, + fsdp_mesh: DeviceMesh, + extra_parallel_fsdp_mesh: DeviceMesh, + ep_shard_dim: int = 0, + fsdp_shard_dim: int = 1, + global_shape: Optional[tuple[int, ...]] = None, +): + """ + Restore ExtraParallel dim so that DCP can be aware about ExtraParallel ranks + + The ExtraParallel (e.g. ``ep``) dim and the FSDP dim shard the tensor along + DIFFERENT axes: ``ep_shard_dim`` (the EP plan's ``Shard.dim``, dim 0 for + experts) and ``fsdp_shard_dim`` (1 by default, 0 for the Muon zero-comm + layout). They must be mapped to the matching mesh dims explicitly. + + args: + orgin_tensor (torch.Tensor): The orgin tensor (FSDP-local shard). + fsdp_mesh (DeviceMesh): The full ExtraParallel mesh, i.e. 2D + ``(ep_fsdp, ep)`` for plain FSDP or 3D + ``(ep_replicate, ep_fsdp, ep)`` for HSDP. + extra_parallel_fsdp_mesh (DeviceMesh): The FSDP sub-mesh (ExtraParallel + dim excluded), used for the pure-ExtraParallel (no FSDP) path. + ep_shard_dim (int): Tensor dim the ExtraParallel dim shards along. + fsdp_shard_dim (int): Tensor dim FSDP shards along. + global_shape (tuple[int, ...], optional): The model parameter's shape + before ExtraParallel and FSDP sharding. It is required for an + FSDP-only DTensor: a local slice alone cannot prove global shape. + + Note: + When ``ep_shard_dim == fsdp_shard_dim`` (e.g. the Muon zero-comm layout + shards experts on dim 0 just like EP), both mesh dims split the SAME + tensor dim. The temporary checkpoint DTensor uses an EP-outer / + FSDP-inner mesh view so DCP metadata matches the optimizer's physical + local shards and global offsets instead of block-transposing the dim. + """ + assert fsdp_mesh.ndim in (2, 3), f"global_mesh.ndim must be 2 or 3, got {fsdp_mesh.ndim}" + ndim = orgin_tensor.ndim + if not (-ndim <= ep_shard_dim < ndim and -ndim <= fsdp_shard_dim < ndim): + raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") + ep_shard_dim %= ndim + fsdp_shard_dim %= ndim + + if isinstance(orgin_tensor, DTensor): + _validate_extra_parallel_tensor_v2( + orgin_tensor, fsdp_mesh, extra_parallel_fsdp_mesh, ep_shard_dim, fsdp_shard_dim, global_shape + ) + # ExtraParallel+FSDP2. mesh order is (ep_fsdp, ep) or, for HSDP, + # (ep_replicate, ep_fsdp, ep): ep_fsdp -> Shard(fsdp_shard_dim), + # ep -> Shard(ep_shard_dim), ep_replicate -> Replicate(). + if ep_shard_dim == fsdp_shard_dim: + checkpoint_mesh = _extra_parallel_first_mesh(fsdp_mesh) + if checkpoint_mesh.ndim == 3: + placements = [Replicate(), Shard(ep_shard_dim), Shard(fsdp_shard_dim)] + else: + placements = [Shard(ep_shard_dim), Shard(fsdp_shard_dim)] + else: + checkpoint_mesh = fsdp_mesh + if fsdp_mesh.ndim == 3: + placements = [Replicate(), Shard(fsdp_shard_dim), Shard(ep_shard_dim)] + else: + placements = [Shard(fsdp_shard_dim), Shard(ep_shard_dim)] + shape = tuple(global_shape) + stride = [1] * len(shape) + for index in range(len(shape) - 2, -1, -1): + stride[index] = stride[index + 1] * shape[index + 1] + dtensor = _from_local_preserve_device( + orgin_tensor._local_tensor, + device_mesh=checkpoint_mesh, + placements=placements, + shape=shape, + stride=tuple(stride), + ) + elif torch.is_tensor(orgin_tensor): + raise RuntimeError("The v2 ExtraParallel checkpoint layout requires an FSDP DTensor") + else: + raise RuntimeError(f"origin_tensor - {orgin_tensor} is not a tensor!") + + return dtensor + + def _local_rank() -> int: """This process's rank within its node, as set by the elastic launcher. @@ -840,6 +1414,7 @@ def save( """ if "model" not in state: raise ValueError("Model must be provided to save a distributed checkpoint.") + _reject_unproven_extra_parallel_optimizer(state, parallel_state, trainable_only=trainable_only) # Rejected up front, before anything reaches disk: an async write is still # running when save() returns and drops the staged copy, so it would write @@ -998,6 +1573,9 @@ def load( if "model" not in state: raise ValueError("Model must be provided to load a distributed checkpoint.") + _reject_unproven_extra_parallel_optimizer( + state, parallel_state, group=process_group, trainable_only=trainable_only + ) def _model_state() -> "ModelState": return ModelState(state["model"], trainable_only=trainable_only, parallel_state=parallel_state) @@ -1013,6 +1591,15 @@ def _optimizer_state() -> "OptimizerState": wants_optimizer = state.get("optimizer") is not None fused_dir = cls._legacy_fused_dir(path, module) + specs = getattr(state["model"], "_fqn2spec_info", None) + if _needs_extra_parallel_layout_marker(specs): + marker_error = None + try: + source_path = fused_dir if fused_dir is not None else weights_dir(path, module) + _read_extra_parallel_marker(cls._create_storage_reader(source_path), _extra_parallel_layout_key(specs)) + except BaseException as exc: + marker_error = exc + raise_if_any_rank_failed(marker_error, "validating ExtraParallel checkpoint marker", group=process_group) if fused_dir is not None: load_state: Dict[str, Any] = {"model": _model_state()} if wants_optimizer: @@ -1275,6 +1862,10 @@ def _normalize_key(key: str) -> Optional[str]: """ if not key.startswith("model."): return None + normalized = _export_extra_parallel_layout_key(key[6:]) + if normalized is None: + return None + key = "model." + normalized if key.startswith("model.model."): # Standard case: model.model.* -> model.* @@ -1319,9 +1910,21 @@ def _get_sharding_plan( if not isinstance(metadata, Metadata): raise ValueError(f"Invalid metadata format in {checkpoint_path}") + model_keys = {key for key in metadata.state_dict_metadata if key.startswith("model.")} + version_prefix = "model." + _EXTRA_PARALLEL_LAYOUT_MARKER + versioned_keys = {key for key in model_keys if key.startswith(version_prefix)} + if versioned_keys: + marker = "model." + _extra_parallel_layout_key({}) + if marker not in versioned_keys: + raise RuntimeError("Versioned ExtraParallel checkpoint marker is missing or unknown") + if any(key != marker and not key.startswith(marker + "__") for key in model_keys): + raise RuntimeError("ExtraParallel checkpoint mixes versioned and legacy tensor keys") + _read_extra_parallel_marker(FileSystemReader(checkpoint_path), marker.removeprefix("model.")) + # Collect model tensors and calculate sizes tensor_infos = [] all_dcp_keys = [] + seen_hf_keys = {} dropped = 0 for key, tensor_meta in metadata.state_dict_metadata.items(): @@ -1330,6 +1933,9 @@ def _get_sharding_plan( if drop_hf_keys and hf_key in drop_hf_keys: dropped += 1 continue + if hf_key in seen_hf_keys: + raise ValueError(f"Multiple DCP tensors map to HF key {hf_key}: {seen_hf_keys[hf_key]} and {key}") + seen_hf_keys[hf_key] = key # Determine dtype for size calculation if not hasattr(tensor_meta.properties, "dtype"): diff --git a/veomni/distributed/parallel_plan.py b/veomni/distributed/parallel_plan.py index 51fca8138e..6c5b54a104 100644 --- a/veomni/distributed/parallel_plan.py +++ b/veomni/distributed/parallel_plan.py @@ -42,6 +42,13 @@ class SpecInfo: # of FSDP2. Its runtime parameter remains a DTensor on the complete # ``(para_fsdp, para)`` mesh for the entire forward/backward lifecycle. persistent_fsdp_shard_dim: Optional[int] = None + # Original model shape before ExtraParallel or FSDP sharding. Checkpoint + # restoration cannot recover an uneven ExtraParallel extent from an + # FSDP-only DTensor because its local storage already lost that slice. + global_shape: Optional[tuple[int, ...]] = None + # Opt in to the versioned EP-outer checkpoint layout. Existing plans and + # checkpoints retain their legacy layout until explicitly migrated. + checkpoint_layout_v2: bool = False @property def para_mesh(self): @@ -149,6 +156,7 @@ def apply(self, model: nn.Module, extra_parallel_fsdp_device_mesh: Dict[str, Dev para_fsdp_mesh=para_fsdp_mesh, placement=shard, fqn=fqn, + global_shape=tuple(param.shape), persistent_fsdp_shard_dim=persistent_fsdp_shard_dim, ) if persistent_fsdp_shard_dim is not None: diff --git a/veomni/models/seed_omni/utils/hf_layout.py b/veomni/models/seed_omni/utils/hf_layout.py index d2374d0ce5..cfaae05361 100644 --- a/veomni/models/seed_omni/utils/hf_layout.py +++ b/veomni/models/seed_omni/utils/hf_layout.py @@ -693,8 +693,13 @@ def save_hf_source_checkpoint( """ from torch.distributed.checkpoint import HuggingFaceStorageWriter - from ....checkpoint.dcp_checkpointer import ModelState + from ....checkpoint.dcp_checkpointer import ( + ModelState, + _export_extra_parallel_layout_key, + _reject_unproven_extra_parallel_optimizer, + ) from ....checkpoint.dcp_consolidation import apply_dcp_consolidation_patch + from ....utils.dist_utils import raise_if_any_rank_failed apply_dcp_consolidation_patch() layout = source.layout @@ -703,31 +708,38 @@ def save_hf_source_checkpoint( save_state: dict[str, torch.Tensor] = {} for name, (model, parallel_state) in modules.items(): - module_state = ModelState(model, parallel_state=parallel_state).state_dict() - tied = _tied_parameter_names(model) - unmapped, reshaped = [], [] - for key, tensor in module_state.items(): - source_key = layout.source_key(name, key) - if source_key is None or source_key not in weight_files: - # A tied alias the source does not store (``tie_word_embeddings`` - # checkpoints omit ``lm_head.weight``) is written through its original. - if not any(layout.source_key(name, alias) in weight_files for alias in tied.get(key, ())): - unmapped.append(key) - continue - if source_key in save_state: - raise ValueError(f"Source key {source_key!r} is produced by more than one module.") - dtype, shape = source_specs[source_key] - if tuple(tensor.shape) != shape: - reshaped.append(f"{key} {tuple(tensor.shape)} vs {source_key} {shape}") - continue - save_state[source_key] = tensor if tensor.dtype == dtype else tensor.to(dtype) - if unmapped or reshaped: - raise ValueError( - f"save_hf_source_checkpoint: module '{name}' cannot be written back in the `{source.model_type}` " - f"source layout. Keys with no source key: {unmapped}; keys whose shape differs from the source: " - f"{reshaped}. A module whose checkpoint tensor converter renames or reshapes tensors cannot export " - "to its upstream layout; disable train.checkpoint.save_hf_weights, or train from a split checkpoint." - ) + _reject_unproven_extra_parallel_optimizer({"model": model}, parallel_state) + error = None + try: + module_state = ModelState(model, parallel_state=parallel_state).state_dict() + tied = _tied_parameter_names(model) + unmapped, reshaped = [], [] + for key, tensor in module_state.items(): + key = _export_extra_parallel_layout_key(key) + if key is None: + continue + source_key = layout.source_key(name, key) + if source_key is None or source_key not in weight_files: + if not any(layout.source_key(name, alias) in weight_files for alias in tied.get(key, ())): + unmapped.append(key) + continue + if source_key in save_state: + raise ValueError(f"Source key {source_key!r} is produced by more than one module.") + dtype, shape = source_specs[source_key] + if tuple(tensor.shape) != shape: + reshaped.append(f"{key} {tuple(tensor.shape)} vs {source_key} {shape}") + continue + save_state[source_key] = tensor if tensor.dtype == dtype else tensor.to(dtype) + if unmapped or reshaped: + raise ValueError( + f"save_hf_source_checkpoint: module '{name}' cannot be written back in the `{source.model_type}` " + f"source layout. Keys with no source key: {unmapped}; keys whose shape differs from the source: " + f"{reshaped}. A module whose checkpoint tensor converter renames or reshapes tensors cannot export " + "to its upstream layout; disable train.checkpoint.save_hf_weights, or train from a split checkpoint." + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, f"preparing SeedOmni HF module {name}") for key, original in layout.tied_source_keys.items(): if key in weight_files and key not in save_state and original in save_state: diff --git a/veomni/utils/save_safetensor_utils.py b/veomni/utils/save_safetensor_utils.py index 24b69b59d4..a487140d12 100644 --- a/veomni/utils/save_safetensor_utils.py +++ b/veomni/utils/save_safetensor_utils.py @@ -53,7 +53,7 @@ def get_model_save_state( """ from contextlib import nullcontext - from veomni.checkpoint.dcp_checkpointer import ModelState + from veomni.checkpoint.dcp_checkpointer import ModelState, _export_extra_parallel_layout_key from veomni.distributed.parallel_state import use_parallel_state from veomni.models.checkpoint import get_checkpoint_tensor_converter @@ -132,6 +132,10 @@ def _should_quantize(origin: str) -> bool: # so the original model parameters remain unchanged. converted_state = {} for k, v in save_state.items(): + # Translate versioned DCP-only names back to their HuggingFace FQNs. + k = _export_extra_parallel_layout_key(k) + if k is None: + continue if v.dtype == torch.float32: logger.info_rank0(f"Converting {k} from {v.dtype} to torch.bfloat16") converted_state[k] = v.to(torch.bfloat16) @@ -171,14 +175,23 @@ def _save_hf_safetensor_distributed( """ from torch.distributed.checkpoint import HuggingFaceStorageWriter + from veomni.checkpoint.dcp_checkpointer import _reject_unproven_extra_parallel_optimizer + # Apply DCP consolidation patch just-in-time for HDFS FUSE compatibility # This patches torch.distributed.checkpoint._consolidate_hf_safetensors._process_output_file # to use append mode instead of r+b mode, which is required for append-only file systems from veomni.checkpoint.dcp_consolidation import apply_dcp_consolidation_patch + from veomni.utils.dist_utils import raise_if_any_rank_failed apply_dcp_consolidation_patch() - save_state = get_model_save_state(model, fqn_to_index_mapping, parallel_state=parallel_state) + _reject_unproven_extra_parallel_optimizer({"model": model}, parallel_state) + save_state, error = None, None + try: + save_state = get_model_save_state(model, fqn_to_index_mapping, parallel_state=parallel_state) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "preparing distributed HF checkpoint") # Filter fqn_to_index_mapping to only include keys that exist in save_state. # This is necessary when training excludes certain modules (e.g., MTP) but the original @@ -357,6 +370,8 @@ def save_lora_adapter_with_dcp( MoE-LoRA adapter instead of every rank reading the whole gathered ``[E, ...]`` tensor. ``load_adapter_state_dict`` already prefers ``adapter_model.safetensors``. """ + from torch.distributed.tensor import Shard + from veomni.lora import is_veomni_lora_model from veomni.lora.state_dict import get_lora_state_dict @@ -366,17 +381,63 @@ def save_lora_adapter_with_dcp( "LoRA is served exclusively by the native veomni.lora stack." ) + from ..checkpoint.dcp_checkpointer import ( + _extra_parallel_layout_key, + _needs_extra_parallel_layout_marker, + _runtime_extra_parallel_layout_signature, + _validate_extra_parallel_tensor_v2, + ) + from .dist_utils import raise_if_any_rank_failed + + fqn2spec_info = getattr(model, "_fqn2spec_info", None) + error, signature = None, None + try: + if _needs_extra_parallel_layout_marker(fqn2spec_info): + _extra_parallel_layout_key(fqn2spec_info) + signature = _runtime_extra_parallel_layout_signature(fqn2spec_info) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking LoRA ExtraParallel layout") + if dist.is_initialized(): + signatures = [None] * dist.get_world_size() + dist.all_gather_object(signatures, signature) + if any(item != signature for item in signatures): + raise RuntimeError("Ranks disagree on LoRA ExtraParallel checkpoint layout") + synchronize() if dist.is_initialized(): dist.barrier() os.makedirs(save_path, exist_ok=True) dcp_save_path = os.path.join(save_path, dcp_subdir) + lora_state, error = None, None + try: + lora_state = get_lora_state_dict(model, adapter_name=adapter_name, config=model.get_lora_config()) + lora_state = {k: v.to(torch.bfloat16) if v.dtype == torch.float32 else v for k, v in lora_state.items()} + for name, value in lora_state.items(): + spec = (fqn2spec_info or {}).get(name) + if spec is None and name.endswith(".weight"): + spec = (fqn2spec_info or {}).get(f"{name[: -len('.weight')]}.{adapter_name}.weight") + if ( + spec is not None + and getattr(spec, "checkpoint_layout_v2", False) + and isinstance(spec.placement, Shard) + and getattr(spec, "persistent_fsdp_shard_dim", None) is None + ): + mesh = spec.para_fsdp_mesh + _validate_extra_parallel_tensor_v2( + value, + mesh, + mesh[mesh.mesh_dim_names[:-1]], + spec.placement.dim, + spec.fsdp_shard_dim, + spec.global_shape, + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking LoRA ExtraParallel runtime tensors") os.makedirs(dcp_save_path, exist_ok=True) - lora_state = get_lora_state_dict(model, adapter_name=adapter_name, config=model.get_lora_config()) - lora_state = {k: v.to(torch.bfloat16) if v.dtype == torch.float32 else v for k, v in lora_state.items()} - # Restore the EP shard dim before DCP sees the LoRA tensors. After # ``ParallelPlan.apply`` the EP-sharded params are stored as plain # local tensors (then wrapped by FSDP only along the FSDP mesh), so @@ -388,11 +449,10 @@ def save_lora_adapter_with_dcp( # as a 2-D (FSDP+EP) DTensor so DCP gathers both dimensions. # Mirrors what ``ModelState.get_state_dict_with_extra_parallel_dim_preprocess`` # does for the regular checkpoint path. - fqn2spec_info = getattr(model, "_fqn2spec_info", None) if fqn2spec_info is not None: from torch.distributed._tensor import Shard as _ShardPlacement - from ..checkpoint.dcp_checkpointer import restore_extra_parallel_dim + from ..checkpoint.dcp_checkpointer import restore_extra_parallel_dim, restore_extra_parallel_dim_v2 # ``get_peft_model_state_dict`` strips the adapter name from keys # (``...gate_proj.lora_A.default.weight`` -> ``...gate_proj.lora_A.weight``) @@ -405,12 +465,27 @@ def save_lora_adapter_with_dcp( if spec is None and name.endswith(".weight"): spec_key = f"{name[: -len('.weight')]}.{adapter_name}.weight" spec = fqn2spec_info.get(spec_key) - if spec is not None and isinstance(spec.placement, _ShardPlacement) and spec.para_fsdp_mesh is not None: - lora_state[name] = restore_extra_parallel_dim( - lora_state[name], - spec.para_fsdp_mesh, - spec.para_fsdp_mesh[f"{spec.para_name}_fsdp"], - ) + if ( + spec is not None + and isinstance(spec.placement, _ShardPlacement) + and spec.para_fsdp_mesh is not None + and getattr(spec, "persistent_fsdp_shard_dim", None) is None + ): + if getattr(spec, "checkpoint_layout_v2", False): + lora_state[name] = restore_extra_parallel_dim_v2( + lora_state[name], + spec.para_fsdp_mesh, + spec.para_fsdp_mesh[spec.para_fsdp_mesh.mesh_dim_names[:-1]], + ep_shard_dim=spec.placement.dim, + fsdp_shard_dim=spec.fsdp_shard_dim, + global_shape=spec.global_shape, + ) + else: + lora_state[name] = restore_extra_parallel_dim( + lora_state[name], + spec.para_fsdp_mesh, + spec.para_fsdp_mesh[f"{spec.para_name}_fsdp"], + ) # ckpt_to_state_dict's DCP conversion path only recognizes keys starting with "model.". # Prefix LoRA keys temporarily for DCP save so consolidation can reuse existing conversion logic. From adf19514b0652f8d49fcce7204b7cc238732fc2c Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sun, 1 Nov 2026 07:55:25 +0800 Subject: [PATCH 012/120] [checkpoint] prove multi-optimizer resume identity --- docs/usage/checkpoint.md | 64 +- tests/checkpoints/test_dcp_checkpointer.py | 51 +- .../test_multi_optimizer_proven_resume.py | 1807 +++++++++++++++++ .../test_optimizer_checkpoint_adapter.py | 200 ++ .../checkpoints/test_rank_local_optimizer.py | 486 +++++ tests/checkpoints/test_rank_local_ple.py | 72 + tests/optim/test_multi_optimizer_sparse_lr.py | 46 +- veomni/arguments/arguments_types.py | 13 + veomni/checkpoint/dcp_checkpointer.py | 665 +++++- .../checkpoint/multi_optimizer_checkpoint.py | 877 ++++++++ veomni/checkpoint/optimizer_resume_layout.py | 469 +++++ veomni/checkpoint/rank_local_optimizer.py | 333 +++ veomni/checkpoint/tensor_bytes.py | 45 + veomni/models/checkpoint/manager.py | 255 ++- veomni/models/model_runtime.py | 49 +- veomni/optim/checkpoint_adapter.py | 57 + veomni/optim/lr_scheduler.py | 35 +- veomni/optim/optimizer.py | 267 ++- 18 files changed, 5633 insertions(+), 158 deletions(-) create mode 100644 tests/checkpoints/test_multi_optimizer_proven_resume.py create mode 100644 tests/checkpoints/test_optimizer_checkpoint_adapter.py create mode 100644 tests/checkpoints/test_rank_local_optimizer.py create mode 100644 tests/checkpoints/test_rank_local_ple.py create mode 100644 veomni/checkpoint/multi_optimizer_checkpoint.py create mode 100644 veomni/checkpoint/optimizer_resume_layout.py create mode 100644 veomni/checkpoint/rank_local_optimizer.py create mode 100644 veomni/checkpoint/tensor_bytes.py create mode 100644 veomni/optim/checkpoint_adapter.py diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 83895d484d..943baf88a8 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -246,14 +246,70 @@ identifies canonical global values independently of source geometry, so DCP can reshard a model checkpoint across supported FSDP shard axes. Use `DistributedCheckpointer` for save and resume: it rejects old/unversioned layouts before copying model tensors. Migrating a legacy checkpoint requires -verified source layout information. Canonical v2 optimizer save and resume are -rejected until an owner-aware codec can preserve rank-specific optimizer -states; this opt-in therefore does not yet support resuming training. The -standard VeOmni HF export strips the namespace. +verified source layout information. Canonical v2 optimizer save and resume use +the rank-local codec described below. The standard VeOmni HF export strips the +namespace. Raw `torch.distributed.checkpoint.load(..., planner=DefaultLoadPlanner(allow_partial_load=True))` is not a supported migration API: it can return success without reading any tensor when the namespaces differ. Replicated values must already agree across ranks. +## MultiOptimizer recovery identity + +Use `DistributedCheckpointer` for training checkpoints containing a +`MultiOptimizer`. The save is synchronous and requires a completed optimizer +step with no pending gradients. Its identity records parameter ownership, +group structure, scheduler recipe and checkpoint generation. Loading rejects an +unversioned checkpoint, a changed optimizer layout, or a model/optimizer/sidecar +from different saves. An old checkpoint whose original group and scheduler +identity cannot be established needs a verified offline migration. + +The generic codec supports AdamW, AnyPrecisionAdamW and DistributedMuon +without ExtraParallel when every checkpoint rank has the same logical parameter +inventory and replicated values already agree. For the ExtraParallel candidate, +optimizer state is stored under rank-specific keys and recovery requires the +identical rank-to-mesh mapping. This avoids merging distinct expert-owner steps +and uninitialized states, but increases checkpoint storage for replicated +optimizer shards and does not support elastic resharding. Model replica shards +are checked for bitwise equality before DCP deduplication. The check streams +bounded chunks to CPU and adds work during checkpoint save. Unknown model +or optimizer hooks, asynchronous saves and unsupported layouts fail closed; +callers should validate these constraints before starting a long run. Synchronous +staged saves are supported. +Partially overlapping shards with different geometries are rejected, even when +their shared values agree; exact replicas and disjoint shards remain supported. +Group/state inventory and generation mismatches are rejected before model +copying. If an I/O or optimizer-application error occurs after model copying, +the load fails and clears its receipt but does not roll back in-memory objects; +discard and rebuild that runtime before retrying. +`MultiOptimizer.state_dict()` returns a detached inspection snapshot for +supported leaves, but `load_state_dict()` alone cannot prove group/scheduler +identity and rejects restoration; use the proven DCP path instead. +The codec also rejects tensor learning rates shared between groups and +overlapping storage between distinct model or optimizer tensors. Its byte-range +check is conservative for non-contiguous interleaved views. +Mutable optimizer options must not be shared between groups or optimizer +defaults: ordinary YAML `betas` are normalized to a tuple; custom shared lists, +dictionaries and tensor-valued learning rates require explicitly independent +copies. This avoids changing aliasing semantics on resume. +The model implementation and training program must also be the same trusted +version on save and resume; custom `state_dict`/`load_state_dict` overrides or +other user code can change training semantics outside this identity check. +Do not use raw DCP `OptimizerState` or `MultiOptimizer.load_state_dict` to +resume: they cannot establish scheduler and group provenance. The standalone +`MultiOptimizer.state_dict()` returns a snapshot and copies tensor states on +their current devices, which can temporarily require as much extra device +memory as the optimizer tensors themselves. That raw snapshot supports the +built-in AdamW, AnyPrecisionAdamW and DistributedMuon leaves, as well as +registered adapters, provided they have no unsupported serialization hooks. +It does not prove the identity needed for raw restoration. Use the high-level +`DistributedCheckpointer` for save and restore; its save path avoids this copy. + +For the base and DPO policy ExtraParallel runtimes, a versioned checkpoint's rank inventory is +used to temporarily include its previously owned frozen parameters while building +the optimizer. Their frozen flags are restored before the first forward and the +complete source group and scheduler identity is still checked on load. Frozen +parameters can be intentionally unfrozen after recovery. Unproven legacy +checkpoints and changed freeze configurations require verified migration. ## Related pages diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py index 1b6a2ba1f9..f4bd4ca234 100644 --- a/tests/checkpoints/test_dcp_checkpointer.py +++ b/tests/checkpoints/test_dcp_checkpointer.py @@ -41,7 +41,7 @@ def _fsdp2_multi_optimizer_worker(rank: int, world_size: int, tmp_path: Path, endpoint): """Worker for the FSDP2 MultiOptimizer DCP round-trip test.""" - from veomni.checkpoint.dcp_checkpointer import ModelState, OptimizerState + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.optim.optimizer import MultiOptimizer from veomni.utils.device import get_device_type, get_dist_comm_backend, get_torch_device @@ -75,25 +75,17 @@ def build_model_and_optimizer(): x = torch.randn(2, 4, device=device) source_model(x).sum().backward() source_optimizer.step() - expected = source_optimizer.state_dict() + source_optimizer.zero_grad() + expected = source_optimizer._sparse_state_dict() checkpoint_dir = tmp_path / "ckpt" - dcp.save( - {"model": ModelState(source_model), "optimizer": OptimizerState(source_model, source_optimizer)}, - checkpoint_id=str(checkpoint_dir), - ) + DistributedCheckpointer.save(str(checkpoint_dir), {"model": source_model, "optimizer": source_optimizer}) dist.barrier() target_model, target_optimizer = build_model_and_optimizer() - dcp.load( - { - "model": ModelState(target_model), - "optimizer": OptimizerState(target_model, target_optimizer, load=True), - }, - checkpoint_id=str(checkpoint_dir), - ) + DistributedCheckpointer.load(str(checkpoint_dir), {"model": target_model, "optimizer": target_optimizer}) - actual = target_optimizer.state_dict() + actual = target_optimizer._sparse_state_dict() assert expected.keys() == actual.keys() for key in expected: expected_value = expected[key] @@ -164,7 +156,7 @@ class TestMultiOptimizerState: """Non-ExtraParallel MultiOptimizer must use its own DCP protocol.""" def test_single_process_dcp_roundtrip(self, tmp_path): - from veomni.checkpoint.dcp_checkpointer import OptimizerState + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.optim.optimizer import MultiOptimizer def build_model_and_optimizer(): @@ -184,21 +176,20 @@ def build_model_and_optimizer(): for param in source_model.parameters(): param.grad = torch.randn_like(param) source_optimizer.step() - expected = source_optimizer.state_dict() + source_optimizer.zero_grad() + expected = source_optimizer._sparse_state_dict() assert any("exp_avg" in key for key in expected) assert any("step" in key for key in expected) target_model, target_optimizer = build_model_and_optimizer() - dcp.save( - {"optimizer": OptimizerState(source_model, source_optimizer, parallel_state=parallel_state)}, - checkpoint_id=tmp_path, + DistributedCheckpointer.save( + str(tmp_path), {"model": source_model, "optimizer": source_optimizer}, parallel_state=parallel_state ) - dcp.load( - {"optimizer": OptimizerState(target_model, target_optimizer, parallel_state=parallel_state, load=True)}, - checkpoint_id=tmp_path, + DistributedCheckpointer.load( + str(tmp_path), {"model": target_model, "optimizer": target_optimizer}, parallel_state=parallel_state ) - actual = target_optimizer.state_dict() + actual = target_optimizer._sparse_state_dict() assert expected.keys() == actual.keys() for key, expected_value in expected.items(): actual_value = actual[key] @@ -219,7 +210,7 @@ def test_multi_optimizer_sparse_state_excludes_synthetic(self, tmp_path): from torch.distributed.checkpoint import FileSystemReader from torch.distributed.checkpoint.metadata import Metadata - from veomni.checkpoint.dcp_checkpointer import OptimizerState + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.optim.optimizer import MultiOptimizer model = nn.Linear(4, 4) @@ -234,16 +225,16 @@ def test_multi_optimizer_sparse_state_excludes_synthetic(self, tmp_path): # Only bias gets a gradient; the adamw_w sub-optimizer remains empty. model.bias.grad = torch.randn_like(model.bias) optimizer.step() + optimizer.zero_grad() assert optimizer.optimizers_dict["adamw_b"].state assert not optimizer.optimizers_dict["adamw_w"].state parallel_state = SimpleNamespace(dp_mode="fsdp2") - dcp.save( - {"optimizer": OptimizerState(model, optimizer, parallel_state=parallel_state)}, - checkpoint_id=tmp_path, + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel_state ) - reader = FileSystemReader(tmp_path) + reader = FileSystemReader(tmp_path / "model" / "optimizer") metadata = reader.read_metadata() assert isinstance(metadata, Metadata) keys = list(metadata.state_dict_metadata.keys()) @@ -2368,6 +2359,10 @@ def test_a_rank_that_cannot_remove_the_markers_fails_the_group(self, tmp_path): patch("veomni.checkpoint.dcp_checkpointer.dist.is_initialized", return_value=True), patch("veomni.checkpoint.dcp_checkpointer.dist.get_rank", return_value=1), patch("veomni.checkpoint.dcp_checkpointer._reject_unproven_extra_parallel_optimizer"), + patch( + "veomni.checkpoint.dcp_checkpointer._coordinated_optimizer_flag", + side_effect=lambda value, *_a, **_k: value, + ), patch("veomni.utils.dist_utils.any_rank_failed", return_value=True), patch.object(DistributedCheckpointer, "execute_save") as execute_save, ): diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py new file mode 100644 index 0000000000..f165d5173d --- /dev/null +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -0,0 +1,1807 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. +"""CPU regressions for proven MultiOptimizer checkpoint recovery.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer +from veomni.optim.lr_scheduler import build_lr_scheduler +from veomni.optim.muon import DistributedMuon +from veomni.optim.optimizer import AnyPrecisionAdamW, MultiOptimizer + + +def _model_and_optimizer(*, split=True): + model = torch.nn.Linear(2, 2) + if split: + weight = torch.optim.AdamW([model.weight], lr=0.01, weight_decay=0.0) + bias = torch.optim.AdamW([model.bias], lr=0.02, weight_decay=0.0) + return model, MultiOptimizer(model, {"weight": weight, "bias": bias}, ["weight", "bias"]) + merged = torch.optim.AdamW([model.weight, model.bias], lr=0.5, weight_decay=0.0) + return model, MultiOptimizer(model, {"merged": merged}, ["merged"]) + + +def test_multi_optimizer_async_checkpoint_rejected_before_writes(tmp_path): + model, optimizer = _model_and_optimizer() + with pytest.raises(RuntimeError, match="require synchronous DCP save"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + save_async=True, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("mode", ["betas", "list", "nested"]) +def test_shared_mutable_group_options_rejected_before_save(tmp_path, mode): + model = torch.nn.Linear(2, 2) + shared = {"values": [0.9, 0.95]} if mode == "nested" else [0.9, 0.95] + groups = [ + {"params": [model.weight], "betas": (0.9, 0.95), "custom_recipe": {"shared": shared}}, + {"params": [model.bias], "betas": (0.9, 0.95), "custom_recipe": {"shared": shared}}, + ] + if mode == "betas": + for group in groups: + group["betas"] = shared + group.pop("custom_recipe") + leaf = torch.optim.AdamW(groups, lr=0.01) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + with pytest.raises(RuntimeError, match="Shared mutable optimizer option"): + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + assert not list(tmp_path.iterdir()) + + +def test_tensor_lr_aliased_to_optimizer_defaults_rejected_before_save(tmp_path): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([model.weight], lr=torch.tensor(0.01)) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + with pytest.raises(RuntimeError, match="Shared mutable optimizer tensor option"): + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("singleton", [False, True]) +def test_noncontiguous_adamw_parameter_recovers_and_updates(tmp_path, singleton): + def build(): + model = torch.nn.Module() + values = torch.arange(3, dtype=torch.float32)[::3] if singleton else torch.arange(12).float().reshape(3, 4).T + model.register_parameter("weight", torch.nn.Parameter(values)) + leaf = torch.optim.AdamW([model.weight], lr=0.01) + return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + source, optimizer = build() + source.weight.grad = torch.ones_like(source.weight) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build() + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + source.weight.grad = torch.full_like(source.weight, 0.125) + target.weight.grad = source.weight.grad.clone() + optimizer.step() + restored.step() + torch.testing.assert_close(target.weight, source.weight, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_tensor_byte_chunks_preserve_logical_order(dtype): + from veomni.checkpoint.tensor_bytes import tensor_byte_chunks + + tensor = torch.arange(60, dtype=dtype).reshape(3, 4, 5).permute(2, 0, 1) + chunks = list(tensor_byte_chunks(tensor, chunk_bytes=8)) + assert b"".join(chunks) == tensor.contiguous().view(torch.uint8).numpy().tobytes() + assert max(map(len, chunks)) <= 8 + + +def test_tensor_byte_chunks_bound_wide_index_iteration(): + from veomni.checkpoint.tensor_bytes import tensor_byte_chunks + + # A narrow transposed tensor must not allocate an index pool per row or + # emit one tiny hash update per row. + tensor = torch.arange(2_000_000, dtype=torch.float32).reshape(2, -1).T + chunks = list(tensor_byte_chunks(tensor, chunk_bytes=4 << 20)) + assert len(chunks) == 2 + assert all(len(chunk) <= 4 << 20 for chunk in chunks) + assert b"".join(chunks) == tensor.contiguous().view(torch.uint8).numpy().tobytes() + + +@pytest.mark.parametrize("chunk_bytes", [4, 8]) +def test_tensor_byte_chunks_accept_singleton_strided_views(chunk_bytes): + from veomni.checkpoint.tensor_bytes import tensor_byte_chunks + + for tensor in (torch.arange(3, dtype=torch.float32)[::3], torch.arange(6, dtype=torch.float32).reshape(2, 3).T): + actual = b"".join(tensor_byte_chunks(tensor, chunk_bytes=chunk_bytes)) + expected = torch.tensor(tensor.tolist()).view(torch.uint8).numpy().tobytes() + assert actual == expected + + +def _leaf_order_mismatch_worker(rank, mode, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + model = torch.nn.Linear(2, 2) + if rank == 0: + leaves = {"weight": torch.optim.AdamW([model.weight]), "bias": torch.optim.AdamW([model.bias])} + else: + leaves = {"both": torch.optim.AdamW(list(model.parameters()))} + optimizer = MultiOptimizer(model, leaves, list(leaves)) + with pytest.raises(RuntimeError, match="optimizer leaf order differs across ranks"): + if mode == "save": + optimizer._sparse_state_dict(coordinate=True) + else: + optimizer._dcp_load_state_dict({}) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("mode", ["save", "load"]) +def test_optimizer_leaf_order_mismatch_stops_all_ranks(mode): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_leaf_order_mismatch_worker, args=(mode,), nprocs=2) + + +def _high_level_leaf_mismatch_worker(rank, checkpoint, mismatch, mode, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + + def build(mismatched): + torch.manual_seed(45) + model = torch.nn.Linear(2, 2) + leaves = {"weight": torch.optim.AdamW([model.weight]), "bias": torch.optim.AdamW([model.bias])} + keys = list(leaves) + if mismatched and rank == 1: + if mismatch == "count": + leaves.pop("bias") + keys.pop() + else: + keys.reverse() + return model, MultiOptimizer(model, leaves, keys) + + parallel = SimpleNamespace(dp_mode="ddp") + if mode == "load": + source, optimizer = build(False) + DistributedCheckpointer.save( + checkpoint, {"model": source, "optimizer": optimizer}, parallel_state=parallel + ) + model, optimizer = build(True) + before = model.weight.detach().clone() + method = DistributedCheckpointer.save if mode == "save" else DistributedCheckpointer.load + with pytest.raises(RuntimeError, match="optimizer leaf order differs across ranks"): + method(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + torch.testing.assert_close(model.weight, before, rtol=0, atol=0) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("mode", ["save", "load"]) +@pytest.mark.parametrize("mismatch", ["count", "order"]) +def test_high_level_leaf_mismatch_rejected_before_model_mutation(tmp_path, mismatch, mode): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_high_level_leaf_mismatch_worker, args=(str(tmp_path), mismatch, mode), nprocs=2) + + +def test_staged_checkpoint_refuses_uncleared_previous_sidecar(tmp_path, monkeypatch): + import veomni.checkpoint.dcp_checkpointer as checkpoint_module + + checkpoint = tmp_path / "final" + stage = tmp_path / "stage" + stale = stage / checkpoint_module._STAGE_ROOT / checkpoint_module._stage_key(str(checkpoint)) + stale_sidecar = stale / "extra_state" / "extra_state_rank_0.pt" + stale_sidecar.parent.mkdir(parents=True) + torch.save({"cursor": "stale"}, stale_sidecar) + original = checkpoint_module.shutil.rmtree + + def fail_to_remove(path, *args, **kwargs): + if str(path) == str(stale): + return + return original(path, *args, **kwargs) + + monkeypatch.setattr(checkpoint_module.shutil, "rmtree", fail_to_remove) + model, optimizer = _model_and_optimizer() + with pytest.raises(FileExistsError): + DistributedCheckpointer.save( + str(checkpoint), + {"model": model, "optimizer": optimizer}, + stage_dir=str(stage), + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert stale_sidecar.exists() + assert not checkpoint.exists() + + +def _trainable_only_mismatch_worker(rank, checkpoint, mode, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + torch.manual_seed(8) + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + if mode == "load": + DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + method = DistributedCheckpointer.save if mode == "save" else DistributedCheckpointer.load + with pytest.raises(RuntimeError, match="trainable-only"): + method( + checkpoint, + {"model": model, "optimizer": optimizer}, + parallel_state=parallel, + trainable_only=bool(rank), + ) + finally: + dist.destroy_process_group() + + +def _proven_staged_resume_worker(rank, checkpoint, stage, endpoint): + import os + from datetime import timedelta + + import torch.distributed as dist + + os.environ["LOCAL_RANK"] = str(rank) + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + torch.manual_seed(9) + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.full_like(model.weight, 0.25) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel, stage_dir=stage + ) + target, restored = _model_and_optimizer() + DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + for source, destination in zip(model.parameters(), target.parameters(), strict=True): + torch.testing.assert_close(source, destination, rtol=0, atol=0) + source.grad = torch.full_like(source, 0.5) + destination.grad = torch.full_like(destination, 0.5) + optimizer.step() + restored.step() + for source, destination in zip(model.parameters(), target.parameters(), strict=True): + torch.testing.assert_close(source, destination, rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +def test_proven_staged_checkpoint_restores_and_updates_on_two_ranks(tmp_path): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous( + _proven_staged_resume_worker, args=(str(tmp_path / "checkpoint"), str(tmp_path / "stage")), nprocs=2 + ) + + +def _writer_error_worker(rank, checkpoint, target, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + original = DistributedCheckpointer._create_storage_writer + try: + torch.manual_seed(12) + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + + def make_writer(cls, path): + if rank == 1 and Path(path).name == target: + raise RuntimeError("injected writer factory failure") + return original(path) + + DistributedCheckpointer._create_storage_writer = classmethod(make_writer) + with pytest.raises(RuntimeError, match="writer"): + DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("target", ["ckpt", "optimizer"]) +def test_writer_factory_failure_propagates_before_dcp_collective(tmp_path, target): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_writer_error_worker, args=(str(tmp_path), target), nprocs=2) + + +def _inventory_reader_failure_worker(rank, checkpoint, rank_local, endpoint): + from datetime import timedelta + + import torch.distributed as dist + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import Shard, distribute_tensor + + from veomni.distributed.parallel_plan import SpecInfo + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + original = DistributedCheckpointer._create_storage_reader + try: + + def build(): + model = torch.nn.Module() + values = torch.arange(16, dtype=torch.float32).reshape(4, 4) + if rank_local: + mesh = init_device_mesh("cpu", (1, 2), mesh_dim_names=("ple_fsdp", "ple")) + model.weight = torch.nn.Parameter(distribute_tensor(values, mesh, [Shard(1), Shard(0)])) + model._fqn2spec_info = {"weight": SpecInfo("ple", Shard(0), "weight", mesh, 1, 1)} + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ple",), extra_parallel_fsdp_device_mesh={"ple": mesh} + ) + else: + model.weight = torch.nn.Parameter(values) + parallel = SimpleNamespace(dp_mode="ddp") + leaf = torch.optim.AdamW([model.weight]) + return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]), parallel + + source, optimizer, parallel = build() + DistributedCheckpointer.save(checkpoint, {"model": source, "optimizer": optimizer}, parallel_state=parallel) + target, restored, parallel = build() + attempts = 0 + + def reader(cls, path): + nonlocal attempts + if Path(path).name == "optimizer": + attempts += 1 + if rank == 1 and attempts == 2: + raise OSError("injected inventory reader failure") + return original(path) + + DistributedCheckpointer._create_storage_reader = classmethod(reader) + with pytest.raises((RuntimeError, OSError), match="inventory reader"): + DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + assert attempts == 2 + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("rank_local", [False, True]) +def test_inventory_reader_failure_propagates_before_collective(tmp_path, rank_local): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_inventory_reader_failure_worker, args=(str(tmp_path), rank_local), nprocs=2) + + +def _checkpoint_cache_failure_worker(rank, fail_at, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + from veomni.models.checkpoint import manager as manager_module + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + original = manager_module.helper.empty_cache + try: + calls = 0 + + def clear_cache(): + nonlocal calls + calls += 1 + if rank == 1 and calls == fail_at: + raise RuntimeError("injected cache failure") + + manager_module.helper.empty_cache = clear_cache + manager = object.__new__(manager_module.ModelCheckpointManager) + manager.config = SimpleNamespace( + save_path="/unused", + save_async=False, + dcp_save_to_lowest_rank=False, + stage_dir=None, + save_timeout_seconds=60, + ) + manager.runtime = SimpleNamespace( + model=torch.nn.Linear(2, 2), optimizer=SimpleNamespace(_is_multi_optimizer=True), lr_scheduler=None + ) + manager._extra_state = lambda _state: {"_veomni_scheduler_recipe": None} + manager._checkpointer_kwargs = lambda: {} + manager.checkpointer = SimpleNamespace(save=lambda *args, **kwargs: "new-generation") + manager._last_saved_step = 3 + manager._last_saved_optimizer_generation = "old-generation" + with pytest.raises(RuntimeError, match="clearing cache|injected cache failure"): + manager.save_dcp(SimpleNamespace(global_step=3)) + assert manager._last_saved_optimizer_generation is None + assert manager._last_saved_step == -1 + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + manager_module.helper.empty_cache = original + dist.destroy_process_group() + + +@pytest.mark.parametrize("fail_at", [1, 2]) +def test_checkpoint_cache_failure_propagates_before_next_collective(fail_at): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_checkpoint_cache_failure_worker, args=(fail_at,), nprocs=2) + + +def _inventory_hash_failure_worker(rank, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + from veomni.checkpoint import multi_optimizer_checkpoint as implementation + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + original = implementation.hashlib.sha256 + try: + if rank == 1: + + def fail_hash(*_args, **_kwargs): + raise MemoryError("injected inventory hash failure") + + implementation.hashlib.sha256 = fail_hash + with pytest.raises((RuntimeError, MemoryError), match="hashing optimizer checkpoint inventory|inventory hash"): + implementation._inventory_digests(["optimizer.state.weight.step"]) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + implementation.hashlib.sha256 = original + dist.destroy_process_group() + + +def test_inventory_hash_failure_propagates_before_collective(): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_inventory_hash_failure_worker, args=(), nprocs=2) + + +def _trainability_mismatch_worker(rank, checkpoint, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + model, optimizer = _model_and_optimizer() + if rank == 1: + model.weight.requires_grad_(False) + with pytest.raises(RuntimeError, match="trainability differs"): + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +def test_rank_asymmetric_trainability_is_rejected_before_save(tmp_path): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_trainability_mismatch_worker, args=(str(tmp_path),), nprocs=2) + + +def _tied_parameter_order_worker(rank, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + from veomni.checkpoint.optimizer_resume_layout import coordinated_multi_layout + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + model = torch.nn.Module() + shared = torch.nn.Parameter(torch.ones(2, 2)) + for name in ("a", "b") if rank == 0 else ("b", "a"): + model.register_parameter(name, shared) + optimizer = MultiOptimizer(model, {"leaf": torch.optim.AdamW([shared])}, ["leaf"]) + coordinated_multi_layout(model, optimizer, writing=True) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +def test_tied_parameter_registration_order_does_not_change_trainability(): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_tied_parameter_order_worker, args=(), nprocs=2) + + +def _save_config_mismatch_worker(rank, checkpoint, mismatch, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + torch.manual_seed(22) + model, optimizer = _model_and_optimizer() + state = {"model": model, "optimizer": optimizer} + kwargs = {"parallel_state": SimpleNamespace(dp_mode="ddp")} + if mismatch == "model" and rank == 1: + state.pop("model") + elif mismatch == "extra" and rank == 1: + state["extra_state"] = {} + elif mismatch == "stage" and rank == 1: + kwargs["stage_dir"] = checkpoint + "-stage" + elif mismatch == "stage_path": + kwargs["stage_dir"] = checkpoint + f"-stage-{rank}" + elif mismatch == "module": + kwargs["module"] = f"model-{rank}" + elif mismatch == "async" and rank == 1: + kwargs["save_async"] = True + kwargs["stage_dir"] = checkpoint + "-stage" + with pytest.raises(RuntimeError, match="differs across ranks"): + DistributedCheckpointer.save(checkpoint, state, **kwargs) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("mismatch", ["model", "extra", "stage", "stage_path", "module", "async"]) +def test_checkpoint_save_rank_config_mismatch_fails_all_ranks(tmp_path, mismatch): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_save_config_mismatch_worker, args=(str(tmp_path), mismatch), nprocs=2) + + +def _finalize_error_worker(rank, checkpoint, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + import veomni.checkpoint.dcp_checkpointer as checkpoint_module + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + torch.manual_seed(27) + model, optimizer = _model_and_optimizer() + + def synchronize(): + if rank == 1: + raise RuntimeError("injected final synchronize failure") + + checkpoint_module.synchronize = synchronize + with pytest.raises(RuntimeError, match="synchronize|finalizing"): + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +def test_save_finalize_error_propagates_to_all_ranks(tmp_path): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_finalize_error_worker, args=(str(tmp_path),), nprocs=2) + + +def _runtime_plan_mismatch_worker(rank, mode, endpoint): + import contextlib + from datetime import timedelta + + import torch.distributed as dist + + import veomni.models.model_runtime as runtime_module + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + + class TestRuntime(runtime_module.VeOmniModelRuntime): + def setup(self): + pass + + def _build_model(self): + self.model = torch.nn.Linear(2, 2) + if mode != "missing" and (mode != "plan" or rank == 0): + self.model._fqn2spec_info = {} + + def _freeze_model_module(self): + pass + + def _build_parallelized_model(self): + pass + + runtime_module.use_parallel_state = lambda *_args, **_kwargs: contextlib.nullcontext() + checkpoint = SimpleNamespace( + extra_parallel_canonical_layout=(mode == "missing" or (mode == "flag" and rank == 1)) + ) + args = SimpleNamespace(optimizer=SimpleNamespace(type="adamw")) + with pytest.raises(RuntimeError, match="differs across ranks|requires an ExtraParallel model plan"): + TestRuntime(args, train=SimpleNamespace(checkpoint=checkpoint)) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("mode", ["plan", "flag", "missing"]) +def test_runtime_checkpoint_plan_mismatch_fails_all_ranks(mode): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_runtime_plan_mismatch_worker, args=(mode,), nprocs=2) + + +@pytest.mark.parametrize("mode", ["save", "load"]) +def test_trainable_only_mismatch_fails_all_ranks(tmp_path, mode): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_trainable_only_mismatch_worker, args=(str(tmp_path), mode), nprocs=2) + + +@pytest.mark.parametrize("same_base", [True, False]) +def test_trainable_only_resume_proves_frozen_base_before_loading(tmp_path, same_base): + def make_model(base): + model = torch.nn.Linear(2, 2) + with torch.no_grad(): + model.weight.fill_(base) + model.bias.fill_(0.5) + model.weight.requires_grad_(False) + leaf = torch.optim.AdamW([model.bias], lr=0.02) + return model, MultiOptimizer(model, {"bias": leaf}, ["bias"]) + + source, optimizer = make_model(1.0) + source.bias.grad = torch.ones_like(source.bias) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel, trainable_only=True + ) + target, restored = make_model(1.0 if same_base else 2.0) + before = target.bias.detach().clone() + if same_base: + DistributedCheckpointer.load( + str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel, trainable_only=True + ) + torch.testing.assert_close(target.bias, source.bias, rtol=0, atol=0) + else: + with pytest.raises(RuntimeError, match="identity|layout|checkpoint"): + DistributedCheckpointer.load( + str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel, trainable_only=True + ) + torch.testing.assert_close(target.bias, before, rtol=0, atol=0) + + +def test_multi_optimizer_rejects_changed_group_identity(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + model.bias.grad = torch.ones_like(model.bias) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + + target, changed = _model_and_optimizer(split=False) + with pytest.raises(RuntimeError, match="(?i)optimizer.*(group|layout|identity)"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": changed}, parallel_state=parallel) + + +def test_unversioned_multi_optimizer_checkpoint_fails_closed(tmp_path): + import torch.distributed.checkpoint as dcp + + from veomni.checkpoint.layout import optimizer_dir + + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + model.bias.grad = torch.ones_like(model.bias) + optimizer.step() + # Historical MultiOptimizer format had flat entries but no group identity. + dcp.save({"optimizer": optimizer._sparse_state_dict()}, checkpoint_id=optimizer_dir(str(tmp_path))) + target, restored = _model_and_optimizer() + with pytest.raises(RuntimeError, match="(?i)group layout is unproven"): + DistributedCheckpointer.load( + str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + + +def test_optimizer_serialization_hooks_fail_before_overwriting_checkpoint(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + model.bias.grad = torch.ones_like(model.bias) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + original = model.weight.detach().clone() + + def change_groups(_optimizer, state): + state["param_groups"][0]["lr"] = 0.5 + return state + + handle = optimizer.optimizers_dict["weight"].register_state_dict_post_hook(change_groups) + with pytest.raises(RuntimeError, match="serialization hooks"): + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + handle.remove() + + target, restored = _model_and_optimizer() + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, original, rtol=0, atol=0) + + +def test_nonfresh_target_is_rejected_before_any_checkpoint_load(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + + target, restored = _model_and_optimizer() + restored.optimizers_dict["weight"].state[target.weight] = { + "step": torch.tensor(1.0), + "exp_avg": torch.zeros_like(target.weight, dtype=torch.bfloat16), + "exp_avg_sq": torch.zeros_like(target.weight, dtype=torch.bfloat16), + } + original = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="fresh optimizer"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, original, rtol=0, atol=0) + + +def test_anyprecision_fp32_moments_with_bf16_parameter_fail_closed(tmp_path): + def build(): + model = torch.nn.Linear(2, 2, bias=False).to(torch.bfloat16) + leaf = AnyPrecisionAdamW( + model.parameters(), + momentum_dtype=torch.float32, + variance_dtype=torch.float32, + compensation_buffer_dtype=torch.float32, + ) + return model, MultiOptimizer(model, {"weight": leaf}, ["weight"]) + + model, optimizer = build() + model.weight.grad = torch.full_like(model.weight, 0.1) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + with pytest.raises(RuntimeError, match="state dtype cannot be safely restored"): + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + assert not list(tmp_path.iterdir()) + + +def test_anyprecision_uninitialized_leaf_restores_and_updates(tmp_path): + def build(): + model = torch.nn.Linear(2, 2) + weight = AnyPrecisionAdamW([model.weight]) + bias = torch.optim.AdamW([model.bias], foreach=False) + return model, MultiOptimizer(model, {"weight": weight, "bias": bias}, ["weight", "bias"]) + + model, optimizer = build() + model.bias.grad = torch.ones_like(model.bias) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build() + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + assert target.weight not in restored.optimizers_dict["weight"].state + model.weight.grad = torch.ones_like(model.weight) + target.weight.grad = torch.ones_like(target.weight) + optimizer.step() + restored.step() + torch.testing.assert_close(model.weight, target.weight, rtol=0, atol=0) + for name in ("exp_avg", "exp_avg_sq"): + torch.testing.assert_close( + optimizer.optimizers_dict["weight"].state[model.weight][name], + restored.optimizers_dict["weight"].state[target.weight][name], + rtol=0, + atol=0, + ) + + +def test_trainable_only_rejects_trained_then_frozen_optimizer_parameter(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + model.weight.requires_grad_(False) + with pytest.raises(RuntimeError, match="trainable-only save would omit a parameter with optimizer history"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + trainable_only=True, + ) + + +def test_trainable_only_resume_rejects_frozen_parameter_with_saved_history(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _model_and_optimizer() + target.weight.requires_grad_(False) + original = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="trainable-only resume cannot prove"): + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored}, + parallel_state=parallel, + trainable_only=True, + ) + torch.testing.assert_close(target.weight, original, rtol=0, atol=0) + + +@pytest.mark.parametrize("mutation", ["freeze", "remove"]) +def test_trainable_only_save_rejects_model_hook_hiding_trained_parameter(tmp_path, mutation): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + + def hide_weight(module, state_dict, prefix, local_metadata): + if mutation == "freeze": + module.weight.requires_grad_(False) + else: + state_dict.pop(prefix + "weight") + + model.register_state_dict_post_hook(hide_weight) + with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + trainable_only=True, + ) + + +def test_multi_optimizer_uninitialized_leaf_uses_its_first_real_step(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + bias_before = model.bias.detach().clone() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + torch.testing.assert_close(model.bias, bias_before, rtol=0, atol=0) + assert model.bias not in optimizer.optimizers_dict["bias"].state + + target, restored = _model_and_optimizer() + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + assert target.bias not in restored.optimizers_dict["bias"].state + model.bias.grad = torch.full_like(model.bias, 0.4) + target.bias.grad = torch.full_like(target.bias, 0.4) + optimizer.step() + restored.step() + torch.testing.assert_close(target.bias, model.bias, rtol=0, atol=0) + + +def test_direct_multi_optimizer_without_provenance_is_rejected_before_mutation(): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + saved = optimizer._sparse_state_dict() + target, restored = _model_and_optimizer() + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="Unverified MultiOptimizer state"): + restored.load_state_dict(saved) + assert torch.equal(before, target.weight) + assert all(not leaf.state for leaf in restored.optimizers_dict.values()) + + +def test_raw_dcp_multi_without_provenance_is_rejected_before_synthetic_step(): + from veomni.checkpoint.dcp_checkpointer import OptimizerState + + target, restored = _model_and_optimizer() + wrapper = OptimizerState(target, restored, parallel_state=SimpleNamespace(dp_mode="ddp"), load=True) + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="Raw DCP MultiOptimizer does not prove"): + wrapper.state_dict() + snapshot = restored.state_dict() + assert snapshot + assert torch.equal(before, target.weight) + assert all(not leaf.state for leaf in restored.optimizers_dict.values()) + + +def test_raw_dcp_load_rejection_preserves_existing_optimizer_state(tmp_path): + import torch.distributed.checkpoint as dcp + + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + dcp.save({"optimizer": optimizer}, checkpoint_id=tmp_path) + target, restored = _model_and_optimizer() + target.weight.grad = torch.full_like(target.weight, 0.4) + restored.step() + restored.zero_grad() + before = {key: value.clone() for key, value in restored.optimizers_dict["weight"].state[target.weight].items()} + with pytest.raises(RuntimeError, match="Unverified MultiOptimizer state"): + dcp.load({"optimizer": restored}, checkpoint_id=tmp_path) + for key, value in before.items(): + torch.testing.assert_close(restored.optimizers_dict["weight"].state[target.weight][key], value, rtol=0, atol=0) + + +def test_raw_snapshot_rejects_live_group_reordering_hook(): + model, optimizer = _model_and_optimizer() + leaf = optimizer.optimizers_dict["weight"] + leaf.register_state_dict_post_hook(lambda opt, saved: (opt.param_groups.clear(), saved)[1]) + with pytest.raises(RuntimeError, match="safe raw MultiOptimizer snapshot"): + optimizer.state_dict() + assert leaf.param_groups[0]["params"][0] is model.weight + + +def test_multi_model_generation_keeps_standard_torch_export_strictly_loadable(tmp_path): + from torch.distributed.checkpoint.format_utils import dcp_to_torch_save + + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + exported = tmp_path / "export.pt" + dcp_to_torch_save(str(tmp_path / "model" / "ckpt"), str(exported)) + state = torch.load(exported, weights_only=False) + target = torch.nn.Linear(2, 2) + target.load_state_dict(state["model"], strict=True) + for actual, expected in zip(target.parameters(), model.parameters(), strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_multi_optimizer_rejects_changed_scheduler_before_loading_weights(tmp_path): + model, optimizer = _model_and_optimizer() + scheduler = build_lr_scheduler(optimizer, train_steps=100, lr_decay_style="linear", lr_warmup_ratio=0.1) + model.weight.grad = torch.ones_like(model.weight) + model.bias.grad = torch.ones_like(model.bias) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": scheduler.state_dict()}, + "_veomni_scheduler_object": scheduler, + }, + parallel_state=parallel, + ) + + target, restored = _model_and_optimizer() + different = build_lr_scheduler(restored, train_steps=200, lr_decay_style="linear", lr_warmup_ratio=0.1) + original = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="(?i)optimizer.*scheduler.*layout"): + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "extra_state": {}, "_veomni_scheduler_object": different}, + parallel_state=parallel, + ) + torch.testing.assert_close(target.weight, original, rtol=0, atol=0) + assert different is not None + + +def test_multi_lr_scheduler_rejects_missing_leaf(): + _, optimizer = _model_and_optimizer() + scheduler = build_lr_scheduler(optimizer, train_steps=100) + state = scheduler.state_dict() + state.pop("bias") + with pytest.raises(RuntimeError, match="leaves"): + scheduler.load_state_dict(state) + + +@pytest.mark.parametrize("override_container", [False, True]) +@pytest.mark.parametrize("method", ["load_state_dict", "step", "get_last_lr"]) +def test_multi_optimizer_rejects_overridden_scheduler_load(tmp_path, override_container, method): + model, optimizer = _model_and_optimizer() + scheduler = build_lr_scheduler(optimizer, train_steps=100, lr_decay_style="linear") + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": scheduler.state_dict()}, + "_veomni_scheduler_object": scheduler, + }, + parallel_state=parallel, + ) + target, restored = _model_and_optimizer() + restored_scheduler = build_lr_scheduler(restored, train_steps=100, lr_decay_style="linear") + owner = restored_scheduler if override_container else restored_scheduler["weight"] + setattr(owner, method, lambda *args, **kwargs: None) + with pytest.raises(RuntimeError, match="overridden scheduler"): + DistributedCheckpointer.load( + str(tmp_path), + { + "model": target, + "optimizer": restored, + "extra_state": {}, + "_veomni_scheduler_object": restored_scheduler, + }, + parallel_state=parallel, + ) + + +def test_multi_optimizer_rejects_snapshot_from_an_old_scheduler_binding(tmp_path): + model, optimizer = _model_and_optimizer() + original = build_lr_scheduler(optimizer, train_steps=100, lr_decay_style="linear") + newer = build_lr_scheduler(optimizer, train_steps=10, lr_decay_style="linear") + state = { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": original.state_dict()}, + "_veomni_scheduler_object": original, + } + with pytest.raises(RuntimeError, match="scheduler object differs"): + DistributedCheckpointer.save(str(tmp_path), state, parallel_state=SimpleNamespace(dp_mode="ddp")) + assert not (tmp_path / "model").exists() + assert newer is not None + + +def test_multi_optimizer_rejects_replaced_framework_scheduler_recipe(tmp_path): + model, optimizer = _model_and_optimizer() + scheduler = build_lr_scheduler(optimizer, train_steps=100, lr_decay_style="linear") + + def custom(step, *, factor=0.5): + return factor + + scheduler["weight"].lr_lambdas[0] = custom + state = { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": scheduler.state_dict()}, + "_veomni_scheduler_object": scheduler, + } + with pytest.raises(RuntimeError, match="scheduler recipe changed after construction"): + DistributedCheckpointer.save(str(tmp_path), state, parallel_state=SimpleNamespace(dp_mode="ddp")) + assert not (tmp_path / "model").exists() + + +def test_multi_optimizer_rejects_modified_scheduler_sidecar_before_loading(tmp_path): + model, optimizer = _model_and_optimizer() + scheduler = build_lr_scheduler(optimizer, train_steps=100) + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": scheduler.state_dict()}, + "_veomni_scheduler_object": scheduler, + }, + parallel_state=parallel, + ) + sidecar = tmp_path / "model" / "extra_state" / "extra_state_rank_0.pt" + stale = torch.load(sidecar, weights_only=False) + stale["lr_scheduler"]["weight"]["last_epoch"] += 1 + torch.save(stale, sidecar) + + target, restored = _model_and_optimizer() + restored_scheduler = build_lr_scheduler(restored, train_steps=100) + original = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="(?i)optimizer.*layout"): + DistributedCheckpointer.load( + str(tmp_path), + { + "model": target, + "optimizer": restored, + "extra_state": {}, + "_veomni_scheduler_object": restored_scheduler, + }, + parallel_state=parallel, + ) + torch.testing.assert_close(target.weight, original, rtol=0, atol=0) + assert restored_scheduler is not None + + +@pytest.mark.parametrize("swapped", ["optimizer", "extra_state"]) +def test_multi_optimizer_rejects_mixed_checkpoint_generations(tmp_path, swapped): + import shutil + + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + first, second = tmp_path / "first", tmp_path / "second" + for index, path in enumerate((first, second), start=1): + model.weight.grad = torch.full_like(model.weight, 0.1 * index) + optimizer.step() + optimizer.zero_grad() + DistributedCheckpointer.save( + str(path), + {"model": model, "optimizer": optimizer, "extra_state": {"condition_rng": torch.tensor(index)}}, + parallel_state=parallel, + ) + source = first / "model" / swapped + destination = second / "model" / swapped + shutil.rmtree(destination) + shutil.copytree(source, destination) + + target, restored = _model_and_optimizer() + original = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="(?i)generation|layout"): + DistributedCheckpointer.load( + str(second), {"model": target, "optimizer": restored, "extra_state": {}}, parallel_state=parallel + ) + torch.testing.assert_close(target.weight, original, rtol=0, atol=0) + + +def test_model_generation_rechecked_by_actual_dcp_load(tmp_path, monkeypatch): + import shutil + + import veomni.checkpoint.dcp_checkpointer as checkpoint + + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + first, second = tmp_path / "first", tmp_path / "second" + for index, path in enumerate((first, second), start=1): + model.weight.grad = torch.full_like(model.weight, 0.1 * index) + optimizer.step() + optimizer.zero_grad() + DistributedCheckpointer.save(str(path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + original_preflight = checkpoint._preflight_multi_optimizer_layout + + def replace_after_preflight(*args, **kwargs): + verified = original_preflight(*args, **kwargs) + shutil.rmtree(second / "model" / "ckpt") + shutil.copytree(first / "model" / "ckpt", second / "model" / "ckpt") + return verified + + monkeypatch.setattr(checkpoint, "_preflight_multi_optimizer_layout", replace_after_preflight) + target, restored = _model_and_optimizer() + with pytest.raises(BaseException, match="(?i)checkpoint.*(changed|generation)"): + DistributedCheckpointer.load(str(second), {"model": target, "optimizer": restored}, parallel_state=parallel) + + +def test_model_load_hook_cannot_change_optimizer_group_layout(tmp_path): + def build(): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) + return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + model, optimizer = build() + model.weight.grad = torch.ones_like(model.weight) + model.bias.grad = torch.ones_like(model.bias) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build() + + def change_groups(_module, _incompatible): + leaf = restored.optimizers_dict["leaf"] + leaf.param_groups[0]["params"].extend(leaf.param_groups.pop()["params"]) + + target.register_load_state_dict_post_hook(change_groups) + with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + + +@pytest.mark.parametrize("mutation", ["gradient", "state"]) +def test_model_load_hook_cannot_populate_optimizer_state_or_gradient(tmp_path, mutation): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _model_and_optimizer() + + def mutate(_module, _incompatible): + if mutation == "gradient": + target.weight.grad = torch.full_like(target.weight, 5) + else: + restored.optimizers_dict["weight"].state[target.weight]["step"] = torch.tensor(1.0) + + target.register_load_state_dict_post_hook(mutate) + with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + + +def test_checkpoint_generation_marker_is_not_a_hf_parameter(): + from veomni.checkpoint.dcp_checkpointer import _MULTI_GENERATION_MARKER, _normalize_key + + assert _normalize_key(_MULTI_GENERATION_MARKER) is None + + +def test_muon_adam_resume_does_not_run_an_optimizer_step(tmp_path): + def build(): + model = torch.nn.Linear(4, 4) + muon = DistributedMuon([model.weight], ns_implementation="std") + adam = torch.optim.AdamW([model.bias]) + return model, MultiOptimizer(model, {"muon": muon, "adam": adam}, ["muon", "adam"]) + + model, optimizer = build() + model.weight.grad = torch.full_like(model.weight, 0.1) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + + target, restored = build() + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, model.weight, rtol=0, atol=0) + torch.testing.assert_close(target.bias, model.bias, rtol=0, atol=0) + assert target.bias not in restored.optimizers_dict["adam"].state + model.weight.grad = torch.full_like(model.weight, 0.2) + target.weight.grad = torch.full_like(target.weight, 0.2) + model.bias.grad = torch.full_like(model.bias, 0.3) + target.bias.grad = torch.full_like(target.bias, 0.3) + optimizer.step() + restored.step() + torch.testing.assert_close(target.weight, model.weight, rtol=0, atol=0) + torch.testing.assert_close(target.bias, model.bias, rtol=0, atol=0) + + +def test_muon_head_partition_change_requires_explicit_migration(tmp_path): + def build(head_blocks): + model = torch.nn.Linear(4, 8, bias=False) + leaf = DistributedMuon([model.weight], ns_implementation="std", head_blocks=head_blocks) + return model, MultiOptimizer(model, {"muon": leaf}, ["muon"]) + + model, optimizer = build(2) + model.weight.grad = torch.full_like(model.weight, 0.1) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, changed = build(4) + with pytest.raises(RuntimeError, match="(?i)optimizer.*layout"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": changed}, parallel_state=parallel) + + +def test_multi_optimizer_layout_binds_extra_parallel_mesh_names(): + from veomni.checkpoint.optimizer_resume_layout import _local_layout + + model, optimizer = _model_and_optimizer() + mesh = SimpleNamespace(mesh=torch.tensor([[0, 1], [2, 3]]), device_type="cpu", mesh_dim_names=("ep_fsdp", "ep")) + spec = SimpleNamespace(para_name="ep", placement="Shard(0)", fsdp_shard_dim=1, para_fsdp_mesh=mesh) + model._fqn2spec_info = {"weight": spec} + before = _local_layout(model, optimizer, None, None, None, writing=False) + mesh.mesh_dim_names = ("other_fsdp", "other") + after = _local_layout(model, optimizer, None, None, None, writing=False) + assert before != after + + +@pytest.mark.parametrize("target_frozen", [False, True]) +def test_frozen_optimizer_history_survives_resume_then_unfreeze(tmp_path, target_frozen): + def build(): + model = torch.nn.Linear(2, 2) + weight = torch.optim.AdamW([model.weight], lr=0.01) + bias = torch.optim.AdamW([model.bias], lr=0.02) + return model, MultiOptimizer(model, {"weight": weight, "bias": bias}, ["weight", "bias"]) + + model, optimizer = build() + model.weight.grad = torch.full_like(model.weight, 0.1) + optimizer.step() + optimizer.zero_grad() + model.weight.requires_grad_(False) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + + target, restored = build() + target.weight.requires_grad_(not target_frozen) + if not target_frozen: + with pytest.raises(RuntimeError, match="(?i)optimizer.*layout"): + DistributedCheckpointer.load( + str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel + ) + return + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + assert target.weight.requires_grad is not target_frozen + model.weight.requires_grad_(True) + target.weight.requires_grad_(True) + model.weight.grad = torch.full_like(model.weight, 0.2) + target.weight.grad = torch.full_like(target.weight, 0.2) + optimizer.step() + restored.step() + torch.testing.assert_close(target.weight, model.weight, rtol=0, atol=0) + source_state = optimizer.optimizers_dict["weight"].state[model.weight] + restored_state = restored.optimizers_dict["weight"].state[target.weight] + for key in ("step", "exp_avg", "exp_avg_sq"): + torch.testing.assert_close(restored_state[key], source_state[key], rtol=0, atol=0) + + +def test_multi_optimizer_rejects_checkpoint_with_pending_gradients(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + before = model.weight.detach().clone() + with pytest.raises(RuntimeError, match="cannot omit pending gradients"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + torch.testing.assert_close(model.weight, before, rtol=0, atol=0) + assert not list(tmp_path.iterdir()) + + +def test_unowned_pending_gradients_reject_save_and_load_before_model_copy(tmp_path): + def build(): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([model.weight]) + return model, MultiOptimizer(model, {"weight": leaf}, ["weight"]) + + parallel = SimpleNamespace(dp_mode="ddp") + source, optimizer = build() + source.weight.grad = torch.ones_like(source.weight) + optimizer.step() + optimizer.zero_grad() + source.bias.grad = torch.ones_like(source.bias) + with pytest.raises(RuntimeError, match="cannot omit pending gradients"): + DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + assert not list(tmp_path.iterdir()) + + source.bias.grad = None + DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build() + target.bias.grad = torch.ones_like(target.bias) + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="cannot omit pending gradients"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +def test_incomplete_optimizer_source_state_cannot_be_saved(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + del optimizer.optimizers_dict["weight"].state[model.weight]["exp_avg_sq"] + with pytest.raises(RuntimeError, match="incomplete or unsupported parameter state"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not list(tmp_path.iterdir()) + + +def test_incomplete_optimizer_checkpoint_rejected_before_model_copy(tmp_path): + import pickle + + from veomni.checkpoint.layout import optimizer_dir + + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + path = Path(optimizer_dir(str(tmp_path))) / ".metadata" + with path.open("rb") as handle: + metadata = pickle.load(handle) + del metadata.state_dict_metadata["optimizer.state.weight.exp_avg_sq"] + with path.open("wb") as handle: + pickle.dump(metadata, handle) + target, restored = _model_and_optimizer() + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="incomplete or unsupported parameter state|inventory"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +@pytest.mark.parametrize("missing", ["group_lr", "entire_history"]) +def test_missing_group_or_all_parameter_history_rejected_before_model_copy(tmp_path, missing): + import pickle + + from veomni.checkpoint.layout import optimizer_dir + + source, optimizer = _model_and_optimizer() + source.weight.grad = torch.ones_like(source.weight) + optimizer.step() + optimizer.zero_grad() + optimizer.optimizers_dict["weight"].param_groups[0]["lr"] = 0.0025 + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + path = Path(optimizer_dir(str(tmp_path))) / ".metadata" + with path.open("rb") as handle: + metadata = pickle.load(handle) + if missing == "group_lr": + keys = ["optimizer.param_groups.weight.lr"] + else: + keys = [key for key in metadata.state_dict_metadata if key.startswith("optimizer.state.weight.")] + assert set(keys) == { + "optimizer.state.weight.step", + "optimizer.state.weight.exp_avg", + "optimizer.state.weight.exp_avg_sq", + } + for key in keys: + del metadata.state_dict_metadata[key] + with path.open("wb") as handle: + pickle.dump(metadata, handle) + target, restored = _model_and_optimizer() + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="inventory"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +def test_model_save_hook_cannot_remove_checkpoint_optimizer(tmp_path): + model, optimizer = _model_and_optimizer() + state = {"model": model, "optimizer": optimizer} + + def remove_optimizer(*_): + state["optimizer"] = None + + model.register_state_dict_post_hook(remove_optimizer) + with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): + DistributedCheckpointer.save(str(tmp_path), state, parallel_state=SimpleNamespace(dp_mode="ddp")) + assert state["optimizer"] is optimizer + + +def test_multi_optimizer_rejects_stale_sidecar_before_overwrite(tmp_path): + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "extra_state": {"cursor": "old"}}, + parallel_state=parallel, + ) + with pytest.raises(RuntimeError, match="sidecar cannot be reused"): + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + + +def test_multi_optimizer_rejects_duplicate_ownership_before_checkpoint_write(tmp_path): + model, optimizer = _model_and_optimizer() + optimizer.optimizers_dict["bias"].param_groups[0]["params"].append(model.weight) + with pytest.raises(RuntimeError, match="ownership is duplicated"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not list(tmp_path.iterdir()) + + +def test_multi_optimizer_rejects_empty_parameter_group_before_checkpoint_write(tmp_path): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([{"params": []}, {"params": [model.weight]}]) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + with pytest.raises(RuntimeError, match="empty parameter group"): + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + assert not list(tmp_path.iterdir()) + + +def test_optimizer_root_names_must_match_checkpoint_model(tmp_path): + weight = torch.nn.Parameter(torch.ones(2)) + model = torch.nn.Module() + model.register_parameter("b", weight) + optimizer_root = torch.nn.Module() + optimizer_root.register_parameter("a", weight) + leaf = torch.optim.AdamW([weight], lr=0.01) + optimizer = MultiOptimizer(optimizer_root, {"leaf": leaf}, ["leaf"]) + with pytest.raises(RuntimeError, match="different canonical parameter identities"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not list(tmp_path.iterdir()) + + +def test_tensor_lr_dtype_change_rejected_before_model_load(tmp_path): + def build(dtype): + model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW([model.weight], lr=torch.tensor(0.03012345, dtype=dtype)) + leaf.param_groups[0]["lr"] = leaf.param_groups[0]["lr"].clone() + return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + model, optimizer = build(torch.float32) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build(torch.bfloat16) + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +def test_shared_tensor_lr_across_groups_is_rejected_before_save(tmp_path): + def build(shared): + model = torch.nn.Linear(2, 2) + lr = torch.tensor(0.01) + other = lr if shared else lr.clone() + leaf = torch.optim.AdamW([{"params": [model.weight], "lr": lr}, {"params": [model.bias], "lr": other}]) + return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + model, optimizer = build(True) + parallel = SimpleNamespace(dp_mode="ddp") + with pytest.raises(RuntimeError, match="Shared mutable optimizer tensor option"): + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + assert not list(tmp_path.iterdir()) + + +def test_model_checkpoint_extra_state_records_optimizer_global_step(): + from veomni.models.checkpoint.manager import ModelCheckpointManager + + _, optimizer = _model_and_optimizer() + manager = object.__new__(ModelCheckpointManager) + manager.runtime = SimpleNamespace(optimizer=optimizer, lr_scheduler=None, extra_state=lambda: {}) + saved = manager._extra_state(SimpleNamespace(global_step=3)) + assert saved["_veomni_optimizer_global_step"] == 3 + + +def test_proven_optimizer_load_sets_identity_receipt(tmp_path): + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + generation = DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + target, restored = _model_and_optimizer() + state = {"model": target, "optimizer": restored} + DistributedCheckpointer.load(str(tmp_path), state, parallel_state=parallel) + assert state["_veomni_optimizer_identity_verified"] is True + assert isinstance(generation, str) and state["_veomni_optimizer_checkpoint_generation"] == generation + + +def test_failed_second_load_clears_old_identity_receipt(tmp_path, monkeypatch): + from veomni.checkpoint import dcp_checkpointer + + parallel = SimpleNamespace(dp_mode="ddp") + first, first_optimizer = _model_and_optimizer() + second, second_optimizer = _model_and_optimizer() + with torch.no_grad(): + first.weight.fill_(1) + second.weight.fill_(9) + DistributedCheckpointer.save( + str(tmp_path / "first"), {"model": first, "optimizer": first_optimizer}, parallel_state=parallel + ) + DistributedCheckpointer.save( + str(tmp_path / "second"), {"model": second, "optimizer": second_optimizer}, parallel_state=parallel + ) + target, restored = _model_and_optimizer() + state = {"model": target, "optimizer": restored} + DistributedCheckpointer.load(str(tmp_path / "first"), state, parallel_state=parallel) + assert state["_veomni_optimizer_identity_verified"] is True + + def fail_optimizer_apply(*_args, **_kwargs): + raise RuntimeError("injected optimizer apply failure") + + monkeypatch.setattr(dcp_checkpointer._ProvenMultiOptimizerState, "load_state_dict", fail_optimizer_apply) + with pytest.raises(BaseException, match="optimizer apply failure"): + DistributedCheckpointer.load(str(tmp_path / "second"), state, parallel_state=parallel) + assert torch.equal(target.weight, second.weight) + assert "_veomni_optimizer_identity_verified" not in state + assert "_veomni_optimizer_checkpoint_generation" not in state + + +def test_model_buffer_dtype_change_rejected_before_model_load(tmp_path): + model, optimizer = _model_and_optimizer() + model.register_buffer("sensitive", torch.tensor([0.1], dtype=torch.float32)) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _model_and_optimizer() + target.register_buffer("sensitive", torch.tensor([0.2], dtype=torch.bfloat16)) + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="tensor dtype or shape differs"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +def test_source_model_extra_buffer_rejected_before_model_load(tmp_path): + model, optimizer = _model_and_optimizer() + model.register_buffer("offset", torch.tensor([0.5])) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _model_and_optimizer() + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="model state inventory differs"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +@pytest.mark.parametrize("buffer_alias", [False, True]) +def test_distinct_model_tensors_with_overlapping_storage_are_rejected(tmp_path, buffer_alias): + model = torch.nn.Module() + model.register_parameter("p", torch.nn.Parameter(torch.ones(3))) + if buffer_alias: + model.register_buffer("q", model.p.detach()) + leaf = torch.optim.AdamW([model.p]) + else: + model.register_parameter("q", torch.nn.Parameter(model.p.detach())) + leaf = torch.optim.AdamW([model.p, model.q]) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + with pytest.raises(RuntimeError, match="checkpoint model storage alias"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("kind", ["parameter_buffer", "buffers"]) +def test_model_object_alias_graph_must_match_on_resume(tmp_path, kind): + def build(shared): + model = torch.nn.Module() + model.register_parameter("p", torch.nn.Parameter(torch.ones(3))) + if kind == "parameter_buffer": + model.register_buffer("view", model.p if shared else model.p.detach().clone()) + else: + model.register_buffer("view1", torch.ones(3)) + model.register_buffer("view2", model.view1 if shared else model.view1.clone()) + leaf = torch.optim.AdamW([model.p]) + return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + model, optimizer = build(True) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build(False) + before = target.p.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.p, before, rtol=0, atol=0) + + +def test_multi_optimizer_rejects_unproven_step_hooks(tmp_path): + model, optimizer = _model_and_optimizer() + optimizer.optimizers_dict["weight"].register_step_pre_hook(lambda opt, *_: opt.param_groups[0].update(lr=0.5)) + with pytest.raises(RuntimeError, match="cannot prove custom optimizer"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not list(tmp_path.iterdir()) + + +def test_multi_optimizer_rejects_overridden_step(tmp_path): + model, optimizer = _model_and_optimizer() + optimizer.step = lambda: None + with pytest.raises(RuntimeError, match="cannot prove overridden optimizer methods"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not list(tmp_path.iterdir()) + + +@pytest.mark.parametrize("old_model_state", [True, False]) +def test_multi_optimizer_save_preserves_completed_checkpoint(tmp_path, old_model_state): + from veomni.checkpoint.layout import checkpoint_is_complete, write_manifest + + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model}, parallel_state=parallel) + write_manifest(str(tmp_path), global_step=0, world_size=1, model_state=old_model_state) + assert checkpoint_is_complete(str(tmp_path)) + with pytest.raises(RuntimeError, match="Cannot overwrite a completed checkpoint"): + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + assert checkpoint_is_complete(str(tmp_path)) + + +def test_failed_multi_optimizer_save_remains_incomplete(tmp_path, monkeypatch): + from veomni.checkpoint.layout import checkpoint_is_complete + + model, optimizer = _model_and_optimizer() + original = DistributedCheckpointer._create_storage_writer + + def failing_writer(path): + if str(path).endswith("optimizer"): + raise OSError("injected optimizer write failure") + return original(path) + + monkeypatch.setattr(DistributedCheckpointer, "_create_storage_writer", failing_writer) + with pytest.raises(OSError, match="injected optimizer write failure"): + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + assert not checkpoint_is_complete(str(tmp_path)) + + +@pytest.mark.parametrize("phase", ["save", "load"]) +def test_model_value_changing_hooks_are_rejected(tmp_path, phase): + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + if phase == "load": + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + model, optimizer = _model_and_optimizer() + + def change_weight(module, _): + with torch.no_grad(): + module.weight.add_(0.5) + + model.register_load_state_dict_post_hook(change_weight) + else: + + def change_weight(module, saved, prefix, metadata): + saved[prefix + "weight"] = saved[prefix + "weight"] + 0.5 + + model.register_state_dict_post_hook(change_weight) + before = model.weight.detach().clone() + with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): + if phase == "save": + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + else: + DistributedCheckpointer.load( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + torch.testing.assert_close(model.weight, before, rtol=0, atol=0) + + +def test_model_hook_cannot_bypass_check_by_copying_fsdp_name(tmp_path): + model, optimizer = _model_and_optimizer() + + def fake_sharded_hook(*_): + with torch.no_grad(): + model.weight.add_(0.5) + + fake_sharded_hook.__module__ = "torch.distributed.fsdp._fully_shard._fsdp_param_group" + fake_sharded_hook.__qualname__ = "FSDPParamGroup._register_state_dict_hooks..to_sharded_hook" + model.register_state_dict_pre_hook(fake_sharded_hook) + before = model.weight.detach().clone() + with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + torch.testing.assert_close(model.weight, before, rtol=0, atol=0) + + +def _replicated_optimizer_metadata_worker(rank, checkpoint, divergent_field, endpoint): + from datetime import timedelta + + import torch.distributed as dist + from torch.distributed.fsdp import fully_shard + + from veomni.distributed.parallel_state import _init_parallel_state, get_parallel_state + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + _init_parallel_state(dp_size=2, dp_mode="fsdp2", device_type="cpu") + parallel = get_parallel_state() + model = torch.nn.Linear(4, 4) + fully_shard(model, mesh=parallel.dp_shard_mesh) + leaf = torch.optim.AdamW(model.parameters(), lr=0.01) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + model(torch.ones(2, 4)).sum().backward() + optimizer.step() + optimizer.zero_grad() + if rank == 1: + if divergent_field == "lr": + leaf.param_groups[0]["lr"] = 0.02 + else: + next(iter(leaf.state.values()))["step"] += 1 + with pytest.raises(RuntimeError, match="Replicated optimizer metadata differs"): + DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("field", ["lr", "step"]) +def test_replicated_optimizer_metadata_must_agree_before_save(tmp_path, field): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_replicated_optimizer_metadata_worker, args=(str(tmp_path), field), nprocs=2) diff --git a/tests/checkpoints/test_optimizer_checkpoint_adapter.py b/tests/checkpoints/test_optimizer_checkpoint_adapter.py new file mode 100644 index 0000000000..c0de985734 --- /dev/null +++ b/tests/checkpoints/test_optimizer_checkpoint_adapter.py @@ -0,0 +1,200 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +"""Exercise the optimizer checkpoint extension without optional optimizer packages.""" + +from contextlib import nullcontext +from types import SimpleNamespace + +import pytest +import torch + +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer, OptimizerState +from veomni.optim.checkpoint_adapter import register_rank_local_adamw +from veomni.optim.optimizer import MultiOptimizer + + +class ToyAdamW(torch.optim.Optimizer): + def __init__(self, params, scale=1.0): + super().__init__(params, {"lr": 0.01}) + self.scale = scale + + @torch.no_grad() + def step(self, closure=None): + for group in self.param_groups: + for param in group["params"]: + if param.grad is None: + continue + state = self.state[param] + if not state: + state.update( + step=torch.tensor(0.0), + master_weight=param.detach().float().clone(), + exp_avg=torch.zeros_like(param, dtype=torch.float32), + exp_avg_sq=torch.zeros_like(param, dtype=torch.float32), + ) + state["step"] += 1 + state["exp_avg"].mul_(0.9).add_(param.grad, alpha=0.1) + state["exp_avg_sq"].mul_(0.95).addcmul_(param.grad, param.grad, value=0.05) + state["master_weight"].add_(state["exp_avg"], alpha=-group["lr"] * self.scale) + param.copy_(state["master_weight"]) + + +register_rank_local_adamw( + ToyAdamW, + lambda _optimizer: nullcontext(), + lambda optimizer: ("toy-v1", optimizer.scale), + lambda _optimizer, _param, _field: torch.float32, + lambda _optimizer, _param: False, +) + + +def _build(scale=1.0): + torch.manual_seed(57) + model = torch.nn.Linear(2, 2) + leaf = ToyAdamW(model.parameters(), scale=scale) + return model, MultiOptimizer(model, {"toy": leaf}, ["toy"]) + + +def test_yaml_beta_sequence_is_immutable(): + from veomni.arguments.arguments_types import OptimizerConfig + from veomni.arguments.parser import _instantiate_recursive + + config = _instantiate_recursive(OptimizerConfig, {"betas": [0.9, 0.95]}) + assert config.betas == (0.9, 0.95) and isinstance(config.betas, tuple) + + +def test_registered_optimizer_preserves_lazy_state_and_next_update(tmp_path): + source, optimizer = _build() + source.weight.grad = torch.ones_like(source.weight) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _build() + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + assert target.bias not in restored.optimizers_dict["toy"].state + for left, right in zip(source.parameters(), target.parameters(), strict=True): + left.grad = torch.full_like(left, 0.125) + right.grad = left.grad.clone() + optimizer.step() + restored.step() + for left, right in zip(source.parameters(), target.parameters(), strict=True): + torch.testing.assert_close(left, right, rtol=0, atol=0) + + +def test_registered_optimizer_raw_snapshot_is_detached_but_restore_requires_proof(): + model, optimizer = _build() + model.weight.grad = torch.ones_like(model.weight) + optimizer.step() + optimizer.zero_grad() + snapshot = optimizer.state_dict() + live = optimizer.optimizers_dict["toy"].state[model.weight]["master_weight"] + saved = snapshot["state.weight.master_weight"] + assert saved.data_ptr() != live.data_ptr() + before = live.detach().clone() + saved.add_(7) + torch.testing.assert_close(live, before, rtol=0, atol=0) + with pytest.raises(RuntimeError, match="Unverified MultiOptimizer state"): + optimizer.load_state_dict(snapshot) + + +def test_registered_optimizer_rejects_changed_algorithm_before_loading(tmp_path): + source, optimizer = _build() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _build(scale=2.0) + with pytest.raises(RuntimeError, match="layout|identity|checkpoint"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + + +@pytest.mark.parametrize("changed_field", ["master_weight", "step"]) +def test_lazy_adapter_state_dtype_policy_is_part_of_checkpoint_identity(tmp_path, changed_field): + class ConfigurableToy(ToyAdamW): + def __init__(self, params, dtype): + super().__init__(params) + self.configured_dtype = dtype + + register_rank_local_adamw( + ConfigurableToy, + lambda _optimizer: nullcontext(), + lambda _optimizer: ("same-algorithm",), + lambda optimizer, _param, field: optimizer.configured_dtype if field == changed_field else torch.float32, + lambda _optimizer, _param: False, + ) + + def build(dtype): + model = torch.nn.Linear(2, 2, dtype=torch.bfloat16) + leaf = ConfigurableToy(model.parameters(), dtype) + return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + + parallel = SimpleNamespace(dp_mode="ddp") + source, optimizer = build(torch.float32) + assert not optimizer.optimizers_dict["leaf"].state + DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build(torch.bfloat16) + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity|checkpoint"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +def test_registered_single_optimizer_is_rejected_before_any_save_or_model_load(tmp_path): + # Registration alone declares exact state precision. In particular, an + # implementation need not also define requires_exact_state_dtypes. + source = torch.nn.Linear(2, 2, dtype=torch.bfloat16) + optimizer = ToyAdamW(source.parameters()) + parallel = SimpleNamespace(dp_mode="ddp") + assert not hasattr(optimizer, "requires_exact_state_dtypes") + source.weight.grad = torch.ones_like(source.weight) + optimizer.step() + assert optimizer.state[source.weight]["master_weight"].dtype == torch.float32 + + invalid_dir = tmp_path / "single-optimizer" + with pytest.raises(RuntimeError, match="require MultiOptimizer"): + DistributedCheckpointer.save( + str(invalid_dir), {"model": source, "optimizer": optimizer}, parallel_state=parallel + ) + assert not invalid_dir.exists() + + raw = OptimizerState(source, optimizer, parallel_state=parallel) + with pytest.raises(RuntimeError, match="require MultiOptimizer"): + raw.state_dict() + with pytest.raises(RuntimeError, match="require MultiOptimizer"): + raw.load_state_dict({}) + + model_dir = tmp_path / "model-only" + DistributedCheckpointer.save(str(model_dir), {"model": source}, parallel_state=parallel) + target = torch.nn.Linear(2, 2, dtype=torch.bfloat16) + with torch.no_grad(): + target.weight.fill_(42) + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="require MultiOptimizer"): + DistributedCheckpointer.load( + str(model_dir), {"model": target, "optimizer": ToyAdamW(target.parameters())}, parallel_state=parallel + ) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +def test_registered_adapter_cannot_fabricate_optimizer_group_options(tmp_path): + class FabricatingToy(ToyAdamW): + def state_dict(self): + state = super().state_dict() + state["param_groups"][0]["factor"] = 2 + return state + + register_rank_local_adamw( + FabricatingToy, + lambda _optimizer: nullcontext(), + lambda _optimizer: ("toy-v1",), + lambda _optimizer, _param, _field: torch.float32, + lambda _optimizer, _param: False, + ) + model = torch.nn.Linear(2, 2) + leaf = FabricatingToy(model.parameters()) + optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) + with pytest.raises(RuntimeError, match="fabricated checkpoint group options"): + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + assert not list(tmp_path.iterdir()) diff --git a/tests/checkpoints/test_rank_local_optimizer.py b/tests/checkpoints/test_rank_local_optimizer.py new file mode 100644 index 0000000000..7ea291fe23 --- /dev/null +++ b/tests/checkpoints/test_rank_local_optimizer.py @@ -0,0 +1,486 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +import torch.multiprocessing as mp +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import Replicate, Shard, distribute_tensor + +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer +from veomni.checkpoint.rank_local_optimizer import ( + local_optimizer_state, + recover_frozen_optimizer_ownership, + restore_optimizer_state, + verify_replicated_model_shards, +) +from veomni.distributed.parallel_plan import SpecInfo +from veomni.optim.muon import DistributedMuon +from veomni.optim.optimizer import MultiOptimizer + + +def _roundtrip_worker(rank, root, inactive_owner): + dist.init_process_group( + "gloo", init_method=f"file://{root}/init", rank=rank, world_size=4, timeout=timedelta(seconds=30) + ) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("fsdp", "ep")) + fsdp, owner = mesh["fsdp"], mesh.get_local_rank("ep") + local = distribute_tensor(torch.arange(16, dtype=torch.float32).reshape(4, 4) + 100 * owner, fsdp, [Shard(0)]) + source = {"param_groups.expert.lr": 0.005 + 0.001 * owner} + if not inactive_owner or owner == 0: + source.update( + { + "state.expert.step": torch.tensor(7 if owner == 0 else 3), + "state.expert.exp_avg": local, + "state.expert.exp_avg_sq": local + 0.5, + } + ) + dcp.save({"optimizer": local_optimizer_state(source, rank)}, checkpoint_id=f"{root}/checkpoint") + target = {"param_groups.expert.lr": 0.0} + if not inactive_owner or owner == 0: + target.update( + { + "state.expert.step": torch.tensor(0), + "state.expert.exp_avg": torch.zeros_like(local), + "state.expert.exp_avg_sq": torch.zeros_like(local), + } + ) + checkpoint = {"optimizer": local_optimizer_state(target, rank)} + dcp.load(checkpoint, checkpoint_id=f"{root}/checkpoint") + restored = restore_optimizer_state(checkpoint["optimizer"], target, rank) + assert set(restored) == set(source) + for name in source: + if torch.is_tensor(source[name]): + actual = restored[name].to_local() if hasattr(restored[name], "to_local") else restored[name] + expected = source[name].to_local() if hasattr(source[name], "to_local") else source[name] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + else: + assert restored[name] == source[name] + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("inactive_owner", [False, True]) +def test_rank_local_state_keeps_owner_steps_and_absence(tmp_path, inactive_owner): + mp.spawn(_roundtrip_worker, args=(str(tmp_path), inactive_owner), nprocs=4, join=True) + + +def _optimizer_resume_worker(rank, root, inactive_owner): + dist.init_process_group( + "gloo", init_method=f"file://{root}/optimizer-init", rank=rank, world_size=4, timeout=timedelta(seconds=40) + ) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ep_fsdp", "ep")) + fsdp, owner = mesh["ep_fsdp"], mesh.get_local_rank("ep") + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + + def build(): + model = torch.nn.Module() + values = torch.arange(16, dtype=torch.float32).reshape(4, 4) + 100 * owner + model.register_parameter("weight", torch.nn.Parameter(distribute_tensor(values, fsdp, [Shard(1)]))) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 1)} + leaf = torch.optim.AdamW([model.weight], lr=0.01) + return model, MultiOptimizer(model, {"expert": leaf}, ["expert"]) + + source, optimizer = build() + if not inactive_owner or owner == 0: + source.weight.grad = torch.ones_like(source.weight) + optimizer.step() + optimizer.zero_grad() + if owner == 0: + source.weight.grad = torch.full_like(source.weight, 0.5) + optimizer.step() + optimizer.zero_grad() + DistributedCheckpointer.save( + f"{root}/training", {"model": source, "optimizer": optimizer}, parallel_state=parallel + ) + target, restored = build() + DistributedCheckpointer.load( + f"{root}/training", {"model": target, "optimizer": restored}, parallel_state=parallel + ) + torch.testing.assert_close(target.weight.to_local(), source.weight.to_local(), rtol=0, atol=0) + source_state = optimizer.optimizers_dict["expert"].state.get(source.weight, {}) + target_state = restored.optimizers_dict["expert"].state.get(target.weight, {}) + assert set(source_state) == set(target_state) + for key, value in source_state.items(): + expected = value.to_local() if hasattr(value, "to_local") else value + actual = target_state[key].to_local() if hasattr(target_state[key], "to_local") else target_state[key] + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + source.weight.grad = torch.full_like(source.weight, 0.2) + target.weight.grad = torch.full_like(target.weight, 0.2) + optimizer.step() + restored.step() + torch.testing.assert_close(target.weight.to_local(), source.weight.to_local(), rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("inactive_owner", [False, True]) +def test_high_level_rank_owned_optimizer_cold_resume(tmp_path, inactive_owner): + mp.spawn(_optimizer_resume_worker, args=(str(tmp_path), inactive_owner), nprocs=4, join=True) + + +def _cold_optimizer_worker(rank, root, inactive_owner, phase, kind, canonical, same_axis): + dist.init_process_group( + "gloo", init_method=f"file://{root}/{phase}-init", rank=rank, world_size=4, timeout=timedelta(seconds=50) + ) + try: + if kind == "muon" and same_axis: + # Gloo lacks list all_to_all; emulate the same rank-to-rank tensor + # exchange with its supported all_to_all_single for this CPU test. + def gloo_all_to_all(outputs, inputs, group=None, async_op=False): + assert not async_op + send_sizes = [tensor.numel() for tensor in inputs] + recv_sizes = [tensor.numel() for tensor in outputs] + send = torch.cat([tensor.reshape(-1) for tensor in inputs]) + recv = torch.empty(sum(recv_sizes), dtype=send.dtype, device=send.device) + dist.all_to_all_single(recv, send, recv_sizes, send_sizes, group=group) + offset = 0 + for tensor, size in zip(outputs, recv_sizes, strict=True): + tensor.copy_(recv[offset : offset + size].reshape(tensor.shape)) + offset += size + + dist.all_to_all = gloo_all_to_all + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ep_fsdp", "ep")) + fsdp, owner = mesh["ep_fsdp"], mesh.get_local_rank("ep") + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + model = torch.nn.Module() + values = torch.arange(16, dtype=torch.float32).reshape(4, 4) + 100 * owner + fsdp_axis = 0 if same_axis else 1 + model.register_parameter("weight", torch.nn.Parameter(distribute_tensor(values, fsdp, [Shard(fsdp_axis)]))) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, fsdp_axis, None, (8, 4), canonical)} + if kind == "muon": + leaf = DistributedMuon([model.weight], lr=0.01, ns_implementation="std") + else: + leaf = torch.optim.AdamW([model.weight], lr=0.01) + leaves = {"expert": leaf} + if kind == "muon": + model.register_parameter("bias", torch.nn.Parameter(torch.ones(2))) + model._fqn2spec_info["bias"] = SpecInfo("ep", Replicate(), "bias", mesh) + leaves["fallback"] = torch.optim.AdamW([model.bias], lr=0.001) + optimizer = MultiOptimizer(model, leaves, list(leaves)) + + def snapshot(): + state = optimizer.optimizers_dict["expert"].state.get(model.weight, {}) + result = { + "weight": model.weight.to_local().detach().clone(), + "state": { + key: (value.to_local() if hasattr(value, "to_local") else value).detach().clone() + for key, value in state.items() + }, + } + if kind == "muon": + result["bias"] = model.bias.detach().clone() + bias_state = optimizer.optimizers_dict["fallback"].state.get(model.bias, {}) + result["bias_state"] = {key: value.detach().clone() for key, value in bias_state.items()} + return result + + def advance(gradient, *, include_bias=True): + model.weight.grad = torch.full_like(model.weight, gradient) + if kind == "muon" and include_bias: + model.bias.grad = torch.full_like(model.bias, 0.03) + optimizer.step() + optimizer.zero_grad() + + if phase == "save": + for _ in range(2 if owner == 0 else (0 if inactive_owner else 1)): + advance(1.0, include_bias=False) + if kind == "muon": + model.bias.grad = torch.full_like(model.bias, 0.03) + optimizer.step() + optimizer.zero_grad() + DistributedCheckpointer.save( + f"{root}/cold", + {"model": model, "optimizer": optimizer}, + parallel_state=parallel, + ) + before = snapshot() + advance(0.2) + torch.save({"before": before, "after": snapshot()}, f"{root}/reference-{rank}.pt") + else: + DistributedCheckpointer.load( + f"{root}/cold", {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + reference = torch.load(f"{root}/reference-{rank}.pt", weights_only=True) + torch.testing.assert_close(snapshot(), reference["before"], rtol=0, atol=0) + advance(0.2) + torch.testing.assert_close(snapshot(), reference["after"], rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("inactive_owner", [False, True]) +@pytest.mark.parametrize("kind", ["adam", "muon"]) +@pytest.mark.parametrize("canonical", [False, True]) +def test_rank_owned_optimizer_fresh_process_resume(tmp_path, inactive_owner, kind, canonical): + mp.spawn( + _cold_optimizer_worker, + args=(str(tmp_path), inactive_owner, "save", kind, canonical, False), + nprocs=4, + join=True, + ) + mp.spawn( + _cold_optimizer_worker, + args=(str(tmp_path), inactive_owner, "load", kind, canonical, False), + nprocs=4, + join=True, + ) + + +def test_canonical_same_axis_muon_fresh_process_resume(tmp_path): + mp.spawn(_cold_optimizer_worker, args=(str(tmp_path), False, "save", "muon", True, True), nprocs=4, join=True) + mp.spawn(_cold_optimizer_worker, args=(str(tmp_path), False, "load", "muon", True, True), nprocs=4, join=True) + + +def _replica_check_worker(rank, root): + dist.init_process_group( + "gloo", init_method=f"file://{root}/replica-init", rank=rank, world_size=2, timeout=timedelta(seconds=20) + ) + try: + with pytest.raises(RuntimeError, match="replicas differ"): + verify_replicated_model_shards({"bias": torch.tensor([float(rank)])}) + with pytest.raises(RuntimeError, match="replicas differ"): + verify_replicated_model_shards({"extra": {"scale": torch.tensor(float(rank))}}) + with pytest.raises(RuntimeError, match="replicas differ"): + verify_replicated_model_shards({"extra": {"scale": float(rank)}}) + finally: + dist.destroy_process_group() + + +def test_divergent_replicas_rejected_before_checkpoint_deduplication(tmp_path): + mp.spawn(_replica_check_worker, args=(str(tmp_path),), nprocs=2, join=True) + + +def _global_tensor_metadata_mismatch_worker(rank, root, mismatch): + from torch.distributed.tensor import DTensor + + dist.init_process_group( + "gloo", + init_method=f"file://{root}/global-metadata-init", + rank=rank, + world_size=2, + timeout=timedelta(seconds=30), + ) + try: + mesh = init_device_mesh("cpu", (2,)) + dtype = torch.bfloat16 if mismatch == "dtype" and rank else torch.float32 + global_size = 6 if mismatch == "shape" and rank else 4 + local = torch.ones(global_size // 2, dtype=dtype) + tensor = DTensor.from_local(local, mesh, [Shard(0)], shape=torch.Size([global_size]), stride=(1,)) + with pytest.raises(RuntimeError, match="global dtype or shape"): + verify_replicated_model_shards({"expert": tensor}) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("mismatch", ["dtype", "shape"]) +def test_checkpoint_global_tensor_metadata_mismatch_fails_all_ranks(tmp_path, mismatch): + mp.spawn(_global_tensor_metadata_mismatch_worker, args=(str(tmp_path), mismatch), nprocs=2, join=True) + + +def _overlapping_tensor_shards_worker(rank, root, same_offset): + from torch.distributed.tensor import DTensor + + dist.init_process_group( + "gloo", + init_method=f"file://{root}/overlap-init", + rank=rank, + world_size=2, + timeout=timedelta(seconds=30), + ) + try: + mesh = init_device_mesh("cpu", (2,)) + shard = rank == (0 if same_offset else 1) + placement = Shard(0) if shard else Replicate() + tensor = DTensor.from_local( + torch.ones(2 if shard else 4), mesh, [placement], shape=torch.Size([4]), stride=(1,), run_check=False + ) + with pytest.raises(RuntimeError, match="overlapping nonidentical shards"): + verify_replicated_model_shards({"expert": tensor}) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("same_offset", [False, True]) +def test_checkpoint_overlapping_shard_geometries_fail_all_ranks(tmp_path, same_offset): + mp.spawn(_overlapping_tensor_shards_worker, args=(str(tmp_path), same_offset), nprocs=2, join=True) + + +def test_checkpoint_rejects_empty_shard_with_same_offset_as_nonempty_shard(): + from veomni.checkpoint.rank_local_optimizer import _reject_overlapping_geometries + + with pytest.raises(RuntimeError, match="overlapping nonidentical shards"): + _reject_overlapping_geometries("expert", [((0,), (0,)), ((0,), (2,))]) + + +def _owner_trainability_worker(rank, root, within_owner_mismatch): + from veomni.checkpoint.optimizer_resume_layout import coordinated_multi_layout + + dist.init_process_group( + "gloo", + init_method=f"file://{root}/owner-trainability-init", + rank=rank, + world_size=4, + timeout=timedelta(seconds=30), + ) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ple_fsdp", "ple")) + model = torch.nn.Module() + values = torch.arange(64, dtype=torch.float32).reshape(8, 8) + model.weight = torch.nn.Parameter(distribute_tensor(values, mesh, [Shard(1), Shard(0)])) + model._fqn2spec_info = {"weight": SpecInfo("ple", Shard(0), "weight", mesh, 1, 1, (8, 8))} + optimizer = MultiOptimizer(model, {"leaf": torch.optim.AdamW([model.weight])}, ["leaf"]) + owner = mesh.get_local_rank("ple") + replica = mesh.get_local_rank("ple_fsdp") + if owner == 1 and (not within_owner_mismatch or replica == 1): + model.weight.requires_grad_(False) + if within_owner_mismatch: + with pytest.raises(RuntimeError, match="trainability differs"): + coordinated_multi_layout(model, optimizer, writing=True) + else: + coordinated_multi_layout(model, optimizer, writing=True) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 4 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("within_owner_mismatch", [False, True]) +def test_trainability_may_differ_between_owners_but_not_replicas(tmp_path, within_owner_mismatch): + mp.spawn(_owner_trainability_worker, args=(str(tmp_path), within_owner_mismatch), nprocs=4, join=True) + + +def _frozen_resume_worker(rank, root, phase): + dist.init_process_group( + "gloo", + init_method=f"file://{root}/frozen-{phase}-init", + rank=rank, + world_size=4, + timeout=timedelta(seconds=40), + ) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ep_fsdp", "ep")) + fsdp, owner = mesh["ep_fsdp"], mesh.get_local_rank("ep") + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + model = torch.nn.Module() + values = torch.arange(16, dtype=torch.float32).reshape(4, 4) + 100 * owner + model.register_parameter("weight", torch.nn.Parameter(distribute_tensor(values, fsdp, [Shard(1)]))) + model.register_parameter("bias", torch.nn.Parameter(torch.ones(2))) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 1)} + + def build(): + leaf = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=0.01) + return MultiOptimizer(model, {"expert": leaf}, ["expert"]) + + if phase == "save": + optimizer = build() + model.weight.grad = torch.ones_like(model.weight) + model.bias.grad = torch.full_like(model.bias, 0.1) + optimizer.step() + optimizer.zero_grad() + model.weight.requires_grad_(False) + DistributedCheckpointer.save( + f"{root}/frozen", {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + model.weight.requires_grad_(True) + model.weight.grad = torch.full_like(model.weight, 0.2) + optimizer.step() + torch.save(model.weight.to_local().detach().clone(), f"{root}/frozen-ref-{rank}.pt") + else: + model.weight.requires_grad_(False) + with recover_frozen_optimizer_ownership(model, f"{root}/frozen"): + assert model.weight.requires_grad + optimizer = build() + assert not model.weight.requires_grad + DistributedCheckpointer.load( + f"{root}/frozen", {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + model.weight.requires_grad_(True) + model.weight.grad = torch.full_like(model.weight, 0.2) + optimizer.step() + expected = torch.load(f"{root}/frozen-ref-{rank}.pt", weights_only=True) + torch.testing.assert_close(model.weight.to_local(), expected, rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +def test_frozen_owned_parameter_rebuilt_on_fresh_process(tmp_path): + mp.spawn(_frozen_resume_worker, args=(str(tmp_path), "save"), nprocs=4, join=True) + mp.spawn(_frozen_resume_worker, args=(str(tmp_path), "load"), nprocs=4, join=True) + + +def _missing_rank_payload_worker(rank, root): + from pathlib import Path + + from veomni.checkpoint.layout import optimizer_dir + from veomni.checkpoint.rank_local_optimizer import rank_scalar_key + + dist.init_process_group( + "gloo", init_method=f"file://{root}/corrupt-init", rank=rank, world_size=2, timeout=timedelta(seconds=25) + ) + try: + mesh = init_device_mesh("cpu", (1, 2), mesh_dim_names=("ep_fsdp", "ep")) + owner = mesh.get_local_rank("ep") + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + + def build(): + model = torch.nn.Module() + tensor = torch.full((2, 2), float(owner)) + model.register_parameter( + "weight", torch.nn.Parameter(distribute_tensor(tensor, mesh["ep_fsdp"], [Shard(1)])) + ) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 1)} + leaf = torch.optim.AdamW([model.weight], lr=0.01) + return model, MultiOptimizer(model, {"expert": leaf}, ["expert"]) + + source, optimizer = build() + source.weight.grad = torch.ones_like(source.weight) + optimizer.step() + optimizer.zero_grad() + checkpoint = f"{root}/corrupt" + DistributedCheckpointer.save(checkpoint, {"model": source, "optimizer": optimizer}, parallel_state=parallel) + if rank == 0: + opt_path = optimizer_dir(checkpoint, "") + metadata = dcp.FileSystemReader(opt_path).read_metadata() + storage = next( + item for index, item in metadata.storage_data.items() if index.fqn == f"optimizer.{rank_scalar_key(0)}" + ) + Path(opt_path, storage.relative_path).unlink() + dist.barrier() + target, restored = build() + before = target.weight.to_local().detach().clone() + from torch.distributed.checkpoint.api import CheckpointException + + with pytest.raises((CheckpointException, RuntimeError)): + DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight.to_local(), before, rtol=0, atol=0) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +def test_one_rank_missing_optimizer_payload_fails_all_before_model_copy(tmp_path): + mp.spawn(_missing_rank_payload_worker, args=(str(tmp_path),), nprocs=2, join=True) diff --git a/tests/checkpoints/test_rank_local_ple.py b/tests/checkpoints/test_rank_local_ple.py new file mode 100644 index 0000000000..05d2f45aa4 --- /dev/null +++ b/tests/checkpoints/test_rank_local_ple.py @@ -0,0 +1,72 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +from datetime import timedelta +from types import SimpleNamespace + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import Shard, distribute_tensor + +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer +from veomni.distributed.parallel_plan import SpecInfo +from veomni.optim.optimizer import MultiOptimizer + + +def _worker(rank, root, phase): + dist.init_process_group( + "gloo", init_method=f"file://{root}/ple-{phase}-init", rank=rank, world_size=4, timeout=timedelta(seconds=40) + ) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ple_fsdp", "ple")) + owner = mesh.get_local_rank("ple") + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ple",), extra_parallel_fsdp_device_mesh={"ple": mesh} + ) + model = torch.nn.Module() + initial = torch.arange(64, dtype=torch.bfloat16).reshape(8, 8) + model.register_parameter("weight", torch.nn.Parameter(distribute_tensor(initial, mesh, [Shard(1), Shard(0)]))) + model._fqn2spec_info = {"weight": SpecInfo("ple", Shard(0), "weight", mesh, 1, 1)} + leaf = torch.optim.AdamW([model.weight], lr=0.01) + optimizer = MultiOptimizer(model, {"ple": leaf}, ["ple"]) + + def snapshot(): + state = leaf.state.get(model.weight, {}) + return { + "weight": model.weight.to_local().detach().clone(), + "state": { + key: (value.to_local() if hasattr(value, "to_local") else value).detach().clone() + for key, value in state.items() + }, + } + + if phase == "save": + if owner == 0: + model.weight.grad = torch.full_like(model.weight, 0.1) + optimizer.step() + optimizer.zero_grad() + DistributedCheckpointer.save( + f"{root}/checkpoint", {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + before = snapshot() + model.weight.grad = torch.full_like(model.weight, 0.2) + optimizer.step() + torch.save({"before": before, "after": snapshot()}, f"{root}/ple-ref-{rank}.pt") + else: + DistributedCheckpointer.load( + f"{root}/checkpoint", {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + reference = torch.load(f"{root}/ple-ref-{rank}.pt", weights_only=True) + torch.testing.assert_close(snapshot(), reference["before"], rtol=0, atol=0) + model.weight.grad = torch.full_like(model.weight, 0.2) + optimizer.step() + torch.testing.assert_close(snapshot(), reference["after"], rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +def test_persistent_ple_owner_local_optimizer_fresh_process(tmp_path): + mp.spawn(_worker, args=(str(tmp_path), "save"), nprocs=4, join=True) + mp.spawn(_worker, args=(str(tmp_path), "load"), nprocs=4, join=True) diff --git a/tests/optim/test_multi_optimizer_sparse_lr.py b/tests/optim/test_multi_optimizer_sparse_lr.py index 669d0cf66f..a2cd9195a5 100644 --- a/tests/optim/test_multi_optimizer_sparse_lr.py +++ b/tests/optim/test_multi_optimizer_sparse_lr.py @@ -166,7 +166,7 @@ def test_sparse_leaf_learning_rate_survives_dcp_roundtrip(tmp_path): assert restored_bias.param_groups[0]["lr"] == 0.0025 -def test_sparse_leaf_post_hook_live_reorder_survives_dcp_roundtrip(tmp_path): +def test_sparse_leaf_post_hook_live_reorder_is_rejected_before_dcp_save(tmp_path): def build(): model = torch.nn.Linear(2, 2) leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) @@ -175,13 +175,11 @@ def build(): model, leaf, optimizer = build() leaf.register_state_dict_post_hook(lambda opt, state: (opt.param_groups.reverse(), state)[1]) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - restored_model, restored_leaf, restored = build() - DistributedCheckpointer.load( - str(tmp_path), {"model": restored_model, "optimizer": restored}, parallel_state=parallel - ) - assert restored_leaf.param_groups[0]["lr"] == 0.01 - assert restored_leaf.param_groups[1]["lr"] == 0.02 + with pytest.raises(RuntimeError, match="cannot prove custom optimizer serialization hooks"): + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + assert leaf.param_groups[0]["params"][0] is model.weight + assert leaf.param_groups[1]["params"][0] is model.bias + assert not list(tmp_path.iterdir()) def _fsdp2_empty_leaf_worker(rank, checkpoint, endpoint): @@ -207,20 +205,32 @@ def build(seed): second = torch.optim.AdamW(model[1].parameters(), lr=0.01) return model, first, second, MultiOptimizer(model, {"first": first, "second": second}, ["first", "second"]) + from veomni.optim.lr_scheduler import build_lr_scheduler + model, first, second, optimizer = build(42) - scheduler = torch.optim.lr_scheduler.StepLR(second, step_size=1, gamma=0.25) + scheduler = build_lr_scheduler(optimizer, train_steps=4, lr=0.01, lr_decay_style="linear", lr_warmup_ratio=0) model(torch.ones(2, 4)).sum().backward() for param in model[1].parameters(): param.grad = None first.step() first.zero_grad() - second.step() - scheduler.step() + for _ in range(3): + second.step() + scheduler.step() assert second.param_groups[0]["lr"] == 0.0025 assert first.state assert not second.state source_shards = [param.to_local().detach().clone() for param in model.parameters()] - DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + checkpoint, + { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": scheduler.state_dict()}, + "_veomni_scheduler_object": scheduler, + }, + parallel_state=parallel, + ) assert not second.state for expected, param in zip(source_shards, model.parameters(), strict=True): torch.testing.assert_close(expected, param.to_local(), rtol=0, atol=0) @@ -230,9 +240,17 @@ def build(seed): not torch.equal(saved, param.to_local()) for saved, param in zip(source_shards, restored_model.parameters(), strict=True) ) - DistributedCheckpointer.load( - checkpoint, {"model": restored_model, "optimizer": restored}, parallel_state=parallel + restored_scheduler = build_lr_scheduler( + restored, train_steps=4, lr=0.01, lr_decay_style="linear", lr_warmup_ratio=0 ) + recovered = { + "model": restored_model, + "optimizer": restored, + "extra_state": {}, + "_veomni_scheduler_object": restored_scheduler, + } + DistributedCheckpointer.load(checkpoint, recovered, parallel_state=parallel) + restored_scheduler.load_state_dict(recovered["extra_state"]["lr_scheduler"]) assert restored_second.param_groups[0]["lr"] == 0.0025 for expected, param in zip(source_shards, restored_model.parameters(), strict=True): torch.testing.assert_close(expected, param.to_local(), rtol=0, atol=0) diff --git a/veomni/arguments/arguments_types.py b/veomni/arguments/arguments_types.py index 982ca75f4c..489cbc38bd 100644 --- a/veomni/arguments/arguments_types.py +++ b/veomni/arguments/arguments_types.py @@ -248,6 +248,12 @@ class OptimizerConfig: }, ) + def __post_init__(self): + # YAML sequences arrive as lists; the declared tuple also prevents + # optimizer defaults and groups from sharing mutable beta options. + if isinstance(self.betas, list): + self.betas = tuple(self.betas) + @dataclass class WandbConfig: @@ -932,6 +938,13 @@ def _validate_init_device(self): class CheckpointConfig: """train.checkpoint.* — Checkpoint saving and loading.""" + extra_parallel_canonical_layout: bool = field( + default=False, + metadata={ + "help": "Use the versioned canonical EP/PLE model checkpoint layout. Legacy checkpoints require migration." + }, + ) + output_dir: str = field( default="output", metadata={"help": "Path to save model checkpoints."}, diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 1f6f8deed0..4a7b7989dd 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -15,8 +15,10 @@ import gc import hashlib +import json import os import shutil +import uuid import weakref from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor @@ -53,11 +55,29 @@ OPTIMIZER_DIRNAME, WEIGHTS_DIRNAME, dcp_markers, + manifest_path, model_dir, optimizer_dir, step_dir, weights_dir, ) +from .multi_optimizer_checkpoint import ( + _MULTI_GENERATION_MARKER, + _MULTI_INVENTORY_MARKER, + _check_model_serialization_hooks, + _check_qwen_omni_checkpoint_keys, + _coordinated_optimizer_flag, + _dcp_parameter_name_map, + _inventory_digests, + _iter_optimizer_leaves, + _multi_extra_snapshot, + _optimizer_payload_digest, + _preflight_multi_optimizer_layout, + _prepare_multi_optimizer_saved_states, + _validate_source_optimizer_states, + _verify_replicated_optimizer_metadata, +) +from .optimizer_resume_layout import MULTI_LAYOUT_MARKER, coordinated_multi_layout, frozen_parameter_digest logger = logging.get_logger(__name__) @@ -70,9 +90,10 @@ class _ModelStrictLoadPlanner(DefaultLoadPlanner): """Allow partial optimizer state while requiring a complete full-model DCP.""" - def __init__(self, strict_model: bool): + def __init__(self, strict_model: bool, expected_save_id=None): super().__init__(allow_partial_load=True) self.strict_model = strict_model + self.expected_save_id = expected_save_id def create_local_plan(self): assert self.metadata is not None @@ -101,6 +122,10 @@ def create_local_plan(self): raise RuntimeError( f"ExtraParallel checkpoint dtype differs for {key}; explicit conversion required" ) + if self.expected_save_id is not None: + actual = getattr(getattr(self.metadata, "storage_meta", None), "save_id", None) + if actual != self.expected_save_id: + raise RuntimeError("Checkpoint changed between identity preflight and tensor load") plan = super().create_local_plan() if not self.strict_model: return plan @@ -206,10 +231,33 @@ def _check_canonical_spec_coverage(model, specs, state_keys, *, trainable_only): raise RuntimeError(f"ExtraParallel checkpoint spec is not present in the canonical state dict: {name}") -def _reject_unproven_extra_parallel_optimizer(state, parallel_state, *, group=None, trainable_only=False): +def _validate_optimizer_model_layout(model): + specs = getattr(model, "_fqn2spec_info", None) + parameters = dict(model.named_parameters(remove_duplicate=False)) + for name, spec in (specs or {}).items(): + if ( + isinstance(spec.placement, Shard) + and getattr(spec, "persistent_fsdp_shard_dim", None) is None + and not getattr(spec, "checkpoint_layout_v2", False) + and spec.para_fsdp_mesh is not None + and spec.para_fsdp_mesh.size(-2) > 1 + and name in parameters + and spec.placement.dim % parameters[name].ndim == spec.fsdp_shard_dim % parameters[name].ndim + ): + raise RuntimeError( + f"ExtraParallel tensor {name} requires train.checkpoint.extra_parallel_canonical_layout=true " + "when EP and FSDP shard the same axis" + ) + + +def _reject_unproven_extra_parallel_optimizer( + state, parallel_state, *, group=None, trainable_only=False, allow_proven_optimizer=False +): error, signature = None, None try: specs = getattr(state.get("model"), "_fqn2spec_info", None) + if allow_proven_optimizer: + _validate_optimizer_model_layout(state["model"]) if _needs_extra_parallel_layout_marker(specs): _extra_parallel_layout_key(specs) signature = _runtime_extra_parallel_layout_signature(specs) @@ -218,7 +266,7 @@ def _reject_unproven_extra_parallel_optimizer(state, parallel_state, *, group=No else: ps = None if state.get("optimizer") is not None and getattr(ps, "dp_mode", None) == "fsdp2": - if _needs_extra_parallel_layout_marker(specs): + if _needs_extra_parallel_layout_marker(specs) and not allow_proven_optimizer: raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") except BaseException as exc: error = exc @@ -645,6 +693,17 @@ def get_state_dict_with_extra_parallel_dim_preprocess(self, state_dict, action): ) +def _requires_rank_local_optimizer_checkpoint(optimizer): + # A registered adapter declares an exact, rank-local state contract even + # when the underlying optimizer does not expose a marker attribute. + from ..optim.checkpoint_adapter import rank_local_adamw_adapter + + return ( + getattr(optimizer, "requires_exact_state_dtypes", False) is True + or rank_local_adamw_adapter(optimizer) is not None + ) + + class OptimizerState(Stateful): """A wrapper around an optimizer to make it stateful. @@ -694,12 +753,19 @@ def __init__(self, model, optimizer, parallel_state=None, *, load: bool = False) self._load = load def state_dict(self): + if _requires_rank_local_optimizer_checkpoint(self.optimizer): + raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) ): raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") if self.should_extra_parallel_aware and self._layout_marker_key is not None: raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") + if getattr(self.optimizer, "_is_multi_optimizer", False): + raise RuntimeError( + "Raw DCP MultiOptimizer does not prove parameter groups or scheduler identity; " + "use DistributedCheckpointer with the live scheduler." + ) if self.should_extra_parallel_aware: logger.info_rank0( "Getting optimizer state_dict from OptimizerState wrapper, would restore ExtraParallel dim for Experts module" @@ -726,6 +792,8 @@ def state_dict(self): def load_state_dict(self, state_dict): """Restore optimizer state from the loaded DCP state dict.""" + if _requires_rank_local_optimizer_checkpoint(self.optimizer): + raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") if self.should_extra_parallel_aware and self._layout_marker_key is not None: raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( @@ -776,6 +844,203 @@ def get_state_dict_with_extra_parallel_dim_preprocess(self, state_dict, action): ) +class _ProvenModelState(ModelState): + def __init__(self, *args, optimizer, group, writing=False, **kwargs): + super().__init__(*args, **kwargs) + self.optimizer = optimizer + self.group = group + self.writing = writing + + def state_dict(self): + result, error = None, None + try: + _check_model_serialization_hooks(self.model) + result = super().state_dict() + _check_model_serialization_hooks(self.model) + logical_keys = {key for raw in result if (key := _export_extra_parallel_layout_key(raw)) is not None} + _check_qwen_omni_checkpoint_keys(self.model, logical_keys) + if self.trainable_only: + names = _dcp_parameter_name_map(self.model) + for leaf in _iter_optimizer_leaves(self.optimizer): + for param, saved in leaf.state.items(): + aliases = [name for name, owned in names.items() if owned is param] + if saved and not any(name in logical_keys for name in aliases): + raise RuntimeError("trainable-only save would omit a parameter with optimizer history") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "constructing proven model state", group=self.group) + if self.writing: + from .rank_local_optimizer import verify_replicated_model_shards + + verify_replicated_model_shards(result, group=self.group) + return result + + def load_state_dict(self, state_dict): + error = None + try: + super().load_state_dict(state_dict) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "applying proven model state", group=self.group) + + +class _ProvenMultiOptimizerState(OptimizerState): + """High-level DCP adapter; direct optimizer APIs retain their own contract.""" + + def __init__( + self, + *args, + generation, + scheduler_snapshot, + scheduler_object, + extra_digest, + expected_layout, + group, + frozen_digest=None, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.generation = generation + self.scheduler_snapshot = scheduler_snapshot + self.scheduler_object = scheduler_object + self.extra_digest = extra_digest + self.expected_layout = expected_layout + self.group = group + self.frozen_digest = frozen_digest + + def state_dict(self): + result, error = None, None + try: + result = self.optimizer._sparse_state_dict(coordinate=True, group=self.group) + if self.should_extra_parallel_aware: + from .rank_local_optimizer import ( + current_rank, + local_optimizer_state, + rank_inventory, + rank_inventory_key, + ) + + inventory = rank_inventory(result, current_rank()) + result = local_optimizer_state(result, current_rank()) + result[rank_inventory_key(current_rank())] = "" if self._load else json.dumps(inventory) + else: + from torch.distributed.checkpoint._nested_dict import flatten_state_dict + + flattened, _ = flatten_state_dict({"optimizer": result}) + inventory = sorted(flattened) + result[_MULTI_INVENTORY_MARKER] = "" if self._load else json.dumps(inventory) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "constructing proven MultiOptimizer state", group=self.group) + if self._load: + result[MULTI_LAYOUT_MARKER] = "" + return result + initial = coordinated_multi_layout( + self.model, + self.optimizer, + self.scheduler_snapshot, + self.scheduler_object, + writing=True, + group=self.group, + frozen_digest=self.frozen_digest, + ) + if initial != self.expected_layout: + raise RuntimeError("MultiOptimizer layout changed while saving checkpoint") + current = coordinated_multi_layout( + self.model, + self.optimizer, + self.scheduler_snapshot, + self.scheduler_object, + self.extra_digest, + writing=True, + group=self.group, + frozen_digest=self.frozen_digest, + ) + result[MULTI_LAYOUT_MARKER] = json.dumps( + { + "generation": self.generation, + "layout": current, + "inventory_digests": _inventory_digests(inventory, group=self.group), + }, + sort_keys=True, + ) + return result + + def load_state_dict(self, state_dict): + error = None + try: + state_dict = dict(state_dict) + marker = json.loads(state_dict.pop(MULTI_LAYOUT_MARKER)) + if ( + not isinstance(marker, dict) + or set(marker) != {"generation", "layout", "inventory_digests"} + or marker["generation"] != self.generation + or marker["layout"] != self.expected_layout + ): + raise RuntimeError("MultiOptimizer identity changed during tensor loading") + if self.should_extra_parallel_aware: + from .rank_local_optimizer import ( + current_rank, + rank_inventory, + rank_inventory_key, + restore_optimizer_state, + ) + + inventory = json.loads(state_dict.pop(rank_inventory_key(current_rank()))) + template = self.optimizer._sparse_state_dict() + if inventory != rank_inventory(template, current_rank()): + raise RuntimeError("rank-local optimizer checkpoint inventory changed during loading") + state_dict = restore_optimizer_state(state_dict, template, current_rank()) + else: + from torch.distributed.checkpoint._nested_dict import flatten_state_dict + + inventory = json.loads(state_dict.pop(_MULTI_INVENTORY_MARKER)) + flattened, _ = flatten_state_dict({"optimizer": state_dict}) + if inventory != sorted(flattened): + raise RuntimeError("MultiOptimizer checkpoint inventory changed during loading") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "validating loaded MultiOptimizer state", group=self.group) + if marker["inventory_digests"] != _inventory_digests(inventory, group=self.group): + raise RuntimeError("MultiOptimizer inventory identity changed during tensor loading") + current = coordinated_multi_layout( + self.model, + self.optimizer, + self.scheduler_snapshot, + self.scheduler_object, + self.extra_digest, + group=self.group, + frozen_digest=self.frozen_digest, + ) + if current != self.expected_layout: + raise RuntimeError("MultiOptimizer layout changed while loading optimizer state") + from ..optim.checkpoint_adapter import rank_local_adamw_adapter + + has_adapter = _coordinated_optimizer_flag( + any(rank_local_adamw_adapter(leaf) is not None for leaf in _iter_optimizer_leaves(self.optimizer)), + "registered optimizer checkpoint adapter presence", + group=self.group, + ) + expected_payload = None + if has_adapter: + error = None + try: + expected_payload = _optimizer_payload_digest(state_dict) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "fingerprinting loaded MultiOptimizer state", group=self.group) + error = None + try: + self.optimizer._dcp_load_state_dict(state_dict, group=self.group) + if has_adapter: + restored_payload = self.optimizer._sparse_state_dict(coordinate=True, group=self.group) + if _optimizer_payload_digest(restored_payload) != expected_payload: + raise RuntimeError("MultiOptimizer load changed or omitted checkpoint state") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "applying proven MultiOptimizer state", group=self.group) + + def drop_extra_parallel_dim(loaded_tensor: torch.Tensor, device_mesh: DeviceMesh, fsdp_shard_dim: int = 1): """ Drop ExtraParallel dims after loading from DCP so that ExtraParallel-FSDP would not be confused. @@ -1135,8 +1400,11 @@ def _prepare_stage_dir(stage_dir: str, path: str, group: Optional[Any] = None) - error: Optional[BaseException] = None if _local_rank() == 0: try: - shutil.rmtree(stage_path, ignore_errors=True) - os.makedirs(stage_path, exist_ok=True) + if os.path.lexists(stage_path): + shutil.rmtree(stage_path) + # Failing rather than reusing an uncleared directory prevents a + # previous checkpoint's sidecars from joining this transaction. + os.makedirs(stage_path, exist_ok=False) except BaseException as e: # noqa: BLE001 - raised once every rank has agreed # ``BaseException`` for the reason ``_promotion_phase`` documents: # anything that escapes here skips the reduction below, and the ranks @@ -1412,20 +1680,105 @@ def save( return: None """ - if "model" not in state: + has_model = _coordinated_optimizer_flag( + isinstance(state, dict) and "model" in state, "checkpoint model presence" + ) + if not has_model: raise ValueError("Model must be provided to save a distributed checkpoint.") - _reject_unproven_extra_parallel_optimizer(state, parallel_state, trainable_only=trainable_only) - - # Rejected up front, before anything reaches disk: an async write is still - # running when save() returns and drops the staged copy, so it would write - # straight to the slow destination staging was meant to avoid. + optimizer = state.get("optimizer") + multi = _coordinated_optimizer_flag( + getattr(optimizer, "_is_multi_optimizer", False) is True, "MultiOptimizer save policy" + ) + if multi: + _coordinated_optimizer_flag(tuple(optimizer.key_names), "optimizer leaf order") + trainable_only = _coordinated_optimizer_flag(trainable_only, "MultiOptimizer trainable-only save") + _coordinated_optimizer_flag(bool(save_async), "MultiOptimizer asynchronous save mode") + _coordinated_optimizer_flag(os.path.abspath(stage_dir) if stage_dir else None, "MultiOptimizer stage path") + _coordinated_optimizer_flag(module, "MultiOptimizer module") + _coordinated_optimizer_flag("extra_state" in state, "MultiOptimizer extra-state presence") + # An async write may outlive a staged copy, so the modes cannot combine. if stage_dir and save_async: raise ValueError("stage_dir cannot be combined with save_async") + single_exact = _coordinated_optimizer_flag( + not multi and _requires_rank_local_optimizer_checkpoint(optimizer), + "exact-precision optimizer save policy", + ) + if single_exact: + raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") + _reject_unproven_extra_parallel_optimizer( + state, parallel_state, trainable_only=trainable_only, allow_proven_optimizer=multi + ) + scheduler_object = state.get("_veomni_scheduler_object") if multi else None + source_model = state.get("model") + extra = state.get("extra_state") if multi else None + snapshot = extra.get("lr_scheduler") if isinstance(extra, dict) else None + initial_layout, generation = None, None + if multi: + error = None + try: + parallel_state = parallel_state if parallel_state is not None else get_parallel_state() + if save_async: + raise RuntimeError("Proven MultiOptimizer checkpoints currently require synchronous DCP save") + _check_model_serialization_hooks(source_model) + _validate_source_optimizer_states(optimizer, source_model) + if "extra_state" in state and not isinstance(extra, dict): + raise RuntimeError("MultiOptimizer extra_state must be a mapping") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "validating MultiOptimizer checkpoint policy") + frozen_digest = frozen_parameter_digest(state["model"], enabled=trainable_only) + initial_layout = coordinated_multi_layout( + state["model"], optimizer, snapshot, scheduler_object, writing=True, frozen_digest=frozen_digest + ) + rank_local = _coordinated_optimizer_flag( + getattr(state["model"], "_fqn2spec_info", None) is not None + and getattr(parallel_state, "dp_mode", None) == "fsdp2", + "rank-local MultiOptimizer checkpoint mode", + ) + if not rank_local: + _verify_replicated_optimizer_metadata(state["model"], optimizer) + from .rank_local_optimizer import verify_replicated_model_shards + + optimizer_snapshot, error = None, None + try: + optimizer_snapshot = optimizer._sparse_state_dict(coordinate=True) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "constructing replicated optimizer snapshot") + verify_replicated_model_shards(optimizer_snapshot) + error, generation_value = None, None + if not dist.is_initialized() or dist.get_rank() == 0: + try: + generation_value = str(uuid.uuid4()) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "generating MultiOptimizer checkpoint identity") + values = [generation_value] + if dist.is_initialized(): + dist.broadcast_object_list(values, src=0) + generation = values[0] # ``is not None`` rather than truthiness: step 0 is a step like any other, and # folding it onto ``path`` would write it over the run's own directory. checkpoint_dir = step_dir(path, global_steps) if global_steps is not None else path model_root = model_dir(checkpoint_dir, module) + if multi: + _coordinated_optimizer_flag(os.path.abspath(checkpoint_dir), "MultiOptimizer checkpoint path") + error = ( + RuntimeError("Cannot overwrite a completed checkpoint transaction with MultiOptimizer state") + if os.path.exists(manifest_path(checkpoint_dir)) + else None + ) + raise_if_any_rank_failed(error, "checking MultiOptimizer checkpoint destination") + if multi and extra is None: + rank = dist.get_rank() if dist.is_initialized() else 0 + extra_path = os.path.join(model_root, _EXTRA_STATE_DIRNAME, _EXTRA_STATE_FORMAT.format(rank)) + error = ( + RuntimeError("Existing checkpoint sidecar cannot be reused when extra_state is omitted") + if os.path.exists(extra_path) + else None + ) + raise_if_any_rank_failed(error, "checking for stale MultiOptimizer extra state") # Keyed on the module's own run-level path, not just ``path``: a # multi-module job calls this once per module, and a single key would have @@ -1438,50 +1791,108 @@ def save( stage_path = _prepare_stage_dir(stage_dir, stage_key_path, group=stage_group) if stage_dir else None write_root = stage_path or model_root - if stage_path is None: - # Nothing stands between this save and the previous checkpoint's - # shards, so its markers have to go before the first byte lands -- - # otherwise a save that dies part-way leaves one describing shards - # that are half this step and half the last. A staged save invalidates - # its destination in ``_promote_staged_checkpoint`` instead, at the - # point where it is about to overwrite it for real. - # - # The manifest is not ours either way; ``GlobalStateCallback`` clears - # the one it writes. - cls._remove_dcp_markers(checkpoint_dir, module) - - cls._create_checkpoint_dir(write_root) - - # Sidecar first, then the DCP directories. Ordering no longer carries the - # completeness guarantee it used to — ``checkpoint_manifest.json`` does, - # 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_extra_state(checkpoint_dir=write_root, state=state) + error = None + try: + if stage_path is None: + # Invalidate stale metadata before writing any new shards. + cls._remove_dcp_markers(checkpoint_dir, module) + cls._create_checkpoint_dir(write_root) + if multi: + os.makedirs(os.path.join(write_root, OPTIMIZER_DIRNAME), exist_ok=True) + cls._save_extra_state(checkpoint_dir=write_root, state=state) + except BaseException as exc: + error = exc + if multi: + raise_if_any_rank_failed(error, "preparing MultiOptimizer checkpoint") + elif error is not None: + raise error + extra_digest = None + if multi: + error = None + try: + extra_digest, _ = _multi_extra_snapshot(write_root, read_state=False) + if (extra is None) != (extra_digest == "absent"): + raise RuntimeError("MultiOptimizer checkpoint sidecar presence does not match this save") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "hashing MultiOptimizer extra state") + + def prepare_writer_state(factory, what): + result, error = None, None + try: + result = factory() + except BaseException as exc: + error = exc + if multi: + raise_if_any_rank_failed(error, what) + elif error is not None: + raise error + return result try: + model_state, model_writer = prepare_writer_state( + lambda: ( + _ProvenModelState( + state["model"], + trainable_only=trainable_only, + parallel_state=parallel_state, + optimizer=optimizer, + group=None, + writing=True, + ) + if multi + else ModelState(state["model"], trainable_only=trainable_only, parallel_state=parallel_state), + cls._create_storage_writer(os.path.join(write_root, WEIGHTS_DIRNAME)), + ), + "preparing model DCP writer", + ) + model_payload = {"model": model_state} + if multi: + model_payload[_MULTI_GENERATION_MARKER] = generation cls.execute_save( - save_state={ - "model": ModelState(state["model"], trainable_only=trainable_only, parallel_state=parallel_state) - }, - storage_writer=cls._create_storage_writer(os.path.join(write_root, WEIGHTS_DIRNAME)), + save_state=model_payload, + storage_writer=model_writer, save_async=save_async, save_to_lowest_rank=save_to_lowest_rank, slot=WEIGHTS_DIRNAME, timeout_seconds=save_timeout_seconds, ) + if multi: + error = None + if state.get("model") is not source_model or state.get("optimizer") is not optimizer: + error = RuntimeError("Model or optimizer object changed while saving model state") + elif ( + state.get("extra_state") is not extra + or state.get("_veomni_scheduler_object") is not scheduler_object + ): + error = RuntimeError("Scheduler or extra-state object changed while saving model state") + raise_if_any_rank_failed(error, "rechecking MultiOptimizer checkpoint ownership") + if "optimizer" in state and state["optimizer"] is not None: - cls.execute_save( - save_state={ - "optimizer": OptimizerState( + optim_state, optim_writer = prepare_writer_state( + lambda: ( + _ProvenMultiOptimizerState( model=state["model"], - optimizer=state["optimizer"], + optimizer=optimizer, parallel_state=parallel_state, - load=False, + generation=generation, + scheduler_snapshot=snapshot, + scheduler_object=scheduler_object, + extra_digest=extra_digest, + expected_layout=initial_layout, + group=None, + frozen_digest=frozen_digest, ) - }, - storage_writer=cls._create_storage_writer(os.path.join(write_root, OPTIMIZER_DIRNAME)), + if multi + else OptimizerState(model=state["model"], optimizer=optimizer, parallel_state=parallel_state), + cls._create_storage_writer(os.path.join(write_root, OPTIMIZER_DIRNAME)), + ), + "preparing optimizer DCP writer", + ) + cls.execute_save( + save_state={"optimizer": optim_state}, + storage_writer=optim_writer, save_async=save_async, save_to_lowest_rank=save_to_lowest_rank, slot=OPTIMIZER_DIRNAME, @@ -1496,6 +1907,7 @@ def save( _promote_staged_checkpoint(stage_path, model_root, step_root=checkpoint_dir, group=stage_group) logger.info_rank0(f"Saved checkpoint to {model_root}") + return generation if multi else None @classmethod def _remove_dcp_markers(cls, checkpoint_dir: str, module: str) -> None: @@ -1568,13 +1980,38 @@ def load( return: state: state loaded """ - if state is None: - raise ValueError("State dict must be provided to load a distributed checkpoint.") - - if "model" not in state: + if isinstance(state, dict): + state.pop("_veomni_optimizer_identity_verified", None) + state.pop("_veomni_optimizer_checkpoint_generation", None) + has_model = _coordinated_optimizer_flag( + isinstance(state, dict) and "model" in state, "checkpoint model presence", group=process_group + ) + if not has_model: + if state is None: + raise ValueError("State dict must be provided to load a distributed checkpoint.") raise ValueError("Model must be provided to load a distributed checkpoint.") + optimizer = state.get("optimizer") + multi = _coordinated_optimizer_flag( + getattr(optimizer, "_is_multi_optimizer", False) is True, "MultiOptimizer load policy", group=process_group + ) + if multi: + _coordinated_optimizer_flag(tuple(optimizer.key_names), "optimizer leaf order", group=process_group) + trainable_only = _coordinated_optimizer_flag( + trainable_only, "MultiOptimizer trainable-only load", group=process_group + ) + _coordinated_optimizer_flag(module, "MultiOptimizer module", group=process_group) + _coordinated_optimizer_flag( + "extra_state" in state, "MultiOptimizer extra-state presence", group=process_group + ) + single_exact = _coordinated_optimizer_flag( + not multi and _requires_rank_local_optimizer_checkpoint(optimizer), + "exact-precision optimizer load policy", + group=process_group, + ) + if single_exact: + raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") _reject_unproven_extra_parallel_optimizer( - state, parallel_state, group=process_group, trainable_only=trainable_only + state, parallel_state, group=process_group, trainable_only=trainable_only, allow_proven_optimizer=multi ) def _model_state() -> "ModelState": @@ -1590,7 +2027,12 @@ def _optimizer_state() -> "OptimizerState": wants_optimizer = state.get("optimizer") is not None - fused_dir = cls._legacy_fused_dir(path, module) + fused_dir, error = None, None + try: + fused_dir = cls._legacy_fused_dir(path, module) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "inspecting checkpoint layout", group=process_group) specs = getattr(state["model"], "_fqn2spec_info", None) if _needs_extra_parallel_layout_marker(specs): marker_error = None @@ -1600,6 +2042,122 @@ def _optimizer_state() -> "OptimizerState": except BaseException as exc: marker_error = exc raise_if_any_rank_failed(marker_error, "validating ExtraParallel checkpoint marker", group=process_group) + if multi: + _coordinated_optimizer_flag(os.path.abspath(path), "MultiOptimizer checkpoint path", group=process_group) + error = None + try: + parallel_state = parallel_state if parallel_state is not None else get_parallel_state() + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "validating MultiOptimizer checkpoint policy", group=process_group) + _coordinated_optimizer_flag( + getattr(state["model"], "_fqn2spec_info", None) is not None + and getattr(parallel_state, "dp_mode", None) == "fsdp2", + "rank-local MultiOptimizer checkpoint mode", + group=process_group, + ) + error = ( + RuntimeError("Fused MultiOptimizer checkpoints require an explicitly verified migration") + if fused_dir is not None + else None + ) + raise_if_any_rank_failed(error, "checking MultiOptimizer checkpoint layout", group=process_group) + scheduler_object = state.get("_veomni_scheduler_object") + optimizer_path = optimizer_dir(path, module) + model_path = weights_dir(path, module) + root = model_dir(path, module) + optimizer_id, extra, model_id, generation, expected_layout, extra_digest, frozen_digest = ( + _preflight_multi_optimizer_layout( + state["model"], + optimizer, + cls._create_storage_reader, + optimizer_path, + model_path, + root, + "extra_state" in state, + scheduler_object, + parallel_state=parallel_state, + trainable_only=trainable_only, + group=process_group, + ) + ) + _prepare_multi_optimizer_saved_states( + state["model"], + optimizer, + cls._create_storage_reader, + optimizer_path, + optimizer_id, + group=process_group, + parallel_state=parallel_state, + validate_only=True, + ) + model_state = _ProvenModelState( + state["model"], + trainable_only=trainable_only, + parallel_state=parallel_state, + optimizer=optimizer, + group=process_group, + ) + dcp.load( + state_dict={"model": model_state}, + storage_reader=cls._create_storage_reader(model_path), + process_group=process_group, + planner=_ModelStrictLoadPlanner( + strict_model=not trainable_only, + expected_save_id=model_id, + ), + ) + error = ( + RuntimeError("Optimizer object changed while loading model state") + if state.get("optimizer") is not optimizer + else None + ) + raise_if_any_rank_failed(error, "rechecking MultiOptimizer identity", group=process_group) + snapshot = extra.get("lr_scheduler") if extra is not None else None + current = coordinated_multi_layout( + state["model"], + optimizer, + snapshot, + scheduler_object, + extra_digest, + group=process_group, + frozen_digest=frozen_digest, + ) + if current != expected_layout: + raise RuntimeError("MultiOptimizer layout changed while loading model state") + _prepare_multi_optimizer_saved_states( + state["model"], + optimizer, + cls._create_storage_reader, + optimizer_path, + optimizer_id, + group=process_group, + parallel_state=parallel_state, + ) + optimizer_state = _ProvenMultiOptimizerState( + model=state["model"], + optimizer=optimizer, + parallel_state=parallel_state, + load=True, + generation=generation, + scheduler_snapshot=snapshot, + scheduler_object=scheduler_object, + extra_digest=extra_digest, + expected_layout=expected_layout, + group=process_group, + frozen_digest=frozen_digest, + ) + dcp.load( + state_dict={"optimizer": optimizer_state}, + storage_reader=cls._create_storage_reader(optimizer_path), + process_group=process_group, + planner=_ModelStrictLoadPlanner(strict_model=False, expected_save_id=optimizer_id), + ) + if extra is not None: + state["extra_state"].update(extra) + state["_veomni_optimizer_identity_verified"] = True + state["_veomni_optimizer_checkpoint_generation"] = generation + return state if fused_dir is not None: load_state: Dict[str, Any] = {"model": _model_state()} if wants_optimizer: @@ -1791,9 +2349,14 @@ def execute_save( ) if dist.is_initialized(): dist.barrier() - gc.collect() - empty_cache() - synchronize() + error = None + try: + gc.collect() + empty_cache() + synchronize() + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "finalizing DCP save") # Private helper methods @classmethod diff --git a/veomni/checkpoint/multi_optimizer_checkpoint.py b/veomni/checkpoint/multi_optimizer_checkpoint.py new file mode 100644 index 0000000000..b0713af172 --- /dev/null +++ b/veomni/checkpoint/multi_optimizer_checkpoint.py @@ -0,0 +1,877 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +"""Versioned identity and lazy-state restoration for MultiOptimizer checkpoints.""" + +import hashlib +import json +import os +import struct +import tempfile +from collections import defaultdict + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed._tensor import DTensor +from torch.distributed.checkpoint import DefaultLoadPlanner +from torch.distributed.checkpoint._nested_dict import flatten_state_dict +from torch.distributed.checkpoint.metadata import BytesStorageMetadata, TensorStorageMetadata +from torch.distributed.checkpoint.state_dict import StateDictOptions, _get_fqns, get_model_state_dict + +from ..utils.dist_utils import raise_if_any_rank_failed +from .optimizer_resume_layout import ( + MULTI_LAYOUT_MARKER, + _group_option_value, + _lr_layout, + coordinated_multi_layout, + frozen_parameter_digest, +) + + +_EXTRA_STATE_DIRNAME = "extra_state" +_EXTRA_STATE_FORMAT = "extra_state_rank_{}.pt" +_MULTI_GENERATION_MARKER = "__veomni_multi_generation_v1__" +_MULTI_INVENTORY_MARKER = "__veomni_multi_inventory_v1__" + + +def _inventory_digests(inventory, *, group=None): + from ..utils.dist_utils import raise_if_any_rank_failed + + local, error = None, None + try: + local = hashlib.sha256(json.dumps(inventory, separators=(",", ":")).encode()).hexdigest() + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "hashing optimizer checkpoint inventory", group=group) + if not dist.is_initialized(): + return [local] + digests = [None] * dist.get_world_size(group) + dist.all_gather_object(digests, local, group=group) + return digests + + +def _coordinated_storage_reader(reader_factory, path, *, group=None): + from ..utils.dist_utils import raise_if_any_rank_failed + + reader, error = None, None + try: + reader = reader_factory(path) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "constructing optimizer inventory reader", group=group) + return reader + + +def _optimizer_payload_digest(payload, *, chunk_bytes=4 << 20): + """Fingerprint a loaded sparse payload before or after optimizer restore.""" + digest = hashlib.sha256() + + def write(blob): + digest.update(len(blob).to_bytes(8, "big")) + digest.update(blob) + + def update(value): + if torch.is_tensor(value): + write(b"tensor") + write(str(value.dtype).encode()) + write(repr(tuple(value.shape)).encode()) + if isinstance(value, DTensor): + mesh = value.device_mesh + write(repr((mesh.device_type, mesh.mesh_dim_names, mesh.mesh.tolist(), value.placements)).encode()) + else: + write(b"local") + local = value.to_local() if hasattr(value, "to_local") else value + write(str(local.device.type).encode()) + write(repr(tuple(local.shape)).encode()) + from .tensor_bytes import tensor_byte_chunks + + digest.update((local.numel() * local.element_size()).to_bytes(8, "big")) + for chunk in tensor_byte_chunks(local, chunk_bytes=chunk_bytes): + digest.update(chunk) + elif type(value) is dict: + write(b"dict") + for key in sorted(value): + update(key) + update(value[key]) + elif type(value) in (tuple, list): + write(type(value).__name__.encode()) + for item in value: + update(item) + elif type(value) is float: + write(b"float") + write(struct.pack("!d", value)) + elif value is None or type(value) in (bool, int, str, bytes, torch.dtype, torch.device): + write(type(value).__name__.encode()) + write(repr(value).encode()) + else: + raise RuntimeError(f"Cannot compare optimizer checkpoint value of type {type(value).__name__}") + + update(payload) + return digest.hexdigest() + + +def _coordinated_optimizer_flag(value, what, group=None): + if not dist.is_initialized(): + return value + gathered = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, value, group=group) + if any(item != gathered[0] for item in gathered): + raise RuntimeError(f"{what} differs across ranks") + return value + + +def _iter_optimizer_leaves(optimizer): + yield from optimizer.optimizers_dict.values() + + +def _expected_optimizer_fields(leaf, param_group): + from ..optim.checkpoint_adapter import rank_local_adamw_adapter + from ..optim.muon import DistributedMuon + from ..optim.optimizer import AnyPrecisionAdamW + + if isinstance(leaf, DistributedMuon): + return {"momentum_buffer"} + if rank_local_adamw_adapter(leaf) is not None: + return {"step", "master_weight", "exp_avg", "exp_avg_sq"} + if type(leaf) is torch.optim.AdamW: + fields = {"step", "exp_avg", "exp_avg_sq"} + if param_group.get("amsgrad"): + fields.add("max_exp_avg_sq") + return fields + if isinstance(leaf, AnyPrecisionAdamW): + fields = {"step", "exp_avg", "exp_avg_sq"} + if param_group.get("use_kahan_summation"): + fields.add("compensation") + return fields + raise RuntimeError("Unsupported MultiOptimizer leaf") + + +def _validate_source_optimizer_states(optimizer, model): + from ..optim.checkpoint_adapter import rank_local_adamw_adapter + from ..optim.optimizer import AnyPrecisionAdamW + + _check_tensor_storage_aliases(model, optimizer) + for leaf in _iter_optimizer_leaves(optimizer): + for group in leaf.param_groups: + expected = _expected_optimizer_fields(leaf, group) + for param in group["params"]: + state = leaf.state.get(param, {}) + if not state: + continue + if set(state) != expected: + raise RuntimeError("MultiOptimizer has incomplete or unsupported parameter state") + for field, value in state.items(): + if not torch.is_tensor(value): + raise RuntimeError(f"MultiOptimizer state {field} must be a tensor") + if field == "step": + if tuple(value.shape) != (): + raise RuntimeError("MultiOptimizer step must be a scalar") + adapter = rank_local_adamw_adapter(leaf) + if adapter is not None and value.dtype != adapter.state_dtype(leaf, param, field): + raise RuntimeError("MultiOptimizer step dtype differs from its registered policy") + else: + if tuple(value.shape) != tuple(param.shape): + raise RuntimeError(f"MultiOptimizer state shape differs for {field}") + adapter = rank_local_adamw_adapter(leaf) + expected_dtype = ( + adapter.state_dtype(leaf, param, field) + if adapter is not None + else group[ + { + "exp_avg": "momentum_dtype", + "exp_avg_sq": "variance_dtype", + "compensation": "compensation_buffer_dtype", + }[field] + ] + if isinstance(leaf, AnyPrecisionAdamW) + else param.dtype + ) + if (adapter is None and expected_dtype != param.dtype) or value.dtype != expected_dtype: + raise RuntimeError(f"MultiOptimizer state dtype cannot be safely restored for {field}") + + +def _qwen_omni_filtered_prefixes(model, hook): + import types + + families = ( + ( + "Qwen3OmniMoeForConditionalGeneration", + "veomni.models.transformers.qwen3_omni_moe.generated.", + "_drop_talker_and_code2wav_keys", + ("talker.", "code2wav."), + ), + ( + "Qwen2_5OmniForConditionalGeneration", + "veomni.models.transformers.qwen2_5_omni.generated.", + "_drop_talker_and_token2wav_keys", + ("talker.", "token2wav."), + ), + ) + for cls_name, cls_module, function_name, prefixes in families: + cls = next( + ( + item + for item in type(model).__mro__ + if item.__name__ == cls_name and item.__module__.startswith(cls_module) + ), + None, + ) + if cls is None: + continue + if type(model) is not cls: + from torch.distributed.fsdp import FSDPModule + + if not isinstance(model, FSDPModule): + return () + code = getattr(hook, "__code__", None) + init_code = getattr(cls.__init__, "__code__", None) + if ( + init_code is None + or getattr(hook, "__closure__", None) + or getattr(hook, "__defaults__", None) + or getattr(hook, "__kwdefaults__", None) + or getattr(hook, "__globals__", None) is not getattr(cls.__init__, "__globals__", None) + ): + return () + if ( + code is not None + and code.co_name == function_name + and any(isinstance(item, types.CodeType) and item is code for item in init_code.co_consts) + ): + return prefixes + return () + + +def _check_qwen_omni_checkpoint_keys(model, keys): + from torch.nn.modules.module import _WrappedHook + + for hook in model._load_state_dict_pre_hooks.values(): + func = hook.hook if isinstance(hook, _WrappedHook) else hook + prefixes = _qwen_omni_filtered_prefixes(model, func) + if prefixes and any(key.startswith(prefixes) for key in keys): + raise RuntimeError("Cannot prove Qwen Omni load hook is a no-op for checkpoint keys") + + +def _check_model_serialization_hooks(model): + """Allow the FSDP2 reshard hook, reject unproven model transforms.""" + import types + + from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam + from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup + from torch.nn.modules.module import _WrappedHook + + def nested_code(parent, name): + return { + item for item in parent.__code__.co_consts if isinstance(item, types.CodeType) and item.co_name == name + } + + post_codes = nested_code(FSDPParam.__init__, "") + pre_codes = nested_code(FSDPParamGroup._register_state_dict_hooks, "to_sharded_hook") + for name, module in model.named_modules(): + for attr in ( + "_state_dict_pre_hooks", + "_state_dict_hooks", + "_load_state_dict_pre_hooks", + "_load_state_dict_post_hooks", + ): + for hook in getattr(module, attr, {}).values(): + func = hook.hook if isinstance(hook, _WrappedHook) else hook + code = getattr(func, "__code__", None) + closure = getattr(func, "__closure__", None) or () + captured = tuple(cell.cell_contents for cell in closure) + builtin_fsdp = ( + attr == "_load_state_dict_post_hooks" + and code in post_codes + and any(isinstance(item, FSDPParam) for item in captured) + ) or ( + attr in ("_state_dict_pre_hooks", "_load_state_dict_pre_hooks") + and code in pre_codes + and any(isinstance(item, FSDPParamGroup) for item in captured) + ) + builtin_qwen = ( + attr == "_load_state_dict_pre_hooks" and not name and _qwen_omni_filtered_prefixes(model, func) + ) + if not builtin_fsdp and not builtin_qwen: + raise RuntimeError(f"Cannot prove model serialization hook {attr} on {name or ''}") + if builtin_qwen and any( + key.startswith(builtin_qwen) + for key, _ in list(model.named_parameters()) + list(model.named_buffers()) + ): + raise RuntimeError("Cannot prove Qwen Omni load hook is a no-op for checkpoint model tensors") + + _check_tensor_storage_aliases(model) + + +def _check_tensor_storage_aliases(model, optimizer=None): + # A single tied model Parameter under multiple FQNs is supported. Distinct + # model/state tensors must not overlap: DCP cannot reconstruct the alias. + storages, seen_model = defaultdict(list), set() + entries = [("model", name, value) for name, value in model.named_parameters(remove_duplicate=False)] + entries += [("model", name, value) for name, value in model.named_buffers(remove_duplicate=False)] + if optimizer is not None: + seen_lr = set() + for leaf_name, leaf in optimizer.optimizers_dict.items(): + for group_index, group in enumerate(leaf.param_groups): + for field in ("lr", "initial_lr"): + value = group.get(field) + if torch.is_tensor(value) and id(value) not in seen_lr: + seen_lr.add(id(value)) + entries.append(("lr", f"{leaf_name}.{group_index}.{field}", value)) + for index, state in enumerate(leaf.state.values()): + entries.extend( + ("optimizer", f"{leaf_name}.{index}.{field}", value) + for field, value in state.items() + if torch.is_tensor(value) + ) + for kind, name, value in entries: + if kind == "model" and id(value) in seen_model: + continue + if kind == "model": + seen_model.add(id(value)) + local = value.to_local() if isinstance(value, DTensor) else value + if local.numel() == 0: + continue + if local.device.type == "meta": + raise RuntimeError("Cannot prove checkpoint storage aliases for a meta tensor") + storage = (str(local.device), local.untyped_storage().data_ptr()) + first = last = local.storage_offset() + for size, stride in zip(local.shape, local.stride(), strict=True): + offset = (size - 1) * stride + first += min(0, offset) + last += max(0, offset) + storages[storage].append((first * local.element_size(), (last + 1) * local.element_size(), f"{kind}:{name}")) + for ranges in storages.values(): + ranges.sort() + for previous, current in zip(ranges, ranges[1:]): + if current[0] < previous[1]: + raise RuntimeError( + f"Cannot prove checkpoint model storage alias between {previous[2]} and {current[2]}" + ) + + +def _dcp_parameter_name_map(model) -> dict[str, torch.nn.Parameter]: + """Map DCP-canonical parameter FQNs back to their runtime parameters. + + ``get_optimizer_state_dict`` strips wrapper prefixes such as DDP's + ``module.`` when it canonicalizes state keys. The absent-state handoff uses + those canonical keys, so resolving them through ``named_parameters()`` + directly breaks every wrapped model. Reuse DCP's own FQN conversion so both + sides of the handoff name the same parameter. + """ + names = {} + for runtime_name, param in model.named_parameters(): + for fqn in _get_fqns(model, runtime_name): + existing = names.get(fqn) + if existing is not None and existing is not param: + raise RuntimeError(f"DCP parameter FQN collision for {fqn}") + names[fqn] = param + return names + + +def _verify_replicated_optimizer_metadata(model, optimizer, group=None): + """DCP deduplicates scalars; compare bounded-size digests on every rank. + + The generic codec requires all DCP ranks to describe the same logical + parameters. Expert/pipeline-specific ownership needs an explicit codec. + """ + if not dist.is_initialized(): + return + error, digest = None, None + try: + names_by_id = defaultdict(list) + for name, param in _dcp_parameter_name_map(model).items(): + names_by_id[id(param)].append(name) + records, step_tensors, lr_aliases = [], defaultdict(list), {} + for leaf_name, leaf in optimizer.optimizers_dict.items(): + for group_index, param_group in enumerate(leaf.param_groups): + options = { + key: _group_option_value(value) for key, value in sorted(param_group.items()) if key != "params" + } + lr_layout = {key: _lr_layout(param_group, key, lr_aliases) for key in ("lr", "initial_lr")} + for param in param_group["params"]: + state = leaf.state.get(param, {}) + record = { + "names": sorted(names_by_id[id(param)]), + "leaf": leaf_name, + "group": group_index, + "present": bool(state), + "step": None, + "options": options, + "lr_layout": lr_layout, + } + records.append(record) + if "step" in state: + step = state["step"] + if isinstance(step, DTensor): + step = step.to_local() + if torch.is_tensor(step): + if step.numel() != 1: + raise RuntimeError("Optimizer step must be a scalar") + step_tensors[(str(step.device), str(step.dtype))].append( + (record, step.detach().reshape(())) + ) + else: + record["step"] = _group_option_value(step) + for (_, dtype), values in step_tensors.items(): + scalars = torch.stack([value for _, value in values]).cpu().tolist() + for (record, _), scalar in zip(values, scalars, strict=True): + record["step"] = {"dtype": dtype, "value": scalar} + records.sort(key=lambda item: (item["names"], item["leaf"], item["group"])) + digest = hashlib.sha256( + json.dumps(records, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking replicated optimizer metadata", group=group) + gathered = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, digest, group=group) + if any(value != gathered[0] for value in gathered): + raise RuntimeError("Replicated optimizer metadata differs across checkpoint ranks") + + +def _multi_extra_snapshot(extra_root, *, read_state): + """Hash the exact sidecar bytes; read through a bounded-memory snapshot.""" + rank = dist.get_rank() if dist.is_initialized() else 0 + extra_path = os.path.join(extra_root, _EXTRA_STATE_DIRNAME, _EXTRA_STATE_FORMAT.format(rank)) + if not os.path.exists(extra_path): + return "absent", None + digest = hashlib.sha256() + with open(extra_path, "rb") as source: + if not read_state: + while chunk := source.read(4 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest(), None + with tempfile.TemporaryFile() as snapshot: + while chunk := source.read(4 * 1024 * 1024): + digest.update(chunk) + snapshot.write(chunk) + snapshot.seek(0) + value = torch.load(snapshot, weights_only=False) + if not isinstance(value, dict): + raise RuntimeError("MultiOptimizer checkpoint extra_state is not a mapping") + return digest.hexdigest(), value + + +def _preflight_multi_optimizer_layout( + model, + optimizer, + reader_factory, + checkpoint_path, + model_checkpoint_path, + extra_root, + load_extra_state, + scheduler_object, + parallel_state=None, + trainable_only=False, + group=None, +): + """Reject an unproven or changed MultiOptimizer group/scheduler layout.""" + multi = _coordinated_optimizer_flag( + getattr(optimizer, "_is_multi_optimizer", False) is True, "MultiOptimizer layout policy", group=group + ) + if not multi: + return None, None, None, None, None, None, None + snapshot, extra, extra_digest, source_metadata, error = None, None, None, None, None + try: + from ..optim.optimizer import AnyPrecisionAdamW + + _check_model_serialization_hooks(model) + source_metadata = reader_factory(checkpoint_path).read_metadata() + parameter_names = _dcp_parameter_name_map(model) + parameter_names_by_id = defaultdict(list) + for name, param in parameter_names.items(): + parameter_names_by_id[id(param)].append(name) + from .rank_local_optimizer import current_rank, load_rank_scalars, rank_prefix + + rank_local = ( + getattr(model, "_fqn2spec_info", None) is not None and getattr(parallel_state, "dp_mode", None) == "fsdp2" + ) + prefix = f"optimizer.{rank_prefix(current_rank())}" if rank_local else "optimizer." + source_bases = { + "optimizer." + key.removeprefix(prefix).rpartition(".")[0] + for key in source_metadata.state_dict_metadata + if key.startswith(prefix + "state.") + } + if trainable_only: + for name, param in parameter_names.items(): + if not param.requires_grad and f"optimizer.state.{name}" in source_bases: + raise RuntimeError( + "trainable-only resume cannot prove the value of a frozen parameter with optimizer history; " + "use a full-model checkpoint or a verified migration" + ) + for leaf in _iter_optimizer_leaves(optimizer): + if leaf.state: + raise RuntimeError("MultiOptimizer checkpoint resume requires a fresh optimizer without state") + if any(param.grad is not None for param_group in leaf.param_groups for param in param_group["params"]): + raise RuntimeError("MultiOptimizer checkpoint resume requires cleared gradients") + if isinstance(leaf, AnyPrecisionAdamW): + for param_group in leaf.param_groups: + fields = ["momentum_dtype", "variance_dtype"] + if param_group.get("use_kahan_summation"): + fields.append("compensation_buffer_dtype") + incompatible = False + for param in param_group["params"]: + if not any(param_group[field] != param.dtype for field in fields): + continue + names = parameter_names_by_id[id(param)] + incompatible |= any(f"optimizer.state.{name}" in source_bases for name in names) + if incompatible: + raise RuntimeError( + "AnyPrecisionAdamW requires explicit precision-preserving migration when optimizer " + "state and parameter dtypes differ" + ) + extra_digest, extra = _multi_extra_snapshot(extra_root, read_state=True) + if extra is not None: + snapshot = extra.get("lr_scheduler") + if extra is not None and not load_extra_state: + raise RuntimeError("MultiOptimizer model extra state cannot be discarded during optimizer resume") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "reading MultiOptimizer scheduler checkpoint", group=group) + if rank_local: + from .rank_local_optimizer import current_rank, rank_inventory_key + + rank = current_rank() + inventory_name = rank_inventory_key(rank) + inventory_path = "optimizer." + inventory_name + inventory_error = ( + RuntimeError("rank-local optimizer checkpoint inventory is missing") + if not isinstance(source_metadata.state_dict_metadata.get(inventory_path), BytesStorageMetadata) + else None + ) + raise_if_any_rank_failed(inventory_error, "inspecting rank-local optimizer inventory", group=group) + inventory_state = {"optimizer": {inventory_name: ""}} + reader = _coordinated_storage_reader(reader_factory, checkpoint_path, group=group) + dcp.load(inventory_state, storage_reader=reader, process_group=group) + inventory_error = None + try: + inventory = json.loads(inventory_state["optimizer"][inventory_name]) + expected = sorted( + key for key in source_metadata.state_dict_metadata if key.startswith(prefix) and key != inventory_path + ) + if inventory != expected: + raise RuntimeError("rank-local optimizer checkpoint inventory differs from DCP metadata") + load_rank_scalars(reader_factory, checkpoint_path, rank) + except BaseException as exc: + inventory_error = exc + raise_if_any_rank_failed(inventory_error, "validating rank-local optimizer inventory", group=group) + else: + inventory_path = f"optimizer.{_MULTI_INVENTORY_MARKER}" + inventory_error = ( + RuntimeError( + "MultiOptimizer checkpoint group layout is unproven or inventory is missing; " + "original grouping and scheduler must be verified and migrated explicitly" + ) + if not isinstance(source_metadata.state_dict_metadata.get(inventory_path), BytesStorageMetadata) + else None + ) + raise_if_any_rank_failed(inventory_error, "inspecting MultiOptimizer inventory", group=group) + inventory_state = {"optimizer": {_MULTI_INVENTORY_MARKER: ""}} + reader = _coordinated_storage_reader(reader_factory, checkpoint_path, group=group) + dcp.load(inventory_state, storage_reader=reader, process_group=group) + inventory_error = None + try: + inventory = json.loads(inventory_state["optimizer"][_MULTI_INVENTORY_MARKER]) + expected = sorted( + key + for key in source_metadata.state_dict_metadata + if key.startswith("optimizer.") and key not in (inventory_path, f"optimizer.{MULTI_LAYOUT_MARKER}") + ) + if inventory != expected: + raise RuntimeError("MultiOptimizer checkpoint inventory differs from DCP metadata") + except BaseException as exc: + inventory_error = exc + raise_if_any_rank_failed(inventory_error, "validating MultiOptimizer inventory", group=group) + inventory_digests = _inventory_digests(inventory, group=group) + frozen_digest = frozen_parameter_digest(model, enabled=trainable_only, group=group) + current = coordinated_multi_layout( + model, optimizer, snapshot, scheduler_object, extra_digest, group=group, frozen_digest=frozen_digest + ) + metadata, reader, error = None, None, None + key = f"optimizer.{MULTI_LAYOUT_MARKER}" + try: + reader = reader_factory(checkpoint_path) + metadata = reader.read_metadata() + first_id = getattr(getattr(source_metadata, "storage_meta", None), "save_id", None) + second_id = getattr(getattr(metadata, "storage_meta", None), "save_id", None) + if first_id and first_id != second_id: + raise RuntimeError("MultiOptimizer checkpoint changed during preflight") + if not isinstance(metadata.state_dict_metadata.get(key), BytesStorageMetadata): + raise RuntimeError( + "MultiOptimizer checkpoint group layout is unproven; original grouping and scheduler " + "must be verified and migrated explicitly" + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "inspecting MultiOptimizer checkpoint layout", group=group) + model_key = _MULTI_GENERATION_MARKER + model_reader, model_error = None, None + try: + model_reader = reader_factory(model_checkpoint_path) + model_metadata = model_reader.read_metadata() + if not isinstance(model_metadata.state_dict_metadata.get(model_key), BytesStorageMetadata): + raise RuntimeError("MultiOptimizer model checkpoint generation is missing") + options = StateDictOptions(ignore_frozen_params=True) if trainable_only else None + if rank_local: + from .dcp_checkpointer import ModelState + + current_model_state = ModelState( + model, trainable_only=trainable_only, parallel_state=parallel_state + ).state_dict() + else: + current_model_state = get_model_state_dict(model, options=options) + from .dcp_checkpointer import _export_extra_parallel_layout_key + + _check_model_serialization_hooks(model) + logical_keys = { + logical_key + for raw in current_model_state + if (logical_key := _export_extra_parallel_layout_key(raw)) is not None + } + _check_qwen_omni_checkpoint_keys(model, logical_keys) + flat_model, _ = flatten_state_dict({"model": current_model_state}) + expected_keys = set(flat_model) + source_keys = {key for key in model_metadata.state_dict_metadata if key.startswith("model.")} + if source_keys != expected_keys: + raise RuntimeError("MultiOptimizer checkpoint model state inventory differs") + for name, tensor in flat_model.items(): + source_entry = model_metadata.state_dict_metadata[name] + if not torch.is_tensor(tensor): + if not isinstance(source_entry, BytesStorageMetadata): + raise RuntimeError(f"MultiOptimizer model checkpoint value type differs for {name}") + continue + if not isinstance(source_entry, TensorStorageMetadata): + raise RuntimeError(f"MultiOptimizer model checkpoint value type differs for {name}") + if source_entry.properties.dtype != tensor.dtype or tuple(source_entry.size) != tuple(tensor.shape): + raise RuntimeError(f"MultiOptimizer model checkpoint tensor dtype or shape differs for {name}") + if trainable_only: + logical_source_keys = { + logical_key + for raw in source_keys + if (logical_key := _export_extra_parallel_layout_key(raw.removeprefix("model."))) is not None + } + for name in parameter_names: + if f"optimizer.state.{name}" in source_bases and name not in logical_source_keys: + raise RuntimeError( + "trainable-only resume cannot prove the model value for a parameter with optimizer history; " + "use a full-model checkpoint or a verified migration" + ) + except BaseException as exc: + model_error = exc + raise_if_any_rank_failed(model_error, "inspecting MultiOptimizer model generation", group=group) + model_scratch = {model_key: ""} + dcp.load( + state_dict=model_scratch, + storage_reader=model_reader, + process_group=group, + planner=DefaultLoadPlanner(flatten_state_dict=False, flatten_sharded_tensors=False, allow_partial_load=False), + ) + scratch = {key: ""} + dcp.load( + state_dict=scratch, + storage_reader=reader, + process_group=group, + planner=DefaultLoadPlanner(flatten_state_dict=False, flatten_sharded_tensors=False, allow_partial_load=False), + ) + mismatch = None + try: + saved = json.loads(scratch[key]) + if ( + not isinstance(saved, dict) + or set(saved) != {"generation", "layout", "inventory_digests"} + or not isinstance(saved["generation"], str) + or not saved["generation"] + or saved["generation"] != model_scratch[model_key] + or saved["layout"] != current + or saved["inventory_digests"] != inventory_digests + ): + raise RuntimeError("MultiOptimizer checkpoint generation, group or scheduler layout differs") + except (ValueError, TypeError, KeyError) as exc: + mismatch = RuntimeError("MultiOptimizer checkpoint layout marker is malformed") + mismatch.__cause__ = exc + except RuntimeError as exc: + mismatch = exc + raise_if_any_rank_failed(mismatch, "validating MultiOptimizer checkpoint layout", group=group) + optimizer_id = getattr(getattr(metadata, "storage_meta", None), "save_id", None) + model_id = getattr(getattr(model_metadata, "storage_meta", None), "save_id", None) + id_error = ( + RuntimeError("MultiOptimizer checkpoint cannot prove model/optimizer save generations") + if not optimizer_id or not model_id + else None + ) + raise_if_any_rank_failed(id_error, "validating MultiOptimizer DCP generation IDs", group=group) + return optimizer_id, extra, model_id, saved["generation"], current, extra_digest, frozen_digest + + +def _prepare_multi_optimizer_saved_states( + model, + optimizer, + reader_factory, + checkpoint_path, + expected_save_id, + group=None, + *, + parallel_state=None, + validate_only=False, +): + """Validate or allocate saved leaf state, without a synthetic optimizer step.""" + multi = _coordinated_optimizer_flag( + getattr(optimizer, "_is_multi_optimizer", False) is True, "MultiOptimizer state preparation", group=group + ) + if not multi: + return + error = None + try: + from ..optim.checkpoint_adapter import rank_local_adamw_adapter + from ..optim.optimizer import AnyPrecisionAdamW + + for leaf in _iter_optimizer_leaves(optimizer): + if leaf.state or any( + param.grad is not None for param_group in leaf.param_groups for param in param_group["params"] + ): + raise RuntimeError("MultiOptimizer state or gradients changed while loading model state") + metadata = reader_factory(checkpoint_path).read_metadata() + save_id = getattr(getattr(metadata, "storage_meta", None), "save_id", None) + if expected_save_id is not None and save_id != expected_save_id: + raise RuntimeError("MultiOptimizer checkpoint changed before preparing its saved state") + names = _dcp_parameter_name_map(model) + names_by_id = defaultdict(list) + for name, param in names.items(): + names_by_id[id(param)].append(name) + saved = defaultdict(dict) + from .rank_local_optimizer import current_rank, load_rank_scalars, rank_prefix + + rank_local = ( + getattr(model, "_fqn2spec_info", None) is not None and getattr(parallel_state, "dp_mode", None) == "fsdp2" + ) + checkpoint_prefix = f"optimizer.{rank_prefix(current_rank())}" if rank_local else "optimizer." + unknown_state_names = set() + for key, value in metadata.state_dict_metadata.items(): + if not key.startswith(checkpoint_prefix): + continue + key = "optimizer." + key.removeprefix(checkpoint_prefix) + prefix, separator, field = key.rpartition(".") + name = prefix.removeprefix("optimizer.state.") + if separator and prefix.startswith("optimizer.state."): + if name not in names: + unknown_state_names.add(name) + else: + saved[name][field] = value + if unknown_state_names: + raise RuntimeError("optimizer checkpoint contains state for unknown model parameters") + expected_groups = {key for key in optimizer._sparse_state_dict() if key.startswith("param_groups.")} + if rank_local: + scalars = load_rank_scalars(reader_factory, checkpoint_path, current_rank()) + actual_groups = {key for key in scalars if key.startswith("param_groups.")} + if expected_groups != actual_groups: + raise RuntimeError("rank-local optimizer parameter group inventory differs") + owned_names = { + name + for leaf in _iter_optimizer_leaves(optimizer) + for param_group in leaf.param_groups + for param in param_group["params"] + for name in names_by_id[id(param)] + } + unknown_steps = { + key + for key in scalars + if key.startswith("state.") and key.removeprefix("state.").removesuffix(".step") not in owned_names + } + if unknown_steps: + raise RuntimeError("rank-local optimizer checkpoint has unowned parameter steps") + if set(saved) - owned_names: + raise RuntimeError("rank-local optimizer checkpoint has unowned parameter state") + for name in names: + step_key = f"state.{name}.step" + if step_key in scalars: + saved[name]["step"] = scalars[step_key] + else: + actual_groups = { + key.removeprefix("optimizer.") + for key in metadata.state_dict_metadata + if key.startswith("optimizer.param_groups.") + } + if expected_groups != actual_groups: + raise RuntimeError("MultiOptimizer checkpoint parameter group inventory differs") + + for leaf in _iter_optimizer_leaves(optimizer): + if any(not param.requires_grad for group in leaf.param_groups for param in group["params"]): + if getattr(leaf, "_optimizer_load_state_dict_pre_hooks", None) or getattr( + leaf, "_optimizer_load_state_dict_post_hooks", None + ): + raise RuntimeError( + "Cannot safely restore frozen optimizer history with custom optimizer load hooks" + ) + for param_group in leaf.param_groups: + for param in param_group["params"]: + candidates = [name for name in names_by_id[id(param)] if name in saved] + if len(candidates) > 1: + raise RuntimeError("Ambiguous shared parameter optimizer history in MultiOptimizer checkpoint") + if not candidates: + continue + fields = saved[candidates[0]] + expected = _expected_optimizer_fields(leaf, param_group) + if set(fields) != expected: + raise RuntimeError("MultiOptimizer checkpoint has incomplete or unsupported parameter state") + state = {} + for field, entry in fields.items(): + if field == "step": + if rank_local: + if not torch.is_tensor(entry): + raise RuntimeError("MultiOptimizer checkpoint step is not a tensor") + step_shape, step_dtype = tuple(entry.shape), entry.dtype + else: + if not isinstance(entry, TensorStorageMetadata): + raise RuntimeError("MultiOptimizer checkpoint state is not a supported tensor") + step_shape, step_dtype = tuple(entry.size), entry.properties.dtype + if step_shape != (): + raise RuntimeError("Non-scalar optimizer steps require an explicit owner-layout codec") + adapter = rank_local_adamw_adapter(leaf) + if adapter is not None and step_dtype != adapter.state_dtype(leaf, param, field): + raise RuntimeError( + "MultiOptimizer checkpoint step dtype differs from its registered policy" + ) + device = ( + param.device if param_group.get("capturable") or param_group.get("fused") else "cpu" + ) + if not validate_only: + state[field] = torch.zeros((), dtype=step_dtype, device=device) + continue + if not isinstance(entry, TensorStorageMetadata): + raise RuntimeError("MultiOptimizer checkpoint state is not a supported tensor") + adapter = rank_local_adamw_adapter(leaf) + if adapter is not None: + dtype = adapter.state_dtype(leaf, param, field) + elif isinstance(leaf, AnyPrecisionAdamW): + dtype = param_group[ + { + "exp_avg": "momentum_dtype", + "exp_avg_sq": "variance_dtype", + "compensation": "compensation_buffer_dtype", + }[field] + ] + if dtype != param.dtype: + raise RuntimeError( + "AnyPrecisionAdamW cannot restore moments at a dtype different from its parameter; " + "explicit precision-preserving migration is required" + ) + else: + dtype = param.dtype + if entry.properties.dtype != dtype: + raise RuntimeError(f"MultiOptimizer checkpoint dtype mismatch for {candidates[0]}.{field}") + local_param = param.to_local() if rank_local and hasattr(param, "to_local") else param + if tuple(entry.size) != tuple(local_param.shape): + raise RuntimeError(f"MultiOptimizer checkpoint shape mismatch for {candidates[0]}.{field}") + if not validate_only: + if adapter is not None and adapter.state_is_offloaded(leaf, param): + from .rank_local_optimizer import zeros_like_parameter_on_cpu + + state[field] = zeros_like_parameter_on_cpu(param, dtype=dtype) + else: + state[field] = torch.zeros_like(param, dtype=dtype, device=param.device) + if not validate_only: + leaf.state[param] = state + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "preparing MultiOptimizer checkpoint state", group=group) diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py new file mode 100644 index 0000000000..e0b785945f --- /dev/null +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -0,0 +1,469 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +"""Identity of a MultiOptimizer's rank-local parameter and scheduler layout.""" + +import hashlib +import json +import marshal +import types + +import torch +import torch.distributed as dist +from torch.distributed.checkpoint.state_dict import _get_fqns + +from ..utils.dist_utils import raise_if_any_rank_failed + + +MULTI_LAYOUT_MARKER = "__veomni_multi_optimizer_layout_v1__" + + +def _immutable_value(value): + if value is None or type(value) in (bool, int, float, str): + return value + if isinstance(value, tuple): + return [_immutable_value(item) for item in value] + raise RuntimeError("MultiOptimizer scheduler captures unsupported mutable state") + + +def _scheduler_snapshot_value(value): + if value is None or type(value) in (bool, int, float, str): + return value + if torch.is_tensor(value) and value.numel() == 1: + return {"dtype": str(value.dtype), "value": value.item()} + if isinstance(value, (list, tuple)): + return [_scheduler_snapshot_value(item) for item in value] + if isinstance(value, dict) and all(isinstance(key, str) for key in value): + return {key: _scheduler_snapshot_value(item) for key, item in sorted(value.items())} + raise RuntimeError("MultiOptimizer scheduler checkpoint contains an unsupported value") + + +def _group_option_value(value): + if isinstance(value, (torch.dtype, torch.device)): + return str(value) + if isinstance(value, (list, tuple)): + return [_group_option_value(item) for item in value] + if isinstance(value, dict) and all(isinstance(key, str) for key in value): + return {key: _group_option_value(item) for key, item in sorted(value.items())} + if torch.is_tensor(value) and value.numel() == 1: + return {"dtype": str(value.dtype), "value": value.item()} + return _immutable_value(value) + + +def _lr_layout(group, key, aliases): + if key not in group: + return None + value = group[key] + if torch.is_tensor(value): + if value.device.type == "meta" or value.numel() != 1: + raise RuntimeError(f"Cannot prove optimizer {key} tensor storage") + storage = (str(value.device), value.untyped_storage().data_ptr()) + alias = aliases.setdefault(storage, len(aliases)) + return {"type": "tensor", "dtype": str(value.dtype), "shape": list(value.shape), "storage_alias": alias} + if type(value) not in (int, float): + raise RuntimeError(f"Unsupported optimizer {key} type") + return {"type": type(value).__name__} + + +def _scheduler_layout(leaf): + import math + + from ..optim.lr_scheduler import checkpoint_scheduler_binding, checkpoint_scheduler_functions + + scheduler = checkpoint_scheduler_binding(leaf) + wrapped = getattr(leaf.step, "_wrapped_by_lr_sched", False) + if scheduler is None: + if wrapped or "step" in vars(leaf): + raise RuntimeError("MultiOptimizer scheduler identity is unknown; explicit migration is required") + return None + import weakref + + def nested_codes(code): + result = {code} + for value in code.co_consts: + if isinstance(value, types.CodeType): + result.update(nested_codes(value)) + return result + + step = vars(leaf).get("step") + permitted = nested_codes(torch.optim.lr_scheduler.LRScheduler.__init__.__code__) + closure = getattr(step, "__closure__", None) or () + captured = tuple(cell.cell_contents for cell in closure) + if ( + getattr(step, "__code__", None) not in permitted + or not any(item is type(leaf).step for item in captured) + or not any(isinstance(item, weakref.ReferenceType) and item() is leaf for item in captured) + ): + raise RuntimeError("MultiOptimizer cannot prove optimizer step implementation") + if type(scheduler) is not torch.optim.lr_scheduler.LambdaLR or any( + name in vars(scheduler) for name in ("state_dict", "load_state_dict", "step", "get_lr", "get_last_lr") + ): + raise RuntimeError("MultiOptimizer cannot prove an overridden scheduler method") + if scheduler.optimizer is not leaf or not wrapped: + raise RuntimeError("MultiOptimizer scheduler binding changed") + original_functions = checkpoint_scheduler_functions(scheduler) + if ( + original_functions is None + or len(original_functions) != len(scheduler.lr_lambdas) + or any( + current is not original for current, original in zip(scheduler.lr_lambdas, original_functions, strict=True) + ) + ): + raise RuntimeError("MultiOptimizer scheduler recipe changed after construction") + functions = [] + for function in scheduler.lr_lambdas: + if not isinstance(function, types.FunctionType): + raise RuntimeError("MultiOptimizer scheduler uses an unsupported callable") + if any( + name in function.__globals__ and (name != "math" or function.__globals__[name] is not math) + for name in function.__code__.co_names + ): + raise RuntimeError("MultiOptimizer scheduler refers to unsupported global state") + + def normalize(code): + constants = tuple(normalize(item) if isinstance(item, types.CodeType) else item for item in code.co_consts) + return code.replace(co_filename="", co_firstlineno=0, co_consts=constants) + + functions.append( + { + "code": hashlib.sha256(marshal.dumps(normalize(function.__code__))).hexdigest(), + "defaults": _immutable_value(function.__defaults__), + "kwdefaults": _immutable_value(function.__kwdefaults__), + "closure": [_immutable_value(cell.cell_contents) for cell in function.__closure__ or ()], + } + ) + return {"type": f"{type(scheduler).__module__}.{type(scheduler).__qualname__}", "functions": functions} + + +def _local_layout( + model, optimizer, scheduler_snapshot, scheduler_object, extra_digest, *, writing, frozen_digest=None +): + from ..optim.checkpoint_adapter import rank_local_adamw_adapter + from ..optim.muon import DistributedMuon + from ..optim.optimizer import AnyPrecisionAdamW, MultiOptimizer + + if type(optimizer) is not MultiOptimizer: + raise RuntimeError("MultiOptimizer checkpoint requires the standard serialization implementation") + seen_mutable_options = {} + seen_tensor_options = {} + + def check_aliases(value, path): + if torch.is_tensor(value): + if value.device.type == "meta": + raise RuntimeError(f"Cannot prove optimizer option tensor storage: {path}") + storage = (str(value.device), value.untyped_storage().data_ptr()) + previous = seen_tensor_options.setdefault(storage, path) + if previous != path: + raise RuntimeError( + f"Shared mutable optimizer tensor option cannot be restored safely: {previous} and {path}" + ) + elif type(value) in (dict, list): + previous = seen_mutable_options.setdefault(id(value), path) + if previous != path: + raise RuntimeError(f"Shared mutable optimizer option cannot be restored safely: {previous} and {path}") + values = value.values() if type(value) is dict else value + for index, item in enumerate(values): + check_aliases(item, f"{path}.{index}") + elif type(value) is tuple: + for index, item in enumerate(value): + check_aliases(item, f"{path}.{index}") + + for leaf_name, leaf in optimizer.optimizers_dict.items(): + for key, value in leaf.defaults.items(): + check_aliases(value, f"{leaf_name}.defaults.{key}") + for index, group in enumerate(leaf.param_groups): + for key, value in group.items(): + if key != "params": + check_aliases(value, f"{leaf_name}.group{index}.{key}") + if any( + name in vars(optimizer) + for name in ("state_dict", "load_state_dict", "_sparse_state_dict", "step", "zero_grad") + ): + raise RuntimeError("MultiOptimizer checkpoint cannot prove overridden optimizer methods") + import torch.optim.optimizer as torch_optimizer + + if torch_optimizer._global_optimizer_pre_hooks or torch_optimizer._global_optimizer_post_hooks: + raise RuntimeError("MultiOptimizer checkpoint cannot prove global optimizer step hooks") + keys = optimizer.key_names + if len(keys) != len(set(keys)) or set(keys) != set(optimizer.optimizers_dict): + raise RuntimeError("MultiOptimizer keys do not uniquely describe its optimizer leaves") + names, model_trainability = {}, {} + for runtime_name, param in model.named_parameters(remove_duplicate=False): + # Even parameters outside the optimizer can affect the next global + # gradient clipping operation. DCP does not persist pending gradients. + if param.grad is not None: + raise RuntimeError("MultiOptimizer checkpoint cannot omit pending gradients") + fqns = _get_fqns(model, runtime_name) + names.setdefault(id(param), set()).update(fqns) + model_trainability.update((name, bool(param.requires_grad)) for name in fqns) + optimizer_names = {} + for runtime_name, param in optimizer.model.named_parameters(remove_duplicate=False): + optimizer_names.setdefault(id(param), set()).update(_get_fqns(optimizer.model, runtime_name)) + if names != optimizer_names: + raise RuntimeError("MultiOptimizer and checkpoint model have different canonical parameter identities") + model_aliases = {} + for kind, items in ( + ("parameter", model.named_parameters(remove_duplicate=False)), + ("buffer", model.named_buffers(remove_duplicate=False)), + ): + for runtime_name, value in items: + model_aliases.setdefault(id(value), set()).update( + f"{kind}:{name}" for name in _get_fqns(model, runtime_name) + ) + alias_graph = sorted(sorted(aliases) for aliases in model_aliases.values() if len(aliases) > 1) + + leaves = [] + seen_param_ids, seen_fqns, lr_aliases, lr_owners = set(), set(), {}, {} + mesh_digests = {} + + def mesh_identity(mesh): + if mesh is None: + return None + mesh_id = id(mesh) + if mesh_id not in mesh_digests: + identity = { + "ranks": mesh.mesh.tolist(), + "device_type": mesh.device_type, + "dimension_names": mesh.mesh_dim_names, + } + mesh_digests[mesh_id] = hashlib.sha256( + json.dumps(identity, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + return mesh_digests[mesh_id] + + specs = getattr(model, "_fqn2spec_info", None) + if specs is not None and not isinstance(specs, dict): + raise RuntimeError("MultiOptimizer checkpoint cannot prove ExtraParallel plan identity") + extra_parallel_specs = [] + for name, spec in sorted((specs or {}).items()): + extra_parallel_specs.append( + { + "name": name, + "parallel_name": spec.para_name, + "placement": repr(spec.placement), + "fsdp_shard_dim": spec.fsdp_shard_dim, + "persistent_fsdp_shard_dim": getattr(spec, "persistent_fsdp_shard_dim", None), + "global_shape": getattr(spec, "global_shape", None), + "checkpoint_layout_v2": getattr(spec, "checkpoint_layout_v2", False), + "mesh": mesh_identity(spec.para_fsdp_mesh), + } + ) + for name in keys: + leaf = optimizer.optimizers_dict[name] + if rank_local_adamw_adapter(leaf) is None and type(leaf) not in ( + torch.optim.AdamW, + AnyPrecisionAdamW, + DistributedMuon, + ): + raise RuntimeError("MultiOptimizer checkpoint requires an explicitly supported optimizer leaf") + if any(name in vars(leaf) for name in ("state_dict", "load_state_dict", "zero_grad")): + raise RuntimeError("MultiOptimizer checkpoint cannot prove overridden leaf methods") + if any( + getattr(leaf, key, None) + for key in ( + "_optimizer_state_dict_pre_hooks", + "_optimizer_state_dict_post_hooks", + "_optimizer_load_state_dict_pre_hooks", + "_optimizer_load_state_dict_post_hooks", + "_optimizer_step_pre_hooks", + "_optimizer_step_post_hooks", + ) + ): + raise RuntimeError("MultiOptimizer checkpoint cannot prove custom optimizer serialization hooks") + groups = [] + for group_index, group in enumerate(leaf.param_groups): + if not group["params"]: + raise RuntimeError("MultiOptimizer cannot checkpoint an empty parameter group") + params = [] + for param in group["params"]: + fqns = names.get(id(param)) + if not fqns: + raise RuntimeError("MultiOptimizer owns a parameter absent from the checkpoint model") + if id(param) in seen_param_ids or seen_fqns.intersection(fqns): + raise RuntimeError("MultiOptimizer parameter ownership is duplicated across groups or leaves") + seen_param_ids.add(id(param)) + seen_fqns.update(fqns) + mesh = getattr(param, "device_mesh", None) + mesh_digest = mesh_identity(mesh) + adapter = rank_local_adamw_adapter(leaf) + params.append( + { + "names": sorted(fqns), + "shape": list(param.shape), + "dtype": str(param.dtype), + "device_type": param.device.type, + "mesh": mesh_digest, + "placements": [repr(place) for place in param.placements] if mesh is not None else None, + "cpu_offload": bool(adapter.state_is_offloaded(leaf, param)) if adapter is not None else False, + "state_dtypes": { + key: str(adapter.state_dtype(leaf, param, key)) + for key in ("step", "master_weight", "exp_avg", "exp_avg_sq") + } + if adapter is not None + else None, + } + ) + lr_layout = _lr_layout(group, "lr", lr_aliases) + initial_lr_layout = _lr_layout(group, "initial_lr", lr_aliases) + for lr_spec in (lr_layout, initial_lr_layout): + if lr_spec is not None and lr_spec["type"] == "tensor": + alias = lr_spec["storage_alias"] + owner = lr_owners.setdefault(alias, (name, group_index)) + if owner != (name, group_index): + raise RuntimeError( + "MultiOptimizer cannot safely restore LR tensor storage shared across groups" + ) + groups.append( + { + "params": params, + "lr_layout": lr_layout, + "initial_lr_layout": initial_lr_layout, + # LR is restored from the checkpoint and may have advanced; + # structural and algorithm options must already agree. + "options": { + key: _group_option_value(value) + for key, value in sorted(group.items()) + if key not in ("params", "lr", "initial_lr") + }, + } + ) + leaves.append( + { + "name": name, + "type": f"{type(leaf).__module__}.{type(leaf).__qualname__}", + "algorithm_identity": ( + _immutable_value(rank_local_adamw_adapter(leaf).algorithm_identity(leaf)) + if rank_local_adamw_adapter(leaf) is not None + else None + ), + "groups": groups, + "scheduler": _scheduler_layout(leaf), + } + ) + has_scheduler = any(leaf["scheduler"] is not None for leaf in leaves) + if has_scheduler: + if not isinstance(scheduler_snapshot, dict) or set(scheduler_snapshot) != set(keys): + raise RuntimeError("MultiOptimizer scheduler checkpoint is missing or has different leaves") + from ..optim.lr_scheduler import MultiLRScheduler, checkpoint_scheduler_binding + + if type(scheduler_object) is not MultiLRScheduler or set(scheduler_object) != set(keys): + raise RuntimeError("MultiOptimizer scheduler checkpoint requires its live scheduler object") + if any(name in vars(scheduler_object) for name in ("state_dict", "load_state_dict", "step", "get_last_lr")): + raise RuntimeError("MultiOptimizer cannot prove an overridden scheduler container method") + + for name in keys: + current_scheduler = scheduler_object[name] + if current_scheduler is not checkpoint_scheduler_binding(optimizer.optimizers_dict[name]): + raise RuntimeError("MultiOptimizer scheduler object differs from its optimizer binding") + if writing and _scheduler_snapshot_value(scheduler_snapshot[name]) != _scheduler_snapshot_value( + current_scheduler.state_dict() + ): + raise RuntimeError("MultiOptimizer scheduler snapshot differs from its live scheduler") + elif scheduler_snapshot is not None: + raise RuntimeError("MultiOptimizer scheduler checkpoint has no matching scheduler binding") + elif scheduler_object is not None: + raise RuntimeError("MultiOptimizer scheduler object has no matching scheduler binding") + layout = { + "leaves": leaves, + "model_aliases": alias_graph, + "model_trainability": sorted(model_trainability.items()), + "extra_parallel_specs": extra_parallel_specs, + "scheduler_snapshot": _scheduler_snapshot_value(scheduler_snapshot), + "extra_digest": extra_digest, + "frozen_digest": frozen_digest, + } + return hashlib.sha256( + json.dumps(layout, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + + +def frozen_parameter_digest(model, *, enabled, group=None, chunk_bytes=4 << 20): + """Bind omitted frozen model values to a trainable-only checkpoint.""" + if not enabled: + return None + digest, error = hashlib.sha256(), None + try: + for name, param in sorted(model.named_parameters(remove_duplicate=False)): + if param.requires_grad: + continue + local = param.to_local() if hasattr(param, "to_local") else param + descriptor = (name, str(param.dtype), tuple(param.shape), tuple(local.shape), str(param.device.type)) + digest.update(json.dumps(descriptor, separators=(",", ":")).encode()) + from .tensor_bytes import tensor_byte_chunks + + for chunk in tensor_byte_chunks(local, chunk_bytes=chunk_bytes): + digest.update(chunk) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "hashing frozen checkpoint parameters", group=group) + return digest.hexdigest() + + +def coordinated_multi_layout( + model, + optimizer, + scheduler_snapshot=None, + scheduler_object=None, + extra_digest=None, + *, + writing=False, + group=None, + frozen_digest=None, +): + """Return the same rank-indexed layout on every participant or fail together.""" + local, trainability = None, None + error = None + try: + local = _local_layout( + model, + optimizer, + scheduler_snapshot, + scheduler_object, + extra_digest, + writing=writing, + frozen_digest=frozen_digest, + ) + from torch.distributed.tensor import Shard + + specs = getattr(model, "_fqn2spec_info", {}) or {} + trainability_by_owner = {} + for runtime_name, param in model.named_parameters(remove_duplicate=False): + for name in _get_fqns(model, runtime_name): + spec = specs.get(runtime_name) or specs.get(name) + owner = None + if spec is not None and isinstance(spec.placement, Shard): + mesh = spec.para_fsdp_mesh + if mesh is None: + raise RuntimeError(f"Cannot prove ExtraParallel owner for trainability of {name}") + owner = (spec.para_name, mesh[spec.para_name].get_local_rank()) + if owner[1] < 0: + raise RuntimeError(f"Current rank is not an ExtraParallel owner of {name}") + values = trainability_by_owner.setdefault(owner, {}) + enabled = bool(param.requires_grad) + previous = values.setdefault(name, enabled) + if previous != enabled: + raise RuntimeError(f"Tied parameter trainability differs for {name}") + trainability = [] + for owner, values in trainability_by_owner.items(): + payload = json.dumps(sorted(values.items()), separators=(",", ":")).encode() + trainability.append((owner, len(values), hashlib.sha256(payload).hexdigest())) + trainability = tuple(sorted(trainability, key=lambda entry: repr(entry[0]))) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "inspecting MultiOptimizer layout", group=group) + if dist.is_initialized(): + world = dist.get_world_size(group) + digests = [None] * world + gathered = [None] * world + dist.all_gather_object(gathered, (local, trainability), group=group) + digests = [item[0] for item in gathered] + owners = {} + for _, entries in gathered: + for owner, count, digest in entries: + previous = owners.setdefault(owner, (count, digest)) + if previous != (count, digest): + raise RuntimeError("Model parameter trainability differs across checkpoint replica ranks") + ranks = list(range(world)) if group is None else dist.get_process_group_ranks(group) + else: + ranks, digests = [0], [local] + return json.dumps({"version": 1, "ranks": ranks, "layouts": digests}, sort_keys=True, separators=(",", ":")) diff --git a/veomni/checkpoint/rank_local_optimizer.py b/veomni/checkpoint/rank_local_optimizer.py new file mode 100644 index 0000000000..665005a26f --- /dev/null +++ b/veomni/checkpoint/rank_local_optimizer.py @@ -0,0 +1,333 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +"""Exact-rank optimizer DCP keys for fixed-geometry recovery experiments. + +Unlike globally named scalar steps, rank-local keys cannot be deduplicated +across different expert owners. This intentionally trades checkpoint storage +for exact state preservation and requires an independently verified rank and +mesh identity before use. It does not support elastic resharding. +""" + +import hashlib +import heapq +import io +import json +import struct +from contextlib import contextmanager + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint._nested_dict import flatten_state_dict +from torch.distributed.checkpoint.metadata import BytesStorageMetadata +from torch.distributed.tensor import DTensor + + +def rank_prefix(rank: int) -> str: + if type(rank) is not int or rank < 0: + raise ValueError("optimizer checkpoint rank must be a non-negative integer") + return f"__veomni_rank_{rank}__" + + +def rank_scalar_key(rank: int) -> str: + return rank_prefix(rank) + "__scalars" + + +def _is_small_optimizer_value(key: str) -> bool: + return key.startswith("param_groups.") or (key.startswith("state.") and key.endswith(".step")) + + +def _decode_scalar_payload(payload) -> dict: + if not isinstance(payload, bytes): + raise RuntimeError("rank-local optimizer scalar payload is not bytes") + values = torch.load(io.BytesIO(payload), map_location="cpu", weights_only=True) + if not isinstance(values, dict) or any( + not isinstance(key, str) or not _is_small_optimizer_value(key) for key in values + ): + raise RuntimeError("rank-local optimizer scalar payload is malformed") + return values + + +def local_optimizer_state(state: dict, rank: int) -> dict: + """Expose physical shards under rank-unique keys; pack steps and group options.""" + prefix = rank_prefix(rank) + local, small, pending = {}, {}, {} + for key, value in state.items(): + if not isinstance(key, str) or key.startswith("__veomni_rank_"): + raise ValueError("invalid optimizer checkpoint key") + if _is_small_optimizer_value(key): + if isinstance(value, DTensor): + raise RuntimeError("rank-local optimizer scalar cannot be a DTensor") + if torch.is_tensor(value) and value.device.type != "cpu": + if value.ndim != 0: + raise RuntimeError("rank-local accelerator optimizer metadata must be scalar") + pending.setdefault((value.device, value.dtype), []).append((key, value.detach())) + else: + small[key] = value.detach() if torch.is_tensor(value) else value + continue + if isinstance(value, DTensor): + value = value.to_local() + local[prefix + key] = value + for entries in pending.values(): + transferred = torch.stack([value for _, value in entries]).cpu() + for index, (key, _) in enumerate(entries): + small[key] = transferred[index] + buffer = io.BytesIO() + torch.save(small, buffer) + local[rank_scalar_key(rank)] = buffer.getvalue() + return local + + +def rank_inventory(state: dict, rank: int) -> list[str]: + flattened, _ = flatten_state_dict({"optimizer": local_optimizer_state(state, rank)}) + return sorted(flattened) + + +def rank_inventory_key(rank: int) -> str: + return rank_prefix(rank) + "__inventory" + + +def load_rank_scalars(reader_factory, checkpoint_path, rank): + key = rank_scalar_key(rank) + reader = reader_factory(checkpoint_path) + metadata = reader.read_metadata() + if not isinstance(metadata.state_dict_metadata.get(f"optimizer.{key}"), BytesStorageMetadata): + raise RuntimeError("rank-local optimizer scalar payload is missing") + scratch = {"optimizer": {key: b""}} + dcp.load(scratch, storage_reader=reader_factory(checkpoint_path), no_dist=True) + return _decode_scalar_payload(scratch["optimizer"][key]) + + +def restore_optimizer_state(local: dict, template: dict, rank: int) -> dict: + """Restore the DTensor wrappers from the already-allocated optimizer state.""" + prefix = rank_prefix(rank) + expected = {prefix + key for key in template if not _is_small_optimizer_value(key)} | {rank_scalar_key(rank)} + if set(local) != expected: + raise RuntimeError("rank-local optimizer checkpoint inventory differs") + small = _decode_scalar_payload(local[rank_scalar_key(rank)]) + if set(small) != {key for key in template if _is_small_optimizer_value(key)}: + raise RuntimeError("rank-local optimizer scalar inventory differs") + result = {} + for key, expected in template.items(): + value = small[key] if _is_small_optimizer_value(key) else local[prefix + key] + if _is_small_optimizer_value(key) and torch.is_tensor(expected): + if ( + not torch.is_tensor(value) + or expected.dtype != value.dtype + or tuple(expected.shape) != tuple(value.shape) + ): + raise RuntimeError("rank-local optimizer scalar shape or dtype differs") + expected.copy_(value.to(device=expected.device)) + result[key] = expected + continue + if isinstance(expected, DTensor): + if not isinstance(value, torch.Tensor): + raise RuntimeError("rank-local optimizer tensor was replaced with a non-tensor") + original = expected.to_local() + if value.dtype != original.dtype or tuple(value.shape) != tuple(original.shape): + raise RuntimeError("rank-local optimizer tensor shape or dtype differs") + if value.data_ptr() != original.data_ptr(): + original.copy_(value) + result[key] = expected + else: + result[key] = value + return result + + +def current_rank() -> int: + return dist.get_rank() if dist.is_initialized() else 0 + + +def wrap_local_like_parameter(local, param): + """Preserve a parameter's distributed geometry for a local state shard.""" + if not isinstance(param, DTensor): + return local + if local.device.type == "cpu" and param.device_mesh.device_type != "cpu": + metadata = DTensor.from_local( + local.to("meta"), + param.device_mesh, + param.placements, + run_check=False, + shape=param.shape, + stride=param.stride(), + ) + return DTensor(local, metadata._spec, requires_grad=False) + return DTensor.from_local( + local, param.device_mesh, param.placements, run_check=False, shape=param.shape, stride=param.stride() + ) + + +def zeros_like_parameter_on_cpu(param, *, dtype): + """Allocate one local optimizer shard on CPU with the parameter's DTensor spec.""" + local = param.to_local() if isinstance(param, DTensor) else param + return wrap_local_like_parameter(torch.zeros_like(local, dtype=dtype, device="cpu"), param) + + +def _reject_overlapping_geometries(name, geometries): + """Reject distinct overlapping boxes before DCP deduplicates by offset.""" + offsets_to_size = {} + nonempty = [] + for offsets, size in geometries: + old_size = offsets_to_size.setdefault(offsets, size) + if old_size != size: + raise RuntimeError(f"Checkpoint tensor {name} has overlapping nonidentical shards") + if all(size): + nonempty.append((offsets, size)) + if len(nonempty) < 2 or not nonempty[0][0]: + return + # Choose a dimension with the most distinct intervals to keep the sweep + # active set small for both row- and column-sharded tensors. + axis = max( + range(len(nonempty[0][0])), + key=lambda dim: len({(offsets[dim], size[dim]) for offsets, size in nonempty}), + ) + ordered = sorted(nonempty, key=lambda item: item[0][axis]) + active, expiration = {}, [] + for index, (offsets, size) in enumerate(ordered): + start = offsets[axis] + while expiration and expiration[0][0] <= start: + _, previous = heapq.heappop(expiration) + del active[previous] + for other_offsets, other_size in active.values(): + if all( + start < other_start + other_length and other_start < start + length + for start, length, other_start, other_length in zip( + offsets, size, other_offsets, other_size, strict=True + ) + ): + raise RuntimeError(f"Checkpoint tensor {name} has overlapping nonidentical shards") + active[index] = (offsets, size) + heapq.heappush(expiration, (start + size[axis], index)) + + +def verify_replicated_model_shards(state: dict, *, group=None, batch_size=128, chunk_bytes=4 << 20): + """Reject divergent model replicas before DCP deduplicates their shards. + + Hashing is chunked to bound temporary CPU/NPU memory, including for + non-contiguous model shards. + """ + from ..utils.dist_utils import raise_if_any_rank_failed + + def scalar_bytes(value): + if type(value) is float: + return b"float:" + struct.pack("!d", value) + if type(value) in (str, int, bool, bytes, type(None), torch.dtype, torch.device): + return type(value).__name__.encode() + b":" + repr(value).encode() + if type(value) in (tuple, list): + parts = [scalar_bytes(item) for item in value] + return ( + type(value).__name__.encode() + b":" + b"".join(len(part).to_bytes(8, "big") + part for part in parts) + ) + raise RuntimeError(f"Cannot prove non-tensor checkpoint state of type {type(value).__name__}") + + error, records = None, [] + try: + flattened, _ = flatten_state_dict({"model": state}) + for name, value in sorted(flattened.items()): + if not torch.is_tensor(value): + payload = scalar_bytes(value) + records.append((name, (), (), type(value).__name__, hashlib.sha256(payload).hexdigest(), None)) + continue + if isinstance(value, DTensor): + chunks = value.__create_chunk_list__() + if len(chunks) != 1: + raise RuntimeError(f"Cannot prove model shard geometry for {name}") + offsets, size = tuple(chunks[0].offsets), tuple(chunks[0].sizes) + local = value.to_local() + else: + local = value + offsets, size = (0,) * value.ndim, tuple(value.shape) + if tuple(local.shape) != size: + raise RuntimeError(f"Cannot prove model shard shape for {name}") + from .tensor_bytes import tensor_byte_chunks + + digest = hashlib.sha256() + for chunk in tensor_byte_chunks(local, chunk_bytes=chunk_bytes): + digest.update(chunk) + records.append((name, offsets, size, str(local.dtype), digest.hexdigest(), tuple(value.shape))) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "hashing replicated model shards", group=group) + if not dist.is_initialized(): + return + names = hashlib.sha256("\0".join(item[0] for item in records).encode()).hexdigest() + identities = [None] * dist.get_world_size(group) + dist.all_gather_object(identities, (len(records), names), group=group) + if any(identity != identities[0] for identity in identities): + raise RuntimeError("Model ranks have different checkpoint tensor inventories") + for start in range(0, len(records), batch_size): + gathered = [None] * len(identities) + dist.all_gather_object(gathered, records[start : start + batch_size], group=group) + seen = {} + geometries = {} + global_signatures = {} + for items in gathered: + for name, offsets, size, dtype, digest, global_shape in items: + signature = (dtype, global_shape) + previous = global_signatures.setdefault(name, signature) + if previous != signature: + raise RuntimeError(f"Checkpoint tensor {name} has different global dtype or shape across ranks") + identity = (name, offsets, size) + if identity not in seen: + geometries.setdefault(name, []).append((offsets, size)) + old = seen.setdefault(identity, (dtype, digest)) + if old != (dtype, digest): + raise RuntimeError(f"Model replicas differ for checkpoint tensor {name} at {offsets}") + for name, shards in geometries.items(): + _reject_overlapping_geometries(name, shards) + + +@contextmanager +def recover_frozen_optimizer_ownership(model, checkpoint_path, *, module="", group=None): + """Include source-recorded frozen parameters while building an optimizer. + + The original flags are restored before the first forward. A later layout + proof still checks the reconstructed optimizer groups against the source. + """ + from torch.distributed.checkpoint.state_dict import _get_fqns + + from ..utils.dist_utils import raise_if_any_rank_failed + from .dcp_checkpointer import DistributedCheckpointer + from .layout import optimizer_dir + from .multi_optimizer_checkpoint import MULTI_LAYOUT_MARKER + + changed, error = [], None + try: + path = optimizer_dir(checkpoint_path, module) + reader = DistributedCheckpointer._create_storage_reader(path) + metadata = reader.read_metadata() + if not isinstance(metadata.state_dict_metadata.get(f"optimizer.{MULTI_LAYOUT_MARKER}"), BytesStorageMetadata): + raise RuntimeError("Cannot recover frozen ownership from an unproven optimizer checkpoint") + rank = current_rank() + prefix = f"optimizer.{rank_prefix(rank)}" + inventory_path = "optimizer." + rank_inventory_key(rank) + if isinstance(metadata.state_dict_metadata.get(inventory_path), BytesStorageMetadata): + scratch = {"optimizer": {rank_inventory_key(rank): ""}} + dcp.load(scratch, storage_reader=DistributedCheckpointer._create_storage_reader(path), no_dist=True) + inventory = json.loads(scratch["optimizer"][rank_inventory_key(rank)]) + actual = sorted(k for k in metadata.state_dict_metadata if k.startswith(prefix) and k != inventory_path) + if inventory != actual: + raise RuntimeError("Optimizer checkpoint ownership inventory differs from DCP metadata") + owned_keys = set(load_rank_scalars(DistributedCheckpointer._create_storage_reader, path, rank)) + else: + owned_keys = { + key.removeprefix("optimizer.") + for key in metadata.state_dict_metadata + if key.startswith("optimizer.param_groups.") + } + for name, param in model.named_parameters(remove_duplicate=False): + if not param.requires_grad and any( + f"param_groups.{fqn}.lr" in owned_keys for fqn in _get_fqns(model, name) + ): + if all(param is not existing for existing in changed): + param.requires_grad_(True) + changed.append(param) + except BaseException as exc: + error = exc + try: + raise_if_any_rank_failed(error, "recovering frozen optimizer ownership", group=group) + yield + finally: + for param in changed: + param.requires_grad_(False) diff --git a/veomni/checkpoint/tensor_bytes.py b/veomni/checkpoint/tensor_bytes.py new file mode 100644 index 0000000000..1d4af81507 --- /dev/null +++ b/veomni/checkpoint/tensor_bytes.py @@ -0,0 +1,45 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +"""Bounded, logical-order tensor bytes for checkpoint identity checks.""" + +import torch + + +def tensor_byte_chunks(tensor, *, chunk_bytes): + """Yield CPU byte chunks without materializing a full non-contiguous tensor.""" + if chunk_bytes <= 0: + raise ValueError("chunk_bytes must be positive") + tensor = tensor.detach() + count = max(1, chunk_bytes // max(1, tensor.element_size())) + + def blocks(view): + if view.numel() == 0: + return + if view.numel() <= count: + yield view + return + if view.is_contiguous(): + flat = view.reshape(-1) + for start in range(0, flat.numel(), count): + yield flat[start : start + count] + return + # If a complete trailing slice fits, copy several leading slices at + # once. Otherwise descend one dimension without materializing a pool + # of indices or a full contiguous copy of the original tensor. + trailing = view.numel() // view.shape[0] + if trailing <= count: + rows = max(1, count // trailing) + for start in range(0, view.shape[0], rows): + yield view[start : start + rows] + else: + for index in range(view.shape[0]): + yield from blocks(view[index]) + + for block in blocks(tensor): + flat = block.contiguous().view(-1) + if flat.stride(0) != 1: + # PyTorch can report a one-element strided view as contiguous, + # while dtype reinterpretation still requires unit stride. + flat = torch.empty(flat.numel(), dtype=flat.dtype, device=flat.device).copy_(flat) + yield flat.view(torch.uint8).cpu().numpy().tobytes() diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 8644cf2c9f..82dd77b543 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -29,12 +29,143 @@ :class:`~veomni.trainer.callbacks.global_state_callback.GlobalStateCallback`. """ +import hashlib +import marshal +import math +import types from typing import TYPE_CHECKING, Any, Dict, Optional import torch.distributed as dist +from torch.optim.lr_scheduler import LambdaLR, StepLR from ...checkpoint import CheckpointerBase, build_checkpointer, layout from ...utils import helper +from ...utils.dist_utils import raise_if_any_rank_failed + + +def _agree_checkpoint_value(value, description): + if dist.is_initialized(): + values = [None] * dist.get_world_size() + dist.all_gather_object(values, value) + if any(item != value for item in values): + raise RuntimeError(f"Ranks disagree on checkpoint {description}") + return value + + +def _agree_multi_optimizer_mode(optimizer): + multi = getattr(optimizer, "_is_multi_optimizer", False) is True + return _agree_checkpoint_value(multi, "optimizer mode") + + +def _scheduler_recipe(scheduler, optimizer): + """Stable schedule identity; LambdaLR state_dict omits its callable.""" + from ...optim.lr_scheduler import MultiLRScheduler + + if scheduler is None: + return None + if type(scheduler) is MultiLRScheduler: + if set(scheduler) != set(optimizer.optimizers_dict): + raise RuntimeError("MultiOptimizer scheduler leaves differ") + return { + name: _scheduler_recipe(scheduler[name], optimizer.optimizers_dict[name]) for name in sorted(scheduler) + } + if getattr(optimizer, "_is_multi_optimizer", False) is True: + raise RuntimeError("MultiOptimizer requires a scheduler for each optimizer leaf") + if type(scheduler) not in (LambdaLR, StepLR) or scheduler.optimizer is not optimizer: + raise RuntimeError("MultiOptimizer scheduler implementation or binding is unsupported") + if any(name in vars(scheduler) for name in ("state_dict", "load_state_dict", "step", "get_lr", "get_last_lr")): + raise RuntimeError("MultiOptimizer scheduler methods are overridden") + if type(scheduler) is StepLR: + return {"type": "StepLR", "step_size": scheduler.step_size, "gamma": scheduler.gamma} + + def immutable(value): + if value is None or type(value) in (bool, int, float, str): + return value + if isinstance(value, tuple): + return [immutable(item) for item in value] + if isinstance(value, dict) and all(isinstance(key, str) for key in value): + return {key: immutable(item) for key, item in sorted(value.items())} + raise RuntimeError("MultiOptimizer scheduler captures unsupported mutable state") + + def normalize(code): + constants = tuple(normalize(item) if isinstance(item, types.CodeType) else item for item in code.co_consts) + return code.replace(co_filename="", co_firstlineno=0, co_consts=constants) + + functions = [] + for function in scheduler.lr_lambdas: + if not isinstance(function, types.FunctionType): + raise RuntimeError("MultiOptimizer scheduler callable is unsupported") + if any( + name in function.__globals__ and (name != "math" or function.__globals__[name] is not math) + for name in function.__code__.co_names + ): + raise RuntimeError("MultiOptimizer scheduler depends on unsupported global state") + functions.append( + { + "code": hashlib.sha256(marshal.dumps(normalize(function.__code__))).hexdigest(), + "defaults": immutable(function.__defaults__), + "kwdefaults": immutable(function.__kwdefaults__), + "closure": [immutable(cell.cell_contents) for cell in function.__closure__ or ()], + } + ) + return {"type": "LambdaLR", "functions": functions} + + +def _validate_scheduler_checkpoint(saved, current, global_step): + """Require every scheduler leaf and its step to match a fresh runtime.""" + + def validate(left, right): + if not isinstance(left, dict) or not isinstance(right, dict) or set(left) != set(right): + raise RuntimeError("MultiOptimizer checkpoint scheduler layout differs") + found = "last_epoch" in right + if found: + initial, restored = right["last_epoch"], left["last_epoch"] + if type(initial) is not int or type(restored) is not int or restored != initial + global_step: + raise RuntimeError("MultiOptimizer checkpoint scheduler step differs from trainer step") + if "_step_count" in right: + initial_count, saved_count = right["_step_count"], left["_step_count"] + if ( + type(initial_count) is not int + or type(saved_count) is not int + or saved_count != initial_count + global_step + ): + raise RuntimeError("MultiOptimizer checkpoint scheduler call count differs from trainer step") + for name, value in right.items(): + if isinstance(value, dict): + child_found = validate(left[name], value) + if "last_epoch" not in right and not child_found: + raise RuntimeError("MultiOptimizer checkpoint scheduler leaf has no provable step") + found = child_found or found + elif isinstance(value, (list, tuple)): + if not isinstance(left[name], (list, tuple)) or len(left[name]) != len(value): + raise RuntimeError("MultiOptimizer checkpoint scheduler state shape differs") + if name != "_last_lr" and left[name] != value: + raise RuntimeError("MultiOptimizer checkpoint scheduler configuration differs") + elif name not in {"last_epoch", "_step_count", "_get_lr_called_within_step"} and left[name] != value: + raise RuntimeError("MultiOptimizer checkpoint scheduler configuration differs") + return found + + if not validate(saved, current): + raise RuntimeError("MultiOptimizer checkpoint scheduler has no provable step") + + +def _validate_scheduler_optimizer_lrs(saved, optimizer): + """A restored scheduler and optimizer must agree on the next-step LR.""" + for name, leaf in optimizer.optimizers_dict.items(): + state = saved.get(name) + if not isinstance(state, dict) or not isinstance(state.get("_last_lr"), (list, tuple)): + raise RuntimeError("MultiOptimizer checkpoint scheduler LR state is missing") + last_lrs = state["_last_lr"] + if len(last_lrs) != len(leaf.param_groups): + raise RuntimeError("MultiOptimizer checkpoint scheduler LR groups differ") + for value, group in zip(last_lrs, leaf.param_groups, strict=True): + actual = group["lr"] + if hasattr(value, "numel") and value.numel() == 1: + value = value.item() + if hasattr(actual, "numel") and actual.numel() == 1: + actual = actual.item() + if type(value) not in (int, float) or type(actual) not in (int, float) or value != actual: + raise RuntimeError("MultiOptimizer checkpoint scheduler and optimizer LRs differ") if TYPE_CHECKING: @@ -85,6 +216,9 @@ def __init__(self, runtime: "VeOmniModelRuntime"): self.runtime = runtime self.config: "CheckpointConfig" = runtime.train_args.checkpoint self._last_saved_step: int = -1 + self._loaded_optimizer_global_step = None + self._loaded_optimizer_generation = None + self._last_saved_optimizer_generation = None self.checkpointer: CheckpointerBase = build_checkpointer( ckpt_manager=self.config.manager, dist_backend=runtime.args.accelerator.fsdp_config.fsdp_mode, @@ -185,6 +319,17 @@ def _extra_state(self, state: "TrainerState") -> Dict[str, Any]: "lr_scheduler": None if lr_scheduler is None else lr_scheduler.state_dict(), } extra_state.update(self.runtime.extra_state()) + if getattr(self.runtime.optimizer, "_is_multi_optimizer", False) is True: + if any(key in extra_state for key in ("_veomni_optimizer_global_step", "_veomni_scheduler_recipe")): + raise RuntimeError("Reserved MultiOptimizer checkpoint step key is already in use") + extra_state["_veomni_optimizer_global_step"] = int(state.global_step) + try: + extra_state["_veomni_scheduler_recipe"] = { + "status": "supported", + "value": _scheduler_recipe(lr_scheduler, self.runtime.optimizer), + } + except RuntimeError as exc: + extra_state["_veomni_scheduler_recipe"] = {"status": "unsupported", "reason": str(exc)} return extra_state def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: @@ -197,7 +342,10 @@ def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: def load(self) -> None: """Restore model, optimizer and model-bound extra state from ``load_path``.""" + self._loaded_optimizer_global_step = None + self._loaded_optimizer_generation = None load_dir = self.load_dir() + _agree_checkpoint_value(load_dir, "load path") if load_dir is None: return @@ -206,13 +354,62 @@ def load(self) -> None: "model": self.runtime.model, "optimizer": self.runtime.optimizer, "extra_state": {}, + "_veomni_scheduler_object": self.runtime.lr_scheduler, } - self.checkpointer.load( - load_dir, - state, - **self._checkpointer_kwargs(), - ) - self._load_extra_state(state["extra_state"]) + multi = _agree_multi_optimizer_mode(self.runtime.optimizer) + error = None + try: + self.checkpointer.load(load_dir, state, **self._checkpointer_kwargs()) + except BaseException as exc: + error = exc + if multi: + raise_if_any_rank_failed(error, "loading MultiOptimizer checkpoint") + elif error is not None: + raise error + error, step = None, None + try: + if multi: + if state.get("_veomni_optimizer_identity_verified") is not True: + raise RuntimeError("MultiOptimizer checkpoint identity has not been proven") + step = state["extra_state"].pop("_veomni_optimizer_global_step", None) + if type(step) is not int or step < 0: + raise RuntimeError("MultiOptimizer checkpoint is missing its proven trainer global step") + if "_veomni_scheduler_recipe" not in state["extra_state"]: + raise RuntimeError("MultiOptimizer checkpoint scheduler identity is missing") + recipe = state["extra_state"].pop("_veomni_scheduler_recipe") + if ( + not isinstance(recipe, dict) + or recipe.get("status") != "supported" + or set(recipe) != {"status", "value"} + ): + raise RuntimeError("MultiOptimizer checkpoint scheduler identity is unsupported") + if recipe["value"] != _scheduler_recipe(self.runtime.lr_scheduler, self.runtime.optimizer): + raise RuntimeError("MultiOptimizer checkpoint scheduler identity differs") + saved_scheduler = state["extra_state"].get("lr_scheduler") + scheduler = self.runtime.lr_scheduler + if (scheduler is None) != (saved_scheduler is None) or ( + scheduler is not None and not isinstance(saved_scheduler, dict) + ): + raise RuntimeError("MultiOptimizer checkpoint scheduler state is missing or incompatible") + if scheduler is not None: + _validate_scheduler_checkpoint(saved_scheduler, scheduler.state_dict(), step) + _validate_scheduler_optimizer_lrs(saved_scheduler, self.runtime.optimizer) + self._load_extra_state(state["extra_state"]) + except BaseException as exc: + error = exc + if multi: + raise_if_any_rank_failed(error, "restoring MultiOptimizer trainer step and scheduler") + self._loaded_optimizer_global_step = step + generation = state.get("_veomni_optimizer_checkpoint_generation") + error = ( + RuntimeError("MultiOptimizer checkpoint generation is missing") + if not isinstance(generation, str) or not generation + else None + ) + raise_if_any_rank_failed(error, "checking restored MultiOptimizer generation") + self._loaded_optimizer_generation = generation + elif error is not None: + raise error dist.barrier() logger.info_rank0(f"Load distributed checkpoint from {load_dir} successfully!") @@ -222,14 +419,43 @@ def save_dcp(self, state: "TrainerState") -> None: 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._last_saved_optimizer_generation = None + if self._last_saved_step == state.global_step: + # A failed rewrite may already have invalidated the old DCP marker. + # A later export must not reuse the previous same-step receipt. + self._last_saved_step = -1 + _agree_checkpoint_value((self.config.save_path, state.global_step), "save path and step") + multi = _agree_multi_optimizer_mode(self.runtime.optimizer) + error, extra_state = None, None + try: + extra_state = self._extra_state(state) + except BaseException as exc: + error = exc + if multi: + raise_if_any_rank_failed(error, "preparing MultiOptimizer model extra state") + _agree_checkpoint_value(extra_state["_veomni_scheduler_recipe"], "scheduler recipe") + elif error is not None: + raise error + + def clear_cache(stage): + error = None + try: + helper.empty_cache() + except BaseException as exc: + error = exc + if multi: + raise_if_any_rank_failed(error, stage) + elif error is not None: + raise error + + clear_cache("clearing cache before MultiOptimizer checkpoint save") + generation = self.checkpointer.save( self.config.save_path, { "model": self.runtime.model, "optimizer": self.runtime.optimizer, "extra_state": extra_state, + "_veomni_scheduler_object": self.runtime.lr_scheduler, }, global_steps=state.global_step, save_async=self.config.save_async, @@ -238,8 +464,17 @@ def save_dcp(self, state: "TrainerState") -> None: save_timeout_seconds=self.config.save_timeout_seconds, **self._checkpointer_kwargs(), ) - helper.empty_cache() + if multi: + error = ( + RuntimeError("MultiOptimizer save did not return its checkpoint generation") + if not isinstance(generation, str) or not generation + else None + ) + raise_if_any_rank_failed(error, "checking saved MultiOptimizer generation") + clear_cache("clearing cache after MultiOptimizer checkpoint save") dist.barrier() + if multi: + self._last_saved_optimizer_generation = generation self._last_saved_step = state.global_step logger.info_rank0(f"Distributed checkpoint saved at {self.save_dir(state)} successfully!") diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index 450dbbdf45..9a2790e20b 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -143,7 +143,54 @@ def __init__( self._build_model() self._freeze_model_module() self._build_parallelized_model() - self._build_optimizer() + from ..checkpoint.multi_optimizer_checkpoint import _coordinated_optimizer_flag + from ..utils.dist_utils import raise_if_any_rank_failed + + specs = getattr(self.model, "_fqn2spec_info", None) + has_specs = _coordinated_optimizer_flag(specs is not None, "ExtraParallel model plan presence") + canonical = _coordinated_optimizer_flag( + getattr(self.train_args.checkpoint, "extra_parallel_canonical_layout", False), + "canonical ExtraParallel checkpoint policy", + ) + layout_error = None + try: + if canonical: + from torch.distributed.tensor import Shard + + if not has_specs: + raise RuntimeError( + "Canonical ExtraParallel checkpointing requires an ExtraParallel model plan" + ) + for spec in specs.values(): + if isinstance(spec.placement, Shard) and spec.persistent_fsdp_shard_dim is None: + spec.checkpoint_layout_v2 = True + except BaseException as exc: + layout_error = exc + raise_if_any_rank_failed(layout_error, "configuring ExtraParallel checkpoint layout") + if has_specs: + from ..checkpoint.dcp_checkpointer import _validate_optimizer_model_layout + + layout_error = None + try: + _validate_optimizer_model_layout(self.model) + except BaseException as exc: + layout_error = exc + raise_if_any_rank_failed(layout_error, "validating optimizer model checkpoint layout") + checkpoint_path = self.train_args.checkpoint.load_path + recover_ownership = bool( + checkpoint_path + and self.model_name in ("base", "policy") + and (getattr(self.model, "_fqn2spec_info", None) is not None or self.args.optimizer.type in ("muon",)) + ) + recover_ownership = _coordinated_optimizer_flag(recover_ownership, "frozen optimizer ownership policy") + if recover_ownership: + _coordinated_optimizer_flag(checkpoint_path, "frozen optimizer checkpoint path") + from ..checkpoint.rank_local_optimizer import recover_frozen_optimizer_ownership + + with recover_frozen_optimizer_ownership(self.model, checkpoint_path): + self._build_optimizer() + else: + self._build_optimizer() self._build_model_assets() self.build_checkpoint() diff --git a/veomni/optim/checkpoint_adapter.py b/veomni/optim/checkpoint_adapter.py new file mode 100644 index 0000000000..8f9b876218 --- /dev/null +++ b/veomni/optim/checkpoint_adapter.py @@ -0,0 +1,57 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +"""Explicit extension point for precision-preserving rank-local optimizer DCP.""" + +from contextlib import nullcontext +from dataclasses import dataclass + +from torch.optim import Optimizer + + +_RANK_LOCAL_ADAMW_ADAPTERS = {} + + +@dataclass(frozen=True) +class RankLocalAdamWAdapter: + context_factory: object + algorithm_identity: object + state_dtype: object + state_is_offloaded: object + + +def register_rank_local_adamw(optimizer_type, context_factory, algorithm_identity, state_dtype, state_is_offloaded): + """Register a trusted exact optimizer type and its sparse checkpoint context. + + The context must expose only actually initialized AdamW master/moment state, + preserve parameter-group semantics, not mutate live training state while + serializing, and accept rank-local scalar steps. Registered implementations + are trusted project code; this is not a sandbox for arbitrary callbacks. + The identity callback must return immutable values covering every + non-group setting that can alter updates or restore behavior. The dtype and + offload callbacks describe each parameter's actual step, master and moment + state; the step dtype must describe its scalar, not the moment dtype. + Registration is an internal compatibility contract, not a class-wide opt-in. + """ + if not isinstance(optimizer_type, type) or not issubclass(optimizer_type, Optimizer): + raise TypeError("Rank-local AdamW adapter requires a torch Optimizer subclass") + if not callable(context_factory): + raise TypeError("Rank-local AdamW adapter context must be callable") + if not callable(algorithm_identity): + raise TypeError("Rank-local AdamW adapter algorithm identity must be callable") + if not callable(state_dtype) or not callable(state_is_offloaded): + raise TypeError("Rank-local AdamW state dtype and offload policies must be callable") + if optimizer_type in _RANK_LOCAL_ADAMW_ADAPTERS: + raise ValueError(f"Optimizer checkpoint adapter already registered: {optimizer_type!r}") + _RANK_LOCAL_ADAMW_ADAPTERS[optimizer_type] = RankLocalAdamWAdapter( + context_factory, algorithm_identity, state_dtype, state_is_offloaded + ) + + +def rank_local_adamw_adapter(optimizer): + return _RANK_LOCAL_ADAMW_ADAPTERS.get(type(optimizer)) + + +def rank_local_adamw_context(optimizer): + adapter = rank_local_adamw_adapter(optimizer) + return adapter.context_factory(optimizer) if adapter is not None else nullcontext() diff --git a/veomni/optim/lr_scheduler.py b/veomni/optim/lr_scheduler.py index 500ed280da..fab230bce7 100644 --- a/veomni/optim/lr_scheduler.py +++ b/veomni/optim/lr_scheduler.py @@ -14,6 +14,7 @@ import math +import weakref from typing import TYPE_CHECKING, Dict, Literal from torch.optim.lr_scheduler import LambdaLR @@ -26,6 +27,19 @@ logger = logging.get_logger(__name__) +_CHECKPOINT_SCHEDULER_BINDINGS = weakref.WeakKeyDictionary() +_CHECKPOINT_SCHEDULER_FUNCTIONS = weakref.WeakKeyDictionary() + + +def checkpoint_scheduler_binding(optimizer): + """Find a framework-built scheduler without mutating the optimizer payload.""" + reference = _CHECKPOINT_SCHEDULER_BINDINGS.get(optimizer) + return None if reference is None else reference() + + +def checkpoint_scheduler_functions(scheduler): + """Identify the immutable function objects installed by the framework builder.""" + return _CHECKPOINT_SCHEDULER_FUNCTIONS.get(scheduler) class MultiLRScheduler(dict): @@ -43,9 +57,10 @@ def state_dict(self) -> Dict[str, any]: return {name: sched.state_dict() for name, sched in self.items()} def load_state_dict(self, state_dict: Dict[str, any]) -> None: + if set(state_dict) != set(self): + raise RuntimeError("MultiLRScheduler checkpoint leaves differ from the current optimizer layout") for name, sched in self.items(): - if name in state_dict: - sched.load_state_dict(state_dict[name]) + sched.load_state_dict(state_dict[name]) def get_last_lr(self): # Return the first scheduler's LR for logging consistency @@ -83,15 +98,15 @@ def build_lr_scheduler( lr_warmup_steps = int(train_steps * lr_warmup_ratio) if lr_decay_style == "constant": - return get_constant_schedule_with_warmup( + scheduler = get_constant_schedule_with_warmup( optimizer=optimizer, num_warmup_steps=lr_warmup_steps, lr_start=lr_start, init_lr=lr, ) - if lr_decay_style == "linear": - return get_linear_schedule_with_warmup( + elif lr_decay_style == "linear": + scheduler = get_linear_schedule_with_warmup( optimizer=optimizer, num_warmup_steps=lr_warmup_steps, num_training_steps=train_steps, @@ -99,8 +114,8 @@ def build_lr_scheduler( lr_start=lr_start, ) - if lr_decay_style == "cosine": - return get_cosine_schedule_with_warmup( + elif lr_decay_style == "cosine": + scheduler = get_cosine_schedule_with_warmup( optimizer=optimizer, num_warmup_steps=lr_warmup_steps, num_training_steps=train_steps, @@ -110,7 +125,11 @@ def build_lr_scheduler( lr_start=lr_start, ) - raise ValueError(f"Unknown learning rate decay style: {lr_decay_style}.") + else: + raise ValueError(f"Unknown learning rate decay style: {lr_decay_style}.") + _CHECKPOINT_SCHEDULER_BINDINGS[optimizer] = weakref.ref(scheduler) + _CHECKPOINT_SCHEDULER_FUNCTIONS[scheduler] = tuple(scheduler.lr_lambdas) + return scheduler def get_constant_schedule_with_warmup( diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index dee3c57933..b77805b1dd 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -15,6 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import contextmanager from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union import torch @@ -257,8 +258,78 @@ def step(self, closure=None): p.data.addcdiv_(exp_avg, centered_variance, value=-step_size) +def _verify_registered_adapter_snapshot(model, optimizer, snapshot): + """Require a registered adapter to serialize each live state without alteration.""" + from ..checkpoint.optimizer_resume_layout import _group_option_value + + names = {} + for runtime_name, param in model.named_parameters(remove_duplicate=False): + names.setdefault(id(param), set()).update(_get_fqns(model, runtime_name)) + owned = {id(param) for group in optimizer.param_groups for param in group["params"]} + expected_state_keys = set() + + def same_view(left, right): + if not torch.is_tensor(left) or not torch.is_tensor(right): + return False + if tuple(left.shape) != tuple(right.shape) or left.dtype != right.dtype: + return False + left = left.to_local() if hasattr(left, "to_local") else left + right = right.to_local() if hasattr(right, "to_local") else right + return ( + left.device == right.device + and tuple(left.shape) == tuple(right.shape) + and tuple(left.stride()) == tuple(right.stride()) + and left.storage_offset() == right.storage_offset() + and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr() + ) + + for param, state in optimizer.state.items(): + if id(param) not in owned or id(param) not in names: + raise RuntimeError("Registered optimizer adapter has unowned live state") + for field, value in state.items(): + candidates = [f"state.{name}.{field}" for name in names[id(param)]] + present = [key for key in candidates if key in snapshot] + if not present or any(not same_view(value, snapshot[key]) for key in present): + raise RuntimeError(f"Registered optimizer adapter omitted or altered live state {field}") + expected_state_keys.update(present) + actual_state_keys = {key for key in snapshot if key.startswith("state.")} + if actual_state_keys != expected_state_keys: + raise RuntimeError("Registered optimizer adapter fabricated checkpoint state") + expected_group_keys = set() + for group in optimizer.param_groups: + for param in group["params"]: + for key, value in group.items(): + if key == "params": + continue + candidates = [f"param_groups.{name}.{key}" for name in names.get(id(param), ())] + present = [candidate for candidate in candidates if candidate in snapshot] + if not present or any( + _group_option_value(snapshot[candidate]) != _group_option_value(value) for candidate in present + ): + raise RuntimeError(f"Registered optimizer adapter omitted or altered group option {key}") + expected_group_keys.update(present) + actual_group_keys = {key for key in snapshot if key.startswith("param_groups.")} + if actual_group_keys != expected_group_keys: + raise RuntimeError("Registered optimizer adapter fabricated checkpoint group options") + + +def _adapter_live_state_versions(optimizer): + versions = [] + for param, state in optimizer.state.items(): + fields = [] + for name, value in sorted(state.items()): + if not torch.is_tensor(value): + raise RuntimeError("Registered optimizer adapter state must contain tensors") + local = value.to_local() if hasattr(value, "to_local") else value + fields.append((name, id(value), value._version, local.untyped_storage().data_ptr(), local._version)) + versions.append((id(param), tuple(fields))) + return tuple(sorted(versions)) + + def _empty_optimizer_state_dict(model, opt): - if getattr(opt.state_dict, "__func__", None) is not Optimizer.state_dict: + from .checkpoint_adapter import rank_local_adamw_adapter + + if getattr(opt.state_dict, "__func__", None) is not Optimizer.state_dict and rank_local_adamw_adapter(opt) is None: raise RuntimeError("Cannot safely serialize an empty optimizer with a custom state_dict") packed_groups = [] id_to_name = {} @@ -304,6 +375,51 @@ def capture_ids(_leaf, packed): return _flatten_optim_state_dict({"state": state, "param_groups": groups}) +@contextmanager +def _dcp_restore_frozen_ownership(optimizer): + frozen = [param for group in optimizer.param_groups for param in group["params"] if not param.requires_grad] + changed = [] + try: + for param in frozen: + param.requires_grad_(True) + changed.append(param) + yield + finally: + for param in changed: + param.requires_grad_(False) + + +@contextmanager +def _dcp_avoid_synthetic_optimizer_step(optimizer): + param = next((param for group in optimizer.param_groups for param in group["params"]), None) + if param is None: + raise RuntimeError("Cannot restore a MultiOptimizer leaf without parameters") + sentinel = {} if not optimizer.state else None + if sentinel is not None: + optimizer.state[param] = sentinel + called = False + + def before_load(current, state_dict): + nonlocal called + called = True + if sentinel is not None: + if current.state.get(param) is not sentinel or sentinel: + raise RuntimeError("Optimizer state changed while preparing checkpoint restore") + current.state.pop(param) + return {**state_dict, "state": {k: v for k, v in state_dict["state"].items() if not isinstance(v, dict) or v}} + + before_load._veomni_internal_dcp_filter = True + handle = optimizer.register_load_state_dict_pre_hook(before_load, prepend=True) + try: + yield + if not called: + raise RuntimeError("Optimizer bypassed checkpoint load-state hooks") + finally: + handle.remove() + if sentinel is not None and optimizer.state.get(param) is sentinel: + optimizer.state.pop(param) + + class MultiOptimizer(Optimizer, Stateful): """ A container that handles multiple optimizers (for extra_parallel and non-extra_parallel parameters when extra_parallel+fsdp2 is enabled) @@ -325,36 +441,57 @@ def __init__( self._is_multi_optimizer: bool = True self.key_names = key_names - def step(self) -> None: + def step(self, closure=None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() for opt in self.optimizers_dict.values(): opt.step() + return loss - def zero_grad(self) -> None: + def zero_grad(self, set_to_none: bool = True) -> None: for opt in self.optimizers_dict.values(): - opt.zero_grad() + opt.zero_grad(set_to_none=set_to_none) def state_dict( self, ) -> Dict[str, Any]: - # get the flatten state dict for multi-optimizer - merged: Dict[str, Any] = {} - for name in self.key_names: - opt = self.optimizers_dict.get(name) - sd = get_optimizer_state_dict(self.model, opt, options=StateDictOptions(flatten_optimizer_state_dict=True)) - # check for key clashes before merging - overlap = set(merged.keys()) & set(sd.keys()) - if overlap: - raise KeyError( - f"Key clash detected while merging state dict for optimizer '{name}': {', '.join(sorted(overlap))}" + # A snapshot must not let a raw DCP load overwrite live optimizer + # tensors before ``load_state_dict`` can reject unproven provenance. + import copy + + from .checkpoint_adapter import rank_local_adamw_adapter + from .muon import DistributedMuon + + for optimizer in self.optimizers_dict.values(): + if ( + ( + type(optimizer) not in (torch.optim.AdamW, AnyPrecisionAdamW, DistributedMuon) + and rank_local_adamw_adapter(optimizer) is None ) - else: - logger.info_rank0("No clashes when merging MultiOptimizer state dicts") - merged.update(sd) - - return merged + or any(name in vars(optimizer) for name in ("state_dict", "load_state_dict")) + or any( + getattr(optimizer, name, None) + for name in ( + "_optimizer_state_dict_pre_hooks", + "_optimizer_state_dict_post_hooks", + "_optimizer_load_state_dict_pre_hooks", + "_optimizer_load_state_dict_post_hooks", + ) + ) + ): + raise RuntimeError("Cannot construct a safe raw MultiOptimizer snapshot with serialization hooks") + return { + key: value.detach().clone() if torch.is_tensor(value) else copy.deepcopy(value) + for key, value in self._sparse_state_dict().items() + } def _sparse_state_dict( self, + *, + coordinate=False, + group=None, ) -> Dict[str, Any]: """Return a flattened optimizer state dict without materializing state. @@ -363,37 +500,79 @@ def _sparse_state_dict( The groups must still be saved because their LR may have been updated by a scheduler even when no parameter has received a gradient. """ + if coordinate: + from ..checkpoint.multi_optimizer_checkpoint import _coordinated_optimizer_flag + + _coordinated_optimizer_flag(tuple(self.key_names), "optimizer leaf order", group=group) merged: Dict[str, Any] = {} for name in self.key_names: - opt = self.optimizers_dict.get(name) - if not opt.state: - sd = _empty_optimizer_state_dict(self.model, opt) - else: - sd = get_optimizer_state_dict( - self.model, opt, options=StateDictOptions(flatten_optimizer_state_dict=True) - ) - overlap = set(merged.keys()) & set(sd.keys()) - if overlap: - raise KeyError( - f"Key clash detected while merging state dict for optimizer '{name}': {', '.join(sorted(overlap))}" - ) - else: - logger.info_rank0("No clashes when merging MultiOptimizer state dicts") + error, sd = None, None + try: + opt = self.optimizers_dict.get(name) + from .checkpoint_adapter import rank_local_adamw_adapter, rank_local_adamw_context + + registered = rank_local_adamw_adapter(opt) is not None + before = _adapter_live_state_versions(opt) if registered else None + with rank_local_adamw_context(opt): + if not opt.state: + sd = _empty_optimizer_state_dict(self.model, opt) + else: + sd = get_optimizer_state_dict( + self.model, opt, options=StateDictOptions(flatten_optimizer_state_dict=True) + ) + if registered: + if _adapter_live_state_versions(opt) != before: + raise RuntimeError("Registered optimizer adapter mutated live state while serializing") + _verify_registered_adapter_snapshot(self.model, opt, sd) + overlap = set(merged) & set(sd) + if overlap: + raise KeyError( + f"Key clash detected while merging state dict for optimizer '{name}': {', '.join(sorted(overlap))}" + ) + except BaseException as exc: + error = exc + if coordinate: + from ..utils.dist_utils import raise_if_any_rank_failed + + raise_if_any_rank_failed(error, f"serializing optimizer leaf {name}", group=group) + elif error is not None: + raise error merged.update(sd) return merged def load_state_dict(self, state_dict: Dict[str, Any]) -> None: - # Feed the same merged flattened dict to each sub-optimizer; PyTorch will - # pick out only the entries for parameters that belong to that optimizer. + raise RuntimeError( + "Unverified MultiOptimizer state cannot establish group and scheduler identity; " + "restore through DistributedCheckpointer or an explicitly verified migration." + ) + + def _dcp_load_state_dict(self, state_dict: Dict[str, Any], *, group=None) -> None: + from ..checkpoint.multi_optimizer_checkpoint import _coordinated_optimizer_flag + + _coordinated_optimizer_flag(tuple(self.key_names), "optimizer leaf order", group=group) for name in self.key_names: - opt = self.optimizers_dict.get(name) - set_optimizer_state_dict( - self.model, - opt, - optim_state_dict=state_dict, - options=StateDictOptions(flatten_optimizer_state_dict=True), - ) + error = None + try: + optimizer = self.optimizers_dict[name] + from .checkpoint_adapter import rank_local_adamw_context + + with ( + rank_local_adamw_context(optimizer), + _dcp_restore_frozen_ownership(optimizer), + _dcp_avoid_synthetic_optimizer_step(optimizer), + ): + set_optimizer_state_dict( + self.model, + optimizer, + optim_state_dict=state_dict, + options=StateDictOptions(flatten_optimizer_state_dict=True), + ) + except BaseException as exc: + error = exc + from ..utils.dist_utils import raise_if_any_rank_failed + + raise_if_any_rank_failed(error, f"loading optimizer leaf {name}", group=group) def register_step_pre_hook(self, hook): return [opt.register_step_pre_hook(hook) for opt in self.optimizers_dict.values()] @@ -501,6 +680,8 @@ def build_optimizer( param groups. ``muon_kwargs`` overrides the config key by key, for callers that build Muon without an ``OptimizerConfig`` at all. """ + if isinstance(betas, list): + betas = tuple(betas) if optimizer_type == "muon": if optimizer_config is not None: muon_kwargs = {**_collect_muon_kwargs(optimizer_config), **(muon_kwargs or {})} @@ -814,6 +995,8 @@ def build_extra_parallel_fsdp2_optimizer( - Each group's params are automatically split into ExtraParallel1, ExtraParallel2, ... and non-ExtraParallel based on DTensor mesh - Custom learning rates and other optimizer settings are preserved per group """ + if isinstance(betas, list): + betas = tuple(betas) parallel_state = get_parallel_state() # Collect all ExtraParallel and non-ExtraParallel parameters across all groups From b47d315650d9001d76db7f91fb5652ab20e8c660 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sun, 1 Nov 2026 09:57:51 +0800 Subject: [PATCH 013/120] [checkpoint] coordinate strict load preflight failures --- .../test_multi_optimizer_proven_resume.py | 92 +++++++++++++++++++ veomni/checkpoint/dcp_checkpointer.py | 84 ++++++++++------- 2 files changed, 145 insertions(+), 31 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index f165d5173d..e59460f0eb 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -24,6 +24,98 @@ def _model_and_optimizer(*, split=True): return model, MultiOptimizer(model, {"merged": merged}, ["merged"]) +def _load_preflight_failure_worker(rank, checkpoint, failure, endpoint): + from datetime import timedelta + from unittest.mock import patch + + import torch.distributed as dist + + import veomni.checkpoint.dcp_checkpointer as checkpoint_module + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + torch.manual_seed(27) + source, optimizer = _model_and_optimizer() + for param in source.parameters(): + param.grad = torch.ones_like(param) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(checkpoint, {"model": source, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _model_and_optimizer() + original_model = checkpoint_module._ProvenModelState + original_optimizer = checkpoint_module._ProvenMultiOptimizerState + original_reader = DistributedCheckpointer._create_storage_reader + original_needs_marker = checkpoint_module._needs_extra_parallel_layout_marker + injected = False + + def model_state(*args, **kwargs): + nonlocal injected + if rank == 1 and failure == "model_state": + injected = True + raise OSError("injected model state failure") + return original_model(*args, **kwargs) + + def optimizer_state(*args, **kwargs): + nonlocal injected + if rank == 1 and failure == "optimizer_state": + injected = True + raise OSError("injected optimizer state failure") + return original_optimizer(*args, **kwargs) + + def reader(path): + import inspect + + nonlocal injected + callers = {frame.function for frame in inspect.stack()} + direct = not callers.intersection( + {"_preflight_multi_optimizer_layout", "_prepare_multi_optimizer_saved_states"} + ) + kind = "optimizer_reader" if str(path).endswith("optimizer") else "model_reader" + if rank == 1 and direct and failure == kind: + injected = True + raise OSError(f"injected {kind} failure") + return original_reader(path) + + def needs_marker(specs): + nonlocal injected + if failure == "layout_policy": + import inspect + + callers = {frame.function for frame in inspect.stack() if frame.filename == checkpoint_module.__file__} + if "_reject_unproven_extra_parallel_optimizer" in callers: + return False + if rank == 1: + injected = True + return rank == 1 + return original_needs_marker(specs) + + with ( + patch.object(checkpoint_module, "_ProvenModelState", side_effect=model_state), + patch.object(checkpoint_module, "_ProvenMultiOptimizerState", side_effect=optimizer_state), + patch.object(checkpoint_module, "_needs_extra_parallel_layout_marker", side_effect=needs_marker), + patch.object(DistributedCheckpointer, "_create_storage_reader", side_effect=reader), + pytest.raises((OSError, RuntimeError), match="injected|failed on another rank|marker policy differs"), + ): + DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + if rank == 1: + assert injected, f"the {failure} fault did not reach its intended load stage" + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "failure", ["model_reader", "optimizer_reader", "model_state", "optimizer_state", "layout_policy"] +) +def test_proven_load_preflight_failure_stops_all_ranks(tmp_path, failure): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_load_preflight_failure_worker, args=(str(tmp_path), failure), nprocs=2) + + def test_multi_optimizer_async_checkpoint_rejected_before_writes(tmp_path): model, optimizer = _model_and_optimizer() with pytest.raises(RuntimeError, match="require synchronous DCP save"): diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 4a7b7989dd..adde74ab12 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -2033,8 +2033,15 @@ def _optimizer_state() -> "OptimizerState": except BaseException as exc: error = exc raise_if_any_rank_failed(error, "inspecting checkpoint layout", group=process_group) - specs = getattr(state["model"], "_fqn2spec_info", None) - if _needs_extra_parallel_layout_marker(specs): + specs, needs_marker, error = None, False, None + try: + specs = getattr(state["model"], "_fqn2spec_info", None) + needs_marker = _needs_extra_parallel_layout_marker(specs) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking ExtraParallel marker policy", group=process_group) + needs_marker = _coordinated_optimizer_flag(needs_marker, "ExtraParallel marker policy", group=process_group) + if needs_marker: marker_error = None try: source_path = fused_dir if fused_dir is not None else weights_dir(path, module) @@ -2091,21 +2098,25 @@ def _optimizer_state() -> "OptimizerState": parallel_state=parallel_state, validate_only=True, ) - model_state = _ProvenModelState( - state["model"], - trainable_only=trainable_only, - parallel_state=parallel_state, - optimizer=optimizer, - group=process_group, - ) + model_state, reader, planner, error = None, None, None, None + try: + model_state = _ProvenModelState( + state["model"], + trainable_only=trainable_only, + parallel_state=parallel_state, + optimizer=optimizer, + group=process_group, + ) + reader = cls._create_storage_reader(model_path) + planner = _ModelStrictLoadPlanner(strict_model=not trainable_only, expected_save_id=model_id) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "preparing proven model checkpoint state", group=process_group) dcp.load( state_dict={"model": model_state}, - storage_reader=cls._create_storage_reader(model_path), + storage_reader=reader, process_group=process_group, - planner=_ModelStrictLoadPlanner( - strict_model=not trainable_only, - expected_save_id=model_id, - ), + planner=planner, ) error = ( RuntimeError("Optimizer object changed while loading model state") @@ -2123,8 +2134,12 @@ def _optimizer_state() -> "OptimizerState": group=process_group, frozen_digest=frozen_digest, ) - if current != expected_layout: - raise RuntimeError("MultiOptimizer layout changed while loading model state") + error = ( + RuntimeError("MultiOptimizer layout changed while loading model state") + if current != expected_layout + else None + ) + raise_if_any_rank_failed(error, "rechecking MultiOptimizer layout", group=process_group) _prepare_multi_optimizer_saved_states( state["model"], optimizer, @@ -2134,24 +2149,31 @@ def _optimizer_state() -> "OptimizerState": group=process_group, parallel_state=parallel_state, ) - optimizer_state = _ProvenMultiOptimizerState( - model=state["model"], - optimizer=optimizer, - parallel_state=parallel_state, - load=True, - generation=generation, - scheduler_snapshot=snapshot, - scheduler_object=scheduler_object, - extra_digest=extra_digest, - expected_layout=expected_layout, - group=process_group, - frozen_digest=frozen_digest, - ) + optimizer_state, reader, planner, error = None, None, None, None + try: + optimizer_state = _ProvenMultiOptimizerState( + model=state["model"], + optimizer=optimizer, + parallel_state=parallel_state, + load=True, + generation=generation, + scheduler_snapshot=snapshot, + scheduler_object=scheduler_object, + extra_digest=extra_digest, + expected_layout=expected_layout, + group=process_group, + frozen_digest=frozen_digest, + ) + reader = cls._create_storage_reader(optimizer_path) + planner = _ModelStrictLoadPlanner(strict_model=False, expected_save_id=optimizer_id) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "preparing proven optimizer checkpoint state", group=process_group) dcp.load( state_dict={"optimizer": optimizer_state}, - storage_reader=cls._create_storage_reader(optimizer_path), + storage_reader=reader, process_group=process_group, - planner=_ModelStrictLoadPlanner(strict_model=False, expected_save_id=optimizer_id), + planner=planner, ) if extra is not None: state["extra_state"].update(extra) From c36b855b39b6e6b0ddf99827c86fbf5575a73b2e Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:34:15 +0800 Subject: [PATCH 014/120] [checkpoint] publish optimizer receipt after barrier --- .../test_multi_optimizer_proven_resume.py | 30 +++++++++++++++++++ veomni/models/checkpoint/manager.py | 5 ++-- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index e59460f0eb..8343e4dfb5 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -128,6 +128,36 @@ def test_multi_optimizer_async_checkpoint_rejected_before_writes(tmp_path): assert not list(tmp_path.iterdir()) +def test_failed_final_load_barrier_does_not_certify_multi_optimizer(): + from unittest.mock import patch + + from veomni.models.checkpoint.manager import ModelCheckpointManager, _scheduler_recipe + + model, optimizer = _model_and_optimizer() + manager = object.__new__(ModelCheckpointManager) + manager.runtime = SimpleNamespace( + model=model, optimizer=optimizer, lr_scheduler=None, load_extra_state=lambda _: None + ) + manager.load_dir = lambda: "/checkpoint/global_step_2" + manager.wait_for_pending_save = lambda: None + manager._checkpointer_kwargs = lambda: {} + + def load(_path, state, **_kwargs): + state["_veomni_optimizer_identity_verified"] = True + state["_veomni_optimizer_checkpoint_generation"] = "generation" + state["extra_state"].update( + _veomni_optimizer_global_step=2, + _veomni_scheduler_recipe={"status": "supported", "value": _scheduler_recipe(None, optimizer)}, + ) + + manager.checkpointer = SimpleNamespace(load=load) + with patch("torch.distributed.barrier", side_effect=RuntimeError("injected final barrier failure")): + with pytest.raises(RuntimeError, match="final barrier failure"): + manager.load() + assert manager._loaded_optimizer_global_step is None + assert manager._loaded_optimizer_generation is None + + @pytest.mark.parametrize("mode", ["betas", "list", "nested"]) def test_shared_mutable_group_options_rejected_before_save(tmp_path, mode): model = torch.nn.Linear(2, 2) diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 82dd77b543..796b464751 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -399,7 +399,6 @@ def load(self) -> None: error = exc if multi: raise_if_any_rank_failed(error, "restoring MultiOptimizer trainer step and scheduler") - self._loaded_optimizer_global_step = step generation = state.get("_veomni_optimizer_checkpoint_generation") error = ( RuntimeError("MultiOptimizer checkpoint generation is missing") @@ -407,10 +406,12 @@ def load(self) -> None: else None ) raise_if_any_rank_failed(error, "checking restored MultiOptimizer generation") - self._loaded_optimizer_generation = generation elif error is not None: raise error dist.barrier() + if multi: + self._loaded_optimizer_global_step = step + self._loaded_optimizer_generation = generation logger.info_rank0(f"Load distributed checkpoint from {load_dir} successfully!") def save_dcp(self, state: "TrainerState") -> None: From 54e0f32c60454b603ca8374418b81101e61abfcf Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:38:30 +0800 Subject: [PATCH 015/120] [checkpoint] invalidate file state on load barrier failure --- .../test_multi_optimizer_proven_resume.py | 5 ++++- veomni/models/checkpoint/manager.py | 15 ++++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 8343e4dfb5..c46fdbea7d 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -129,11 +129,13 @@ def test_multi_optimizer_async_checkpoint_rejected_before_writes(tmp_path): def test_failed_final_load_barrier_does_not_certify_multi_optimizer(): - from unittest.mock import patch + from unittest.mock import MagicMock, patch from veomni.models.checkpoint.manager import ModelCheckpointManager, _scheduler_recipe model, optimizer = _model_and_optimizer() + leaf = next(iter(optimizer.optimizers_dict.values())) + leaf.checkpoint_abort_file_backed_state = MagicMock() manager = object.__new__(ModelCheckpointManager) manager.runtime = SimpleNamespace( model=model, optimizer=optimizer, lr_scheduler=None, load_extra_state=lambda _: None @@ -156,6 +158,7 @@ def load(_path, state, **_kwargs): manager.load() assert manager._loaded_optimizer_global_step is None assert manager._loaded_optimizer_generation is None + leaf.checkpoint_abort_file_backed_state.assert_called_once() @pytest.mark.parametrize("mode", ["betas", "list", "nested"]) diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 796b464751..0b9a498723 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -408,7 +408,20 @@ def load(self) -> None: raise_if_any_rank_failed(error, "checking restored MultiOptimizer generation") elif error is not None: raise error - dist.barrier() + try: + dist.barrier() + except BaseException: + optimizer = state.get("optimizer") + leaves = getattr(optimizer, "optimizers_dict", {}) + targets = [optimizer, *(leaves.values() if isinstance(leaves, dict) else ())] + for leaf in targets: + abort = getattr(leaf, "checkpoint_abort_file_backed_state", None) + if callable(abort): + try: + abort() + except Exception: + pass # Keep the original barrier error. + raise if multi: self._loaded_optimizer_global_step = step self._loaded_optimizer_generation = generation From f7831f6fc5e1fe2d74dfd326f3142135ec63c7e1 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:47:43 +0800 Subject: [PATCH 016/120] [ci] declare CPU plans for rank-local checkpoint tests --- .../test_multi_optimizer_proven_resume.py | 19 +++++++++++++++++++ .../checkpoints/test_rank_local_optimizer.py | 14 ++++++++++++++ tests/checkpoints/test_rank_local_ple.py | 6 ++++++ 3 files changed, 39 insertions(+) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index c46fdbea7d..11290ddc2f 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -8,12 +8,16 @@ import pytest import torch +from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.optim.lr_scheduler import build_lr_scheduler from veomni.optim.muon import DistributedMuon from veomni.optim.optimizer import AnyPrecisionAdamW, MultiOptimizer +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + def _model_and_optimizer(*, split=True): model = torch.nn.Linear(2, 2) if split: @@ -107,6 +111,7 @@ def needs_marker(specs): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize( "failure", ["model_reader", "optimizer_reader", "model_state", "optimizer_state", "layout_policy"] ) @@ -274,6 +279,7 @@ def _leaf_order_mismatch_worker(rank, mode, endpoint): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("mode", ["save", "load"]) def test_optimizer_leaf_order_mismatch_stops_all_ranks(mode): from tests.tools.launch_utils import spawn_with_rendezvous @@ -321,6 +327,7 @@ def build(mismatched): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("mode", ["save", "load"]) @pytest.mark.parametrize("mismatch", ["count", "order"]) def test_high_level_leaf_mismatch_rejected_before_model_mutation(tmp_path, mismatch, mode): @@ -414,6 +421,7 @@ def _proven_staged_resume_worker(rank, checkpoint, stage, endpoint): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_proven_staged_checkpoint_restores_and_updates_on_two_ranks(tmp_path): from tests.tools.launch_utils import spawn_with_rendezvous @@ -449,6 +457,7 @@ def make_writer(cls, path): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("target", ["ckpt", "optimizer"]) def test_writer_factory_failure_propagates_before_dcp_collective(tmp_path, target): from tests.tools.launch_utils import spawn_with_rendezvous @@ -509,6 +518,7 @@ def reader(cls, path): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("rank_local", [False, True]) def test_inventory_reader_failure_propagates_before_collective(tmp_path, rank_local): from tests.tools.launch_utils import spawn_with_rendezvous @@ -563,6 +573,7 @@ def clear_cache(): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("fail_at", [1, 2]) def test_checkpoint_cache_failure_propagates_before_next_collective(fail_at): from tests.tools.launch_utils import spawn_with_rendezvous @@ -596,6 +607,7 @@ def fail_hash(*_args, **_kwargs): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_inventory_hash_failure_propagates_before_collective(): from tests.tools.launch_utils import spawn_with_rendezvous @@ -623,6 +635,7 @@ def _trainability_mismatch_worker(rank, checkpoint, endpoint): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_rank_asymmetric_trainability_is_rejected_before_save(tmp_path): from tests.tools.launch_utils import spawn_with_rendezvous @@ -651,6 +664,7 @@ def _tied_parameter_order_worker(rank, endpoint): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_tied_parameter_registration_order_does_not_change_trainability(): from tests.tools.launch_utils import spawn_with_rendezvous @@ -690,6 +704,7 @@ def _save_config_mismatch_worker(rank, checkpoint, mismatch, endpoint): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("mismatch", ["model", "extra", "stage", "stage_path", "module", "async"]) def test_checkpoint_save_rank_config_mismatch_fails_all_ranks(tmp_path, mismatch): from tests.tools.launch_utils import spawn_with_rendezvous @@ -725,6 +740,7 @@ def synchronize(): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_save_finalize_error_propagates_to_all_ranks(tmp_path): from tests.tools.launch_utils import spawn_with_rendezvous @@ -771,6 +787,7 @@ def _build_parallelized_model(self): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("mode", ["plan", "flag", "missing"]) def test_runtime_checkpoint_plan_mismatch_fails_all_ranks(mode): from tests.tools.launch_utils import spawn_with_rendezvous @@ -778,6 +795,7 @@ def test_runtime_checkpoint_plan_mismatch_fails_all_ranks(mode): spawn_with_rendezvous(_runtime_plan_mismatch_worker, args=(mode,), nprocs=2) +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("mode", ["save", "load"]) def test_trainable_only_mismatch_fails_all_ranks(tmp_path, mode): from tests.tools.launch_utils import spawn_with_rendezvous @@ -1925,6 +1943,7 @@ def _replicated_optimizer_metadata_worker(rank, checkpoint, divergent_field, end dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("field", ["lr", "step"]) def test_replicated_optimizer_metadata_must_agree_before_save(tmp_path, field): from tests.tools.launch_utils import spawn_with_rendezvous diff --git a/tests/checkpoints/test_rank_local_optimizer.py b/tests/checkpoints/test_rank_local_optimizer.py index 7ea291fe23..b63d70ef03 100644 --- a/tests/checkpoints/test_rank_local_optimizer.py +++ b/tests/checkpoints/test_rank_local_optimizer.py @@ -12,6 +12,7 @@ from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import Replicate, Shard, distribute_tensor +from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.checkpoint.rank_local_optimizer import ( local_optimizer_state, @@ -24,6 +25,9 @@ from veomni.optim.optimizer import MultiOptimizer +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + def _roundtrip_worker(rank, root, inactive_owner): dist.init_process_group( "gloo", init_method=f"file://{root}/init", rank=rank, world_size=4, timeout=timedelta(seconds=30) @@ -66,6 +70,7 @@ def _roundtrip_worker(rank, root, inactive_owner): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) @pytest.mark.parametrize("inactive_owner", [False, True]) def test_rank_local_state_keeps_owner_steps_and_absence(tmp_path, inactive_owner): mp.spawn(_roundtrip_worker, args=(str(tmp_path), inactive_owner), nprocs=4, join=True) @@ -123,6 +128,7 @@ def build(): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) @pytest.mark.parametrize("inactive_owner", [False, True]) def test_high_level_rank_owned_optimizer_cold_resume(tmp_path, inactive_owner): mp.spawn(_optimizer_resume_worker, args=(str(tmp_path), inactive_owner), nprocs=4, join=True) @@ -219,6 +225,7 @@ def advance(gradient, *, include_bias=True): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) @pytest.mark.parametrize("inactive_owner", [False, True]) @pytest.mark.parametrize("kind", ["adam", "muon"]) @pytest.mark.parametrize("canonical", [False, True]) @@ -237,6 +244,7 @@ def test_rank_owned_optimizer_fresh_process_resume(tmp_path, inactive_owner, kin ) +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) def test_canonical_same_axis_muon_fresh_process_resume(tmp_path): mp.spawn(_cold_optimizer_worker, args=(str(tmp_path), False, "save", "muon", True, True), nprocs=4, join=True) mp.spawn(_cold_optimizer_worker, args=(str(tmp_path), False, "load", "muon", True, True), nprocs=4, join=True) @@ -257,6 +265,7 @@ def _replica_check_worker(rank, root): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_divergent_replicas_rejected_before_checkpoint_deduplication(tmp_path): mp.spawn(_replica_check_worker, args=(str(tmp_path),), nprocs=2, join=True) @@ -286,6 +295,7 @@ def _global_tensor_metadata_mismatch_worker(rank, root, mismatch): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("mismatch", ["dtype", "shape"]) def test_checkpoint_global_tensor_metadata_mismatch_fails_all_ranks(tmp_path, mismatch): mp.spawn(_global_tensor_metadata_mismatch_worker, args=(str(tmp_path), mismatch), nprocs=2, join=True) @@ -317,6 +327,7 @@ def _overlapping_tensor_shards_worker(rank, root, same_offset): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) @pytest.mark.parametrize("same_offset", [False, True]) def test_checkpoint_overlapping_shard_geometries_fail_all_ranks(tmp_path, same_offset): mp.spawn(_overlapping_tensor_shards_worker, args=(str(tmp_path), same_offset), nprocs=2, join=True) @@ -362,6 +373,7 @@ def _owner_trainability_worker(rank, root, within_owner_mismatch): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) @pytest.mark.parametrize("within_owner_mismatch", [False, True]) def test_trainability_may_differ_between_owners_but_not_replicas(tmp_path, within_owner_mismatch): mp.spawn(_owner_trainability_worker, args=(str(tmp_path), within_owner_mismatch), nprocs=4, join=True) @@ -423,6 +435,7 @@ def build(): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) def test_frozen_owned_parameter_rebuilt_on_fresh_process(tmp_path): mp.spawn(_frozen_resume_worker, args=(str(tmp_path), "save"), nprocs=4, join=True) mp.spawn(_frozen_resume_worker, args=(str(tmp_path), "load"), nprocs=4, join=True) @@ -482,5 +495,6 @@ def build(): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_one_rank_missing_optimizer_payload_fails_all_before_model_copy(tmp_path): mp.spawn(_missing_rank_payload_worker, args=(str(tmp_path),), nprocs=2, join=True) diff --git a/tests/checkpoints/test_rank_local_ple.py b/tests/checkpoints/test_rank_local_ple.py index 05d2f45aa4..ae97e4ba3b 100644 --- a/tests/checkpoints/test_rank_local_ple.py +++ b/tests/checkpoints/test_rank_local_ple.py @@ -4,17 +4,22 @@ from datetime import timedelta from types import SimpleNamespace +import pytest import torch import torch.distributed as dist import torch.multiprocessing as mp from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import Shard, distribute_tensor +from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.distributed.parallel_plan import SpecInfo from veomni.optim.optimizer import MultiOptimizer +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + def _worker(rank, root, phase): dist.init_process_group( "gloo", init_method=f"file://{root}/ple-{phase}-init", rank=rank, world_size=4, timeout=timedelta(seconds=40) @@ -67,6 +72,7 @@ def snapshot(): dist.destroy_process_group() +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) def test_persistent_ple_owner_local_optimizer_fresh_process(tmp_path): mp.spawn(_worker, args=(str(tmp_path), "save"), nprocs=4, join=True) mp.spawn(_worker, args=(str(tmp_path), "load"), nprocs=4, join=True) From 139299ca0bd093bf4a3d9761d42e0673a90c4980 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sun, 1 Nov 2026 23:36:32 +0800 Subject: [PATCH 017/120] [ci] allow CUDA fixture in process-group test --- tests/special_sanity/check_device_api_usage.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/special_sanity/check_device_api_usage.py b/tests/special_sanity/check_device_api_usage.py index bcc5d46182..a515185c04 100644 --- a/tests/special_sanity/check_device_api_usage.py +++ b/tests/special_sanity/check_device_api_usage.py @@ -61,6 +61,9 @@ "tests/ci/test_accelerator.py", # EP option tests use literal backend names to verify cross-device rejection. "tests/ci/test_ep_hccl_options.py", + # Gloo failure-propagation tests mock an available CUDA device to verify + # that the process-group backend still selects CPU tensors. + "tests/utils/test_dist_utils_group.py", # Implicit-CUDA-sync gate. Calls ``torch.cuda.{get,set}_sync_debug_mode`` # directly because the API is intrinsically CUDA-only and has no # ``veomni.utils.device`` equivalent; the test is gated on From 3f1a11449a63865c2033fb2d605bc5192b5ceec4 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 02:15:27 +0800 Subject: [PATCH 018/120] [checkpoint] preserve positional configuration compatibility --- tests/arguments/test_checkpoint_config.py | 34 +++++++++++++++++++++++ veomni/arguments/arguments_types.py | 13 ++++----- 2 files changed, 40 insertions(+), 7 deletions(-) create mode 100644 tests/arguments/test_checkpoint_config.py diff --git a/tests/arguments/test_checkpoint_config.py b/tests/arguments/test_checkpoint_config.py new file mode 100644 index 0000000000..e9f1c18bb5 --- /dev/null +++ b/tests/arguments/test_checkpoint_config.py @@ -0,0 +1,34 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0. + +import pytest + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.arguments.arguments_types import CheckpointConfig + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +def test_checkpoint_config_preserves_legacy_positional_arguments(): + config = CheckpointConfig("checkpoint-dir", "dcp", False) + assert config.output_dir == "checkpoint-dir" + assert config.manager == "dcp" + assert config.save_async is False + assert config.extra_parallel_canonical_layout is False + + +def test_checkpoint_config_accepts_explicit_extra_parallel_layout(): + config = CheckpointConfig("checkpoint-dir", extra_parallel_canonical_layout=True) + assert config.output_dir == "checkpoint-dir" + assert config.extra_parallel_canonical_layout is True + + +def test_checkpoint_config_parser_preserves_explicit_layout(): + from veomni.arguments.parser import _instantiate_recursive + + config = _instantiate_recursive( + CheckpointConfig, {"output_dir": "checkpoint-dir", "extra_parallel_canonical_layout": True} + ) + assert config.output_dir == "checkpoint-dir" + assert config.extra_parallel_canonical_layout is True diff --git a/veomni/arguments/arguments_types.py b/veomni/arguments/arguments_types.py index 489cbc38bd..4acff1601a 100644 --- a/veomni/arguments/arguments_types.py +++ b/veomni/arguments/arguments_types.py @@ -938,13 +938,6 @@ def _validate_init_device(self): class CheckpointConfig: """train.checkpoint.* — Checkpoint saving and loading.""" - extra_parallel_canonical_layout: bool = field( - default=False, - metadata={ - "help": "Use the versioned canonical EP/PLE model checkpoint layout. Legacy checkpoints require migration." - }, - ) - output_dir: str = field( default="output", metadata={"help": "Path to save model checkpoints."}, @@ -1031,6 +1024,12 @@ class CheckpointConfig: default=True, metadata={"help": "Save the huggingface format weights to the last checkpoint dir."}, ) + extra_parallel_canonical_layout: bool = field( + default=False, + metadata={ + "help": "Use the versioned canonical EP/PLE model checkpoint layout. Legacy checkpoints require migration." + }, + ) def __post_init__(self): """Reject a save configuration that cannot do what it says. From 29797c9601737d414ee83adc54cdb11248fcec42 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:17:57 +0800 Subject: [PATCH 019/120] [checkpoint] reject nonboolean canonical layout policy --- tests/arguments/test_checkpoint_config.py | 18 ++++++++++++++++++ veomni/arguments/arguments_types.py | 6 ++++++ 2 files changed, 24 insertions(+) diff --git a/tests/arguments/test_checkpoint_config.py b/tests/arguments/test_checkpoint_config.py index e9f1c18bb5..4c7570f253 100644 --- a/tests/arguments/test_checkpoint_config.py +++ b/tests/arguments/test_checkpoint_config.py @@ -2,6 +2,7 @@ # Licensed under the Apache License, Version 2.0. import pytest +import yaml from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.arguments.arguments_types import CheckpointConfig @@ -32,3 +33,20 @@ def test_checkpoint_config_parser_preserves_explicit_layout(): ) assert config.output_dir == "checkpoint-dir" assert config.extra_parallel_canonical_layout is True + + +@pytest.mark.parametrize("value", ["'false'", "'true'", "0", "1", "null"]) +def test_checkpoint_config_parser_rejects_non_boolean_layout(value): + from veomni.arguments.parser import _instantiate_recursive + + config = yaml.safe_load(f"extra_parallel_canonical_layout: {value}") + with pytest.raises(ValueError, match="extra_parallel_canonical_layout must be a boolean"): + _instantiate_recursive(CheckpointConfig, config) + + +@pytest.mark.parametrize("value", [False, True]) +def test_checkpoint_config_cli_converts_layout_boolean(monkeypatch, value): + from veomni.arguments.parser import parse_args + + monkeypatch.setattr("sys.argv", ["test", "--extra_parallel_canonical_layout", str(value).lower()]) + assert parse_args(CheckpointConfig).extra_parallel_canonical_layout is value diff --git a/veomni/arguments/arguments_types.py b/veomni/arguments/arguments_types.py index 4acff1601a..4d323ac5c8 100644 --- a/veomni/arguments/arguments_types.py +++ b/veomni/arguments/arguments_types.py @@ -1038,6 +1038,12 @@ def __post_init__(self): ``save_async`` too, but not until the first checkpoint is due -- a ``save_steps``-long wait to be told the configuration was never valid. """ + if not isinstance(self.extra_parallel_canonical_layout, bool): + raise ValueError( + "extra_parallel_canonical_layout must be a boolean; use unquoted true or false in YAML, " + f"got {self.extra_parallel_canonical_layout!r}." + ) + if self.stage_dir and self.save_async: raise ValueError( "stage_dir cannot be combined with save_async: the staged copy is dropped when the save " From 1b81246a9db704b9eac1a3d285e210020ec34b1c Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 16:07:41 +0800 Subject: [PATCH 020/120] [test] schedule sparse optimizer state regressions --- scripts/ci/accelerator/suites/gpu.json | 10 ++++++++++ scripts/ci/accelerator/suites/npu.json | 10 ++++++++++ tests/optim/test_multi_optimizer_sparse_lr.py | 3 +++ 3 files changed, 23 insertions(+) diff --git a/scripts/ci/accelerator/suites/gpu.json b/scripts/ci/accelerator/suites/gpu.json index 422591c7bb..b302bae557 100644 --- a/scripts/ci/accelerator/suites/gpu.json +++ b/scripts/ci/accelerator/suites/gpu.json @@ -454,6 +454,16 @@ "tests/optim/test_muon_unit.py" ] }, + { + "name": "sparse-optimizer-state-unit", + "suite": "unit", + "conditional": false, + "args": [ + "-s", + "-x", + "tests/optim/test_multi_optimizer_sparse_lr.py" + ] + }, { "name": "43-unit", "suite": "unit", diff --git a/scripts/ci/accelerator/suites/npu.json b/scripts/ci/accelerator/suites/npu.json index e8644c5baf..6739733e6a 100644 --- a/scripts/ci/accelerator/suites/npu.json +++ b/scripts/ci/accelerator/suites/npu.json @@ -210,6 +210,16 @@ "tests/checkpoints/" ] }, + { + "name": "sparse-optimizer-state-unit", + "suite": "unit", + "conditional": false, + "args": [ + "-s", + "-x", + "tests/optim/test_multi_optimizer_sparse_lr.py" + ] + }, { "name": "npu-unit-20", "suite": "unit", diff --git a/tests/optim/test_multi_optimizer_sparse_lr.py b/tests/optim/test_multi_optimizer_sparse_lr.py index 669d0cf66f..785f47867b 100644 --- a/tests/optim/test_multi_optimizer_sparse_lr.py +++ b/tests/optim/test_multi_optimizer_sparse_lr.py @@ -10,6 +10,9 @@ from veomni.optim.optimizer import MultiOptimizer +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + def _build(): model = torch.nn.Linear(2, 2) weight = torch.optim.AdamW([model.weight], lr=0.01) From ce5fecd75e276d25b4a0317c58e0cc02356a6185 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 16:10:48 +0800 Subject: [PATCH 021/120] [checkpoint] avoid redundant layout agreement gathers --- docs/usage/checkpoint.md | 7 +++++++ tests/checkpoints/test_extra_parallel_layout_v2.py | 12 +++++------- veomni/checkpoint/dcp_checkpointer.py | 5 ----- veomni/utils/save_safetensor_utils.py | 9 +-------- 4 files changed, 13 insertions(+), 20 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 83895d484d..b719563785 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -251,6 +251,13 @@ rejected until an owner-aware codec can preserve rank-specific optimizer states; this opt-in therefore does not yet support resuming training. The standard VeOmni HF export strips the namespace. +All ranks in a checkpoint operation must use the same logical model and +checkpoint configuration, including the layout version and corresponding +sharding specifications. Normal rank-local shard contents and mesh coordinates +may differ. Cross-rank agreement is an input contract, not a runtime preflight; +divergent configurations are unsupported and can hang or produce an invalid +checkpoint. Rank-local layout validation errors are still propagated to peers. + Raw `torch.distributed.checkpoint.load(..., planner=DefaultLoadPlanner(allow_partial_load=True))` is not a supported migration API: it can return success without reading any tensor when the namespaces differ. Replicated values must already agree across ranks. diff --git a/tests/checkpoints/test_extra_parallel_layout_v2.py b/tests/checkpoints/test_extra_parallel_layout_v2.py index 8eab059ca9..2ca2a56656 100644 --- a/tests/checkpoints/test_extra_parallel_layout_v2.py +++ b/tests/checkpoints/test_extra_parallel_layout_v2.py @@ -190,7 +190,7 @@ def test_v2_optimizer_fails_before_model_or_checkpoint_mutation(tmp_path): assert not list(tmp_path.iterdir()) -def _rank_local_invalid_layout_worker(rank, root, failure): +def _rank_local_invalid_layout_worker(rank, root): from datetime import timedelta dist.init_process_group( @@ -199,9 +199,8 @@ def _rank_local_invalid_layout_worker(rank, root, failure): try: model = torch.nn.Linear(2, 2, bias=False) mesh = SimpleNamespace(shape=(2, 2)) - shape = None if failure == "shape" and rank == 0 else (2, 2) - v2 = not (failure == "version" and rank == 0) - model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, shape, v2)} + shape = None if rank == 0 else (2, 2) + model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, shape, True)} with pytest.raises(RuntimeError): DistributedCheckpointer.save( f"{root}/checkpoint", {"model": model}, parallel_state=SimpleNamespace(dp_mode="fsdp2") @@ -211,9 +210,8 @@ def _rank_local_invalid_layout_worker(rank, root, failure): @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -@pytest.mark.parametrize("failure", ["shape", "version"]) -def test_rank_local_invalid_layout_stops_all_ranks_before_write(tmp_path, failure): - torch.multiprocessing.spawn(_rank_local_invalid_layout_worker, args=(str(tmp_path), failure), nprocs=2, join=True) +def test_rank_local_invalid_layout_stops_all_ranks_before_write(tmp_path): + torch.multiprocessing.spawn(_rank_local_invalid_layout_worker, args=(str(tmp_path),), nprocs=2, join=True) assert not (tmp_path / "checkpoint").exists() diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 1f6f8deed0..f6eba2de6e 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -223,11 +223,6 @@ def _reject_unproven_extra_parallel_optimizer(state, parallel_state, *, group=No except BaseException as exc: error = exc raise_if_any_rank_failed(error, "checking canonical ExtraParallel optimizer state", group=group) - if dist.is_initialized(): - signatures = [None] * dist.get_world_size(group) - dist.all_gather_object(signatures, signature, group=group) - if any(item != signature for item in signatures): - raise RuntimeError("Ranks disagree on canonical ExtraParallel checkpoint layout") error = None try: if signature is not None: diff --git a/veomni/utils/save_safetensor_utils.py b/veomni/utils/save_safetensor_utils.py index a487140d12..14633e6dd2 100644 --- a/veomni/utils/save_safetensor_utils.py +++ b/veomni/utils/save_safetensor_utils.py @@ -384,25 +384,18 @@ def save_lora_adapter_with_dcp( from ..checkpoint.dcp_checkpointer import ( _extra_parallel_layout_key, _needs_extra_parallel_layout_marker, - _runtime_extra_parallel_layout_signature, _validate_extra_parallel_tensor_v2, ) from .dist_utils import raise_if_any_rank_failed fqn2spec_info = getattr(model, "_fqn2spec_info", None) - error, signature = None, None + error = None try: if _needs_extra_parallel_layout_marker(fqn2spec_info): _extra_parallel_layout_key(fqn2spec_info) - signature = _runtime_extra_parallel_layout_signature(fqn2spec_info) except BaseException as exc: error = exc raise_if_any_rank_failed(error, "checking LoRA ExtraParallel layout") - if dist.is_initialized(): - signatures = [None] * dist.get_world_size() - dist.all_gather_object(signatures, signature) - if any(item != signature for item in signatures): - raise RuntimeError("Ranks disagree on LoRA ExtraParallel checkpoint layout") synchronize() if dist.is_initialized(): From c61c7334a650984c3dac4a41b07bb503ddf7b259 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 17:18:58 +0800 Subject: [PATCH 022/120] [checkpoint] remove unsupported divergent-layout test --- .../test_extra_parallel_layout_v2.py | 25 ------------------- 1 file changed, 25 deletions(-) diff --git a/tests/checkpoints/test_extra_parallel_layout_v2.py b/tests/checkpoints/test_extra_parallel_layout_v2.py index 585ab1e9f3..b268d4003f 100644 --- a/tests/checkpoints/test_extra_parallel_layout_v2.py +++ b/tests/checkpoints/test_extra_parallel_layout_v2.py @@ -248,31 +248,6 @@ def test_v2_optimizer_fails_before_model_or_checkpoint_mutation(tmp_path): assert not list(tmp_path.iterdir()) -def _rank_local_invalid_layout_worker(rank, root): - from datetime import timedelta - - dist.init_process_group( - "gloo", init_method=f"file://{root}/init", rank=rank, world_size=2, timeout=timedelta(seconds=20) - ) - try: - model = torch.nn.Linear(2, 2, bias=False) - mesh = SimpleNamespace(shape=(2, 2)) - shape = None if rank == 0 else (2, 2) - model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, shape, True)} - with pytest.raises(RuntimeError): - DistributedCheckpointer.save( - f"{root}/checkpoint", {"model": model}, parallel_state=SimpleNamespace(dp_mode="fsdp2") - ) - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_rank_local_invalid_layout_stops_all_ranks_before_write(tmp_path): - torch.multiprocessing.spawn(_rank_local_invalid_layout_worker, args=(str(tmp_path),), nprocs=2, join=True) - assert not (tmp_path / "checkpoint").exists() - - def test_same_v2_wrapper_can_save_and_load(tmp_path, monkeypatch): model = torch.nn.Linear(2, 2, bias=False) mesh = SimpleNamespace(shape=(2, 2)) From 17f5903f52c2191cb0568f543042e00d61038b88 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 18:00:15 +0800 Subject: [PATCH 023/120] [docs] fix SeedOmni runtime source anchors --- docs/seed_omni/mixins/offline_encoding.md | 4 ++-- docs/seed_omni/usage/hf_layout.md | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/seed_omni/mixins/offline_encoding.md b/docs/seed_omni/mixins/offline_encoding.md index 03c59adce1..10e7f1b790 100644 --- a/docs/seed_omni/mixins/offline_encoding.md +++ b/docs/seed_omni/mixins/offline_encoding.md @@ -7,7 +7,7 @@ it declares no method. - Mixin: [`OfflineEncodingMixin`](../../../veomni/models/seed_omni/mixins/offline_encoding_mixin.py#L20) - Run selector: [`train.training_task`](../../../veomni/arguments/omni_arguments_types.py#L702) -- Meta build: [`ModuleRuntime.reads_offline_cache`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L216) +- Meta build: [`ModuleRuntime.reads_offline_cache`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L219) ## Why cache a module's encoder @@ -120,7 +120,7 @@ This follows the DiT trainer's condition model: either the module is loaded in full, or it is built on meta. - **`offline_training`.** - [`ModuleRuntime.reads_offline_cache`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L216) + [`ModuleRuntime.reads_offline_cache`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L219) is true when the module's class mixes in `OfflineEncodingMixin`. The class is resolved from the module config with `get_model_class`, because the init device has to be chosen before the module exists. Then diff --git a/docs/seed_omni/usage/hf_layout.md b/docs/seed_omni/usage/hf_layout.md index 2579c4c39c..be36b703e8 100644 --- a/docs/seed_omni/usage/hf_layout.md +++ b/docs/seed_omni/usage/hf_layout.md @@ -150,7 +150,7 @@ Training from an upstream root: - [`OmniConfig.load_checkpoint_sidecars`](../../../veomni/models/seed_omni/configuration_omni.py#L428) records the `HFSource`, which each `ModuleRuntime` receives. - `ModuleRuntime._build_model` builds the module on meta and [attaches the key converter](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L347); - [`weights_path`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L233) is the upstream root. + [`weights_path`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L236) is the upstream root. - The FSDP2 / DDP wrap (per module, or [composed](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L306)) loads weights: - [`load_model_weights`](../../../veomni/models/checkpoint/weights.py#L410) and the [rank-0 broadcast loader](../../../veomni/models/checkpoint/weights.py#L1207) skip keys the converter declares, without reading them; - [`HFSourceKeyConverter`](../../../veomni/models/seed_omni/utils/hf_layout.py#L331) renames the module's own keys and hands them to the module's converter. From 88849502d219ce1bc286957abc7a4f00e1a7224c Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 18:14:09 +0800 Subject: [PATCH 024/120] [checkpoint] verify legacy expert provenance across targets --- docs/usage/checkpoint.md | 5 ++- tests/checkpoints/test_expert_shard_layout.py | 43 +++++++++++++++++++ veomni/checkpoint/dcp_checkpointer.py | 18 ++++++-- 3 files changed, 62 insertions(+), 4 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 9f9301a524..8d30a15b38 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -257,7 +257,10 @@ Older readers may ignore that metadata and interpret these shards using the old FSDP-first order; do not use them for these checkpoints without a verified migration. Conversely, an old unversioned checkpoint with expert-only chunks may not prove its source order even when resumed with the same topology. The -distributed reader rejects ambiguous sources when loading into EP/FSDP DTensors. +distributed reader rejects ambiguous sources for EP/FSDP as well as ordinary +FSDP or unsharded targets. An old expert-only FSDP checkpoint without layout +metadata also needs verified source-layout migration, since it cannot be +distinguished safely from an old EP checkpoint. Offline export also rejects unversioned expert-only chunks for stacked tensors under an `experts` module when their source order cannot be established. Other legacy tensor naming schemes may still need explicit source-layout verification diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 7d0a93944b..151b205bb9 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -210,3 +210,46 @@ def test_lora_adapter_export_preserves_expert_shard_layout(world, shard_dim, neg spawn_with_rendezvous( _lora_export_worker, args=(world, shard_dim, negative_ep, use_v2, str(tmp_path / "adapter")), nprocs=world ) + + +def _ordinary_fsdp_expert_export_worker(rank, root, ndim, endpoint): + from torch.distributed.tensor import distribute_tensor + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("fsdp",)) + shape = (8,) + (4,) * (ndim - 1) + expected = torch.arange(8 * 4 ** (ndim - 1), dtype=torch.float32).reshape(shape) + weight = distribute_tensor(expected, mesh, [Shard(0)]) + dcp.save({"model": {"experts": {"weight": weight}}}, checkpoint_id=root, planner=_ExtraParallelSavePlanner()) + dist.barrier() + result = [None] + if rank == 0: + try: + metadata = dcp.FileSystemReader(root).read_metadata() + assert metadata.planner_data[_EXPERT_LAYOUTS_KEY] == {"version": 1, "tensors": {}} + exported = dcp_to_torch_state_dict(root) + torch.testing.assert_close(exported["experts.weight"], expected, rtol=0, atol=0) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + + # An old expert-only chunk has no provenance, even when the target is + # ordinary FSDP. It must not be loaded and re-saved as canonical. + legacy = root + "-legacy" + dcp.save({"model": {"experts": {"weight": weight}}}, checkpoint_id=legacy) + target = {"model": {"experts": {"weight": torch.zeros_like(weight)}}} + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + dcp.load(target, checkpoint_id=legacy, planner=_ModelStrictLoadPlanner(strict_model=True)) + assert target["model"]["experts"]["weight"].to_local().count_nonzero() == 0 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [2]) +@pytest.mark.parametrize("ndim", [2, 3]) +def test_recorded_non_ep_fsdp_expert_export(world, ndim, tmp_path): + assert world == 2 + spawn_with_rendezvous(_ordinary_fsdp_expert_export_worker, args=(str(tmp_path / "dcp"), ndim), nprocs=world) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 17478e839a..818ebc17e9 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -209,13 +209,22 @@ def _validate_expert_layouts(self): for key, value in self.state_dict.items(): layout = _expert_layout(value) saved = self.metadata.state_dict_metadata.get(key) - if layout is None or saved is None: + if saved is None: continue # For legacy VeOmni EP weights, the faulty same-axis layout only # split the expert dimension. A source chunk that splits matrix # rows (old dim-1 FSDP), or spans all experts, cannot have that # permutation. This uses saved chunks, never the target FSDP dim. - ep_dim = layout["ep_shard_dim"] + if layout is None: + # A non-EP target can also load an ambiguous old EP source; + # letting it re-save would incorrectly mark that data as + # canonical. The stacked VeOmni expert axis is dimension 0. + shape = getattr(saved, "size", ()) + if "experts" not in key.split(".") or len(shape) < 2: + continue + ep_dim = 0 + else: + ep_dim = layout["ep_shard_dim"] if all( chunk.sizes[ep_dim] == saved.size[ep_dim] or any(size < saved.size[dim] for dim, size in enumerate(chunk.sizes) if dim != ep_dim) @@ -2014,6 +2023,9 @@ def _validate_expert_export_provenance(metadata: Metadata, model_keys) -> None: raise RuntimeError( f"Checkpoint tensor {key!r} has noncanonical expert order; migrate it before export" ) + # The writer records every EP tensor; omissions are ordinary tensors, + # including FSDP-only expert shards, whose coordinates are canonical. + return marker_key = "model." + _extra_parallel_layout_key({}) if marker_key in metadata.state_dict_metadata or planner_data.get("__veomni_ep_first_layout__") == 1: @@ -2024,7 +2036,7 @@ def _validate_expert_export_provenance(metadata: Metadata, model_keys) -> None: for key in model_keys: tensor_metadata = metadata.state_dict_metadata[key] shape = getattr(tensor_metadata, "size", ()) - if "experts" not in key.split(".") or len(shape) < 3 or key in layouts: + if "experts" not in key.split(".") or len(shape) < 2 or key in layouts: continue if any( chunk.sizes[0] < shape[0] and all(chunk.sizes[dim] == shape[dim] for dim in range(1, len(shape))) From 79309efb5d3f420a925001a833a2cca41370c8d7 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 19:25:33 +0800 Subject: [PATCH 025/120] Refresh remaining SeedOmni documentation links --- docs/seed_omni/mixins/module_preprocessor.md | 2 +- docs/seed_omni/mixins/offline_encoding.md | 6 +++--- docs/seed_omni/usage/hf_layout.md | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/seed_omni/mixins/module_preprocessor.md b/docs/seed_omni/mixins/module_preprocessor.md index dfa241f85d..7df4830363 100644 --- a/docs/seed_omni/mixins/module_preprocessor.md +++ b/docs/seed_omni/mixins/module_preprocessor.md @@ -119,7 +119,7 @@ class VisionPreprocessor(ModulePreprocessorBase): | Function | Called by | Effect | |----------|-----------|--------| | [`ModulePreprocessorBase.__call__(batch, inference=False, **kwargs)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L85) | [`OmniProcessor.preprocess_batch`](../../../veomni/models/seed_omni/processing_omni.py#L228) | Delegates `batch["conversation_list"]` to `preprocess_conversations` in `config.module_names` order. | -| [`bind_module_assets(model, *, checkpoint_path=None, preprocessor=None, config_overrides=None)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L128) | [`PretrainedOmniModule.from_pretrained`](../../../veomni/models/seed_omni/modules/module_modeling_base.py#L32) and [`ModuleRuntime._build_model_assets`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L354) | Copies CPU asset attributes (`_processor`, `_image_processor`, `_video_processor`, `_tokenizer`, `_chat_template`) from the preprocessor onto `model` if not already set, so `generate()` and HF checkpoint export can access them. | +| [`bind_module_assets(model, *, checkpoint_path=None, preprocessor=None, config_overrides=None)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L128) | [`PretrainedOmniModule.from_pretrained`](../../../veomni/models/seed_omni/modules/module_modeling_base.py#L32) and [`ModuleRuntime._build_model_assets`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L357) | Copies CPU asset attributes (`_processor`, `_image_processor`, `_video_processor`, `_tokenizer`, `_chat_template`) from the preprocessor onto `model` if not already set, so `generate()` and HF checkpoint export can access them. | ## Call flow diff --git a/docs/seed_omni/mixins/offline_encoding.md b/docs/seed_omni/mixins/offline_encoding.md index 10e7f1b790..b6c76962e1 100644 --- a/docs/seed_omni/mixins/offline_encoding.md +++ b/docs/seed_omni/mixins/offline_encoding.md @@ -124,7 +124,7 @@ full, or it is built on meta. is true when the module's class mixes in `OfflineEncodingMixin`. The class is resolved from the module config with `get_model_class`, because the init device has to be chosen before the module exists. Then - [`_build_model`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L318) + [`_build_model`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L321) passes `init_device="meta"`, and the [training build stops](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L160) right after `requires_grad_(False)`: no parallel wrap, no weight load, no @@ -136,12 +136,12 @@ full, or it is built on meta. detaches cache readers before the composed wrap so they stay on meta and are not materialised with the rest. - **`offline_embedding`.** Every module, cached or not, is built and - [frozen](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L171), + [frozen](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L176), then loaded and wrapped like any other: the run trains nothing, so no module gets an optimizer or a checkpoint manager, and [`OmniTrainer`](../../../veomni/trainer/omni/omni_trainer.py#L276) allows the empty optimizer. A `fsdp_mode: ddp` module is - [loaded but left unwrapped](../../../veomni/distributed/torch_parallelize.py#L1013), + [loaded but left unwrapped](../../../veomni/distributed/torch_parallelize.py#L1071), because DDP rejects a module with no parameter that requires grad. - **`online_training`.** The module is built like any other. diff --git a/docs/seed_omni/usage/hf_layout.md b/docs/seed_omni/usage/hf_layout.md index be36b703e8..4d662999d4 100644 --- a/docs/seed_omni/usage/hf_layout.md +++ b/docs/seed_omni/usage/hf_layout.md @@ -149,7 +149,7 @@ Training from an upstream root: (skipped for `train.training_task: offline_embedding`, which exports nothing). - [`OmniConfig.load_checkpoint_sidecars`](../../../veomni/models/seed_omni/configuration_omni.py#L428) records the `HFSource`, which each `ModuleRuntime` receives. - `ModuleRuntime._build_model` builds the module on meta and - [attaches the key converter](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L347); + [attaches the key converter](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L350); [`weights_path`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L236) is the upstream root. - The FSDP2 / DDP wrap (per module, or [composed](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L306)) loads weights: - [`load_model_weights`](../../../veomni/models/checkpoint/weights.py#L410) and the [rank-0 broadcast loader](../../../veomni/models/checkpoint/weights.py#L1207) skip keys the converter declares, without reading them; From a2ca5ebe3be9d688343170fe9e6bdca4e6391b5d Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 20:26:53 +0800 Subject: [PATCH 026/120] Fix AnyPrecision checkpoint state dtype and ownership --- .../test_multi_optimizer_proven_resume.py | 186 +++++++++++++++++- .../checkpoints/test_rank_local_optimizer.py | 26 +++ .../checkpoint/multi_optimizer_checkpoint.py | 42 ++-- veomni/optim/optimizer.py | 140 +++++++++++++ 4 files changed, 359 insertions(+), 35 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 11290ddc2f..26d9eb32c3 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -913,14 +913,18 @@ def test_nonfresh_target_is_rejected_before_any_checkpoint_load(tmp_path): torch.testing.assert_close(target.weight, original, rtol=0, atol=0) -def test_anyprecision_fp32_moments_with_bf16_parameter_fail_closed(tmp_path): +@pytest.mark.parametrize( + ("param_dtype", "moment_dtype"), + [(torch.float32, torch.bfloat16), (torch.bfloat16, torch.float32)], +) +def test_anyprecision_mixed_state_dtype_save_resume_and_update(tmp_path, monkeypatch, param_dtype, moment_dtype): def build(): - model = torch.nn.Linear(2, 2, bias=False).to(torch.bfloat16) + model = torch.nn.Linear(2, 2, bias=False).to(param_dtype) leaf = AnyPrecisionAdamW( model.parameters(), - momentum_dtype=torch.float32, - variance_dtype=torch.float32, - compensation_buffer_dtype=torch.float32, + momentum_dtype=moment_dtype, + variance_dtype=moment_dtype, + compensation_buffer_dtype=moment_dtype, ) return model, MultiOptimizer(model, {"weight": leaf}, ["weight"]) @@ -929,9 +933,175 @@ def build(): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - with pytest.raises(RuntimeError, match="state dtype cannot be safely restored"): - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - assert not list(tmp_path.iterdir()) + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build() + target_leaf = restored.optimizers_dict["weight"] + original_load = AnyPrecisionAdamW.load_state_dict + owned_pointers = [] + + def tracked_load(self, state_dict): + if self is target_leaf: + before = {name: value.data_ptr() for name, value in self.state[target.weight].items()} + result = original_load(self, state_dict) + if self is target_leaf: + after = {name: value.data_ptr() for name, value in self.state[target.weight].items()} + owned_pointers.append((before, after)) + return result + + monkeypatch.setattr(AnyPrecisionAdamW, "load_state_dict", tracked_load) + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + assert len(owned_pointers) == 1 + before, after = owned_pointers[0] + assert before == after + assert set(before) == {"step", "exp_avg", "exp_avg_sq", "compensation"} + source_state = optimizer.optimizers_dict["weight"].state[model.weight] + target_state = restored.optimizers_dict["weight"].state[target.weight] + for field in ("exp_avg", "exp_avg_sq", "compensation"): + assert source_state[field].dtype == target_state[field].dtype == moment_dtype + torch.testing.assert_close(source_state, target_state, rtol=0, atol=0) + torch.testing.assert_close(model.weight, target.weight, rtol=0, atol=0) + model.weight.grad = torch.full_like(model.weight, 0.2) + target.weight.grad = torch.full_like(target.weight, 0.2) + optimizer.step() + restored.step() + torch.testing.assert_close(model.weight, target.weight, rtol=0, atol=0) + torch.testing.assert_close(source_state, target_state, rtol=0, atol=0) + + +def test_anyprecision_direct_load_preserves_hooks_and_state_dtype(): + source = torch.nn.Parameter(torch.ones(2, dtype=torch.float32)) + first = AnyPrecisionAdamW([source], momentum_dtype=torch.bfloat16, variance_dtype=torch.bfloat16) + source.grad = torch.full_like(source, 0.2) + first.step() + target = torch.nn.Parameter(source.detach().clone()) + second = AnyPrecisionAdamW([target], momentum_dtype=torch.bfloat16, variance_dtype=torch.bfloat16) + calls = [] + second.register_load_state_dict_pre_hook(lambda opt, state: calls.append("pre")) + second.register_load_state_dict_post_hook(lambda opt: calls.append("post")) + second.load_state_dict(first.state_dict()) + assert calls == ["pre", "post"] + assert second.state[target]["exp_avg"].dtype == torch.bfloat16 + assert second.state[target]["exp_avg_sq"].dtype == torch.bfloat16 + for field in ("step", "exp_avg", "exp_avg_sq", "compensation"): + assert first.state[source][field].data_ptr() != second.state[target][field].data_ptr() + source.grad = torch.full_like(source, 0.3) + target.grad = torch.full_like(target, 0.3) + first.step() + second.step() + torch.testing.assert_close(source, target, rtol=0, atol=0) + + +def test_anyprecision_direct_load_does_not_reuse_preexisting_alias(): + source = torch.nn.Parameter(torch.ones(2)) + target = torch.nn.Parameter(torch.ones(2)) + first = AnyPrecisionAdamW([source]) + second = AnyPrecisionAdamW([target]) + source.grad = torch.ones_like(source) + first.step() + second.state[target] = dict(first.state[source]) + second.load_state_dict(first.state_dict()) + for field in ("step", "exp_avg", "exp_avg_sq", "compensation"): + assert first.state[source][field].data_ptr() != second.state[target][field].data_ptr() + + +def test_anyprecision_owned_buffer_reuse_requires_exact_view(): + from veomni.optim.optimizer import _same_owned_optimizer_tensor + + storage = torch.arange(8) + assert _same_owned_optimizer_tensor(storage[:4], storage[:4]) + assert not _same_owned_optimizer_tensor(storage[:4], storage[1:5]) + assert not _same_owned_optimizer_tensor(storage[:4], storage[:8:2]) + assert not _same_owned_optimizer_tensor(storage[:0], storage[:0]) + + +@pytest.mark.parametrize("bad_field", ["exp_avg", "exp_avg_sq", "compensation", "step"]) +def test_anyprecision_direct_load_rejects_dtensor_state_for_plain_parameter(tmp_path, bad_field): + import torch.distributed as dist + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, Replicate + + dist.init_process_group("gloo", init_method=f"file://{tmp_path / 'rdzv'}", rank=0, world_size=1) + try: + param = torch.nn.Parameter(torch.ones(2)) + optimizer = AnyPrecisionAdamW([param]) + param.grad = torch.full_like(param, 0.2) + optimizer.step() + bad = optimizer.state_dict() + mesh = init_device_mesh("cpu", (1,)) + value = bad["state"][0][bad_field] + bad["state"][0][bad_field] = DTensor.from_local(value, mesh, [Replicate()]) + target = torch.nn.Parameter(param.detach().clone()) + restored = AnyPrecisionAdamW([target]) + with pytest.raises(RuntimeError, match="incompatible|scalar tensor"): + restored.load_state_dict(bad) + assert not restored.state + finally: + dist.destroy_process_group() + + +def test_anyprecision_direct_load_rejects_invalid_local_shard_shape(tmp_path): + import torch.distributed as dist + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, Shard + + dist.init_process_group("gloo", init_method=f"file://{tmp_path / 'rdzv'}", rank=0, world_size=1) + try: + mesh = init_device_mesh("cpu", (1,)) + param = torch.nn.Parameter(DTensor.from_local(torch.ones(4, 4), mesh, [Shard(0)])) + optimizer = AnyPrecisionAdamW([param]) + param.grad = torch.ones_like(param) + optimizer.step() + bad = optimizer.state_dict() + bad["state"][0]["exp_avg"] = DTensor.from_local( + torch.zeros(2, 4, dtype=torch.bfloat16), mesh, [Shard(0)], shape=torch.Size((4, 4)), stride=(4, 1) + ) + restored = AnyPrecisionAdamW([torch.nn.Parameter(param.detach().clone())]) + with pytest.raises(RuntimeError, match="incompatible exp_avg layout"): + restored.load_state_dict(bad) + assert not restored.state + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("bad_field", ["replicated_moment", "wrong_local_shape", "plain_moment", "dtensor_step"]) +def test_anyprecision_source_rejects_incompatible_state_layout(tmp_path, bad_field): + import torch.distributed as dist + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor, Replicate, Shard + + from veomni.checkpoint.multi_optimizer_checkpoint import _validate_source_optimizer_states + + dist.init_process_group("gloo", init_method=f"file://{tmp_path / 'rdzv'}", rank=0, world_size=1) + try: + mesh = init_device_mesh("cpu", (1,)) + param = torch.nn.Parameter(DTensor.from_local(torch.ones(2, 2), mesh, [Shard(0)])) + model = torch.nn.Linear(2, 2, bias=False) + model.weight = param + leaf = AnyPrecisionAdamW( + [param], momentum_dtype=torch.bfloat16, variance_dtype=torch.bfloat16, use_kahan_summation=False + ) + optimizer = MultiOptimizer(model, {"weight": leaf}, ["weight"]) + state = { + "step": torch.tensor(1.0), + "exp_avg": torch.zeros_like(param, dtype=torch.bfloat16), + "exp_avg_sq": torch.zeros_like(param, dtype=torch.bfloat16), + } + if bad_field == "replicated_moment": + state["exp_avg"] = DTensor.from_local(torch.zeros(2, 2, dtype=torch.bfloat16), mesh, [Replicate()]) + elif bad_field == "wrong_local_shape": + state["exp_avg"] = DTensor.from_local( + torch.zeros(1, 2, dtype=torch.bfloat16), mesh, [Shard(0)], shape=torch.Size((2, 2)), stride=(2, 1) + ) + elif bad_field == "plain_moment": + state["exp_avg"] = torch.zeros(2, 2, dtype=torch.bfloat16) + else: + state["step"] = DTensor.from_local(torch.tensor(1.0), mesh, [Replicate()]) + leaf.state[param] = state + with pytest.raises(RuntimeError, match="layout differs|ordinary scalar tensor"): + _validate_source_optimizer_states(optimizer, model) + assert leaf.state[param] is state + finally: + dist.destroy_process_group() def test_anyprecision_uninitialized_leaf_restores_and_updates(tmp_path): diff --git a/tests/checkpoints/test_rank_local_optimizer.py b/tests/checkpoints/test_rank_local_optimizer.py index b63d70ef03..74dbc7900a 100644 --- a/tests/checkpoints/test_rank_local_optimizer.py +++ b/tests/checkpoints/test_rank_local_optimizer.py @@ -167,6 +167,10 @@ def gloo_all_to_all(outputs, inputs, group=None, async_op=False): model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, fsdp_axis, None, (8, 4), canonical)} if kind == "muon": leaf = DistributedMuon([model.weight], lr=0.01, ns_implementation="std") + elif kind == "anyprecision": + from veomni.optim.optimizer import AnyPrecisionAdamW + + leaf = AnyPrecisionAdamW([model.weight], lr=0.01) else: leaf = torch.optim.AdamW([model.weight], lr=0.01) leaves = {"expert": leaf} @@ -178,6 +182,10 @@ def gloo_all_to_all(outputs, inputs, group=None, async_op=False): def snapshot(): state = optimizer.optimizers_dict["expert"].state.get(model.weight, {}) + if kind == "anyprecision" and state: + assert set(state) == {"step", "exp_avg", "exp_avg_sq", "compensation"} + assert state["step"].dtype == torch.float32 + assert all(state[key].dtype == torch.bfloat16 for key in ("exp_avg", "exp_avg_sq", "compensation")) result = { "weight": model.weight.to_local().detach().clone(), "state": { @@ -250,6 +258,24 @@ def test_canonical_same_axis_muon_fresh_process_resume(tmp_path): mp.spawn(_cold_optimizer_worker, args=(str(tmp_path), False, "load", "muon", True, True), nprocs=4, join=True) +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) +@pytest.mark.parametrize("same_axis", [False, True]) +@pytest.mark.parametrize("inactive_owner", [False, True]) +def test_anyprecision_bf16_state_extra_parallel_fresh_process_resume(tmp_path, same_axis, inactive_owner): + mp.spawn( + _cold_optimizer_worker, + args=(str(tmp_path), inactive_owner, "save", "anyprecision", same_axis, same_axis), + nprocs=4, + join=True, + ) + mp.spawn( + _cold_optimizer_worker, + args=(str(tmp_path), inactive_owner, "load", "anyprecision", same_axis, same_axis), + nprocs=4, + join=True, + ) + + def _replica_check_worker(rank, root): dist.init_process_group( "gloo", init_method=f"file://{root}/replica-init", rank=rank, world_size=2, timeout=timedelta(seconds=20) diff --git a/veomni/checkpoint/multi_optimizer_checkpoint.py b/veomni/checkpoint/multi_optimizer_checkpoint.py index b0713af172..45d438f03c 100644 --- a/veomni/checkpoint/multi_optimizer_checkpoint.py +++ b/veomni/checkpoint/multi_optimizer_checkpoint.py @@ -167,12 +167,26 @@ def _validate_source_optimizer_states(optimizer, model): if field == "step": if tuple(value.shape) != (): raise RuntimeError("MultiOptimizer step must be a scalar") + if isinstance(leaf, AnyPrecisionAdamW) and isinstance(value, DTensor): + raise RuntimeError("AnyPrecisionAdamW step must be an ordinary scalar tensor") adapter = rank_local_adamw_adapter(leaf) if adapter is not None and value.dtype != adapter.state_dtype(leaf, param, field): raise RuntimeError("MultiOptimizer step dtype differs from its registered policy") else: if tuple(value.shape) != tuple(param.shape): raise RuntimeError(f"MultiOptimizer state shape differs for {field}") + if isinstance(leaf, AnyPrecisionAdamW) and ( + isinstance(param, DTensor) != isinstance(value, DTensor) + or ( + isinstance(param, DTensor) + and ( + value.device_mesh != param.device_mesh + or value.placements != param.placements + or value.to_local().shape != param.to_local().shape + ) + ) + ): + raise RuntimeError(f"AnyPrecisionAdamW state layout differs for {field}") adapter = rank_local_adamw_adapter(leaf) expected_dtype = ( adapter.state_dtype(leaf, param, field) @@ -187,7 +201,7 @@ def _validate_source_optimizer_states(optimizer, model): if isinstance(leaf, AnyPrecisionAdamW) else param.dtype ) - if (adapter is None and expected_dtype != param.dtype) or value.dtype != expected_dtype: + if value.dtype != expected_dtype: raise RuntimeError(f"MultiOptimizer state dtype cannot be safely restored for {field}") @@ -474,14 +488,9 @@ def _preflight_multi_optimizer_layout( return None, None, None, None, None, None, None snapshot, extra, extra_digest, source_metadata, error = None, None, None, None, None try: - from ..optim.optimizer import AnyPrecisionAdamW - _check_model_serialization_hooks(model) source_metadata = reader_factory(checkpoint_path).read_metadata() parameter_names = _dcp_parameter_name_map(model) - parameter_names_by_id = defaultdict(list) - for name, param in parameter_names.items(): - parameter_names_by_id[id(param)].append(name) from .rank_local_optimizer import current_rank, load_rank_scalars, rank_prefix rank_local = ( @@ -505,22 +514,6 @@ def _preflight_multi_optimizer_layout( raise RuntimeError("MultiOptimizer checkpoint resume requires a fresh optimizer without state") if any(param.grad is not None for param_group in leaf.param_groups for param in param_group["params"]): raise RuntimeError("MultiOptimizer checkpoint resume requires cleared gradients") - if isinstance(leaf, AnyPrecisionAdamW): - for param_group in leaf.param_groups: - fields = ["momentum_dtype", "variance_dtype"] - if param_group.get("use_kahan_summation"): - fields.append("compensation_buffer_dtype") - incompatible = False - for param in param_group["params"]: - if not any(param_group[field] != param.dtype for field in fields): - continue - names = parameter_names_by_id[id(param)] - incompatible |= any(f"optimizer.state.{name}" in source_bases for name in names) - if incompatible: - raise RuntimeError( - "AnyPrecisionAdamW requires explicit precision-preserving migration when optimizer " - "state and parameter dtypes differ" - ) extra_digest, extra = _multi_extra_snapshot(extra_root, read_state=True) if extra is not None: snapshot = extra.get("lr_scheduler") @@ -851,11 +844,6 @@ def _prepare_multi_optimizer_saved_states( "compensation": "compensation_buffer_dtype", }[field] ] - if dtype != param.dtype: - raise RuntimeError( - "AnyPrecisionAdamW cannot restore moments at a dtype different from its parameter; " - "explicit precision-preserving migration is required" - ) else: dtype = param.dtype if entry.properties.dtype != dtype: diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index b77805b1dd..04762ce998 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -15,7 +15,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections import defaultdict from contextlib import contextmanager +from copy import deepcopy +from itertools import chain from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union import torch @@ -153,6 +156,58 @@ def filter_empty_param_groups( return non_empty +def _same_owned_optimizer_tensor(current, incoming): + """Recognize a DCP-filled target tensor without taking a second copy.""" + if current is incoming: + return True + if not torch.is_tensor(current) or not torch.is_tensor(incoming): + return False + if isinstance(current, DTensor) != isinstance(incoming, DTensor): + return False + if current.layout != torch.strided or incoming.layout != torch.strided: + return False + if ( + current.device != incoming.device + or current.dtype != incoming.dtype + or current.shape != incoming.shape + or current.stride() != incoming.stride() + ): + return False + if isinstance(current, DTensor): + if current.device_mesh != incoming.device_mesh or current.placements != incoming.placements: + return False + current, incoming = current.to_local(), incoming.to_local() + if current.layout != torch.strided or incoming.layout != torch.strided or current.numel() == 0: + return False + return ( + current.device == incoming.device + and current.dtype == incoming.dtype + and current.shape == incoming.shape + and current.stride() == incoming.stride() + and current.storage_offset() == incoming.storage_offset() + and current.data_ptr() == incoming.data_ptr() + ) + + +@contextmanager +def _dcp_reuse_owned_anyprecision_state(optimizer): + """Permit reuse only for fresh state buffers created by the proven DCP loader.""" + if type(optimizer) is not AnyPrecisionAdamW or any( + getattr(optimizer, name, None) + for name in ("_optimizer_load_state_dict_pre_hooks", "_optimizer_load_state_dict_post_hooks") + ): + yield + return + marker = "_veomni_dcp_reuse_owned_state" + if marker in vars(optimizer): + raise RuntimeError("AnyPrecisionAdamW DCP state reuse is already active") + setattr(optimizer, marker, True) + try: + yield + finally: + delattr(optimizer, marker) + + # https://github.com/meta-llama/llama-recipes/blob/v0.0.4/src/llama_recipes/policies/anyprecision_optimizer.py class AnyPrecisionAdamW(Optimizer): def __init__( @@ -179,6 +234,90 @@ def __init__( } super().__init__(params, defaults) + def load_state_dict(self, state_dict): + """Restore each moment at its configured dtype without a parameter-dtype cast.""" + state_dict = state_dict.copy() + for hook in self._optimizer_load_state_dict_pre_hooks.values(): + result = hook(self, state_dict) + if result is not None: + state_dict = result + groups = self.param_groups + saved_groups = deepcopy(state_dict["param_groups"]) + if len(groups) != len(saved_groups): + raise ValueError("loaded state dict has a different number of parameter groups") + if any( + len(group["params"]) != len(saved["params"]) for group, saved in zip(groups, saved_groups, strict=True) + ): + raise ValueError("loaded state dict contains a parameter group that doesn't match the optimizer's group") + + id_map = dict( + zip( + chain.from_iterable(group["params"] for group in saved_groups), + chain.from_iterable(group["params"] for group in groups), + strict=True, + ) + ) + group_by_id = {key: group for group in saved_groups for key in group["params"]} + restored = defaultdict(dict) + for key, saved in state_dict["state"].items(): + if key not in id_map: + restored[key] = deepcopy(saved) + continue + param = id_map[key] + group = group_by_id[key] + previous = self.state.get(param, {}) + reuse_owned = getattr(self, "_veomni_dcp_reuse_owned_state", False) is True + fields = {"step", "exp_avg", "exp_avg_sq"} + if group.get("use_kahan_summation"): + fields.add("compensation") + if saved and set(saved) != fields: + raise RuntimeError("AnyPrecisionAdamW checkpoint contains incomplete or unsupported state") + result = dict(saved) + if result and ( + not torch.is_tensor(result["step"]) + or isinstance(result["step"], DTensor) + or result["step"].shape != () + ): + raise RuntimeError("AnyPrecisionAdamW checkpoint step must be a scalar tensor") + if result: + value = result["step"] + if reuse_owned and _same_owned_optimizer_tensor(previous.get("step"), value): + result["step"] = previous["step"] + else: + result["step"] = value.detach().clone() + for field, dtype_key in ( + ("exp_avg", "momentum_dtype"), + ("exp_avg_sq", "variance_dtype"), + ("compensation", "compensation_buffer_dtype"), + ): + if field not in result: + continue + value = result[field] + if not torch.is_tensor(value) or value.dtype != group[dtype_key] or value.shape != param.shape: + raise RuntimeError(f"AnyPrecisionAdamW checkpoint has an incompatible {field} tensor") + if isinstance(param, DTensor) != isinstance(value, DTensor) or ( + isinstance(param, DTensor) + and ( + value.device_mesh != param.device_mesh + or value.placements != param.placements + or value.to_local().shape != param.to_local().shape + ) + ): + raise RuntimeError(f"AnyPrecisionAdamW checkpoint has an incompatible {field} layout") + if reuse_owned and _same_owned_optimizer_tensor(previous.get(field), value): + result[field] = previous[field] + else: + result[field] = value.to(device=param.device).detach().clone() + restored[param] = result + + for group, saved in zip(groups, saved_groups, strict=True): + saved["params"] = group["params"] + if "param_names" in group and "param_names" not in saved: + saved["param_names"] = group["param_names"] + self.__setstate__({"state": restored, "param_groups": saved_groups}) + for hook in self._optimizer_load_state_dict_post_hooks.values(): + hook(self) + @torch.no_grad() def step(self, closure=None): """ @@ -560,6 +699,7 @@ def _dcp_load_state_dict(self, state_dict: Dict[str, Any], *, group=None) -> Non with ( rank_local_adamw_context(optimizer), _dcp_restore_frozen_ownership(optimizer), + _dcp_reuse_owned_anyprecision_state(optimizer), _dcp_avoid_synthetic_optimizer_step(optimizer), ): set_optimizer_state_dict( From dcc195da263e9cceead80783bdece97da82ff466 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 20:44:57 +0800 Subject: [PATCH 027/120] Correct remaining SeedOmni documentation anchors --- docs/seed_omni/mixins/metric_meter.md | 2 +- docs/seed_omni/mixins/offline_encoding.md | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/seed_omni/mixins/metric_meter.md b/docs/seed_omni/mixins/metric_meter.md index 8c5d1edf4e..1427cc70c7 100644 --- a/docs/seed_omni/mixins/metric_meter.md +++ b/docs/seed_omni/mixins/metric_meter.md @@ -220,7 +220,7 @@ warning. To make a backbone eligible, stash exactly one length per conversation. - [`tests/seed_omni/mixins/test_metric_meter_mixin.py`](../../../tests/seed_omni/mixins/test_metric_meter_mixin.py): the mixin's stash/accumulate/reset contract and the `OmniEnvironMeter` roll-up, including a two-rank gloo run. -- [`test_metric_meter_collect_drains_only_the_metered_modules`](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L399): +- [`test_metric_meter_collect_drains_only_the_metered_modules`](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L400): the runtime drain. - [`tests/seed_omni/trainer/test_step_metrics_callback.py`](../../../tests/seed_omni/trainer/test_step_metrics_callback.py): the callback wiring. diff --git a/docs/seed_omni/mixins/offline_encoding.md b/docs/seed_omni/mixins/offline_encoding.md index b6c76962e1..b9f86bdf29 100644 --- a/docs/seed_omni/mixins/offline_encoding.md +++ b/docs/seed_omni/mixins/offline_encoding.md @@ -126,7 +126,7 @@ full, or it is built on meta. device has to be chosen before the module exists. Then [`_build_model`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L321) passes `init_device="meta"`, and the - [training build stops](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L160) + [training build stops](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L163) right after `requires_grad_(False)`: no parallel wrap, no weight load, no optimizer, no checkpoint manager. The offline-training method therefore must read only the config, never a parameter or buffer: parameters are meta @@ -323,7 +323,7 @@ A missing module or method still fails before the first request runs. ## Current scope - **Cached modules are not checkpointed.** They are frozen in both cache tasks, - and [fully frozen modules have no checkpoint manager](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L624), + and [fully frozen modules have no checkpoint manager](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L627), so a checkpoint never holds a meta module. - **Every cacheable module is read from the cache.** `offline_training` builds every module that mixes in `OfflineEncodingMixin` on meta; one run cannot @@ -353,8 +353,8 @@ A missing module or method still fails before the first request runs. [it is frozen and never wrapped, trained or saved](../../../tests/seed_omni/runtime/test_module_runtime.py#L209), and [an `offline_embedding` run freezes every module](../../../tests/seed_omni/runtime/test_module_runtime.py#L241). - [`tests/seed_omni/runtime/test_omni_model_runtime.py`](../../../tests/seed_omni/runtime/test_omni_model_runtime.py): - [under `fsdp_scope: model`, cache readers stay on meta outside the composed wrap](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L417), and - [under a module-level input cast, a node's item edits reach the batch only from its hooks or its return](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L304). + [under `fsdp_scope: model`, cache readers stay on meta outside the composed wrap](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L418), and + [under a module-level input cast, a node's item edits reach the batch only from its hooks or its return](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L305). - [`tests/seed_omni/test_offline_cache_writer.py`](../../../tests/seed_omni/test_offline_cache_writer.py): [the `seedomni_cached` transform](../../../tests/seed_omni/test_offline_cache_writer.py#L33), the [writer round trip](../../../tests/seed_omni/test_offline_cache_writer.py#L51), @@ -371,7 +371,7 @@ A missing module or method still fails before the first request runs. - [`tests/checkpoints/test_checkpoint_callback.py`](../../../tests/checkpoints/test_checkpoint_callback.py): [the writer's position is saved and restored with the cursor, and the cursor-only step is complete](../../../tests/checkpoints/test_checkpoint_callback.py#L632). - [`tests/checkpoints/test_dcp_checkpointer.py`](../../../tests/checkpoints/test_dcp_checkpointer.py): - [a step without `model/` is found only when its manifest says it carries no model state](../../../tests/checkpoints/test_dcp_checkpointer.py#L2551). + [a step without `model/` is found only when its manifest says it carries no model state](../../../tests/checkpoints/test_dcp_checkpointer.py#L2552). - [`tests/seed_omni/trainer/test_step_metrics_callback.py`](../../../tests/seed_omni/trainer/test_step_metrics_callback.py): [a step without an lr scheduler logs no lr](../../../tests/seed_omni/trainer/test_step_metrics_callback.py#L82). - [`tests/seed_omni/model/test_graph.py`](../../../tests/seed_omni/model/test_graph.py): From 5b11c0243694b17e9d27ce515e231cd83cfb17cb Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:14:11 +0800 Subject: [PATCH 028/120] Test asymmetric EP expert checkpoint recovery --- tests/ci/_npu_moe_fp32_cold_worker.py | 47 ++++++-- tests/ci/test_npu_moe_fp32_cold.py | 150 ++++++++++++++++++++++++++ 2 files changed, 186 insertions(+), 11 deletions(-) diff --git a/tests/ci/_npu_moe_fp32_cold_worker.py b/tests/ci/_npu_moe_fp32_cold_worker.py index 2b166826aa..ee1e67993f 100644 --- a/tests/ci/_npu_moe_fp32_cold_worker.py +++ b/tests/ci/_npu_moe_fp32_cold_worker.py @@ -21,16 +21,16 @@ import torch import torch.distributed as dist -import torch.distributed.checkpoint as dcp from torch import nn from torch.distributed._tensor import DTensor, Shard -from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict from veomni.arguments.arguments_types import MixedPrecisionConfig +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.distributed.parallel_plan import ParallelPlan -from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state +from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state, get_parallel_state from veomni.distributed.torch_parallelize import build_parallelize_model from veomni.models.transformers.qwen4_exp.generated.patched_modeling_qwen4_exp_npu import Qwen4ExpTextExperts +from veomni.optim.optimizer import build_optimizer class ToyDecoder(nn.Module): @@ -77,10 +77,17 @@ def _local(tensor): return tensor.to_local() if isinstance(tensor, DTensor) else tensor +def _optimizer_state(optimizer, param): + leaves = getattr(optimizer, "optimizers_dict", None) + owners = [leaf for leaf in leaves.values() if param in leaf.state] if leaves is not None else [optimizer] + assert len(owners) == 1 and param in owners[0].state + return owners[0].state[param] + + def _snapshot(model, optimizer): result = {} for name, param in model.named_parameters(): - state = optimizer.state[param] + state = _optimizer_state(optimizer, param) result[name] = { "param": _local(param.detach()).cpu().clone(), "m": _local(state["exp_avg"]).detach().cpu().clone(), @@ -97,7 +104,7 @@ def _step(model, optimizer, rank, ep, step): # Each EP rank owns a zero-token expert while experts 0 and 2 receive tokens. indices = torch.tensor([[0, 2]] * 8, device="npu", dtype=torch.int64) weights = torch.full((8, 2), 0.5, device="npu", dtype=torch.float32) - optimizer.zero_grad(set_to_none=True) + optimizer.zero_grad() loss = model(x, indices, weights).float().square().mean() loss.backward() grads = {} @@ -149,16 +156,34 @@ def run(rank, phase, output_dir, ep, endpoint=None): enable_forward_prefetch=False, ) assert model.layer.experts._veomni_fp32_expert_wgrad - optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.0) + optimizer = build_optimizer(model, lr=1e-3, betas=(0.9, 0.999), weight_decay=0.0, fused=False) ckpt = str(Path(output_dir) / f"ep{ep}-dcp") if phase == "fresh": + # The locally unused expert is different on each EP owner, so a + # duplicated or swapped shard cannot pass the resume comparison. + with torch.no_grad(): + for param in model.parameters(): + _local(param)[0].add_(rank / 2048) + _local(param)[1].add_(rank / 1024) _step(model, optimizer, rank, ep, 1) - model_state, optim_state = get_state_dict(model, optimizer) - dcp.save({"model": model_state, "optimizer": optim_state}, checkpoint_id=ckpt) + # An expert can have historical moments even when this batch gave + # it no tokens. Make those moments distinct across EP owners too. + with torch.no_grad(): + for param in model.parameters(): + state = _optimizer_state(optimizer, param) + _local(state["exp_avg"])[1].fill_((rank + 1) / 1024) + _local(state["exp_avg_sq"])[1].fill_((rank + 1) / 2048) + optimizer.zero_grad() + DistributedCheckpointer.save( + ckpt, {"model": model, "optimizer": optimizer}, parallel_state=get_parallel_state() + ) else: - model_state, optim_state = get_state_dict(model, optimizer) - dcp.load({"model": model_state, "optimizer": optim_state}, checkpoint_id=ckpt) - set_state_dict(model, optimizer, model_state_dict=model_state, optim_state_dict=optim_state) + with torch.no_grad(): + for param in model.parameters(): + _local(param).fill_(-0.05 - rank * 0.01) + DistributedCheckpointer.load( + ckpt, {"model": model, "optimizer": optimizer}, parallel_state=get_parallel_state() + ) before = _snapshot(model, optimizer) after = _step(model, optimizer, rank, ep, 2) for name in before: diff --git a/tests/ci/test_npu_moe_fp32_cold.py b/tests/ci/test_npu_moe_fp32_cold.py index 08a4813de5..c3607fad00 100644 --- a/tests/ci/test_npu_moe_fp32_cold.py +++ b/tests/ci/test_npu_moe_fp32_cold.py @@ -14,13 +14,23 @@ """FSDP2 + expert-parallel optimizer/DCP cold-recovery regression on Ascend.""" +import hashlib +import json +import math import os import pytest import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint import FileSystemReader +from torch.distributed.checkpoint.metadata import BytesStorageMetadata +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import Shard, distribute_tensor from scripts.ci.accelerator.launch_plan import LaunchSpec from tests.ci._npu_moe_fp32_cold_worker import run +from veomni.checkpoint.layout import optimizer_dir, weights_dir from veomni.utils.device import IS_NPU_AVAILABLE @@ -35,6 +45,133 @@ def _assert_equal(left, right): assert left == right +def _assert_ep2_checkpoint_layout(checkpoint_dir): + weights = FileSystemReader(weights_dir(str(checkpoint_dir))).read_metadata().state_dict_metadata + optimizer = FileSystemReader(optimizer_dir(str(checkpoint_dir))).read_metadata().state_dict_metadata + marker = "__veomni_multi_optimizer_layout_v1__" + rank_local = f"optimizer.{marker}" in optimizer + if rank_local: + keys = [marker, *(f"__veomni_rank_{rank}____inventory" for rank in range(2))] + assert all(isinstance(optimizer[f"optimizer.{key}"], BytesStorageMetadata) for key in keys) + scratch = {"optimizer": dict.fromkeys(keys, "")} + dcp.load(scratch, checkpoint_id=optimizer_dir(str(checkpoint_dir)), no_dist=True) + layout = json.loads(scratch["optimizer"][marker]) + assert set(layout) == {"generation", "layout", "inventory_digests"} + assert isinstance(layout["generation"], str) and layout["generation"] + geometry = json.loads(layout["layout"]) + assert set(geometry) == {"version", "ranks", "layouts"} + assert geometry["version"] == 1 and geometry["ranks"] == [0, 1] + assert len(geometry["layouts"]) == 2 + assert all( + isinstance(item, str) and len(item) == 64 and all(char in "0123456789abcdef" for char in item) + for item in geometry["layouts"] + ) + generation_key = "__veomni_multi_generation_v1__" + assert isinstance(weights[generation_key], BytesStorageMetadata) + model_state = {generation_key: ""} + dcp.load(model_state, checkpoint_id=weights_dir(str(checkpoint_dir)), no_dist=True) + assert model_state[generation_key] == layout["generation"] + assert len(layout["inventory_digests"]) == 2 + for rank, key in enumerate(keys[1:]): + inventory = json.loads(scratch["optimizer"][key]) + prefix = f"optimizer.__veomni_rank_{rank}__" + assert inventory == sorted( + name for name in optimizer if name.startswith(prefix) and name != f"optimizer.{key}" + ) + digest = hashlib.sha256(json.dumps(inventory, separators=(",", ":")).encode()).hexdigest() + assert digest == layout["inventory_digests"][rank] + for name, shape in ( + ("layer.experts.gate_up_proj", (4, 64, 256)), + ("layer.experts.down_proj", (4, 256, 32)), + ): + weight = weights[f"model.{name}"] + assert tuple(weight.size) == shape + assert weight.properties.dtype == torch.float32 + local_shape = (2, *shape[1:]) + canonical_chunks = {((0, 0, 0), local_shape), ((2, 0, 0), local_shape)} + assert len(weight.chunks) == len(canonical_chunks) + assert {(tuple(chunk.offsets), tuple(chunk.sizes)) for chunk in weight.chunks} == canonical_chunks + for field in ("exp_avg", "exp_avg_sq"): + owners = range(2) if rank_local else (None,) + for rank in owners: + prefix = f"__veomni_rank_{rank}__" if rank is not None else "" + entry = optimizer[f"optimizer.{prefix}state.{name}.{field}"] + expected_shape = local_shape if rank_local else shape + expected_chunks = {((0, 0, 0), local_shape)} if rank_local else canonical_chunks + assert tuple(entry.size) == expected_shape + assert entry.properties.dtype == torch.float32 + assert len(entry.chunks) == len(expected_chunks) + assert {(tuple(chunk.offsets), tuple(chunk.sizes)) for chunk in entry.chunks} == expected_chunks + + +def _save_cpu_expert_metadata(rank, checkpoint_dir, checkpoint_format, endpoint): + from datetime import timedelta + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + mesh = init_device_mesh("cpu", (2,)) + model, optimizer = {}, {} + for name, shape in ( + ("layer.experts.gate_up_proj", (4, 64, 256)), + ("layer.experts.down_proj", (4, 256, 32)), + ): + tensor = torch.arange(math.prod(shape), dtype=torch.float32).reshape(shape) + tensor = ( + tensor[2 * rank : 2 * rank + 2] + if checkpoint_format == "bare" + else distribute_tensor(tensor, mesh, [Shard(0)]) + ) + model[name] = tensor + rank_local = checkpoint_format in ("rank_local", "corrupt_inventory", "wrong_generation", "wrong_layout") + prefix = f"__veomni_rank_{rank}__" if rank_local else "" + value = tensor.to_local() if rank_local else tensor + optimizer[f"{prefix}state.{name}.exp_avg"] = value + optimizer[f"{prefix}state.{name}.exp_avg_sq"] = value + if checkpoint_format in ("rank_local", "corrupt_inventory", "wrong_generation", "wrong_layout"): + prefix = f"__veomni_rank_{rank}__" + optimizer[f"{prefix}__scalars"] = b"cpu-test" + inventory = sorted(f"optimizer.{key}" for key in optimizer) + local_digest = hashlib.sha256(json.dumps(inventory, separators=(",", ":")).encode()).hexdigest() + digests = [None, None] + dist.all_gather_object(digests, local_digest) + if checkpoint_format == "corrupt_inventory": + digests[1] = "0" * 64 + optimizer[f"{prefix}__inventory"] = json.dumps(inventory) + geometry = { + "version": 1, + "ranks": [1, 0] if checkpoint_format == "wrong_layout" else [0, 1], + "layouts": ["a" * 64, "b" * 64], + } + optimizer["__veomni_multi_optimizer_layout_v1__"] = json.dumps( + {"generation": "cpu-test", "layout": json.dumps(geometry), "inventory_digests": digests} + ) + model_state = {"model": model} + if checkpoint_format in ("rank_local", "corrupt_inventory", "wrong_generation", "wrong_layout"): + model_state["__veomni_multi_generation_v1__"] = ( + "other-generation" if checkpoint_format == "wrong_generation" else "cpu-test" + ) + dcp.save(model_state, checkpoint_id=weights_dir(checkpoint_dir)) + dcp.save({"optimizer": optimizer}, checkpoint_id=optimizer_dir(checkpoint_dir)) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +@pytest.mark.parametrize( + "checkpoint_format", ["bare", "canonical", "rank_local", "corrupt_inventory", "wrong_generation", "wrong_layout"] +) +def test_ep2_metadata_requires_proven_expert_shards(tmp_path, checkpoint_format): + from tests.tools.launch_utils import spawn_with_rendezvous + + checkpoint_dir = str(tmp_path / "ep2-dcp") + spawn_with_rendezvous(_save_cpu_expert_metadata, args=(checkpoint_dir, checkpoint_format), nprocs=2) + if checkpoint_format in ("canonical", "rank_local"): + _assert_ep2_checkpoint_layout(checkpoint_dir) + else: + with pytest.raises(AssertionError): + _assert_ep2_checkpoint_layout(checkpoint_dir) + + @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2)) @pytest.mark.parametrize("ep", [1, 2]) def test_native_fsdp_expert_fp32_dcp_cold_resume(tmp_path, ep): @@ -46,7 +183,20 @@ def test_native_fsdp_expert_fp32_dcp_cold_resume(tmp_path, ep): for phase in ("fresh", "resume"): spawn_with_rendezvous(run, args=(phase, str(tmp_path), ep), nprocs=2) + fresh_by_rank = [] for rank in range(2): fresh = torch.load(tmp_path / f"ep{ep}-fresh-rank{rank}.pt", map_location="cpu", weights_only=True) resumed = torch.load(tmp_path / f"ep{ep}-resume-rank{rank}.pt", map_location="cpu", weights_only=True) _assert_equal(fresh, resumed) + fresh_by_rank.append(fresh) + if ep == 2: + for name in fresh_by_rank[0]["before"]: + for row in (0, 1): + left = fresh_by_rank[0]["before"][name]["param"][row] + right = fresh_by_rank[1]["before"][name]["param"][row] + assert not torch.equal(left, right), f"symmetric EP fixture for {name}: param row {row}" + for field in ("m", "v"): + left = fresh_by_rank[0]["before"][name][field][1] + right = fresh_by_rank[1]["before"][name][field][1] + assert not torch.equal(left, right), f"symmetric EP fixture for {name}: {field}" + _assert_ep2_checkpoint_layout(tmp_path / "ep2-dcp") From 924e8ecdb3173b9d34ac2da53bebceecbce49b58 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:16:29 +0800 Subject: [PATCH 029/120] Reject mismatched trainer and optimizer checkpoints --- docs/usage/checkpoint.md | 8 + tests/checkpoints/test_checkpoint_callback.py | 5 +- .../test_multi_optimizer_trainer_step.py | 270 ++++++++++++++++++ .../callbacks/global_state_callback.py | 109 ++++++- 4 files changed, 383 insertions(+), 9 deletions(-) create mode 100644 tests/checkpoints/test_multi_optimizer_trainer_step.py diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 1b4a000fdd..6d45ad98d2 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -153,6 +153,14 @@ earlier checkpoint, write the new run to a different output directory to avoid colliding with checkpoints from the previous run. A same-step export can reuse an existing checkpoint only when its model, optimizer and trainer state have been proven to belong to the same completed save. +Trainer sidecars record the generation of each `MultiOptimizer` checkpoint. +An absent sidecar, a missing generation, or a trainer step/generation that +differs from the restored optimizer stops resume before training. An older +unversioned sidecar requires a verified migration or a fresh weights-only run; +it is not silently combined with a restored optimizer and scheduler. Matching +step and generation alone do not prove exact dataloader or RNG recovery. +If a restored run has not established a same-step save receipt, an immediate +export to the already completed step fails rather than overwriting that DCP. With `stage_dir` the model half waits: a staged save has overwritten nothing until the copy to the destination, so it invalidates the destination there diff --git a/tests/checkpoints/test_checkpoint_callback.py b/tests/checkpoints/test_checkpoint_callback.py index b848273845..2da910ddf0 100644 --- a/tests/checkpoints/test_checkpoint_callback.py +++ b/tests/checkpoints/test_checkpoint_callback.py @@ -695,8 +695,9 @@ def drop_presence(flag, op=None): assert cb.load_global_state() is None trainer.channel_loss_callback.load_state_dict.assert_not_called() assert trainer.state.global_step == 0 - mock_dist.all_reduce.assert_called_once() - assert mock_dist.all_reduce.call_args.kwargs["op"] is torch.distributed.ReduceOp.MIN + assert mock_dist.all_reduce.call_count == 2 + assert mock_dist.all_reduce.call_args_list[0].kwargs["op"] is torch.distributed.ReduceOp.MIN + assert mock_dist.all_reduce.call_args_list[1].kwargs["op"] is torch.distributed.ReduceOp.MAX def test_load_restores_0_1_12_extra_state_job_cursor(self, mock_dist, tmp_path): mock_dist.is_initialized.return_value = False diff --git a/tests/checkpoints/test_multi_optimizer_trainer_step.py b/tests/checkpoints/test_multi_optimizer_trainer_step.py new file mode 100644 index 0000000000..7eb3cf0e74 --- /dev/null +++ b/tests/checkpoints/test_multi_optimizer_trainer_step.py @@ -0,0 +1,270 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 + +"""Trainer state must match the optimizer step restored from DCP.""" + +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch.distributed as dist + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from tests.checkpoints.test_checkpoint_callback import _make_mock_trainer +from tests.tools.launch_utils import spawn_with_rendezvous +from veomni.trainer.callbacks.global_state_callback import GlobalStateCallback + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +def _trainer(tmp_path, *, optimizer_step=3, strict=True): + trainer = _make_mock_trainer(save_path=str(tmp_path)) + trainer.args.train.checkpoint.load_path = str(tmp_path) + trainer.state = SimpleNamespace(global_step=0) + trainer.train_dataloader = None + trainer.data_iterator = SimpleNamespace(state_dict=lambda: {}) + trainer.environ_meter.state_dict.return_value = {} + trainer.channel_loss_callback = None + trainer.offline_cache_writer = None + runtime = SimpleNamespace(optimizer=SimpleNamespace(_is_multi_optimizer=strict)) + runtime.checkpoint = SimpleNamespace( + runtime=runtime, + _loaded_optimizer_global_step=optimizer_step if strict else None, + _loaded_optimizer_generation="generation" if strict else None, + ) + trainer.model = runtime + return trainer + + +@pytest.mark.parametrize( + "mode", + [ + "missing", + "wrong_step", + "invalid_step", + "optimizer_step_missing", + "generation_missing", + "cursor_generation_missing", + "cursor_generation_wrong", + "matching", + "legacy_missing", + ], +) +def test_trainer_state_cannot_silently_diverge_from_restored_optimizer(tmp_path, mode): + trainer = _trainer(tmp_path, strict=mode != "legacy_missing") + if mode == "optimizer_step_missing": + trainer.model.checkpoint._loaded_optimizer_global_step = None + if mode == "generation_missing": + trainer.model.checkpoint._loaded_optimizer_generation = None + callback = GlobalStateCallback(trainer) + state = ( + None + if mode in ("missing", "legacy_missing") + else { + "global_step": 4 if mode == "wrong_step" else 3, + "_veomni_optimizer_generations": {"model": "generation"}, + } + ) + if mode == "invalid_step": + state["global_step"] = True + if mode == "cursor_generation_missing": + state.pop("_veomni_optimizer_generations") + if mode == "cursor_generation_wrong": + state["_veomni_optimizer_generations"]["model"] = "older-generation" + with ( + patch.object(callback, "_read_current", return_value=state), + patch("veomni.checkpoint.legacy_v0_1_12.apply_legacy_global_state", return_value=None), + patch("veomni.trainer.callbacks.global_state_callback.set_device_rng_state"), + ): + if mode in ( + "missing", + "wrong_step", + "invalid_step", + "optimizer_step_missing", + "generation_missing", + "cursor_generation_missing", + "cursor_generation_wrong", + ): + with pytest.raises( + RuntimeError, + match="Trainer state is missing|checkpoint steps differ|step is missing|step or generation|generations differ", + ): + callback.load_global_state() + assert trainer.state.global_step == 0 + else: + callback.load_global_state() + assert trainer.state.global_step == (3 if mode == "matching" else 0) + + +@pytest.mark.parametrize("mode", ["missing", "stale_generation", "matching"]) +def test_real_dcp_with_scheduler_requires_matching_trainer_generation(tmp_path, mode): + import torch + + from tests.checkpoints.test_checkpoint_callback import _make_mock_runtime + from veomni.checkpoint import layout + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer + from veomni.models.checkpoint.manager import ModelCheckpointManager + from veomni.optim.lr_scheduler import build_lr_scheduler + from veomni.optim.optimizer import MultiOptimizer + from veomni.trainer.callbacks.base import TrainerState + + def build(): + runtime = _make_mock_runtime(save_path=str(tmp_path)) + runtime.model = torch.nn.Linear(2, 2) + leaf = torch.optim.AdamW(runtime.model.parameters(), lr=0.01) + runtime.optimizer = MultiOptimizer(runtime.model, {"model": leaf}, ["model"]) + runtime.lr_scheduler = build_lr_scheduler(runtime.optimizer, train_steps=100, lr=0.01) + runtime.parallel_state = SimpleNamespace(global_rank=0, dp_mode="ddp", extra_parallel_names=[]) + return runtime + + with ( + patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer), + patch("torch.distributed.barrier"), + ): + source = build() + manager = ModelCheckpointManager(source) + for _ in range(3): + source.model(torch.ones(1, 2)).sum().backward() + source.optimizer.step() + source.optimizer.zero_grad() + source.lr_scheduler.step() + manager.save_dcp(TrainerState(global_step=3)) + source.checkpoint = manager + if mode == "matching": + source_trainer = _make_mock_trainer(save_path=str(tmp_path)) + source_trainer.model = source + source_trainer.train_dataloader = None + source_trainer.data_iterator = SimpleNamespace(state_dict=lambda: {}) + source_trainer.environ_meter.state_dict.return_value = {} + source_trainer.channel_loss_callback = None + source_trainer.offline_cache_writer = None + GlobalStateCallback(source_trainer).save_global_state(TrainerState(global_step=3)) + elif mode == "stale_generation": + sidecar = layout.extra_state_path(str(tmp_path / "global_step_3"), 0) + import os + + os.makedirs(os.path.dirname(sidecar), exist_ok=True) + torch.save({"global_step": 3, "_veomni_optimizer_generations": {"model": "stale"}}, sidecar) + target = build() + target.train_args.checkpoint.load_path = str(tmp_path / "global_step_3") + target.checkpoint = ModelCheckpointManager(target) + target.checkpoint.load() + assert target.checkpoint._loaded_optimizer_global_step == 3 + assert target.lr_scheduler["model"].last_epoch == 3 + trainer = _make_mock_trainer(save_path=str(tmp_path)) + trainer.args.train.checkpoint.load_path = target.train_args.checkpoint.load_path + trainer.model = target + trainer.state = SimpleNamespace(global_step=0) + trainer.train_dataloader = None + if mode == "matching": + GlobalStateCallback(trainer).load_global_state() + assert trainer.state.global_step == 3 + else: + with pytest.raises(RuntimeError, match="Trainer state is missing|generations differ"): + GlobalStateCallback(trainer).load_global_state() + assert trainer.state.global_step == 0 + assert target.lr_scheduler["model"].last_epoch == 3 + + +def test_train_end_sidecar_keeps_generation_after_optimizer_release(tmp_path): + import torch + + from veomni.checkpoint import layout + from veomni.trainer.callbacks.base import TrainerState + + trainer = _trainer(tmp_path) + trainer.model.optimizer = None + trainer.model.checkpoint._last_saved_step = 3 + trainer.model.checkpoint._last_saved_optimizer_generation = "saved-generation" + trainer.environ_meter.state_dict.return_value = {} + callback = GlobalStateCallback(trainer) + callback.save_global_state(TrainerState(global_step=3)) + sidecar = layout.extra_state_path(str(tmp_path / "global_step_3"), 0) + saved = torch.load(sidecar, map_location="cpu", weights_only=False) + assert saved["_veomni_optimizer_generations"] == {"model": "saved-generation"} + + +def _resume_worker(rank, tmp_path, mode, endpoint): + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + _init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cpu") + step = 4 if mode == "optimizer_step" and rank == 1 else 3 + trainer = _trainer(tmp_path, optimizer_step=step, strict=mode != "mixed_optimizer" or rank == 0) + trainer.args.train.global_rank = rank + trainer.args.train.world_size = 2 + callback = GlobalStateCallback(trainer) + state = ( + None + if mode == "missing" and rank == 1 + else { + "global_step": 4 if mode == "trainer_step" and rank == 1 else 3, + "_veomni_optimizer_generations": { + "model": "older" if mode == "trainer_generation" and rank == 1 else "generation" + }, + } + ) + with ( + patch.object(callback, "_read_current", return_value=state), + patch("veomni.checkpoint.legacy_v0_1_12.apply_legacy_global_state", return_value=None), + patch("veomni.trainer.callbacks.global_state_callback.set_device_rng_state"), + ): + if mode == "matching": + callback.load_global_state() + assert trainer.state.global_step == 3 + else: + with pytest.raises( + RuntimeError, + match="Trainer state is missing|checkpoint steps differ|steps differ across ranks|" + "generations differ|checking the trainer checkpoint step failed on another rank", + ): + callback.load_global_state() + assert trainer.state.global_step == 0 + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +@pytest.mark.parametrize( + "mode", ["missing", "optimizer_step", "trainer_step", "trainer_generation", "mixed_optimizer", "matching"] +) +def test_multi_optimizer_trainer_resume_is_rank_coordinated(tmp_path, mode): + spawn_with_rendezvous(_resume_worker, args=(str(tmp_path), mode), nprocs=2) + + +def _save_failure_worker(rank, tmp_path, mode, endpoint): + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state + from veomni.trainer.callbacks.base import TrainerState + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + _init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cpu") + trainer = _trainer(tmp_path) + trainer.args.train.global_rank = rank + trainer.args.train.world_size = 2 + trainer.model.checkpoint._last_saved_step = 3 + trainer.model.checkpoint._last_saved_optimizer_generation = ( + None if mode == "generation" and rank == 1 else "gen" + ) + callback = GlobalStateCallback(trainer) + if mode == "cursor" and rank == 1: + callback.state_dict = lambda _: (_ for _ in ()).throw(ValueError("injected cursor failure")) + with pytest.raises( + (RuntimeError, ValueError), match="trainer checkpoint state|injected cursor failure|generation" + ): + callback.save_global_state(TrainerState(global_step=3)) + dist.barrier() + assert not (tmp_path / "global_step_3").exists() + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +@pytest.mark.parametrize("mode", ["cursor", "generation"]) +def test_trainer_save_failure_is_coordinated_before_write(tmp_path, mode): + spawn_with_rendezvous(_save_failure_worker, args=(tmp_path, mode), nprocs=2) diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index e972b672f6..daf3e5e1a4 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -37,6 +37,7 @@ """ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Dict, Optional import torch @@ -59,6 +60,63 @@ _LOADER_KEYS = ("train_dataloader",) +def _multi_optimizer_checkpoints(trainer): + """Find actual model checkpoint owners, including Omni module runtimes.""" + root = getattr(trainer, "model", None) + runtimes = [("model", root)] + modules = getattr(root, "module_runtimes", None) + if isinstance(modules, Mapping): + runtimes.extend((f"module:{name}", runtime) for name, runtime in sorted(modules.items())) + result = [] + seen = set() + for name, runtime in runtimes: + if runtime is None: + continue + optimizer = getattr(runtime, "optimizer", None) + checkpoint = getattr(runtime, "checkpoint", None) + loaded_step = getattr(checkpoint, "_loaded_optimizer_global_step", None) + saved_generation = getattr(checkpoint, "_last_saved_optimizer_generation", None) + if ( + getattr(optimizer, "_is_multi_optimizer", False) is not True + and type(loaded_step) is not int + and not isinstance(saved_generation, str) + ): + continue + if checkpoint is None or id(checkpoint) in seen: + raise RuntimeError("Restored MultiOptimizer has no unique checkpoint owner") + owner = getattr(checkpoint, "runtime", None) + if owner is not None and owner is not runtime: + raise RuntimeError("Restored MultiOptimizer checkpoint belongs to another runtime") + seen.add(id(checkpoint)) + result.append((name, checkpoint)) + return result + + +def _restored_multi_optimizer_state(trainer): + """Read the proven optimizer step and generations loaded by the runtimes.""" + steps, generations = [], {} + for name, checkpoint in _multi_optimizer_checkpoints(trainer): + step = getattr(checkpoint, "_loaded_optimizer_global_step", None) + generation = getattr(checkpoint, "_loaded_optimizer_generation", None) + if type(step) is not int or step < 0 or not isinstance(generation, str) or not generation: + raise RuntimeError("Restored MultiOptimizer step or generation is missing") + steps.append(step) + generations[name] = generation + if len(set(steps)) > 1: + raise RuntimeError("Model runtimes restored different MultiOptimizer steps") + return (steps[0] if steps else None), generations + + +def _saved_multi_optimizer_generations(trainer, step): + generations = {} + for name, checkpoint in _multi_optimizer_checkpoints(trainer): + generation = getattr(checkpoint, "_last_saved_optimizer_generation", None) + if getattr(checkpoint, "_last_saved_step", None) != step or not isinstance(generation, str) or not generation: + raise RuntimeError("Trainer state has no matching MultiOptimizer checkpoint generation") + generations[name] = generation + return generations + + class GlobalStateCallback(Callback): """Save and resume the state that belongs to the job rather than to a model. @@ -160,8 +218,16 @@ def save_global_state(self, state: TrainerState) -> None: it clears. """ args: "VeOmniArguments" = self.trainer.args - step_root = layout.step_dir(args.train.checkpoint.save_path, state.global_step) - payload = self.state_dict(state) + step_root, payload, generations, error = None, None, {}, None + try: + step_root = layout.step_dir(args.train.checkpoint.save_path, state.global_step) + payload = self.state_dict(state) + generations = _saved_multi_optimizer_generations(self.trainer, state.global_step) + except Exception as exc: # noqa: BLE001 - all ranks must fail before writing + error = exc + raise_if_any_rank_failed(error, "preparing trainer checkpoint state") + if generations: + payload["_veomni_optimizer_generations"] = generations loader_payload = {key: payload[key] for key in _LOADER_KEYS if key in payload} extra_payload = {key: value for key, value in payload.items() if key not in _LOADER_KEYS} @@ -247,6 +313,13 @@ def load_global_state(self) -> Optional[Dict[str, Any]]: if load_path is None: return None + restored_step, generations, optimizer_error = None, {}, None + try: + restored_step, generations = _restored_multi_optimizer_state(self.trainer) + except Exception as exc: # noqa: BLE001 - coordinate before any rank reads state + optimizer_error = exc + raise_if_any_rank_failed(optimizer_error, "checking the restored MultiOptimizer step") + state_path = layout.extra_state_path(load_path, self.rank) # A file that is present but unreadable is not the same as an absent one: # the reduction below treats absence as "resume weights only", which would @@ -269,17 +342,39 @@ def load_global_state(self) -> Optional[Dict[str, Any]]: found = current_state is not None or legacy_state is not None if dist.is_initialized(): - flag = torch.tensor([int(found)], dtype=torch.int32, device=get_device_type()) + device = "cpu" if dist.get_backend() == "gloo" else get_device_type() + flag = torch.tensor([int(found)], dtype=torch.int32, device=device) dist.all_reduce(flag, op=torch.distributed.ReduceOp.MIN) found = bool(flag.item()) - if not found: + strict = torch.tensor([int(restored_step is not None)], dtype=torch.int32, device=device) + dist.all_reduce(strict, op=torch.distributed.ReduceOp.MAX) + if strict.item(): + status = torch.tensor( + [restored_step if restored_step is not None else -1], dtype=torch.int64, device=device + ) + statuses = [torch.empty_like(status) for _ in range(dist.get_world_size())] + dist.all_gather(statuses, status) + if len({int(item.item()) for item in statuses}) != 1: + raise RuntimeError("Restored MultiOptimizer steps differ across ranks") + if not found: + if restored_step is not None: + raise RuntimeError("Trainer state is missing; exact MultiOptimizer resume is unsafe") + if dist.is_initialized(): logger.warning_rank0("Trainer state missing on at least one rank; resuming weights only.") - return None - elif not found: - logger.warning(f"No trainer state at {state_path}; resuming weights only.") + else: + logger.warning(f"No trainer state at {state_path}; resuming weights only.") return None global_state = current_state if current_state is not None else legacy_state + error = None + if restored_step is not None: + if not isinstance(global_state, dict) or type(global_state.get("global_step")) is not int: + error = RuntimeError("Trainer checkpoint step is missing or invalid") + elif global_state["global_step"] != restored_step: + error = RuntimeError("Trainer and MultiOptimizer checkpoint steps differ") + elif global_state.get("_veomni_optimizer_generations") != generations: + error = RuntimeError("Trainer and MultiOptimizer checkpoint generations differ or are missing") + raise_if_any_rank_failed(error, "checking the trainer checkpoint step") self.trainer.state.global_step = global_state["global_step"] self._restore_position(global_state) From c25cf3ff3aa617c084771bb1dce9a4913b720b47 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:36:37 +0800 Subject: [PATCH 030/120] Document rank-local adapter and support nested Omni hooks --- .../rank_local_adamw_checkpoint_adapter.md | 95 +++++++++++ docs/index.md | 1 + docs/usage/checkpoint.md | 1 + .../test_qwen_omni_nested_checkpoint_hooks.py | 154 ++++++++++++++++++ .../checkpoint/multi_optimizer_checkpoint.py | 33 +++- 5 files changed, 278 insertions(+), 6 deletions(-) create mode 100644 docs/design/rank_local_adamw_checkpoint_adapter.md create mode 100644 tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py diff --git a/docs/design/rank_local_adamw_checkpoint_adapter.md b/docs/design/rank_local_adamw_checkpoint_adapter.md new file mode 100644 index 0000000000..1b365d1507 --- /dev/null +++ b/docs/design/rank_local_adamw_checkpoint_adapter.md @@ -0,0 +1,95 @@ +# Rank-local AdamW checkpoint adapters + +`register_rank_local_adamw` is an internal extension point for an AdamW-like +optimizer whose master weights or moments must retain their own dtype or CPU +placement when restored through `MultiOptimizer` and `DistributedCheckpointer`. +It is not a generic optimizer plugin, a security boundary, or a conversion tool +for an incompatible checkpoint. Register trusted code only. + +## Registration and ownership + +```python +from contextlib import nullcontext + +import torch + +from veomni.optim.checkpoint_adapter import register_rank_local_adamw + +# Illustrative callbacks for an AdamW-like optimizer that already implements +# precision-preserving state loading; MyAdamW is the application's own class. +register_rank_local_adamw( + MyAdamW, + context_factory=lambda optimizer: nullcontext(), + algorithm_identity=lambda optimizer: ("my-adamw-v1", str(optimizer.moment_dtype)), + state_dtype=lambda optimizer, param, field: ( + torch.float32 if field in ("step", "master_weight") else optimizer.moment_dtype + ), + state_is_offloaded=lambda optimizer, param: optimizer.is_offloaded(param), +) +``` + +The registry is process-local and matches `type(optimizer)` **exactly**; +register each supported concrete optimizer type before saving or loading on +every rank. Duplicate registration raises an error. Registration does not make +an arbitrary optimizer compatible or automatically preserve its precision: the +optimizer and its adapter own the correctness of their serialization and load +behavior. The in-tree checkpoint adapter tests provide an illustrative +registered implementation. + +An initialized parameter entry must contain exactly `step`, `master_weight`, +`exp_avg`, and `exp_avg_sq`. The step is a rank-local plain scalar tensor, not +a scalar DTensor; the other fields have the parameter's shape. A parameter with +no optimizer state must stay absent, rather than receiving synthetic moments +or a synthetic step. The current adapter contract does not support additional +per-parameter fields. + +## Callback contract + +| Callback | Signature and result | Responsibility | +| --- | --- | --- | +| `context_factory` | `context_factory(optimizer) -> context manager` | Entered while the framework obtains or applies that leaf optimizer's state. It may temporarily select an optimizer's supported serialization mode, but must expose the live initialized fields and group settings without mutating training state. On restore, it must allow the optimizer's loader to install the saved precision and placement without an implicit parameter-dtype cast. `nullcontext()` is sufficient only when the optimizer already has that behavior. | +| `algorithm_identity` | `algorithm_identity(optimizer) -> immutable value` | Describes non-group settings that change updates or restore semantics. Use `None`, booleans, integers, floats, strings, or tuples recursively composed of those types; encode dtypes as strings. Include an algorithm/schema version and all such settings. The checkpoint records this value and rejects a different identity on resume. Group settings and scheduler identity are checked separately. | +| `state_dtype` | `state_dtype(optimizer, param, field) -> torch.dtype` | Returns the expected stored dtype for each of `step`, `master_weight`, `exp_avg`, and `exp_avg_sq`, including when the state is not initialized yet. `step` describes the scalar dtype, not the moment dtype. The framework records the policy and validates existing and restored tensors against it; this callback does not cast tensors. | +| `state_is_offloaded` | `state_is_offloaded(optimizer, param) -> bool` | Reports whether that parameter's master and moment state belongs on CPU. The per-parameter policy is part of checkpoint identity and selects CPU versus parameter-device allocation during restoration. The scalar step follows the optimizer group's capturable/fused placement rules separately. | + +Callbacks must be deterministic for the same training configuration and +stable across the save and load. `optimizer` is the live concrete leaf, and +`param` is a parameter owned by that leaf, possibly a DTensor. A fresh process +constructs new optimizer and parameter objects; callbacks must reconstruct the +same policies for them. Parameter ownership, optimizer group +membership, scheduler identity, state dtypes and offload selection must match +the saved checkpoint. Change these only with an explicitly verified migration. + +## Framework call sites and restore behavior + +1. `MultiOptimizer._sparse_state_dict()` enters the context around PyTorch's + optimizer-state serialization. For registered optimizers it checks that + live tensor versions/storage were not changed, that each initialized state + field and group option is represented faithfully, and that no state was + fabricated. Each serialized live state tensor must be the same view/storage + as its live tensor, including device, dtype, shape, stride and offset; an + equal-valued clone is not enough. State-empty leaves retain their parameter + groups without creating optimizer moments. +2. Checkpoint layout construction records the algorithm identity, parameter + ownership, dtype policy and offload policy. Save validates the source state; + load checks the checkpoint's identity, field names, dtypes and shapes before + preparing initialized state for the saved owners. Rank-local ownership and + the supported parallel geometry are not automatically resharded. +3. `MultiOptimizer._dcp_load_state_dict()` enters the context again while + PyTorch applies the optimizer state. The framework checks the resulting + registered-optimizer payload for omissions or changes. An optimizer may + implement its own precision-preserving loader; the adapter callbacks alone + do not repair a lossy one. + +The currently supported exact-precision training-checkpoint path requires +`MultiOptimizer` and synchronous `DistributedCheckpointer` save/load. +`MultiOptimizer.load_state_dict()` and raw DCP `OptimizerState` reject +unproven recovery; an individual leaf's direct `load_state_dict()` cannot prove +the MultiOptimizer group, scheduler and generation identity. Unverified older +group layouts, unsupported optimizer hooks, and changing the rank-local layout +require an explicit verified migration or fail closed. See +[MultiOptimizer recovery identity](../usage/checkpoint.md#multioptimizer-recovery-identity) +for the complete recovery constraints. Test an extension with initialized and +unused parameters, save, fresh-process restore, and the next update at its +actual dtype and placement; an adapter registration or a single forward pass +is not an acceptance test. diff --git a/docs/index.md b/docs/index.md index ad69bbb332..586cf9bba3 100644 --- a/docs/index.md +++ b/docs/index.md @@ -120,6 +120,7 @@ seed_omni/mixins/offline_encoding.md design/ci_case_scheduling.md design/op_selection.md design/module_config.md +design/rank_local_adamw_checkpoint_adapter.md design/fused_moe_ops.md design/deepgemm_fp8_moe.md design/local_parallel_state.md diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 6d45ad98d2..0f37819ac1 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -338,5 +338,6 @@ checkpoints and changed freeze configurations require verified migration. ## Related pages - [Checkpoint conversion](checkpoint_conversion.md) — DCP shards → HuggingFace safetensors (`scripts/merge_dcp_to_hf.py`). +- [Rank-local AdamW checkpoint adapters](../design/rank_local_adamw_checkpoint_adapter.md) — registration, ownership and save/load contract. - [Trainer callbacks](trainer.md#callbacks) — `CheckpointCallback` (when) vs `GlobalStateCallback` (job cursor). - [LoRA checkpoints](../key_features/lora.md#4-checkpoint-saving) — adapter export under `lora_ckpt/`. diff --git a/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py b/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py new file mode 100644 index 0000000000..fcac9db656 --- /dev/null +++ b/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py @@ -0,0 +1,154 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 + +"""Known Omni load hooks must be safe when PEFT nests the model.""" + +import pytest +import torch + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.checkpoint.multi_optimizer_checkpoint import _check_model_serialization_hooks + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +class Qwen3OmniMoeForConditionalGeneration(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(2, 2) + + def _drop_talker_and_code2wav_keys( + module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + if prefix: + return + for key in list(state_dict): + if key.startswith(("talker.", "code2wav.")): + del state_dict[key] + + self.register_load_state_dict_pre_hook(_drop_talker_and_code2wav_keys) + + +Qwen3OmniMoeForConditionalGeneration.__module__ = "veomni.models.transformers.qwen3_omni_moe.generated.test" + + +class Qwen2_5OmniForConditionalGeneration(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(2, 2) + + def _drop_talker_and_token2wav_keys( + module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + if prefix: + return + for key in list(state_dict): + if key.startswith(("talker.", "token2wav.")): + del state_dict[key] + + self.register_load_state_dict_pre_hook(_drop_talker_and_token2wav_keys) + + +Qwen2_5OmniForConditionalGeneration.__module__ = "veomni.models.transformers.qwen2_5_omni.generated.test" + + +@pytest.mark.parametrize("model_class", [Qwen3OmniMoeForConditionalGeneration, Qwen2_5OmniForConditionalGeneration]) +def test_nested_builtin_omni_hook_allows_peft_prefix(model_class): + root = torch.nn.Module() + root.base_model = torch.nn.Module() + root.base_model.model = model_class() + _check_model_serialization_hooks(root) + + +def test_builtin_omni_hook_allows_actual_peft_wrapper(): + peft = pytest.importorskip("peft") + model = peft.get_peft_model( + Qwen3OmniMoeForConditionalGeneration(), peft.LoraConfig(r=2, lora_alpha=4, target_modules=["linear"]) + ) + assert "base_model.model" in dict(model.named_modules()) + _check_model_serialization_hooks(model) + saved = {name: value.clone() for name, value in model.state_dict().items()} + with torch.no_grad(): + for parameter in model.parameters(): + parameter.fill_(42) + model.load_state_dict(saved, strict=True) + for name, value in model.state_dict().items(): + torch.testing.assert_close(value, saved[name], rtol=0, atol=0) + + +def test_root_builtin_omni_hook_still_rejects_filtered_model_tensors(): + root = Qwen3OmniMoeForConditionalGeneration() + root.talker = torch.nn.Linear(2, 2) + with pytest.raises(RuntimeError, match="load hook is a no-op"): + _check_model_serialization_hooks(root) + + +def test_unrecognized_nested_hook_still_fails_closed(): + root = torch.nn.Module() + root.base_model = torch.nn.Module() + root.base_model.model = Qwen3OmniMoeForConditionalGeneration() + root.base_model.model.register_load_state_dict_pre_hook(lambda *args: None) + with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): + _check_model_serialization_hooks(root) + + +def test_known_name_nested_hook_that_changes_state_still_fails_closed(): + class Qwen3OmniMoeForConditionalGeneration(torch.nn.Module): + def __init__(self): + super().__init__() + + def _drop_talker_and_code2wav_keys( + module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + state_dict.clear() + + self.register_load_state_dict_pre_hook(_drop_talker_and_code2wav_keys) + + Qwen3OmniMoeForConditionalGeneration.__module__ = "veomni.models.transformers.qwen3_omni_moe.generated.test" + root = torch.nn.Module() + root.base_model = Qwen3OmniMoeForConditionalGeneration() + with pytest.raises(RuntimeError, match="nested Qwen Omni load hook"): + _check_model_serialization_hooks(root) + + +def test_known_name_nested_hook_that_changes_real_weight_is_rejected(): + class Qwen3OmniMoeForConditionalGeneration(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(2, 2) + + def _drop_talker_and_code2wav_keys( + module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + if prefix + "linear.weight" in state_dict: + state_dict[prefix + "linear.weight"].zero_() + + self.register_load_state_dict_pre_hook(_drop_talker_and_code2wav_keys) + + Qwen3OmniMoeForConditionalGeneration.__module__ = "veomni.models.transformers.qwen3_omni_moe.generated.test" + root = torch.nn.Module() + root.base_model = Qwen3OmniMoeForConditionalGeneration() + with pytest.raises(RuntimeError, match="nested Qwen Omni load hook"): + _check_model_serialization_hooks(root) + + +def test_known_qwen25_nested_hook_that_changes_errors_is_rejected(): + class Qwen2_5OmniForConditionalGeneration(torch.nn.Module): + def __init__(self): + super().__init__() + + def _drop_talker_and_token2wav_keys( + module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + if prefix: + error_msgs.append("unexpected mutation") + return + + self.register_load_state_dict_pre_hook(_drop_talker_and_token2wav_keys) + + Qwen2_5OmniForConditionalGeneration.__module__ = "veomni.models.transformers.qwen2_5_omni.generated.test" + root = torch.nn.Module() + root.base_model = Qwen2_5OmniForConditionalGeneration() + with pytest.raises(RuntimeError, match="nested Qwen Omni load hook"): + _check_model_serialization_hooks(root) diff --git a/veomni/checkpoint/multi_optimizer_checkpoint.py b/veomni/checkpoint/multi_optimizer_checkpoint.py index 45d438f03c..810e3d660c 100644 --- a/veomni/checkpoint/multi_optimizer_checkpoint.py +++ b/veomni/checkpoint/multi_optimizer_checkpoint.py @@ -269,6 +269,7 @@ def _check_qwen_omni_checkpoint_keys(model, keys): def _check_model_serialization_hooks(model): """Allow the FSDP2 reshard hook, reject unproven model transforms.""" + import dis import types from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam @@ -280,6 +281,22 @@ def nested_code(parent, name): item for item in parent.__code__.co_consts if isinstance(item, types.CodeType) and item.co_name == name } + def returns_before_nested_load(hook): + # A sample call cannot prove this: the hook could modify other real + # keys, the module or error lists without changing the probe. + ops = [item for item in dis.get_instructions(hook) if item.opname not in ("RESUME", "NOP", "CACHE")] + if len(ops) < 3 or (ops[0].opname, ops[0].argval) != ("LOAD_FAST", "prefix"): + return False + if ops[1].opname not in ("POP_JUMP_IF_FALSE", "POP_JUMP_FORWARD_IF_FALSE"): + return False + if (ops[2].opname, ops[2].argval) == ("RETURN_CONST", None): + return True + return ( + len(ops) >= 4 + and (ops[2].opname, ops[2].argval) == ("LOAD_CONST", None) + and ops[3].opname == "RETURN_VALUE" + ) + post_codes = nested_code(FSDPParam.__init__, "") pre_codes = nested_code(FSDPParamGroup._register_state_dict_hooks, "to_sharded_hook") for name, module in model.named_modules(): @@ -303,14 +320,18 @@ def nested_code(parent, name): and code in pre_codes and any(isinstance(item, FSDPParamGroup) for item in captured) ) - builtin_qwen = ( - attr == "_load_state_dict_pre_hooks" and not name and _qwen_omni_filtered_prefixes(model, func) - ) + builtin_qwen = attr == "_load_state_dict_pre_hooks" and _qwen_omni_filtered_prefixes(module, func) if not builtin_fsdp and not builtin_qwen: raise RuntimeError(f"Cannot prove model serialization hook {attr} on {name or ''}") - if builtin_qwen and any( - key.startswith(builtin_qwen) - for key, _ in list(model.named_parameters()) + list(model.named_buffers()) + if builtin_qwen and name and not returns_before_nested_load(func): + raise RuntimeError("Cannot prove nested Qwen Omni load hook preserves checkpoint keys") + if ( + builtin_qwen + and not name + and any( + key.startswith(builtin_qwen) + for key, _ in list(model.named_parameters()) + list(model.named_buffers()) + ) ): raise RuntimeError("Cannot prove Qwen Omni load hook is a no-op for checkpoint model tensors") From 13745a36b047a608cd89076deff19b66b161e1d9 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:46:18 +0800 Subject: [PATCH 031/120] Ignore transient-only cache aliases in checkpoint identity --- docs/usage/checkpoint.md | 10 ++ .../test_multi_optimizer_rope_cache.py | 154 ++++++++++++++++++ veomni/checkpoint/optimizer_resume_layout.py | 19 ++- 3 files changed, 182 insertions(+), 1 deletion(-) create mode 100644 tests/checkpoints/test_multi_optimizer_rope_cache.py diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 0f37819ac1..d45113605a 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -311,6 +311,16 @@ identity and rejects restoration; use the proven DCP path instead. The codec also rejects tensor learning rates shared between groups and overlapping storage between distinct model or optimizer tensors. Its byte-range check is conservative for non-contiguous interleaved views. +The saved alias identity covers parameters and persistent buffers, including +their aliases to nonpersistent buffers. Aliases solely among transient buffers, +such as dynamic RoPE caches, are not part of that identity. The separate +conservative storage-overlap check still applies to distinct buffer views. +Older checkpoints whose identity included such aliases can require an +independently verified migration. +Dynamic RoPE frequencies and cached maximum lengths are not checkpointed. If +the cache grew beyond its original context and has not reset, a fresh process +can compute different later outputs; exact resume for that case is not yet +supported. Mutable optimizer options must not be shared between groups or optimizer defaults: ordinary YAML `betas` are normalized to a tuple; custom shared lists, dictionaries and tensor-valued learning rates require explicitly independent diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py new file mode 100644 index 0000000000..2d6d5da90a --- /dev/null +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -0,0 +1,154 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 + +"""Transient RoPE buffers are not part of checkpoint ownership identity.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.multiprocessing as mp +from transformers import Qwen3Config +from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer +from veomni.optim.optimizer import MultiOptimizer + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +class _TinyQwenRope(torch.nn.Module): + def __init__(self): + super().__init__() + self.projection = torch.nn.Linear(8, 8) + config = Qwen3Config( + hidden_size=32, + head_dim=8, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=16, + rope_parameters={"rope_type": "dynamic", "factor": 2.0, "rope_theta": 10000}, + ) + self.rope = Qwen3RotaryEmbedding(config) + + +def _build(): + model = _TinyQwenRope() + leaf = torch.optim.AdamW(model.parameters(), lr=0.01) + return model, MultiOptimizer(model, {"model": leaf}, ["model"]) + + +def _update(model, optimizer, value): + for parameter in model.parameters(): + parameter.grad = torch.full_like(parameter, value) + optimizer.step() + optimizer.zero_grad() + + +def _rope_update(model, optimizer): + inputs = torch.arange(64, dtype=torch.float32).reshape(1, 8, 8) / 64 + projected = model.projection(inputs) + cos, sin = model.rope(projected, torch.arange(8).reshape(1, -1)) + loss = (projected * cos + sin).square().mean() + loss.backward() + trace = {"cos": cos.detach().clone(), "sin": sin.detach().clone(), "loss": loss.detach().clone()} + trace.update({f"grad.{name}": value.grad.detach().clone() for name, value in model.named_parameters()}) + optimizer.step() + optimizer.zero_grad() + return trace + + +def _fresh_resume(_rank, checkpoint, expected_path): + model, optimizer = _build() + assert model.rope.inv_freq is not model.rope.original_inv_freq + DistributedCheckpointer.load( + checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + assert model.rope.inv_freq is not model.rope.original_inv_freq + trace = _rope_update(model, optimizer) + expected = torch.load(expected_path, map_location="cpu", weights_only=True) + actual = {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace} + for group in ("model", "optimizer", "trace"): + assert actual[group].keys() == expected[group].keys() + for key, value in actual[group].items(): + if torch.is_tensor(value): + torch.testing.assert_close(value, expected[group][key], rtol=0, atol=0) + else: + assert value == expected[group][key] + + +def test_old_all_buffer_identity_rejects_real_rope_cache_reset(tmp_path, monkeypatch): + import veomni.checkpoint.optimizer_resume_layout as layout + + monkeypatch.setattr(layout, "persistent_named_buffers", lambda model: model.named_buffers(remove_duplicate=False)) + model, optimizer = _build() + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) + assert model.rope.inv_freq is model.rope.original_inv_freq + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _build() + assert target.rope.inv_freq is not target.rope.original_inv_freq + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + + +def test_real_qwen3_dynamic_rope_reset_saves_and_fresh_process_resumes(tmp_path): + model, optimizer = _build() + _update(model, optimizer, 1.0) + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) + assert model.rope.inv_freq is model.rope.original_inv_freq + assert not any("inv_freq" in key for key in model.state_dict()) + checkpoint = str(tmp_path / "checkpoint") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + trace = _rope_update(model, optimizer) + expected = tmp_path / "expected.pt" + torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace}, expected) + mp.spawn(_fresh_resume, args=(checkpoint, str(expected)), nprocs=1, join=True) + + +def test_distinct_overlapping_nonpersistent_buffer_views_still_fail_closed(tmp_path): + model, optimizer = _build() + storage = torch.zeros(8) + model.register_buffer("cache_view_a", storage[:6], persistent=False) + model.register_buffer("cache_view_b", storage[2:], persistent=False) + with pytest.raises(RuntimeError, match="checkpoint model storage alias"): + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + + +@pytest.mark.parametrize("source_alias", [False, True]) +@pytest.mark.parametrize("anchor_kind", ["parameter", "persistent_buffer"]) +def test_nonpersistent_buffer_persistent_alias_mismatch_still_fails_before_load(tmp_path, source_alias, anchor_kind): + class Model(torch.nn.Module): + def __init__(self, alias): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2)) + self.register_buffer("persistent", torch.ones(2)) + anchor = self.weight if anchor_kind == "parameter" else self.persistent + self.register_buffer("cache", anchor if alias else anchor.detach().clone(), persistent=False) + + def build(alias): + model = Model(alias) + leaf = torch.optim.AdamW(model.parameters(), lr=0.01) + return model, MultiOptimizer(model, {"model": leaf}, ["model"]) + + model, optimizer = build(source_alias) + with torch.no_grad(): + model.weight.fill_(5) + model.persistent.fill_(7) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build(not source_alias) + before = target.weight.detach().clone() + persistent_before = target.persistent.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + torch.testing.assert_close(target.persistent, persistent_before, rtol=0, atol=0) diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index e0b785945f..69ff658242 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -18,6 +18,14 @@ MULTI_LAYOUT_MARKER = "__veomni_multi_optimizer_layout_v1__" +def persistent_named_buffers(model): + """Yield all persistent buffer aliases using their runtime module paths.""" + for prefix, module in model.named_modules(remove_duplicate=False): + for name, value in module._buffers.items(): + if value is not None and name not in module._non_persistent_buffers_set: + yield f"{prefix}.{name}" if prefix else name, value + + def _immutable_value(value): if value is None or type(value) in (bool, int, float, str): return value @@ -202,6 +210,8 @@ def check_aliases(value, path): if names != optimizer_names: raise RuntimeError("MultiOptimizer and checkpoint model have different canonical parameter identities") model_aliases = {} + persistent_objects = set(names) + persistent_objects.update(id(value) for _, value in persistent_named_buffers(model)) for kind, items in ( ("parameter", model.named_parameters(remove_duplicate=False)), ("buffer", model.named_buffers(remove_duplicate=False)), @@ -210,7 +220,14 @@ def check_aliases(value, path): model_aliases.setdefault(id(value), set()).update( f"{kind}:{name}" for name in _get_fqns(model, runtime_name) ) - alias_graph = sorted(sorted(aliases) for aliases in model_aliases.values() if len(aliases) > 1) + # Ignore aliases consisting exclusively of transient caches. A transient + # buffer that aliases a parameter or persistent buffer still affects model + # behavior and must remain bound to the checkpoint's ownership identity. + alias_graph = sorted( + sorted(aliases) + for object_id, aliases in model_aliases.items() + if object_id in persistent_objects and len(aliases) > 1 + ) leaves = [] seen_param_ids, seen_fqns, lr_aliases, lr_owners = set(), set(), {}, {} From a38c191cb3781d8ddc2c8e09ce19e4c5b6604341 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:07:52 +0800 Subject: [PATCH 032/120] Reject unproven dynamic RoPE checkpoint recovery --- docs/usage/checkpoint.md | 20 +- .../test_multi_optimizer_rope_cache.py | 215 ++++++++++++++++-- veomni/checkpoint/dynamic_rope.py | 68 ++++++ veomni/checkpoint/optimizer_resume_layout.py | 26 +-- 4 files changed, 288 insertions(+), 41 deletions(-) create mode 100644 veomni/checkpoint/dynamic_rope.py diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index d45113605a..c45eab8faa 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -311,16 +311,20 @@ identity and rejects restoration; use the proven DCP path instead. The codec also rejects tensor learning rates shared between groups and overlapping storage between distinct model or optimizer tensors. Its byte-range check is conservative for non-contiguous interleaved views. -The saved alias identity covers parameters and persistent buffers, including -their aliases to nonpersistent buffers. Aliases solely among transient buffers, -such as dynamic RoPE caches, are not part of that identity. The separate -conservative storage-overlap check still applies to distinct buffer views. -Older checkpoints whose identity included such aliases can require an -independently verified migration. +The saved alias identity covers parameters and buffers, including +nonpersistent buffer aliases; the separate conservative storage-overlap check +still applies to distinct buffer views. A reset dynamic RoPE cache can have a +different alias graph from a newly constructed model and is conservatively +rejected on resume until that cache state can be restored safely. Dynamic RoPE frequencies and cached maximum lengths are not checkpointed. If the cache grew beyond its original context and has not reset, a fresh process -can compute different later outputs; exact resume for that case is not yet -supported. +can compute different later outputs. MultiOptimizer checkpoints therefore +reject a non-reconstructible dynamic or longrope cache at save. With a dynamic +or longrope target, resume requires a matching cache contract; older checkpoints +without that proof need a verified migration. Older checkpoints do not identify +the source's RoPE configuration, so the operator must also independently verify +that the source and target configurations match. Saving an expanded cache is not +yet supported. Mutable optimizer options must not be shared between groups or optimizer defaults: ordinary YAML `betas` are normalized to a tuple; custom shared lists, dictionaries and tensor-valued learning rates require explicitly independent diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index 2d6d5da90a..fa9d8ba1d3 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -1,17 +1,20 @@ # Copyright 2026 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: Apache-2.0 -"""Transient RoPE buffers are not part of checkpoint ownership identity.""" +"""Dynamic RoPE cache and alias state must not silently change on resume.""" +from datetime import timedelta from types import SimpleNamespace import pytest import torch +import torch.distributed as dist import torch.multiprocessing as mp from transformers import Qwen3Config from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding from scripts.ci.accelerator.launch_plan import LaunchSpec +from tests.tools.launch_utils import spawn_with_rendezvous from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.optim.optimizer import MultiOptimizer @@ -20,7 +23,7 @@ class _TinyQwenRope(torch.nn.Module): - def __init__(self): + def __init__(self, theta=10000): super().__init__() self.projection = torch.nn.Linear(8, 8) config = Qwen3Config( @@ -29,13 +32,13 @@ def __init__(self): num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=16, - rope_parameters={"rope_type": "dynamic", "factor": 2.0, "rope_theta": 10000}, + rope_parameters={"rope_type": "dynamic", "factor": 2.0, "rope_theta": theta}, ) self.rope = Qwen3RotaryEmbedding(config) -def _build(): - model = _TinyQwenRope() +def _build(theta=10000): + model = _TinyQwenRope(theta) leaf = torch.optim.AdamW(model.parameters(), lr=0.01) return model, MultiOptimizer(model, {"model": leaf}, ["model"]) @@ -62,11 +65,9 @@ def _rope_update(model, optimizer): def _fresh_resume(_rank, checkpoint, expected_path): model, optimizer = _build() - assert model.rope.inv_freq is not model.rope.original_inv_freq DistributedCheckpointer.load( checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") ) - assert model.rope.inv_freq is not model.rope.original_inv_freq trace = _rope_update(model, optimizer) expected = torch.load(expected_path, map_location="cpu", weights_only=True) actual = {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace} @@ -79,23 +80,95 @@ def _fresh_resume(_rank, checkpoint, expected_path): assert value == expected[group][key] -def test_old_all_buffer_identity_rejects_real_rope_cache_reset(tmp_path, monkeypatch): - import veomni.checkpoint.optimizer_resume_layout as layout +def test_standard_qwen4_rope_128k_does_not_require_dynamic_cache_contract(): + from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpTextConfig + from transformers.models.qwen4_exp.modeling_qwen4_exp import Qwen4ExpTextRotaryEmbedding - monkeypatch.setattr(layout, "persistent_named_buffers", lambda model: model.named_buffers(remove_duplicate=False)) + from veomni.checkpoint.dynamic_rope import dynamic_rope_checkpoint_contract + + config = Qwen4ExpTextConfig( + hidden_size=32, + head_dim=8, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=262144, + rope_parameters={ + "rope_type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 1.0, + "mrope_section": [1, 1, 2], + }, + ) + rope = Qwen4ExpTextRotaryEmbedding(config) + positions = torch.arange(131072).reshape(1, 1, -1).expand(3, 1, -1) + rope(torch.zeros((1, 131072, 8)), positions) + assert rope.max_seq_len_cached == 262144 + assert dynamic_rope_checkpoint_contract(rope) == [] + + +def test_expanded_dynamic_rope_cache_is_rejected_before_checkpoint_write(tmp_path): model, optimizer = _build() model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) - model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) - assert model.rope.inv_freq is model.rope.original_inv_freq + assert int(model.rope.max_seq_len_cached) == 32 + checkpoint = tmp_path / "checkpoint" + with pytest.raises(RuntimeError, match="Dynamic RoPE cache.*depends on prior sequence lengths"): + DistributedCheckpointer.save( + str(checkpoint), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + assert not checkpoint.exists() + + +def test_old_unproven_expanded_rope_checkpoint_is_rejected_before_model_load(tmp_path, monkeypatch): + import veomni.checkpoint.dynamic_rope as dynamic_rope + + model, optimizer = _build() + with torch.no_grad(): + model.projection.weight.fill_(7) + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + with monkeypatch.context() as old: + old.setattr(dynamic_rope, "dynamic_rope_checkpoint_contract", lambda _model: []) + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) target, restored = _build() - assert target.rope.inv_freq is not target.rope.original_inv_freq + before = target.projection.weight.detach().clone() with pytest.raises(RuntimeError, match="layout|identity"): DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) -def test_real_qwen3_dynamic_rope_reset_saves_and_fresh_process_resumes(tmp_path): +def test_changed_dynamic_rope_config_is_rejected_before_model_load(tmp_path): + model, optimizer = _build() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _build(theta=5000) + before = target.projection.weight.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) + + +def _rank_local_rope_error(rank, checkpoint, endpoint): + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + model, optimizer = _build() + if rank == 1: + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + error = "Dynamic RoPE cache.*prior sequence lengths" if rank == 1 else "layout failed on another rank" + with pytest.raises(RuntimeError, match=error): + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_rank_local_expanded_rope_error_is_coordinated(tmp_path): + spawn_with_rendezvous(_rank_local_rope_error, args=(str(tmp_path / "checkpoint"),), nprocs=2) + assert not (tmp_path / "checkpoint").exists() + + +def test_real_qwen3_dynamic_rope_reset_alias_mismatch_fails_closed(tmp_path): model, optimizer = _build() _update(model, optimizer, 1.0) model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) @@ -106,6 +179,23 @@ def test_real_qwen3_dynamic_rope_reset_saves_and_fresh_process_resumes(tmp_path) DistributedCheckpointer.save( checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") ) + target, restored = _build() + before = target.projection.weight.detach().clone() + assert target.rope.inv_freq is not target.rope.original_inv_freq + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load( + checkpoint, {"model": target, "optimizer": restored}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) + + +def test_unexpanded_dynamic_rope_fresh_process_matches_training_update(tmp_path): + model, optimizer = _build() + _update(model, optimizer, 1.0) + checkpoint = str(tmp_path / "checkpoint") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) trace = _rope_update(model, optimizer) expected = tmp_path / "expected.pt" torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace}, expected) @@ -152,3 +242,98 @@ def build(alias): DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) torch.testing.assert_close(target.persistent, persistent_before, rtol=0, atol=0) + + +@pytest.mark.parametrize("source_alias", [False, True]) +def test_non_rope_transient_buffer_alias_mismatch_still_fails_before_load(tmp_path, source_alias): + class Model(torch.nn.Module): + def __init__(self, alias): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2)) + first = torch.ones(2) + self.register_buffer("cache_a", first, persistent=False) + self.register_buffer("cache_b", first if alias else first.clone(), persistent=False) + + def build(alias): + model = Model(alias) + leaf = torch.optim.AdamW(model.parameters(), lr=0.01) + return model, MultiOptimizer(model, {"model": leaf}, ["model"]) + + model, optimizer = build(source_alias) + with torch.no_grad(): + model.weight.fill_(5) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build(not source_alias) + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + +def test_custom_dynamic_rope_alias_mismatch_still_fails_before_load(tmp_path): + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + class CustomRope(torch.nn.Module): + def __init__(self, alias): + super().__init__() + self.config = Qwen3Config( + hidden_size=32, + head_dim=8, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=16, + rope_parameters={"rope_type": "dynamic", "factor": 2.0, "rope_theta": 10000}, + ) + self.rope_type = "dynamic" + self.original_max_seq_len = self.max_seq_len_cached = 16 + inv, self.attention_scaling = ROPE_INIT_FUNCTIONS["dynamic"](self.config, device="cpu") + self.register_buffer("inv_freq", inv, persistent=False) + self.register_buffer("original_inv_freq", inv if alias else inv.clone(), persistent=False) + + def forward(self): + self.inv_freq.add_(1) + return self.original_inv_freq.clone() + + def build(alias): + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.ones(2)) + model.rope = CustomRope(alias) + leaf = torch.optim.AdamW(model.parameters(), lr=0.01) + return model, MultiOptimizer(model, {"model": leaf}, ["model"]) + + model, optimizer = build(True) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build(False) + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + + +@pytest.mark.parametrize("hook_scope", ["local", "global"]) +def test_rope_forward_hook_disables_alias_exemption(tmp_path, hook_scope): + model, optimizer = _build() + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) + assert model.rope.inv_freq is model.rope.original_inv_freq + + def alter_rope(module, _inputs): + if isinstance(module, Qwen3RotaryEmbedding): + module.original_inv_freq.add_(1) + + if hook_scope == "local": + handle = model.rope.register_forward_pre_hook(alter_rope) + else: + handle = torch.nn.modules.module.register_module_forward_pre_hook(alter_rope) + try: + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = _build() + if hook_scope == "local": + target.rope.register_forward_pre_hook(alter_rope) + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load( + str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel + ) + finally: + handle.remove() diff --git a/veomni/checkpoint/dynamic_rope.py b/veomni/checkpoint/dynamic_rope.py new file mode 100644 index 0000000000..d7fab24fae --- /dev/null +++ b/veomni/checkpoint/dynamic_rope.py @@ -0,0 +1,68 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Checkpoint preflight for dynamic RoPE's nonpersistent runtime state.""" + +import hashlib +import json +import math + +import torch + + +def dynamic_rope_checkpoint_contract(model): + """Reject dynamic caches that a fresh runtime cannot reconstruct exactly. + + HF dynamic RoPE can retain frequencies based on a previous longer input. + DCP does not save these nonpersistent buffers, so an apparently identical + model/optimizer restore can otherwise run the next input differently. + """ + descriptors = [] + for name, module in model.named_modules(): + rope_type = getattr(module, "rope_type", None) + layers = rope_type.items() if isinstance(rope_type, dict) else [(None, rope_type)] + for layer_type, kind in layers: + if not isinstance(kind, str) or ("dynamic" not in kind and kind != "longrope"): + continue + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + prefix = "" if layer_type is None else f"{layer_type}_" + try: + original = int(module.original_max_seq_len) + cached = int(getattr(module, f"{prefix}max_seq_len_cached", module.max_seq_len_cached)) + inv = getattr(module, f"{prefix}inv_freq") + initial_inv = getattr(module, f"{prefix}original_inv_freq") + kwargs = {} if layer_type is None else {"layer_type": layer_type} + expected_inv, expected_scale = ROPE_INIT_FUNCTIONS[kind]( + module.config, device=inv.device, seq_len=None if kind == "longrope" else original, **kwargs + ) + scale = float( + getattr(module, f"{prefix}attention_scaling", getattr(module, "attention_scaling", None)) + ) + config = module.config.to_dict() + # These identify the config source, not its RoPE mathematics. + config.pop("_name_or_path", None) + config.pop("transformers_version", None) + config_hash = hashlib.sha256( + json.dumps(config, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + valid = ( + original > 0 + and cached == original + and not inv.requires_grad + and not initial_inv.requires_grad + and inv.dtype == initial_inv.dtype + and torch.equal(inv, initial_inv.to(device=inv.device)) + and torch.equal(inv, expected_inv.to(device=inv.device, dtype=inv.dtype)) + and math.isfinite(scale) + and scale == float(expected_scale) + and (kind != "longrope" or not hasattr(module, f"{prefix}long_inv_freq")) + ) + except Exception as exc: + raise RuntimeError(f"Cannot prove dynamic RoPE checkpoint cache for {name or ''}") from exc + if not valid: + raise RuntimeError( + f"Dynamic RoPE cache for {name or ''} depends on prior sequence lengths; " + "its nonpersistent state cannot yet be restored exactly" + ) + descriptors.append((name, layer_type, kind, original, str(inv.dtype), config_hash)) + return descriptors diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index 69ff658242..507e1d2c99 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -18,14 +18,6 @@ MULTI_LAYOUT_MARKER = "__veomni_multi_optimizer_layout_v1__" -def persistent_named_buffers(model): - """Yield all persistent buffer aliases using their runtime module paths.""" - for prefix, module in model.named_modules(remove_duplicate=False): - for name, value in module._buffers.items(): - if value is not None and name not in module._non_persistent_buffers_set: - yield f"{prefix}.{name}" if prefix else name, value - - def _immutable_value(value): if value is None or type(value) in (bool, int, float, str): return value @@ -209,9 +201,9 @@ def check_aliases(value, path): optimizer_names.setdefault(id(param), set()).update(_get_fqns(optimizer.model, runtime_name)) if names != optimizer_names: raise RuntimeError("MultiOptimizer and checkpoint model have different canonical parameter identities") + from .dynamic_rope import dynamic_rope_checkpoint_contract + model_aliases = {} - persistent_objects = set(names) - persistent_objects.update(id(value) for _, value in persistent_named_buffers(model)) for kind, items in ( ("parameter", model.named_parameters(remove_duplicate=False)), ("buffer", model.named_buffers(remove_duplicate=False)), @@ -220,14 +212,7 @@ def check_aliases(value, path): model_aliases.setdefault(id(value), set()).update( f"{kind}:{name}" for name in _get_fqns(model, runtime_name) ) - # Ignore aliases consisting exclusively of transient caches. A transient - # buffer that aliases a parameter or persistent buffer still affects model - # behavior and must remain bound to the checkpoint's ownership identity. - alias_graph = sorted( - sorted(aliases) - for object_id, aliases in model_aliases.items() - if object_id in persistent_objects and len(aliases) > 1 - ) + alias_graph = sorted(sorted(aliases) for aliases in model_aliases.values() if len(aliases) > 1) leaves = [] seen_param_ids, seen_fqns, lr_aliases, lr_owners = set(), set(), {}, {} @@ -389,6 +374,11 @@ def mesh_identity(mesh): "extra_digest": extra_digest, "frozen_digest": frozen_digest, } + rope_contract = dynamic_rope_checkpoint_contract(model) + if rope_contract: + # Earlier checkpoints did not prove that transient cache history was + # reconstructible by a fresh process. + layout["dynamic_rope_cache_v1"] = rope_contract return hashlib.sha256( json.dumps(layout, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() ).hexdigest() From 6669054d1a4c114a5000512dc034ffba28853dd6 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:19:35 +0800 Subject: [PATCH 033/120] Apply canonical checkpoint layout to Omni modules --- docs/usage/checkpoint.md | 5 ++ .../seed_omni/runtime/test_module_runtime.py | 88 +++++++++++++++++++ .../runtime/test_omni_model_runtime.py | 71 +++++++++++++++ veomni/models/model_runtime.py | 27 ++++-- .../omni_model/omni_model_runtime.py | 18 ++++ .../omni_module/omni_module_runtime.py | 20 ++++- 6 files changed, 219 insertions(+), 10 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index c45eab8faa..cf3b6ad4d8 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -252,6 +252,11 @@ resume instead of falling back. ## Versioned ExtraParallel layout +For SeedOmni, `extra_parallel_canonical_layout: true` is applied after each +module-scope wrap. Combined `fsdp_scope: model` and active ExtraParallel is not +supported with this option and fails before wrapping; use module scope. All +ranks participating in a module or model scope must use the same layout flag. + The EP-outer layout is opt-in through `SpecInfo.checkpoint_layout_v2`; existing plans remain on their original layout. A plan enabling it must retain the original `global_shape` and must use FSDP DTensor parameters. Every diff --git a/tests/seed_omni/runtime/test_module_runtime.py b/tests/seed_omni/runtime/test_module_runtime.py index 737389356c..4d41444aa6 100644 --- a/tests/seed_omni/runtime/test_module_runtime.py +++ b/tests/seed_omni/runtime/test_module_runtime.py @@ -62,6 +62,94 @@ def test_module_runtime_is_a_model_runtime(): assert issubclass(ModuleRuntime, VeOmniModelRuntime) +@pytest.mark.parametrize( + "enabled,persistent_dim,expected", [(False, None, False), (True, None, True), (True, 1, False)] +) +def test_module_scope_applies_canonical_checkpoint_plan_after_wrap(monkeypatch, enabled, persistent_dim, expected): + from torch.distributed.tensor import Shard + + spec = SimpleNamespace(placement=Shard(0), persistent_fsdp_shard_dim=persistent_dim, checkpoint_layout_v2=False) + model = nn.Linear(2, 2) + model._fqn2spec_info = {"weight": spec} + for step in ( + "setup", + "_build_model_assets", + "_freeze_model_module", + "_build_parallelized_model", + "build_checkpoint", + ): + monkeypatch.setattr(ModuleRuntime, step, lambda self, *a, **k: None) + monkeypatch.setattr(ModuleRuntime, "_build_model", lambda self: setattr(self, "model", model)) + monkeypatch.setattr(ModuleRuntime, "_build_optimizer_checked", lambda self: None) + monkeypatch.setattr(ModuleRuntime, "_scoped", lambda self: nullcontext()) + monkeypatch.setattr( + "veomni.models.model_runtime.get_parallel_state_by_name", + lambda name: SimpleNamespace(any_extra_parallel_enabled=True), + ) + train = SimpleNamespace( + training_task="online_training", + checkpoint=SimpleNamespace(load_path=None, extra_parallel_canonical_layout=enabled), + ) + ModuleRuntime( + SimpleNamespace(accelerator=_fsdp("module")), + "vision_encoder", + module_config=SimpleNamespace(), + global_accelerator=_fsdp("module"), + train_args=train, + ) + assert spec.checkpoint_layout_v2 is expected + + +def test_canonical_checkpoint_layout_rejects_missing_required_plan(): + from veomni.models.model_runtime import _configure_extra_parallel_checkpoint_layout + + checkpoint = SimpleNamespace(extra_parallel_canonical_layout=True) + with pytest.raises(RuntimeError, match="requires an ExtraParallel model plan"): + _configure_extra_parallel_checkpoint_layout(nn.Linear(2, 2), checkpoint, require_plan=True) + _configure_extra_parallel_checkpoint_layout(nn.Linear(2, 2), checkpoint, require_plan=False) + + +def test_module_scope_active_plan_without_wrapped_specs_fails_before_optimizer(monkeypatch): + from torch.distributed.tensor import Shard + + from veomni.distributed.parallel_plan import ParallelPlan + + model = nn.Linear(2, 2) + model.get_parallel_plan = lambda: ParallelPlan(extra_parallel_plan={"ep": {"weight": Shard(0)}}) + for step in ( + "setup", + "_build_model_assets", + "_freeze_model_module", + "_build_parallelized_model", + "build_checkpoint", + ): + monkeypatch.setattr(ModuleRuntime, step, lambda self, *a, **k: None) + monkeypatch.setattr(ModuleRuntime, "_build_model", lambda self: setattr(self, "model", model)) + called = [] + monkeypatch.setattr(ModuleRuntime, "_build_optimizer_checked", lambda self: called.append(True)) + monkeypatch.setattr(ModuleRuntime, "_scoped", lambda self: nullcontext()) + state = SimpleNamespace( + any_extra_parallel_enabled=True, + extra_parallel_sizes=("ep",), + extra_parallel_enabled=lambda name: name == "ep", + ) + monkeypatch.setattr("veomni.models.model_runtime.get_parallel_state_by_name", lambda _name: state) + monkeypatch.setattr("veomni.distributed.torch_parallelize.get_parallel_state", lambda: state) + train = SimpleNamespace( + training_task="online_training", + checkpoint=SimpleNamespace(load_path=None, extra_parallel_canonical_layout=True), + ) + with pytest.raises(RuntimeError, match="requires an ExtraParallel model plan"): + ModuleRuntime( + SimpleNamespace(accelerator=_fsdp("module")), + "vision_encoder", + module_config=SimpleNamespace(), + global_accelerator=_fsdp("module"), + train_args=train, + ) + assert not called + + def test_omni_runtime_arguments_are_model_arguments(): from veomni.arguments import ModelArguments from veomni.arguments.omni_arguments_types import OmniModelRuntimeArguments, OmniModuleRuntimeArguments diff --git a/tests/seed_omni/runtime/test_omni_model_runtime.py b/tests/seed_omni/runtime/test_omni_model_runtime.py index 7b61916772..71d0f7a1d6 100644 --- a/tests/seed_omni/runtime/test_omni_model_runtime.py +++ b/tests/seed_omni/runtime/test_omni_model_runtime.py @@ -414,6 +414,77 @@ def estimate_flops(self, seqlens): assert runtime.metric_meter_collect() == {"metered": (0.0, [])} +def test_composed_canonical_extra_parallel_rejected_before_wrap(monkeypatch): + model = _model() + children = tuple(model._modules) + train = SimpleNamespace(checkpoint=SimpleNamespace(extra_parallel_canonical_layout=True)) + runtime = OmniModelRuntime(model, train_args=train) + runtime.omni_model_runtime_args = SimpleNamespace(accelerator=SimpleNamespace()) + monkeypatch.setattr( + "veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime.composed_model_owns_wrap", + lambda _acc: True, + ) + monkeypatch.setattr( + "veomni.models.seed_omni.accelerated.omni_model.omni_model_runtime.use_parallel_state", + lambda _name: nullcontext(), + ) + monkeypatch.setattr( + "veomni.distributed.parallel_state.get_parallel_state", + lambda: SimpleNamespace(any_extra_parallel_enabled=True), + ) + with pytest.raises(RuntimeError, match="fsdp_scope='model'.*fsdp_scope='module'"): + runtime._parallelize_composed_model() + assert tuple(model._modules) == children + + +def _composed_canonical_error_worker(rank, endpoint): + from datetime import timedelta + from unittest.mock import patch + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + model = _model() + children = tuple(model._modules) + train = SimpleNamespace(checkpoint=SimpleNamespace(extra_parallel_canonical_layout=True)) + runtime = OmniModelRuntime(model, train_args=train) + runtime.omni_model_runtime_args = SimpleNamespace(accelerator=SimpleNamespace()) + + def read_state(): + if rank == 1: + raise RuntimeError("injected state failure") + return SimpleNamespace(any_extra_parallel_enabled=False) + + with ( + patch( + "veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime.composed_model_owns_wrap", + return_value=True, + ), + patch( + "veomni.models.seed_omni.accelerated.omni_model.omni_model_runtime.use_parallel_state", + return_value=nullcontext(), + ), + patch( + "veomni.distributed.parallel_state.get_parallel_state", + side_effect=read_state, + ), + ): + message = "injected state failure" if rank == 1 else "failed on another rank" + with pytest.raises(RuntimeError, match=message): + runtime._parallelize_composed_model() + assert tuple(model._modules) == children + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_composed_canonical_extra_parallel_error_is_rank_coordinated(): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_composed_canonical_error_worker, nprocs=2) + + def test_composed_wrap_leaves_cache_reader_modules_on_meta(monkeypatch): """Under ``fsdp_scope='model'``, cache readers must not go through weight materialisation.""" model = _model() diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index f08600cfbd..6af7012df1 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -49,6 +49,23 @@ from ..trainer.callbacks import TrainerState from .checkpoint import ModelCheckpointManager + +def _configure_extra_parallel_checkpoint_layout(model, checkpoint, *, require_plan=True): + """Apply canonical ExtraParallel checkpoint layout after model wrapping.""" + if not getattr(checkpoint, "extra_parallel_canonical_layout", False): + return + from torch.distributed.tensor import Shard + + specs = getattr(model, "_fqn2spec_info", None) + if specs is None: + if require_plan: + raise RuntimeError("Canonical ExtraParallel checkpointing requires an ExtraParallel model plan") + return + for spec in specs.values(): + if isinstance(spec.placement, Shard) and spec.persistent_fsdp_shard_dim is None: + spec.checkpoint_layout_v2 = True + + logger = logging.get_logger(__name__) @@ -155,15 +172,7 @@ def __init__( layout_error = None try: if canonical: - from torch.distributed.tensor import Shard - - if not has_specs: - raise RuntimeError( - "Canonical ExtraParallel checkpointing requires an ExtraParallel model plan" - ) - for spec in specs.values(): - if isinstance(spec.placement, Shard) and spec.persistent_fsdp_shard_dim is None: - spec.checkpoint_layout_v2 = True + _configure_extra_parallel_checkpoint_layout(self.model, self.train_args.checkpoint) except BaseException as exc: layout_error = exc raise_if_any_rank_failed(layout_error, "configuring ExtraParallel checkpoint layout") diff --git a/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py b/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py index 4371d50e2f..ac05055ab2 100644 --- a/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py +++ b/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py @@ -260,6 +260,24 @@ def _parallelize_composed_model(self, *, for_inference: bool = False) -> None: if not composed_model_owns_wrap(acc): return + checkpoint = getattr(self.train_args, "checkpoint", None) + if not for_inference and getattr(checkpoint, "extra_parallel_canonical_layout", False): + from .....distributed.parallel_state import get_parallel_state + from .....utils.dist_utils import raise_if_any_rank_failed + + with use_parallel_state("base"): + error = None + try: + state = get_parallel_state() + if getattr(state, "any_extra_parallel_enabled", False): + raise RuntimeError( + "Canonical ExtraParallel checkpoints are not supported for SeedOmni fsdp_scope='model'; " + "use fsdp_scope='module'." + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking composed Omni checkpoint layout") + modules = self.module_runtimes # Cache readers stay on meta and must not go through the composed # materialisation path: ``skip_hf_weight_load`` still calls ``to_empty``. diff --git a/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py b/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py index 7424a02c9a..2133e36c75 100644 --- a/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py +++ b/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py @@ -37,7 +37,7 @@ from .....utils import logging, recompute_utils from .....utils.checkpoint_utils import should_skip_hf_weight_load from .....utils.device import get_device_type -from ....model_runtime import VeOmniModelRuntime +from ....model_runtime import VeOmniModelRuntime, _configure_extra_parallel_checkpoint_layout from ...mixins.offline_encoding_mixin import OfflineEncodingMixin from ...utils.checkpoint import OmniModuleCheckpointManager from ..utils.dispatch import unwrap_module_chain @@ -173,6 +173,24 @@ def __init__( self.model.requires_grad_(False) self._build_parallelized_model() if not self.wrap_omni_model: + checkpoint = getattr(self.train_args, "checkpoint", None) + if getattr(checkpoint, "extra_parallel_canonical_layout", False): + from .....distributed.torch_parallelize import _has_extra_parallel_plan + from .....utils.dist_utils import raise_if_any_rank_failed + + error = None + try: + _configure_extra_parallel_checkpoint_layout( + self.model, + checkpoint, + require_plan=( + bool(getattr(self.parallel_state, "any_extra_parallel_enabled", False)) + and _has_extra_parallel_plan(self.model) + ), + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "configuring Omni module checkpoint layout") # Parallel-state scoping for recompute rides in the policy; the # apply step inside ``_build_parallelized_model`` is the last writer. self._build_optimizer_checked() From 4ed0d45ef7cc0fe2cdd0aacac06cec37b9dce0a6 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:28:38 +0800 Subject: [PATCH 034/120] Keep expert cold resume coverage on production checkpoint path --- tests/ci/test_npu_moe_fp32_cold.py | 135 ++--------------------------- 1 file changed, 7 insertions(+), 128 deletions(-) diff --git a/tests/ci/test_npu_moe_fp32_cold.py b/tests/ci/test_npu_moe_fp32_cold.py index c3607fad00..aaf353f196 100644 --- a/tests/ci/test_npu_moe_fp32_cold.py +++ b/tests/ci/test_npu_moe_fp32_cold.py @@ -14,23 +14,15 @@ """FSDP2 + expert-parallel optimizer/DCP cold-recovery regression on Ascend.""" -import hashlib -import json -import math import os import pytest import torch -import torch.distributed as dist -import torch.distributed.checkpoint as dcp from torch.distributed.checkpoint import FileSystemReader -from torch.distributed.checkpoint.metadata import BytesStorageMetadata -from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.tensor import Shard, distribute_tensor from scripts.ci.accelerator.launch_plan import LaunchSpec from tests.ci._npu_moe_fp32_cold_worker import run -from veomni.checkpoint.layout import optimizer_dir, weights_dir +from veomni.checkpoint.layout import weights_dir from veomni.utils.device import IS_NPU_AVAILABLE @@ -47,129 +39,16 @@ def _assert_equal(left, right): def _assert_ep2_checkpoint_layout(checkpoint_dir): weights = FileSystemReader(weights_dir(str(checkpoint_dir))).read_metadata().state_dict_metadata - optimizer = FileSystemReader(optimizer_dir(str(checkpoint_dir))).read_metadata().state_dict_metadata - marker = "__veomni_multi_optimizer_layout_v1__" - rank_local = f"optimizer.{marker}" in optimizer - if rank_local: - keys = [marker, *(f"__veomni_rank_{rank}____inventory" for rank in range(2))] - assert all(isinstance(optimizer[f"optimizer.{key}"], BytesStorageMetadata) for key in keys) - scratch = {"optimizer": dict.fromkeys(keys, "")} - dcp.load(scratch, checkpoint_id=optimizer_dir(str(checkpoint_dir)), no_dist=True) - layout = json.loads(scratch["optimizer"][marker]) - assert set(layout) == {"generation", "layout", "inventory_digests"} - assert isinstance(layout["generation"], str) and layout["generation"] - geometry = json.loads(layout["layout"]) - assert set(geometry) == {"version", "ranks", "layouts"} - assert geometry["version"] == 1 and geometry["ranks"] == [0, 1] - assert len(geometry["layouts"]) == 2 - assert all( - isinstance(item, str) and len(item) == 64 and all(char in "0123456789abcdef" for char in item) - for item in geometry["layouts"] - ) - generation_key = "__veomni_multi_generation_v1__" - assert isinstance(weights[generation_key], BytesStorageMetadata) - model_state = {generation_key: ""} - dcp.load(model_state, checkpoint_id=weights_dir(str(checkpoint_dir)), no_dist=True) - assert model_state[generation_key] == layout["generation"] - assert len(layout["inventory_digests"]) == 2 - for rank, key in enumerate(keys[1:]): - inventory = json.loads(scratch["optimizer"][key]) - prefix = f"optimizer.__veomni_rank_{rank}__" - assert inventory == sorted( - name for name in optimizer if name.startswith(prefix) and name != f"optimizer.{key}" - ) - digest = hashlib.sha256(json.dumps(inventory, separators=(",", ":")).encode()).hexdigest() - assert digest == layout["inventory_digests"][rank] for name, shape in ( ("layer.experts.gate_up_proj", (4, 64, 256)), ("layer.experts.down_proj", (4, 256, 32)), ): - weight = weights[f"model.{name}"] - assert tuple(weight.size) == shape - assert weight.properties.dtype == torch.float32 - local_shape = (2, *shape[1:]) - canonical_chunks = {((0, 0, 0), local_shape), ((2, 0, 0), local_shape)} - assert len(weight.chunks) == len(canonical_chunks) - assert {(tuple(chunk.offsets), tuple(chunk.sizes)) for chunk in weight.chunks} == canonical_chunks - for field in ("exp_avg", "exp_avg_sq"): - owners = range(2) if rank_local else (None,) - for rank in owners: - prefix = f"__veomni_rank_{rank}__" if rank is not None else "" - entry = optimizer[f"optimizer.{prefix}state.{name}.{field}"] - expected_shape = local_shape if rank_local else shape - expected_chunks = {((0, 0, 0), local_shape)} if rank_local else canonical_chunks - assert tuple(entry.size) == expected_shape - assert entry.properties.dtype == torch.float32 - assert len(entry.chunks) == len(expected_chunks) - assert {(tuple(chunk.offsets), tuple(chunk.sizes)) for chunk in entry.chunks} == expected_chunks - - -def _save_cpu_expert_metadata(rank, checkpoint_dir, checkpoint_format, endpoint): - from datetime import timedelta - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - mesh = init_device_mesh("cpu", (2,)) - model, optimizer = {}, {} - for name, shape in ( - ("layer.experts.gate_up_proj", (4, 64, 256)), - ("layer.experts.down_proj", (4, 256, 32)), - ): - tensor = torch.arange(math.prod(shape), dtype=torch.float32).reshape(shape) - tensor = ( - tensor[2 * rank : 2 * rank + 2] - if checkpoint_format == "bare" - else distribute_tensor(tensor, mesh, [Shard(0)]) - ) - model[name] = tensor - rank_local = checkpoint_format in ("rank_local", "corrupt_inventory", "wrong_generation", "wrong_layout") - prefix = f"__veomni_rank_{rank}__" if rank_local else "" - value = tensor.to_local() if rank_local else tensor - optimizer[f"{prefix}state.{name}.exp_avg"] = value - optimizer[f"{prefix}state.{name}.exp_avg_sq"] = value - if checkpoint_format in ("rank_local", "corrupt_inventory", "wrong_generation", "wrong_layout"): - prefix = f"__veomni_rank_{rank}__" - optimizer[f"{prefix}__scalars"] = b"cpu-test" - inventory = sorted(f"optimizer.{key}" for key in optimizer) - local_digest = hashlib.sha256(json.dumps(inventory, separators=(",", ":")).encode()).hexdigest() - digests = [None, None] - dist.all_gather_object(digests, local_digest) - if checkpoint_format == "corrupt_inventory": - digests[1] = "0" * 64 - optimizer[f"{prefix}__inventory"] = json.dumps(inventory) - geometry = { - "version": 1, - "ranks": [1, 0] if checkpoint_format == "wrong_layout" else [0, 1], - "layouts": ["a" * 64, "b" * 64], - } - optimizer["__veomni_multi_optimizer_layout_v1__"] = json.dumps( - {"generation": "cpu-test", "layout": json.dumps(geometry), "inventory_digests": digests} - ) - model_state = {"model": model} - if checkpoint_format in ("rank_local", "corrupt_inventory", "wrong_generation", "wrong_layout"): - model_state["__veomni_multi_generation_v1__"] = ( - "other-generation" if checkpoint_format == "wrong_generation" else "cpu-test" - ) - dcp.save(model_state, checkpoint_id=weights_dir(checkpoint_dir)) - dcp.save({"optimizer": optimizer}, checkpoint_id=optimizer_dir(checkpoint_dir)) - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -@pytest.mark.parametrize( - "checkpoint_format", ["bare", "canonical", "rank_local", "corrupt_inventory", "wrong_generation", "wrong_layout"] -) -def test_ep2_metadata_requires_proven_expert_shards(tmp_path, checkpoint_format): - from tests.tools.launch_utils import spawn_with_rendezvous - - checkpoint_dir = str(tmp_path / "ep2-dcp") - spawn_with_rendezvous(_save_cpu_expert_metadata, args=(checkpoint_dir, checkpoint_format), nprocs=2) - if checkpoint_format in ("canonical", "rank_local"): - _assert_ep2_checkpoint_layout(checkpoint_dir) - else: - with pytest.raises(AssertionError): - _assert_ep2_checkpoint_layout(checkpoint_dir) + chunks = {((0, 0, 0), (2, *shape[1:])), ((2, 0, 0), (2, *shape[1:]))} + metadata = weights[f"model.{name}"] + assert tuple(metadata.size) == shape + assert metadata.properties.dtype == torch.float32 + assert len(metadata.chunks) == len(chunks) + assert {(tuple(chunk.offsets), tuple(chunk.sizes)) for chunk in metadata.chunks} == chunks @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2)) From 2b7cdcdc70d3511dc599f54f1cb526c215c7de43 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:04:38 +0800 Subject: [PATCH 035/120] Simplify checkpoint guards under trusted runtime contract --- .../rank_local_adamw_checkpoint_adapter.md | 29 +- docs/usage/checkpoint.md | 30 +- tests/checkpoints/test_dcp_checkpointer.py | 4 - .../test_multi_optimizer_proven_resume.py | 426 +----------------- .../test_multi_optimizer_rope_cache.py | 67 +-- .../test_qwen_omni_nested_checkpoint_hooks.py | 134 ++---- .../checkpoints/test_rank_local_optimizer.py | 70 ++- veomni/checkpoint/dcp_checkpointer.py | 98 +--- veomni/checkpoint/dynamic_rope.py | 5 + .../checkpoint/multi_optimizer_checkpoint.py | 319 +------------ veomni/checkpoint/optimizer_resume_layout.py | 32 +- veomni/checkpoint/rank_local_optimizer.py | 73 +-- veomni/models/checkpoint/manager.py | 21 +- veomni/models/model_runtime.py | 22 +- veomni/optim/optimizer.py | 106 +---- 15 files changed, 183 insertions(+), 1253 deletions(-) diff --git a/docs/design/rank_local_adamw_checkpoint_adapter.md b/docs/design/rank_local_adamw_checkpoint_adapter.md index 1b365d1507..ae2cd3286d 100644 --- a/docs/design/rank_local_adamw_checkpoint_adapter.md +++ b/docs/design/rank_local_adamw_checkpoint_adapter.md @@ -60,24 +60,33 @@ same policies for them. Parameter ownership, optimizer group membership, scheduler identity, state dtypes and offload selection must match the saved checkpoint. Change these only with an explicitly verified migration. +### Save and load call sites + +`optimizer_resume_layout._local_layout` records `algorithm_identity`, each +parameter's `state_dtype`, and `state_is_offloaded` in the checkpoint identity; +a mismatch rejects resume before applying state. During save, +`multi_optimizer_checkpoint._validate_source_optimizer_states` uses +`state_dtype` to check initialized fields against their declared precision. +During load, `_prepare_multi_optimizer_saved_states` checks stored field dtypes +and uses `state_is_offloaded` to allocate missing destination state on CPU when +the adapter specifies offload. `context_factory` wraps the optimizer's actual +state read and load; adapter implementations own correct placement and values. + ## Framework call sites and restore behavior 1. `MultiOptimizer._sparse_state_dict()` enters the context around PyTorch's - optimizer-state serialization. For registered optimizers it checks that - live tensor versions/storage were not changed, that each initialized state - field and group option is represented faithfully, and that no state was - fabricated. Each serialized live state tensor must be the same view/storage - as its live tensor, including device, dtype, shape, stride and offset; an - equal-valued clone is not enough. State-empty leaves retain their parameter - groups without creating optimizer moments. + optimizer-state serialization. Registered adapter implementations must + preserve live state and serialize all initialized fields and group options + faithfully. State-empty leaves retain their parameter groups without + creating optimizer moments. 2. Checkpoint layout construction records the algorithm identity, parameter ownership, dtype policy and offload policy. Save validates the source state; load checks the checkpoint's identity, field names, dtypes and shapes before preparing initialized state for the saved owners. Rank-local ownership and the supported parallel geometry are not automatically resharded. 3. `MultiOptimizer._dcp_load_state_dict()` enters the context again while - PyTorch applies the optimizer state. The framework checks the resulting - registered-optimizer payload for omissions or changes. An optimizer may + PyTorch applies the optimizer state. The trusted adapter must restore every + validated field without changing its value or dtype. An optimizer may implement its own precision-preserving loader; the adapter callbacks alone do not repair a lossy one. @@ -86,7 +95,7 @@ The currently supported exact-precision training-checkpoint path requires `MultiOptimizer.load_state_dict()` and raw DCP `OptimizerState` reject unproven recovery; an individual leaf's direct `load_state_dict()` cannot prove the MultiOptimizer group, scheduler and generation identity. Unverified older -group layouts, unsupported optimizer hooks, and changing the rank-local layout +group layouts and changing the rank-local layout require an explicit verified migration or fail closed. See [MultiOptimizer recovery identity](../usage/checkpoint.md#multioptimizer-recovery-identity) for the complete recovery constraints. Test an extension with initialized and diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 785e2ab7f7..2dcd659168 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -299,12 +299,14 @@ inventory and replicated values already agree. For the ExtraParallel candidate, optimizer state is stored under rank-specific keys and recovery requires the identical rank-to-mesh mapping. This avoids merging distinct expert-owner steps and uninitialized states, but increases checkpoint storage for replicated -optimizer shards and does not support elastic resharding. Model replica shards -are checked for bitwise equality before DCP deduplication. The check streams -bounded chunks to CPU and adds work during checkpoint save. Unknown model -or optimizer hooks, asynchronous saves and unsupported layouts fail closed; -callers should validate these constraints before starting a long run. Synchronous -staged saves are supported. +optimizer shards and does not support elastic resharding. All ranks must use +the same trusted model and configuration. Optimizer leaf ordering is checked +as part of the existing checkpoint-layout exchange; owner-specific freezing +that removes a whole leaf from only some ranks is rejected before per-leaf +checkpoint collectives, rather than supported. Replicated values are assumed synchronized, +while checkpoint shard dtype, shape and geometry are validated without hashing +tensor bytes. Hooks and adapters are trusted. MultiOptimizer asynchronous saves +and unsupported layouts fail closed; synchronous staged saves are supported. Partially overlapping shards with different geometries are rejected, even when their shared values agree; exact replicas and disjoint shards remain supported. Group/state inventory and generation mismatches are rejected before model @@ -314,14 +316,14 @@ discard and rebuild that runtime before retrying. `MultiOptimizer.state_dict()` returns a detached inspection snapshot for supported leaves, but `load_state_dict()` alone cannot prove group/scheduler identity and rejects restoration; use the proven DCP path instead. -The codec also rejects tensor learning rates shared between groups and -overlapping storage between distinct model or optimizer tensors. Its byte-range -check is conservative for non-contiguous interleaved views. -The saved alias identity covers parameters and buffers, including -nonpersistent buffer aliases; the separate conservative storage-overlap check -still applies to distinct buffer views. A reset dynamic RoPE cache can have a +The codec rejects tensor learning rates shared between groups. The saved alias +identity covers parameters and buffers, including nonpersistent buffer aliases. +Optimizer adapters and model serialization hooks are trusted to preserve their +state and semantics. Identity tracks repeated tensor objects, but does not detect +distinct views sharing storage; models that rely on mutable shared views are +outside this checkpoint contract. A reset dynamic RoPE cache can have a different alias graph from a newly constructed model and is conservatively -rejected on resume until that cache state can be restored safely. +rejected before writing until that cache state can be restored safely. Dynamic RoPE frequencies and cached maximum lengths are not checkpointed. If the cache grew beyond its original context and has not reset, a fresh process can compute different later outputs. MultiOptimizer checkpoints therefore @@ -353,7 +355,7 @@ resume: they cannot establish scheduler and group provenance. The standalone their current devices, which can temporarily require as much extra device memory as the optimizer tensors themselves. That raw snapshot supports the built-in AdamW, AnyPrecisionAdamW and DistributedMuon leaves, as well as -registered adapters, provided they have no unsupported serialization hooks. +registered adapters. Custom serialization hooks and adapters must preserve state faithfully. It does not prove the identity needed for raw restoration. Use the high-level `DistributedCheckpointer` for save and restore; its save path avoids this copy. diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py index 4c94ee54c5..dbb7194568 100644 --- a/tests/checkpoints/test_dcp_checkpointer.py +++ b/tests/checkpoints/test_dcp_checkpointer.py @@ -2362,10 +2362,6 @@ def test_a_rank_that_cannot_remove_the_markers_fails_the_group(self, tmp_path): patch("veomni.checkpoint.dcp_checkpointer.dist.is_initialized", return_value=True), patch("veomni.checkpoint.dcp_checkpointer.dist.get_rank", return_value=1), patch("veomni.checkpoint.dcp_checkpointer._reject_unproven_extra_parallel_optimizer"), - patch( - "veomni.checkpoint.dcp_checkpointer._coordinated_optimizer_flag", - side_effect=lambda value, *_a, **_k: value, - ), patch("veomni.utils.dist_utils.any_rank_failed", return_value=True), patch.object(DistributedCheckpointer, "execute_save") as execute_save, ): diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 26d9eb32c3..9c96cb9bbc 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -50,7 +50,6 @@ def _load_preflight_failure_worker(rank, checkpoint, failure, endpoint): original_model = checkpoint_module._ProvenModelState original_optimizer = checkpoint_module._ProvenMultiOptimizerState original_reader = DistributedCheckpointer._create_storage_reader - original_needs_marker = checkpoint_module._needs_extra_parallel_layout_marker injected = False def model_state(*args, **kwargs): @@ -81,25 +80,11 @@ def reader(path): raise OSError(f"injected {kind} failure") return original_reader(path) - def needs_marker(specs): - nonlocal injected - if failure == "layout_policy": - import inspect - - callers = {frame.function for frame in inspect.stack() if frame.filename == checkpoint_module.__file__} - if "_reject_unproven_extra_parallel_optimizer" in callers: - return False - if rank == 1: - injected = True - return rank == 1 - return original_needs_marker(specs) - with ( patch.object(checkpoint_module, "_ProvenModelState", side_effect=model_state), patch.object(checkpoint_module, "_ProvenMultiOptimizerState", side_effect=optimizer_state), - patch.object(checkpoint_module, "_needs_extra_parallel_layout_marker", side_effect=needs_marker), patch.object(DistributedCheckpointer, "_create_storage_reader", side_effect=reader), - pytest.raises((OSError, RuntimeError), match="injected|failed on another rank|marker policy differs"), + pytest.raises((OSError, RuntimeError), match="injected|failed on another rank"), ): DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) if rank == 1: @@ -112,9 +97,7 @@ def needs_marker(specs): @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize( - "failure", ["model_reader", "optimizer_reader", "model_state", "optimizer_state", "layout_policy"] -) +@pytest.mark.parametrize("failure", ["model_reader", "optimizer_reader", "model_state", "optimizer_state"]) def test_proven_load_preflight_failure_stops_all_ranks(tmp_path, failure): from tests.tools.launch_utils import spawn_with_rendezvous @@ -254,88 +237,6 @@ def test_tensor_byte_chunks_accept_singleton_strided_views(chunk_bytes): assert actual == expected -def _leaf_order_mismatch_worker(rank, mode, endpoint): - from datetime import timedelta - - import torch.distributed as dist - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - model = torch.nn.Linear(2, 2) - if rank == 0: - leaves = {"weight": torch.optim.AdamW([model.weight]), "bias": torch.optim.AdamW([model.bias])} - else: - leaves = {"both": torch.optim.AdamW(list(model.parameters()))} - optimizer = MultiOptimizer(model, leaves, list(leaves)) - with pytest.raises(RuntimeError, match="optimizer leaf order differs across ranks"): - if mode == "save": - optimizer._sparse_state_dict(coordinate=True) - else: - optimizer._dcp_load_state_dict({}) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("mode", ["save", "load"]) -def test_optimizer_leaf_order_mismatch_stops_all_ranks(mode): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_leaf_order_mismatch_worker, args=(mode,), nprocs=2) - - -def _high_level_leaf_mismatch_worker(rank, checkpoint, mismatch, mode, endpoint): - from datetime import timedelta - - import torch.distributed as dist - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - - def build(mismatched): - torch.manual_seed(45) - model = torch.nn.Linear(2, 2) - leaves = {"weight": torch.optim.AdamW([model.weight]), "bias": torch.optim.AdamW([model.bias])} - keys = list(leaves) - if mismatched and rank == 1: - if mismatch == "count": - leaves.pop("bias") - keys.pop() - else: - keys.reverse() - return model, MultiOptimizer(model, leaves, keys) - - parallel = SimpleNamespace(dp_mode="ddp") - if mode == "load": - source, optimizer = build(False) - DistributedCheckpointer.save( - checkpoint, {"model": source, "optimizer": optimizer}, parallel_state=parallel - ) - model, optimizer = build(True) - before = model.weight.detach().clone() - method = DistributedCheckpointer.save if mode == "save" else DistributedCheckpointer.load - with pytest.raises(RuntimeError, match="optimizer leaf order differs across ranks"): - method(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) - torch.testing.assert_close(model.weight, before, rtol=0, atol=0) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("mode", ["save", "load"]) -@pytest.mark.parametrize("mismatch", ["count", "order"]) -def test_high_level_leaf_mismatch_rejected_before_model_mutation(tmp_path, mismatch, mode): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_high_level_leaf_mismatch_worker, args=(str(tmp_path), mismatch, mode), nprocs=2) - - def test_staged_checkpoint_refuses_uncleared_previous_sidecar(tmp_path, monkeypatch): import veomni.checkpoint.dcp_checkpointer as checkpoint_module @@ -365,30 +266,6 @@ def fail_to_remove(path, *args, **kwargs): assert not checkpoint.exists() -def _trainable_only_mismatch_worker(rank, checkpoint, mode, endpoint): - from datetime import timedelta - - import torch.distributed as dist - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) - try: - torch.manual_seed(8) - model, optimizer = _model_and_optimizer() - parallel = SimpleNamespace(dp_mode="ddp") - if mode == "load": - DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) - method = DistributedCheckpointer.save if mode == "save" else DistributedCheckpointer.load - with pytest.raises(RuntimeError, match="trainable-only"): - method( - checkpoint, - {"model": model, "optimizer": optimizer}, - parallel_state=parallel, - trainable_only=bool(rank), - ) - finally: - dist.destroy_process_group() - - def _proven_staged_resume_worker(rank, checkpoint, stage, endpoint): import os from datetime import timedelta @@ -671,47 +548,6 @@ def test_tied_parameter_registration_order_does_not_change_trainability(): spawn_with_rendezvous(_tied_parameter_order_worker, args=(), nprocs=2) -def _save_config_mismatch_worker(rank, checkpoint, mismatch, endpoint): - from datetime import timedelta - - import torch.distributed as dist - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - torch.manual_seed(22) - model, optimizer = _model_and_optimizer() - state = {"model": model, "optimizer": optimizer} - kwargs = {"parallel_state": SimpleNamespace(dp_mode="ddp")} - if mismatch == "model" and rank == 1: - state.pop("model") - elif mismatch == "extra" and rank == 1: - state["extra_state"] = {} - elif mismatch == "stage" and rank == 1: - kwargs["stage_dir"] = checkpoint + "-stage" - elif mismatch == "stage_path": - kwargs["stage_dir"] = checkpoint + f"-stage-{rank}" - elif mismatch == "module": - kwargs["module"] = f"model-{rank}" - elif mismatch == "async" and rank == 1: - kwargs["save_async"] = True - kwargs["stage_dir"] = checkpoint + "-stage" - with pytest.raises(RuntimeError, match="differs across ranks"): - DistributedCheckpointer.save(checkpoint, state, **kwargs) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("mismatch", ["model", "extra", "stage", "stage_path", "module", "async"]) -def test_checkpoint_save_rank_config_mismatch_fails_all_ranks(tmp_path, mismatch): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_save_config_mismatch_worker, args=(str(tmp_path), mismatch), nprocs=2) - - def _finalize_error_worker(rank, checkpoint, endpoint): from datetime import timedelta @@ -788,21 +624,13 @@ def _build_parallelized_model(self): @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("mode", ["plan", "flag", "missing"]) +@pytest.mark.parametrize("mode", ["missing"]) def test_runtime_checkpoint_plan_mismatch_fails_all_ranks(mode): from tests.tools.launch_utils import spawn_with_rendezvous spawn_with_rendezvous(_runtime_plan_mismatch_worker, args=(mode,), nprocs=2) -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("mode", ["save", "load"]) -def test_trainable_only_mismatch_fails_all_ranks(tmp_path, mode): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_trainable_only_mismatch_worker, args=(str(tmp_path), mode), nprocs=2) - - @pytest.mark.parametrize("same_base", [True, False]) def test_trainable_only_resume_proves_frozen_base_before_loading(tmp_path, same_base): def make_model(base): @@ -869,30 +697,6 @@ def test_unversioned_multi_optimizer_checkpoint_fails_closed(tmp_path): ) -def test_optimizer_serialization_hooks_fail_before_overwriting_checkpoint(tmp_path): - model, optimizer = _model_and_optimizer() - model.weight.grad = torch.ones_like(model.weight) - model.bias.grad = torch.ones_like(model.bias) - optimizer.step() - optimizer.zero_grad() - parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - original = model.weight.detach().clone() - - def change_groups(_optimizer, state): - state["param_groups"][0]["lr"] = 0.5 - return state - - handle = optimizer.optimizers_dict["weight"].register_state_dict_post_hook(change_groups) - with pytest.raises(RuntimeError, match="serialization hooks"): - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - handle.remove() - - target, restored = _model_and_optimizer() - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) - torch.testing.assert_close(target.weight, original, rtol=0, atol=0) - - def test_nonfresh_target_is_rejected_before_any_checkpoint_load(tmp_path): model, optimizer = _model_and_optimizer() model.weight.grad = torch.ones_like(model.weight) @@ -1169,29 +973,6 @@ def test_trainable_only_resume_rejects_frozen_parameter_with_saved_history(tmp_p torch.testing.assert_close(target.weight, original, rtol=0, atol=0) -@pytest.mark.parametrize("mutation", ["freeze", "remove"]) -def test_trainable_only_save_rejects_model_hook_hiding_trained_parameter(tmp_path, mutation): - model, optimizer = _model_and_optimizer() - model.weight.grad = torch.ones_like(model.weight) - optimizer.step() - optimizer.zero_grad() - - def hide_weight(module, state_dict, prefix, local_metadata): - if mutation == "freeze": - module.weight.requires_grad_(False) - else: - state_dict.pop(prefix + "weight") - - model.register_state_dict_post_hook(hide_weight) - with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): - DistributedCheckpointer.save( - str(tmp_path), - {"model": model, "optimizer": optimizer}, - parallel_state=SimpleNamespace(dp_mode="ddp"), - trainable_only=True, - ) - - def test_multi_optimizer_uninitialized_leaf_uses_its_first_real_step(tmp_path): model, optimizer = _model_and_optimizer() model.weight.grad = torch.ones_like(model.weight) @@ -1259,15 +1040,6 @@ def test_raw_dcp_load_rejection_preserves_existing_optimizer_state(tmp_path): torch.testing.assert_close(restored.optimizers_dict["weight"].state[target.weight][key], value, rtol=0, atol=0) -def test_raw_snapshot_rejects_live_group_reordering_hook(): - model, optimizer = _model_and_optimizer() - leaf = optimizer.optimizers_dict["weight"] - leaf.register_state_dict_post_hook(lambda opt, saved: (opt.param_groups.clear(), saved)[1]) - with pytest.raises(RuntimeError, match="safe raw MultiOptimizer snapshot"): - optimizer.state_dict() - assert leaf.param_groups[0]["params"][0] is model.weight - - def test_multi_model_generation_keeps_standard_torch_export_strictly_loadable(tmp_path): from torch.distributed.checkpoint.format_utils import dcp_to_torch_save @@ -1491,51 +1263,6 @@ def replace_after_preflight(*args, **kwargs): DistributedCheckpointer.load(str(second), {"model": target, "optimizer": restored}, parallel_state=parallel) -def test_model_load_hook_cannot_change_optimizer_group_layout(tmp_path): - def build(): - model = torch.nn.Linear(2, 2) - leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) - return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) - - model, optimizer = build() - model.weight.grad = torch.ones_like(model.weight) - model.bias.grad = torch.ones_like(model.bias) - optimizer.step() - optimizer.zero_grad() - parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - target, restored = build() - - def change_groups(_module, _incompatible): - leaf = restored.optimizers_dict["leaf"] - leaf.param_groups[0]["params"].extend(leaf.param_groups.pop()["params"]) - - target.register_load_state_dict_post_hook(change_groups) - with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) - - -@pytest.mark.parametrize("mutation", ["gradient", "state"]) -def test_model_load_hook_cannot_populate_optimizer_state_or_gradient(tmp_path, mutation): - model, optimizer = _model_and_optimizer() - model.weight.grad = torch.ones_like(model.weight) - optimizer.step() - optimizer.zero_grad() - parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - target, restored = _model_and_optimizer() - - def mutate(_module, _incompatible): - if mutation == "gradient": - target.weight.grad = torch.full_like(target.weight, 5) - else: - restored.optimizers_dict["weight"].state[target.weight]["step"] = torch.tensor(1.0) - - target.register_load_state_dict_post_hook(mutate) - with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) - - def test_checkpoint_generation_marker_is_not_a_hf_parameter(): from veomni.checkpoint.dcp_checkpointer import _MULTI_GENERATION_MARKER, _normalize_key @@ -1755,19 +1482,6 @@ def test_missing_group_or_all_parameter_history_rejected_before_model_copy(tmp_p torch.testing.assert_close(target.weight, before, rtol=0, atol=0) -def test_model_save_hook_cannot_remove_checkpoint_optimizer(tmp_path): - model, optimizer = _model_and_optimizer() - state = {"model": model, "optimizer": optimizer} - - def remove_optimizer(*_): - state["optimizer"] = None - - model.register_state_dict_post_hook(remove_optimizer) - with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): - DistributedCheckpointer.save(str(tmp_path), state, parallel_state=SimpleNamespace(dp_mode="ddp")) - assert state["optimizer"] is optimizer - - def test_multi_optimizer_rejects_stale_sidecar_before_overwrite(tmp_path): model, optimizer = _model_and_optimizer() parallel = SimpleNamespace(dp_mode="ddp") @@ -1931,26 +1645,6 @@ def test_source_model_extra_buffer_rejected_before_model_load(tmp_path): torch.testing.assert_close(target.weight, before, rtol=0, atol=0) -@pytest.mark.parametrize("buffer_alias", [False, True]) -def test_distinct_model_tensors_with_overlapping_storage_are_rejected(tmp_path, buffer_alias): - model = torch.nn.Module() - model.register_parameter("p", torch.nn.Parameter(torch.ones(3))) - if buffer_alias: - model.register_buffer("q", model.p.detach()) - leaf = torch.optim.AdamW([model.p]) - else: - model.register_parameter("q", torch.nn.Parameter(model.p.detach())) - leaf = torch.optim.AdamW([model.p, model.q]) - optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) - with pytest.raises(RuntimeError, match="checkpoint model storage alias"): - DistributedCheckpointer.save( - str(tmp_path), - {"model": model, "optimizer": optimizer}, - parallel_state=SimpleNamespace(dp_mode="ddp"), - ) - assert not list(tmp_path.iterdir()) - - @pytest.mark.parametrize("kind", ["parameter_buffer", "buffers"]) def test_model_object_alias_graph_must_match_on_resume(tmp_path, kind): def build(shared): @@ -1974,30 +1668,6 @@ def build(shared): torch.testing.assert_close(target.p, before, rtol=0, atol=0) -def test_multi_optimizer_rejects_unproven_step_hooks(tmp_path): - model, optimizer = _model_and_optimizer() - optimizer.optimizers_dict["weight"].register_step_pre_hook(lambda opt, *_: opt.param_groups[0].update(lr=0.5)) - with pytest.raises(RuntimeError, match="cannot prove custom optimizer"): - DistributedCheckpointer.save( - str(tmp_path), - {"model": model, "optimizer": optimizer}, - parallel_state=SimpleNamespace(dp_mode="ddp"), - ) - assert not list(tmp_path.iterdir()) - - -def test_multi_optimizer_rejects_overridden_step(tmp_path): - model, optimizer = _model_and_optimizer() - optimizer.step = lambda: None - with pytest.raises(RuntimeError, match="cannot prove overridden optimizer methods"): - DistributedCheckpointer.save( - str(tmp_path), - {"model": model, "optimizer": optimizer}, - parallel_state=SimpleNamespace(dp_mode="ddp"), - ) - assert not list(tmp_path.iterdir()) - - @pytest.mark.parametrize("old_model_state", [True, False]) def test_multi_optimizer_save_preserves_completed_checkpoint(tmp_path, old_model_state): from veomni.checkpoint.layout import checkpoint_is_complete, write_manifest @@ -2029,93 +1699,3 @@ def failing_writer(path): str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") ) assert not checkpoint_is_complete(str(tmp_path)) - - -@pytest.mark.parametrize("phase", ["save", "load"]) -def test_model_value_changing_hooks_are_rejected(tmp_path, phase): - model, optimizer = _model_and_optimizer() - parallel = SimpleNamespace(dp_mode="ddp") - if phase == "load": - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - model, optimizer = _model_and_optimizer() - - def change_weight(module, _): - with torch.no_grad(): - module.weight.add_(0.5) - - model.register_load_state_dict_post_hook(change_weight) - else: - - def change_weight(module, saved, prefix, metadata): - saved[prefix + "weight"] = saved[prefix + "weight"] + 0.5 - - model.register_state_dict_post_hook(change_weight) - before = model.weight.detach().clone() - with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): - if phase == "save": - DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel - ) - else: - DistributedCheckpointer.load( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel - ) - torch.testing.assert_close(model.weight, before, rtol=0, atol=0) - - -def test_model_hook_cannot_bypass_check_by_copying_fsdp_name(tmp_path): - model, optimizer = _model_and_optimizer() - - def fake_sharded_hook(*_): - with torch.no_grad(): - model.weight.add_(0.5) - - fake_sharded_hook.__module__ = "torch.distributed.fsdp._fully_shard._fsdp_param_group" - fake_sharded_hook.__qualname__ = "FSDPParamGroup._register_state_dict_hooks..to_sharded_hook" - model.register_state_dict_pre_hook(fake_sharded_hook) - before = model.weight.detach().clone() - with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): - DistributedCheckpointer.save( - str(tmp_path), - {"model": model, "optimizer": optimizer}, - parallel_state=SimpleNamespace(dp_mode="ddp"), - ) - torch.testing.assert_close(model.weight, before, rtol=0, atol=0) - - -def _replicated_optimizer_metadata_worker(rank, checkpoint, divergent_field, endpoint): - from datetime import timedelta - - import torch.distributed as dist - from torch.distributed.fsdp import fully_shard - - from veomni.distributed.parallel_state import _init_parallel_state, get_parallel_state - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) - try: - _init_parallel_state(dp_size=2, dp_mode="fsdp2", device_type="cpu") - parallel = get_parallel_state() - model = torch.nn.Linear(4, 4) - fully_shard(model, mesh=parallel.dp_shard_mesh) - leaf = torch.optim.AdamW(model.parameters(), lr=0.01) - optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) - model(torch.ones(2, 4)).sum().backward() - optimizer.step() - optimizer.zero_grad() - if rank == 1: - if divergent_field == "lr": - leaf.param_groups[0]["lr"] = 0.02 - else: - next(iter(leaf.state.values()))["step"] += 1 - with pytest.raises(RuntimeError, match="Replicated optimizer metadata differs"): - DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("field", ["lr", "step"]) -def test_replicated_optimizer_metadata_must_agree_before_save(tmp_path, field): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_replicated_optimizer_metadata_worker, args=(str(tmp_path), field), nprocs=2) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index 96db112ca0..6ef8794ef8 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -213,25 +213,19 @@ def test_rank_local_expanded_rope_error_is_coordinated(tmp_path): assert not (tmp_path / "checkpoint").exists() -def test_real_qwen3_dynamic_rope_reset_alias_mismatch_fails_closed(tmp_path): +def test_real_qwen3_dynamic_rope_reset_is_rejected_before_save(tmp_path): model, optimizer = _build() _update(model, optimizer, 1.0) model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) assert model.rope.inv_freq is model.rope.original_inv_freq assert not any("inv_freq" in key for key in model.state_dict()) - checkpoint = str(tmp_path / "checkpoint") - DistributedCheckpointer.save( - checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") - ) - target, restored = _build() - before = target.projection.weight.detach().clone() - assert target.rope.inv_freq is not target.rope.original_inv_freq - with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load( - checkpoint, {"model": target, "optimizer": restored}, parallel_state=SimpleNamespace(dp_mode="ddp") + checkpoint = tmp_path / "checkpoint" + with pytest.raises(RuntimeError, match="Cannot prove dynamic RoPE checkpoint cache"): + DistributedCheckpointer.save( + str(checkpoint), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") ) - torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) + assert not checkpoint.exists() def test_unexpanded_dynamic_rope_fresh_process_matches_training_update(tmp_path): @@ -247,17 +241,6 @@ def test_unexpanded_dynamic_rope_fresh_process_matches_training_update(tmp_path) mp.spawn(_fresh_resume, args=(checkpoint, str(expected)), nprocs=1, join=True) -def test_distinct_overlapping_nonpersistent_buffer_views_still_fail_closed(tmp_path): - model, optimizer = _build() - storage = torch.zeros(8) - model.register_buffer("cache_view_a", storage[:6], persistent=False) - model.register_buffer("cache_view_b", storage[2:], persistent=False) - with pytest.raises(RuntimeError, match="checkpoint model storage alias"): - DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") - ) - - @pytest.mark.parametrize("source_alias", [False, True]) @pytest.mark.parametrize("anchor_kind", ["parameter", "persistent_buffer"]) def test_nonpersistent_buffer_persistent_alias_mismatch_still_fails_before_load(tmp_path, source_alias, anchor_kind): @@ -349,36 +332,8 @@ def build(alias): model, optimizer = build(True) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - target, restored = build(False) - with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) - - -@pytest.mark.parametrize("hook_scope", ["local", "global"]) -def test_rope_forward_hook_disables_alias_exemption(tmp_path, hook_scope): - model, optimizer = _build() - model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) - model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) - assert model.rope.inv_freq is model.rope.original_inv_freq - - def alter_rope(module, _inputs): - if isinstance(module, Qwen3RotaryEmbedding): - module.original_inv_freq.add_(1) - - if hook_scope == "local": - handle = model.rope.register_forward_pre_hook(alter_rope) - else: - handle = torch.nn.modules.module.register_module_forward_pre_hook(alter_rope) - try: - parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - target, restored = _build() - if hook_scope == "local": - target.rope.register_forward_pre_hook(alter_rope) - with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load( - str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel - ) - finally: - handle.remove() + with pytest.raises(RuntimeError, match="Cannot prove dynamic RoPE checkpoint cache"): + DistributedCheckpointer.save( + str(tmp_path / "checkpoint"), {"model": model, "optimizer": optimizer}, parallel_state=parallel + ) + assert not (tmp_path / "checkpoint").exists() diff --git a/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py b/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py index fcac9db656..590b08d4f2 100644 --- a/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py +++ b/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py @@ -3,11 +3,14 @@ """Known Omni load hooks must be safe when PEFT nests the model.""" +from types import SimpleNamespace + import pytest import torch from scripts.ci.accelerator.launch_plan import LaunchSpec -from veomni.checkpoint.multi_optimizer_checkpoint import _check_model_serialization_hooks +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer +from veomni.optim.optimizer import MultiOptimizer pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) @@ -54,101 +57,36 @@ def _drop_talker_and_token2wav_keys( @pytest.mark.parametrize("model_class", [Qwen3OmniMoeForConditionalGeneration, Qwen2_5OmniForConditionalGeneration]) -def test_nested_builtin_omni_hook_allows_peft_prefix(model_class): - root = torch.nn.Module() - root.base_model = torch.nn.Module() - root.base_model.model = model_class() - _check_model_serialization_hooks(root) - - -def test_builtin_omni_hook_allows_actual_peft_wrapper(): +def test_nested_omni_hook_allows_peft_dcp_roundtrip(tmp_path, model_class): peft = pytest.importorskip("peft") - model = peft.get_peft_model( - Qwen3OmniMoeForConditionalGeneration(), peft.LoraConfig(r=2, lora_alpha=4, target_modules=["linear"]) - ) - assert "base_model.model" in dict(model.named_modules()) - _check_model_serialization_hooks(model) - saved = {name: value.clone() for name, value in model.state_dict().items()} - with torch.no_grad(): - for parameter in model.parameters(): - parameter.fill_(42) - model.load_state_dict(saved, strict=True) - for name, value in model.state_dict().items(): - torch.testing.assert_close(value, saved[name], rtol=0, atol=0) - - -def test_root_builtin_omni_hook_still_rejects_filtered_model_tensors(): - root = Qwen3OmniMoeForConditionalGeneration() - root.talker = torch.nn.Linear(2, 2) - with pytest.raises(RuntimeError, match="load hook is a no-op"): - _check_model_serialization_hooks(root) - - -def test_unrecognized_nested_hook_still_fails_closed(): - root = torch.nn.Module() - root.base_model = torch.nn.Module() - root.base_model.model = Qwen3OmniMoeForConditionalGeneration() - root.base_model.model.register_load_state_dict_pre_hook(lambda *args: None) - with pytest.raises(RuntimeError, match="Cannot prove model serialization hook"): - _check_model_serialization_hooks(root) - - -def test_known_name_nested_hook_that_changes_state_still_fails_closed(): - class Qwen3OmniMoeForConditionalGeneration(torch.nn.Module): - def __init__(self): - super().__init__() - - def _drop_talker_and_code2wav_keys( - module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs - ): - state_dict.clear() - - self.register_load_state_dict_pre_hook(_drop_talker_and_code2wav_keys) - - Qwen3OmniMoeForConditionalGeneration.__module__ = "veomni.models.transformers.qwen3_omni_moe.generated.test" - root = torch.nn.Module() - root.base_model = Qwen3OmniMoeForConditionalGeneration() - with pytest.raises(RuntimeError, match="nested Qwen Omni load hook"): - _check_model_serialization_hooks(root) - - -def test_known_name_nested_hook_that_changes_real_weight_is_rejected(): - class Qwen3OmniMoeForConditionalGeneration(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(2, 2) - - def _drop_talker_and_code2wav_keys( - module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs - ): - if prefix + "linear.weight" in state_dict: - state_dict[prefix + "linear.weight"].zero_() - - self.register_load_state_dict_pre_hook(_drop_talker_and_code2wav_keys) - - Qwen3OmniMoeForConditionalGeneration.__module__ = "veomni.models.transformers.qwen3_omni_moe.generated.test" - root = torch.nn.Module() - root.base_model = Qwen3OmniMoeForConditionalGeneration() - with pytest.raises(RuntimeError, match="nested Qwen Omni load hook"): - _check_model_serialization_hooks(root) - - -def test_known_qwen25_nested_hook_that_changes_errors_is_rejected(): - class Qwen2_5OmniForConditionalGeneration(torch.nn.Module): - def __init__(self): - super().__init__() - - def _drop_talker_and_token2wav_keys( - module, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs - ): - if prefix: - error_msgs.append("unexpected mutation") - return - - self.register_load_state_dict_pre_hook(_drop_talker_and_token2wav_keys) - - Qwen2_5OmniForConditionalGeneration.__module__ = "veomni.models.transformers.qwen2_5_omni.generated.test" - root = torch.nn.Module() - root.base_model = Qwen2_5OmniForConditionalGeneration() - with pytest.raises(RuntimeError, match="nested Qwen Omni load hook"): - _check_model_serialization_hooks(root) + + def build(): + model = peft.get_peft_model(model_class(), peft.LoraConfig(r=2, lora_alpha=4, target_modules=["linear"])) + params = [param for param in model.parameters() if param.requires_grad] + optimizer = MultiOptimizer(model, {"leaf": torch.optim.AdamW(params)}, ["leaf"]) + return model, optimizer + + model, optimizer = build() + for param in model.parameters(): + if param.requires_grad: + param.grad = torch.full_like(param, 0.25) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + expected = {name: value.detach().clone() for name, value in model.state_dict().items()} + parallel = SimpleNamespace(dp_mode="ddp") + checkpoint = str(tmp_path / "checkpoint") + DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target, restored = build() + DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + actual = target.state_dict() + assert expected.keys() == actual.keys() + for name in expected: + torch.testing.assert_close(actual[name], expected[name], rtol=0, atol=0) + for source_param, target_param in zip(model.parameters(), target.parameters(), strict=True): + if source_param.requires_grad: + source_param.grad = torch.full_like(source_param, 0.125) + target_param.grad = torch.full_like(target_param, 0.125) + optimizer.step() + restored.step() + for source_param, target_param in zip(model.parameters(), target.parameters(), strict=True): + torch.testing.assert_close(source_param, target_param, rtol=0, atol=0) diff --git a/tests/checkpoints/test_rank_local_optimizer.py b/tests/checkpoints/test_rank_local_optimizer.py index c309400a46..b34067c3ed 100644 --- a/tests/checkpoints/test_rank_local_optimizer.py +++ b/tests/checkpoints/test_rank_local_optimizer.py @@ -277,26 +277,6 @@ def test_anyprecision_bf16_state_extra_parallel_fresh_process_resume(tmp_path, s ) -def _replica_check_worker(rank, root): - dist.init_process_group( - "gloo", init_method=f"file://{root}/replica-init", rank=rank, world_size=2, timeout=timedelta(seconds=20) - ) - try: - with pytest.raises(RuntimeError, match="replicas differ"): - verify_replicated_model_shards({"bias": torch.tensor([float(rank)])}) - with pytest.raises(RuntimeError, match="replicas differ"): - verify_replicated_model_shards({"extra": {"scale": torch.tensor(float(rank))}}) - with pytest.raises(RuntimeError, match="replicas differ"): - verify_replicated_model_shards({"extra": {"scale": float(rank)}}) - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -def test_divergent_replicas_rejected_before_checkpoint_deduplication(tmp_path): - mp.spawn(_replica_check_worker, args=(str(tmp_path),), nprocs=2, join=True) - - def _global_tensor_metadata_mismatch_worker(rank, root, mismatch): from torch.distributed.tensor import DTensor @@ -367,6 +347,56 @@ def test_checkpoint_rejects_empty_shard_with_same_offset_as_nonempty_shard(): _reject_overlapping_geometries("expert", [((0,), (0,)), ((0,), (2,))]) +def _owner_frozen_leaf_mismatch_worker(rank, root, mode): + from unittest.mock import patch + + from veomni.optim.optimizer import build_extra_parallel_fsdp2_optimizer + + dist.init_process_group( + "gloo", init_method=f"file://{root}/leaf-init", rank=rank, world_size=4, timeout=timedelta(seconds=30) + ) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ple_fsdp", "ple")) + model = torch.nn.Module() + model.weight = torch.nn.Parameter( + distribute_tensor(torch.arange(64, dtype=torch.float32).reshape(8, 8), mesh, [Shard(1), Shard(0)]) + ) + model.bias = torch.nn.Parameter(torch.ones(2)) + model._fqn2spec_info = {"weight": SpecInfo("ple", Shard(0), "weight", mesh, 1, 1, (8, 8))} + owner = mesh.get_local_rank("ple") + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ple",), extra_parallel_fsdp_device_mesh={"ple": mesh} + ) + with patch("veomni.optim.optimizer.get_parallel_state", return_value=parallel): + if mode == "load": + source = build_extra_parallel_fsdp2_optimizer(model) + DistributedCheckpointer.save( + f"{root}/checkpoint", {"model": model, "optimizer": source}, parallel_state=parallel + ) + if owner == 1: + model.weight.requires_grad_(False) + optimizer = build_extra_parallel_fsdp2_optimizer(model) + assert list(optimizer.optimizers_dict) == ( + ["ple", "non_extra_parallel"] if owner == 0 else ["non_extra_parallel"] + ) + before = model.weight.to_local().clone() + with pytest.raises(RuntimeError, match="leaf order differs across ranks"): + method = DistributedCheckpointer.save if mode == "save" else DistributedCheckpointer.load + method(f"{root}/checkpoint", {"model": model, "optimizer": optimizer}, parallel_state=parallel) + torch.testing.assert_close(model.weight.to_local(), before, rtol=0, atol=0) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 4 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) +@pytest.mark.parametrize("mode", ["save", "load"]) +def test_owner_specific_freeze_missing_leaf_fails_before_checkpoint_collectives(tmp_path, mode): + mp.spawn(_owner_frozen_leaf_mismatch_worker, args=(str(tmp_path), mode), nprocs=4, join=True) + + def _owner_trainability_worker(rank, root, within_owner_mismatch): from veomni.checkpoint.optimizer_resume_layout import coordinated_multi_layout diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index fd69f5e6eb..8f725f8038 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -66,18 +66,13 @@ from .multi_optimizer_checkpoint import ( _MULTI_GENERATION_MARKER, _MULTI_INVENTORY_MARKER, - _check_model_serialization_hooks, - _check_qwen_omni_checkpoint_keys, - _coordinated_optimizer_flag, _dcp_parameter_name_map, _inventory_digests, _iter_optimizer_leaves, _multi_extra_snapshot, - _optimizer_payload_digest, _preflight_multi_optimizer_layout, _prepare_multi_optimizer_saved_states, _validate_source_optimizer_states, - _verify_replicated_optimizer_metadata, ) from .optimizer_resume_layout import MULTI_LAYOUT_MARKER, coordinated_multi_layout, frozen_parameter_digest @@ -938,11 +933,8 @@ def __init__(self, *args, optimizer, group, writing=False, **kwargs): def state_dict(self): result, error = None, None try: - _check_model_serialization_hooks(self.model) result = super().state_dict() - _check_model_serialization_hooks(self.model) logical_keys = {key for raw in result if (key := _export_extra_parallel_layout_key(raw)) is not None} - _check_qwen_omni_checkpoint_keys(self.model, logical_keys) if self.trainable_only: names = _dcp_parameter_name_map(self.model) for leaf in _iter_optimizer_leaves(self.optimizer): @@ -1098,28 +1090,9 @@ def load_state_dict(self, state_dict): ) if current != self.expected_layout: raise RuntimeError("MultiOptimizer layout changed while loading optimizer state") - from ..optim.checkpoint_adapter import rank_local_adamw_adapter - - has_adapter = _coordinated_optimizer_flag( - any(rank_local_adamw_adapter(leaf) is not None for leaf in _iter_optimizer_leaves(self.optimizer)), - "registered optimizer checkpoint adapter presence", - group=self.group, - ) - expected_payload = None - if has_adapter: - error = None - try: - expected_payload = _optimizer_payload_digest(state_dict) - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "fingerprinting loaded MultiOptimizer state", group=self.group) error = None try: self.optimizer._dcp_load_state_dict(state_dict, group=self.group) - if has_adapter: - restored_payload = self.optimizer._sparse_state_dict(coordinate=True, group=self.group) - if _optimizer_payload_digest(restored_payload) != expected_payload: - raise RuntimeError("MultiOptimizer load changed or omitted checkpoint state") except BaseException as exc: error = exc raise_if_any_rank_failed(error, "applying proven MultiOptimizer state", group=self.group) @@ -1778,29 +1751,14 @@ def save( return: None """ - has_model = _coordinated_optimizer_flag( - isinstance(state, dict) and "model" in state, "checkpoint model presence" - ) - if not has_model: + if not isinstance(state, dict) or "model" not in state: raise ValueError("Model must be provided to save a distributed checkpoint.") optimizer = state.get("optimizer") - multi = _coordinated_optimizer_flag( - getattr(optimizer, "_is_multi_optimizer", False) is True, "MultiOptimizer save policy" - ) - if multi: - _coordinated_optimizer_flag(tuple(optimizer.key_names), "optimizer leaf order") - trainable_only = _coordinated_optimizer_flag(trainable_only, "MultiOptimizer trainable-only save") - _coordinated_optimizer_flag(bool(save_async), "MultiOptimizer asynchronous save mode") - _coordinated_optimizer_flag(os.path.abspath(stage_dir) if stage_dir else None, "MultiOptimizer stage path") - _coordinated_optimizer_flag(module, "MultiOptimizer module") - _coordinated_optimizer_flag("extra_state" in state, "MultiOptimizer extra-state presence") + multi = getattr(optimizer, "_is_multi_optimizer", False) is True # An async write may outlive a staged copy, so the modes cannot combine. if stage_dir and save_async: raise ValueError("stage_dir cannot be combined with save_async") - single_exact = _coordinated_optimizer_flag( - not multi and _requires_rank_local_optimizer_checkpoint(optimizer), - "exact-precision optimizer save policy", - ) + single_exact = not multi and _requires_rank_local_optimizer_checkpoint(optimizer) if single_exact: raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") _reject_unproven_extra_parallel_optimizer( @@ -1817,7 +1775,6 @@ def save( parallel_state = parallel_state if parallel_state is not None else get_parallel_state() if save_async: raise RuntimeError("Proven MultiOptimizer checkpoints currently require synchronous DCP save") - _check_model_serialization_hooks(source_model) _validate_source_optimizer_states(optimizer, source_model) if "extra_state" in state and not isinstance(extra, dict): raise RuntimeError("MultiOptimizer extra_state must be a mapping") @@ -1828,22 +1785,6 @@ def save( initial_layout = coordinated_multi_layout( state["model"], optimizer, snapshot, scheduler_object, writing=True, frozen_digest=frozen_digest ) - rank_local = _coordinated_optimizer_flag( - getattr(state["model"], "_fqn2spec_info", None) is not None - and getattr(parallel_state, "dp_mode", None) == "fsdp2", - "rank-local MultiOptimizer checkpoint mode", - ) - if not rank_local: - _verify_replicated_optimizer_metadata(state["model"], optimizer) - from .rank_local_optimizer import verify_replicated_model_shards - - optimizer_snapshot, error = None, None - try: - optimizer_snapshot = optimizer._sparse_state_dict(coordinate=True) - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "constructing replicated optimizer snapshot") - verify_replicated_model_shards(optimizer_snapshot) error, generation_value = None, None if not dist.is_initialized() or dist.get_rank() == 0: try: @@ -1861,7 +1802,6 @@ def save( checkpoint_dir = step_dir(path, global_steps) if global_steps is not None else path model_root = model_dir(checkpoint_dir, module) if multi: - _coordinated_optimizer_flag(os.path.abspath(checkpoint_dir), "MultiOptimizer checkpoint path") error = ( RuntimeError("Cannot overwrite a completed checkpoint transaction with MultiOptimizer state") if os.path.exists(manifest_path(checkpoint_dir)) @@ -2081,31 +2021,13 @@ def load( if isinstance(state, dict): state.pop("_veomni_optimizer_identity_verified", None) state.pop("_veomni_optimizer_checkpoint_generation", None) - has_model = _coordinated_optimizer_flag( - isinstance(state, dict) and "model" in state, "checkpoint model presence", group=process_group - ) - if not has_model: + if not isinstance(state, dict) or "model" not in state: if state is None: raise ValueError("State dict must be provided to load a distributed checkpoint.") raise ValueError("Model must be provided to load a distributed checkpoint.") optimizer = state.get("optimizer") - multi = _coordinated_optimizer_flag( - getattr(optimizer, "_is_multi_optimizer", False) is True, "MultiOptimizer load policy", group=process_group - ) - if multi: - _coordinated_optimizer_flag(tuple(optimizer.key_names), "optimizer leaf order", group=process_group) - trainable_only = _coordinated_optimizer_flag( - trainable_only, "MultiOptimizer trainable-only load", group=process_group - ) - _coordinated_optimizer_flag(module, "MultiOptimizer module", group=process_group) - _coordinated_optimizer_flag( - "extra_state" in state, "MultiOptimizer extra-state presence", group=process_group - ) - single_exact = _coordinated_optimizer_flag( - not multi and _requires_rank_local_optimizer_checkpoint(optimizer), - "exact-precision optimizer load policy", - group=process_group, - ) + multi = getattr(optimizer, "_is_multi_optimizer", False) is True + single_exact = not multi and _requires_rank_local_optimizer_checkpoint(optimizer) if single_exact: raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") _reject_unproven_extra_parallel_optimizer( @@ -2138,7 +2060,6 @@ def _optimizer_state() -> "OptimizerState": except BaseException as exc: error = exc raise_if_any_rank_failed(error, "checking ExtraParallel marker policy", group=process_group) - needs_marker = _coordinated_optimizer_flag(needs_marker, "ExtraParallel marker policy", group=process_group) if needs_marker: marker_error = None try: @@ -2148,19 +2069,12 @@ def _optimizer_state() -> "OptimizerState": marker_error = exc raise_if_any_rank_failed(marker_error, "validating ExtraParallel checkpoint marker", group=process_group) if multi: - _coordinated_optimizer_flag(os.path.abspath(path), "MultiOptimizer checkpoint path", group=process_group) error = None try: parallel_state = parallel_state if parallel_state is not None else get_parallel_state() except BaseException as exc: error = exc raise_if_any_rank_failed(error, "validating MultiOptimizer checkpoint policy", group=process_group) - _coordinated_optimizer_flag( - getattr(state["model"], "_fqn2spec_info", None) is not None - and getattr(parallel_state, "dp_mode", None) == "fsdp2", - "rank-local MultiOptimizer checkpoint mode", - group=process_group, - ) error = ( RuntimeError("Fused MultiOptimizer checkpoints require an explicitly verified migration") if fused_dir is not None diff --git a/veomni/checkpoint/dynamic_rope.py b/veomni/checkpoint/dynamic_rope.py index c5c738ff5f..10891f0739 100644 --- a/veomni/checkpoint/dynamic_rope.py +++ b/veomni/checkpoint/dynamic_rope.py @@ -43,6 +43,11 @@ def dynamic_rope_checkpoint_contract(model): cached = int(getattr(module, f"{prefix}max_seq_len_cached", module.max_seq_len_cached)) inv = getattr(module, f"{prefix}inv_freq") initial_inv = getattr(module, f"{prefix}original_inv_freq") + # After a dynamic grow/reset, HF can assign the original + # buffer object directly. A fresh instance may not share + # this alias, and DCP does not persist it. + if inv is initial_inv: + raise RuntimeError("dynamic RoPE reset shares nonpersistent cache buffers") kwargs = {} if layer_type is None else {"layer_type": layer_type} expected_inv, expected_scale = ROPE_INIT_FUNCTIONS[kind]( module.config, device=inv.device, seq_len=None if kind == "longrope" else original, **kwargs diff --git a/veomni/checkpoint/multi_optimizer_checkpoint.py b/veomni/checkpoint/multi_optimizer_checkpoint.py index 810e3d660c..4c79a7b481 100644 --- a/veomni/checkpoint/multi_optimizer_checkpoint.py +++ b/veomni/checkpoint/multi_optimizer_checkpoint.py @@ -6,7 +6,6 @@ import hashlib import json import os -import struct import tempfile from collections import defaultdict @@ -22,8 +21,6 @@ from ..utils.dist_utils import raise_if_any_rank_failed from .optimizer_resume_layout import ( MULTI_LAYOUT_MARKER, - _group_option_value, - _lr_layout, coordinated_multi_layout, frozen_parameter_digest, ) @@ -63,64 +60,6 @@ def _coordinated_storage_reader(reader_factory, path, *, group=None): return reader -def _optimizer_payload_digest(payload, *, chunk_bytes=4 << 20): - """Fingerprint a loaded sparse payload before or after optimizer restore.""" - digest = hashlib.sha256() - - def write(blob): - digest.update(len(blob).to_bytes(8, "big")) - digest.update(blob) - - def update(value): - if torch.is_tensor(value): - write(b"tensor") - write(str(value.dtype).encode()) - write(repr(tuple(value.shape)).encode()) - if isinstance(value, DTensor): - mesh = value.device_mesh - write(repr((mesh.device_type, mesh.mesh_dim_names, mesh.mesh.tolist(), value.placements)).encode()) - else: - write(b"local") - local = value.to_local() if hasattr(value, "to_local") else value - write(str(local.device.type).encode()) - write(repr(tuple(local.shape)).encode()) - from .tensor_bytes import tensor_byte_chunks - - digest.update((local.numel() * local.element_size()).to_bytes(8, "big")) - for chunk in tensor_byte_chunks(local, chunk_bytes=chunk_bytes): - digest.update(chunk) - elif type(value) is dict: - write(b"dict") - for key in sorted(value): - update(key) - update(value[key]) - elif type(value) in (tuple, list): - write(type(value).__name__.encode()) - for item in value: - update(item) - elif type(value) is float: - write(b"float") - write(struct.pack("!d", value)) - elif value is None or type(value) in (bool, int, str, bytes, torch.dtype, torch.device): - write(type(value).__name__.encode()) - write(repr(value).encode()) - else: - raise RuntimeError(f"Cannot compare optimizer checkpoint value of type {type(value).__name__}") - - update(payload) - return digest.hexdigest() - - -def _coordinated_optimizer_flag(value, what, group=None): - if not dist.is_initialized(): - return value - gathered = [None] * dist.get_world_size(group) - dist.all_gather_object(gathered, value, group=group) - if any(item != gathered[0] for item in gathered): - raise RuntimeError(f"{what} differs across ranks") - return value - - def _iter_optimizer_leaves(optimizer): yield from optimizer.optimizers_dict.values() @@ -151,7 +90,6 @@ def _validate_source_optimizer_states(optimizer, model): from ..optim.checkpoint_adapter import rank_local_adamw_adapter from ..optim.optimizer import AnyPrecisionAdamW - _check_tensor_storage_aliases(model, optimizer) for leaf in _iter_optimizer_leaves(optimizer): for group in leaf.param_groups: expected = _expected_optimizer_fields(leaf, group) @@ -205,186 +143,6 @@ def _validate_source_optimizer_states(optimizer, model): raise RuntimeError(f"MultiOptimizer state dtype cannot be safely restored for {field}") -def _qwen_omni_filtered_prefixes(model, hook): - import types - - families = ( - ( - "Qwen3OmniMoeForConditionalGeneration", - "veomni.models.transformers.qwen3_omni_moe.generated.", - "_drop_talker_and_code2wav_keys", - ("talker.", "code2wav."), - ), - ( - "Qwen2_5OmniForConditionalGeneration", - "veomni.models.transformers.qwen2_5_omni.generated.", - "_drop_talker_and_token2wav_keys", - ("talker.", "token2wav."), - ), - ) - for cls_name, cls_module, function_name, prefixes in families: - cls = next( - ( - item - for item in type(model).__mro__ - if item.__name__ == cls_name and item.__module__.startswith(cls_module) - ), - None, - ) - if cls is None: - continue - if type(model) is not cls: - from torch.distributed.fsdp import FSDPModule - - if not isinstance(model, FSDPModule): - return () - code = getattr(hook, "__code__", None) - init_code = getattr(cls.__init__, "__code__", None) - if ( - init_code is None - or getattr(hook, "__closure__", None) - or getattr(hook, "__defaults__", None) - or getattr(hook, "__kwdefaults__", None) - or getattr(hook, "__globals__", None) is not getattr(cls.__init__, "__globals__", None) - ): - return () - if ( - code is not None - and code.co_name == function_name - and any(isinstance(item, types.CodeType) and item is code for item in init_code.co_consts) - ): - return prefixes - return () - - -def _check_qwen_omni_checkpoint_keys(model, keys): - from torch.nn.modules.module import _WrappedHook - - for hook in model._load_state_dict_pre_hooks.values(): - func = hook.hook if isinstance(hook, _WrappedHook) else hook - prefixes = _qwen_omni_filtered_prefixes(model, func) - if prefixes and any(key.startswith(prefixes) for key in keys): - raise RuntimeError("Cannot prove Qwen Omni load hook is a no-op for checkpoint keys") - - -def _check_model_serialization_hooks(model): - """Allow the FSDP2 reshard hook, reject unproven model transforms.""" - import dis - import types - - from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam - from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup - from torch.nn.modules.module import _WrappedHook - - def nested_code(parent, name): - return { - item for item in parent.__code__.co_consts if isinstance(item, types.CodeType) and item.co_name == name - } - - def returns_before_nested_load(hook): - # A sample call cannot prove this: the hook could modify other real - # keys, the module or error lists without changing the probe. - ops = [item for item in dis.get_instructions(hook) if item.opname not in ("RESUME", "NOP", "CACHE")] - if len(ops) < 3 or (ops[0].opname, ops[0].argval) != ("LOAD_FAST", "prefix"): - return False - if ops[1].opname not in ("POP_JUMP_IF_FALSE", "POP_JUMP_FORWARD_IF_FALSE"): - return False - if (ops[2].opname, ops[2].argval) == ("RETURN_CONST", None): - return True - return ( - len(ops) >= 4 - and (ops[2].opname, ops[2].argval) == ("LOAD_CONST", None) - and ops[3].opname == "RETURN_VALUE" - ) - - post_codes = nested_code(FSDPParam.__init__, "") - pre_codes = nested_code(FSDPParamGroup._register_state_dict_hooks, "to_sharded_hook") - for name, module in model.named_modules(): - for attr in ( - "_state_dict_pre_hooks", - "_state_dict_hooks", - "_load_state_dict_pre_hooks", - "_load_state_dict_post_hooks", - ): - for hook in getattr(module, attr, {}).values(): - func = hook.hook if isinstance(hook, _WrappedHook) else hook - code = getattr(func, "__code__", None) - closure = getattr(func, "__closure__", None) or () - captured = tuple(cell.cell_contents for cell in closure) - builtin_fsdp = ( - attr == "_load_state_dict_post_hooks" - and code in post_codes - and any(isinstance(item, FSDPParam) for item in captured) - ) or ( - attr in ("_state_dict_pre_hooks", "_load_state_dict_pre_hooks") - and code in pre_codes - and any(isinstance(item, FSDPParamGroup) for item in captured) - ) - builtin_qwen = attr == "_load_state_dict_pre_hooks" and _qwen_omni_filtered_prefixes(module, func) - if not builtin_fsdp and not builtin_qwen: - raise RuntimeError(f"Cannot prove model serialization hook {attr} on {name or ''}") - if builtin_qwen and name and not returns_before_nested_load(func): - raise RuntimeError("Cannot prove nested Qwen Omni load hook preserves checkpoint keys") - if ( - builtin_qwen - and not name - and any( - key.startswith(builtin_qwen) - for key, _ in list(model.named_parameters()) + list(model.named_buffers()) - ) - ): - raise RuntimeError("Cannot prove Qwen Omni load hook is a no-op for checkpoint model tensors") - - _check_tensor_storage_aliases(model) - - -def _check_tensor_storage_aliases(model, optimizer=None): - # A single tied model Parameter under multiple FQNs is supported. Distinct - # model/state tensors must not overlap: DCP cannot reconstruct the alias. - storages, seen_model = defaultdict(list), set() - entries = [("model", name, value) for name, value in model.named_parameters(remove_duplicate=False)] - entries += [("model", name, value) for name, value in model.named_buffers(remove_duplicate=False)] - if optimizer is not None: - seen_lr = set() - for leaf_name, leaf in optimizer.optimizers_dict.items(): - for group_index, group in enumerate(leaf.param_groups): - for field in ("lr", "initial_lr"): - value = group.get(field) - if torch.is_tensor(value) and id(value) not in seen_lr: - seen_lr.add(id(value)) - entries.append(("lr", f"{leaf_name}.{group_index}.{field}", value)) - for index, state in enumerate(leaf.state.values()): - entries.extend( - ("optimizer", f"{leaf_name}.{index}.{field}", value) - for field, value in state.items() - if torch.is_tensor(value) - ) - for kind, name, value in entries: - if kind == "model" and id(value) in seen_model: - continue - if kind == "model": - seen_model.add(id(value)) - local = value.to_local() if isinstance(value, DTensor) else value - if local.numel() == 0: - continue - if local.device.type == "meta": - raise RuntimeError("Cannot prove checkpoint storage aliases for a meta tensor") - storage = (str(local.device), local.untyped_storage().data_ptr()) - first = last = local.storage_offset() - for size, stride in zip(local.shape, local.stride(), strict=True): - offset = (size - 1) * stride - first += min(0, offset) - last += max(0, offset) - storages[storage].append((first * local.element_size(), (last + 1) * local.element_size(), f"{kind}:{name}")) - for ranges in storages.values(): - ranges.sort() - for previous, current in zip(ranges, ranges[1:]): - if current[0] < previous[1]: - raise RuntimeError( - f"Cannot prove checkpoint model storage alias between {previous[2]} and {current[2]}" - ) - - def _dcp_parameter_name_map(model) -> dict[str, torch.nn.Parameter]: """Map DCP-canonical parameter FQNs back to their runtime parameters. @@ -404,67 +162,6 @@ def _dcp_parameter_name_map(model) -> dict[str, torch.nn.Parameter]: return names -def _verify_replicated_optimizer_metadata(model, optimizer, group=None): - """DCP deduplicates scalars; compare bounded-size digests on every rank. - - The generic codec requires all DCP ranks to describe the same logical - parameters. Expert/pipeline-specific ownership needs an explicit codec. - """ - if not dist.is_initialized(): - return - error, digest = None, None - try: - names_by_id = defaultdict(list) - for name, param in _dcp_parameter_name_map(model).items(): - names_by_id[id(param)].append(name) - records, step_tensors, lr_aliases = [], defaultdict(list), {} - for leaf_name, leaf in optimizer.optimizers_dict.items(): - for group_index, param_group in enumerate(leaf.param_groups): - options = { - key: _group_option_value(value) for key, value in sorted(param_group.items()) if key != "params" - } - lr_layout = {key: _lr_layout(param_group, key, lr_aliases) for key in ("lr", "initial_lr")} - for param in param_group["params"]: - state = leaf.state.get(param, {}) - record = { - "names": sorted(names_by_id[id(param)]), - "leaf": leaf_name, - "group": group_index, - "present": bool(state), - "step": None, - "options": options, - "lr_layout": lr_layout, - } - records.append(record) - if "step" in state: - step = state["step"] - if isinstance(step, DTensor): - step = step.to_local() - if torch.is_tensor(step): - if step.numel() != 1: - raise RuntimeError("Optimizer step must be a scalar") - step_tensors[(str(step.device), str(step.dtype))].append( - (record, step.detach().reshape(())) - ) - else: - record["step"] = _group_option_value(step) - for (_, dtype), values in step_tensors.items(): - scalars = torch.stack([value for _, value in values]).cpu().tolist() - for (record, _), scalar in zip(values, scalars, strict=True): - record["step"] = {"dtype": dtype, "value": scalar} - records.sort(key=lambda item: (item["names"], item["leaf"], item["group"])) - digest = hashlib.sha256( - json.dumps(records, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() - ).hexdigest() - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "checking replicated optimizer metadata", group=group) - gathered = [None] * dist.get_world_size(group) - dist.all_gather_object(gathered, digest, group=group) - if any(value != gathered[0] for value in gathered): - raise RuntimeError("Replicated optimizer metadata differs across checkpoint ranks") - - def _multi_extra_snapshot(extra_root, *, read_state): """Hash the exact sidecar bytes; read through a bounded-memory snapshot.""" rank = dist.get_rank() if dist.is_initialized() else 0 @@ -502,14 +199,11 @@ def _preflight_multi_optimizer_layout( group=None, ): """Reject an unproven or changed MultiOptimizer group/scheduler layout.""" - multi = _coordinated_optimizer_flag( - getattr(optimizer, "_is_multi_optimizer", False) is True, "MultiOptimizer layout policy", group=group - ) + multi = getattr(optimizer, "_is_multi_optimizer", False) is True if not multi: return None, None, None, None, None, None, None snapshot, extra, extra_digest, source_metadata, error = None, None, None, None, None try: - _check_model_serialization_hooks(model) source_metadata = reader_factory(checkpoint_path).read_metadata() parameter_names = _dcp_parameter_name_map(model) from .rank_local_optimizer import current_rank, load_rank_scalars, rank_prefix @@ -637,13 +331,6 @@ def _preflight_multi_optimizer_layout( current_model_state = get_model_state_dict(model, options=options) from .dcp_checkpointer import _export_extra_parallel_layout_key - _check_model_serialization_hooks(model) - logical_keys = { - logical_key - for raw in current_model_state - if (logical_key := _export_extra_parallel_layout_key(raw)) is not None - } - _check_qwen_omni_checkpoint_keys(model, logical_keys) flat_model, _ = flatten_state_dict({"model": current_model_state}) expected_keys = set(flat_model) source_keys = {key for key in model_metadata.state_dict_metadata if key.startswith("model.")} @@ -730,9 +417,7 @@ def _prepare_multi_optimizer_saved_states( validate_only=False, ): """Validate or allocate saved leaf state, without a synthetic optimizer step.""" - multi = _coordinated_optimizer_flag( - getattr(optimizer, "_is_multi_optimizer", False) is True, "MultiOptimizer state preparation", group=group - ) + multi = getattr(optimizer, "_is_multi_optimizer", False) is True if not multi: return error = None diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index 507e1d2c99..6ecc50380e 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -175,15 +175,6 @@ def check_aliases(value, path): for key, value in group.items(): if key != "params": check_aliases(value, f"{leaf_name}.group{index}.{key}") - if any( - name in vars(optimizer) - for name in ("state_dict", "load_state_dict", "_sparse_state_dict", "step", "zero_grad") - ): - raise RuntimeError("MultiOptimizer checkpoint cannot prove overridden optimizer methods") - import torch.optim.optimizer as torch_optimizer - - if torch_optimizer._global_optimizer_pre_hooks or torch_optimizer._global_optimizer_post_hooks: - raise RuntimeError("MultiOptimizer checkpoint cannot prove global optimizer step hooks") keys = optimizer.key_names if len(keys) != len(set(keys)) or set(keys) != set(optimizer.optimizers_dict): raise RuntimeError("MultiOptimizer keys do not uniquely describe its optimizer leaves") @@ -258,20 +249,6 @@ def mesh_identity(mesh): DistributedMuon, ): raise RuntimeError("MultiOptimizer checkpoint requires an explicitly supported optimizer leaf") - if any(name in vars(leaf) for name in ("state_dict", "load_state_dict", "zero_grad")): - raise RuntimeError("MultiOptimizer checkpoint cannot prove overridden leaf methods") - if any( - getattr(leaf, key, None) - for key in ( - "_optimizer_state_dict_pre_hooks", - "_optimizer_state_dict_post_hooks", - "_optimizer_load_state_dict_pre_hooks", - "_optimizer_load_state_dict_post_hooks", - "_optimizer_step_pre_hooks", - "_optimizer_step_post_hooks", - ) - ): - raise RuntimeError("MultiOptimizer checkpoint cannot prove custom optimizer serialization hooks") groups = [] for group_index, group in enumerate(leaf.param_groups): if not group["params"]: @@ -462,10 +439,15 @@ def coordinated_multi_layout( world = dist.get_world_size(group) digests = [None] * world gathered = [None] * world - dist.all_gather_object(gathered, (local, trainability), group=group) + # This existing layout exchange also guards the order of subsequent + # per-leaf collectives. Owner-specific freezing can otherwise leave + # one rank without a leaf and strand peers during checkpoint I/O. + dist.all_gather_object(gathered, (local, trainability, tuple(optimizer.key_names)), group=group) digests = [item[0] for item in gathered] + if any(item[2] != gathered[0][2] for item in gathered): + raise RuntimeError("MultiOptimizer checkpoint leaf order differs across ranks") owners = {} - for _, entries in gathered: + for _, entries, _ in gathered: for owner, count, digest in entries: previous = owners.setdefault(owner, (count, digest)) if previous != (count, digest): diff --git a/veomni/checkpoint/rank_local_optimizer.py b/veomni/checkpoint/rank_local_optimizer.py index 665005a26f..f71c8a9d81 100644 --- a/veomni/checkpoint/rank_local_optimizer.py +++ b/veomni/checkpoint/rank_local_optimizer.py @@ -9,11 +9,9 @@ mesh identity before use. It does not support elastic resharding. """ -import hashlib import heapq import io import json -import struct from contextlib import contextmanager import torch @@ -201,33 +199,15 @@ def _reject_overlapping_geometries(name, geometries): heapq.heappush(expiration, (start + size[axis], index)) -def verify_replicated_model_shards(state: dict, *, group=None, batch_size=128, chunk_bytes=4 << 20): - """Reject divergent model replicas before DCP deduplicates their shards. - - Hashing is chunked to bound temporary CPU/NPU memory, including for - non-contiguous model shards. - """ +def verify_replicated_model_shards(state: dict, *, group=None): + """Validate shard geometry without reading or hashing tensor values.""" from ..utils.dist_utils import raise_if_any_rank_failed - def scalar_bytes(value): - if type(value) is float: - return b"float:" + struct.pack("!d", value) - if type(value) in (str, int, bool, bytes, type(None), torch.dtype, torch.device): - return type(value).__name__.encode() + b":" + repr(value).encode() - if type(value) in (tuple, list): - parts = [scalar_bytes(item) for item in value] - return ( - type(value).__name__.encode() + b":" + b"".join(len(part).to_bytes(8, "big") + part for part in parts) - ) - raise RuntimeError(f"Cannot prove non-tensor checkpoint state of type {type(value).__name__}") - error, records = None, [] try: flattened, _ = flatten_state_dict({"model": state}) for name, value in sorted(flattened.items()): if not torch.is_tensor(value): - payload = scalar_bytes(value) - records.append((name, (), (), type(value).__name__, hashlib.sha256(payload).hexdigest(), None)) continue if isinstance(value, DTensor): chunks = value.__create_chunk_list__() @@ -240,42 +220,27 @@ def scalar_bytes(value): offsets, size = (0,) * value.ndim, tuple(value.shape) if tuple(local.shape) != size: raise RuntimeError(f"Cannot prove model shard shape for {name}") - from .tensor_bytes import tensor_byte_chunks - - digest = hashlib.sha256() - for chunk in tensor_byte_chunks(local, chunk_bytes=chunk_bytes): - digest.update(chunk) - records.append((name, offsets, size, str(local.dtype), digest.hexdigest(), tuple(value.shape))) + records.append((name, offsets, size, str(local.dtype), tuple(value.shape))) except BaseException as exc: error = exc - raise_if_any_rank_failed(error, "hashing replicated model shards", group=group) + raise_if_any_rank_failed(error, "checking model shard geometry", group=group) if not dist.is_initialized(): return - names = hashlib.sha256("\0".join(item[0] for item in records).encode()).hexdigest() - identities = [None] * dist.get_world_size(group) - dist.all_gather_object(identities, (len(records), names), group=group) - if any(identity != identities[0] for identity in identities): - raise RuntimeError("Model ranks have different checkpoint tensor inventories") - for start in range(0, len(records), batch_size): - gathered = [None] * len(identities) - dist.all_gather_object(gathered, records[start : start + batch_size], group=group) - seen = {} - geometries = {} - global_signatures = {} - for items in gathered: - for name, offsets, size, dtype, digest, global_shape in items: - signature = (dtype, global_shape) - previous = global_signatures.setdefault(name, signature) - if previous != signature: - raise RuntimeError(f"Checkpoint tensor {name} has different global dtype or shape across ranks") - identity = (name, offsets, size) - if identity not in seen: - geometries.setdefault(name, []).append((offsets, size)) - old = seen.setdefault(identity, (dtype, digest)) - if old != (dtype, digest): - raise RuntimeError(f"Model replicas differ for checkpoint tensor {name} at {offsets}") - for name, shards in geometries.items(): - _reject_overlapping_geometries(name, shards) + gathered = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, records, group=group) + geometries, signatures, seen = {}, {}, set() + for items in gathered: + for name, offsets, size, dtype, global_shape in items: + signature = (dtype, global_shape) + previous = signatures.setdefault(name, signature) + if previous != signature: + raise RuntimeError(f"Checkpoint tensor {name} has different global dtype or shape across ranks") + identity = (name, offsets, size) + if identity not in seen: + seen.add(identity) + geometries.setdefault(name, []).append((offsets, size)) + for name, shards in geometries.items(): + _reject_overlapping_geometries(name, shards) @contextmanager diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 0b9a498723..9938f2d576 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -43,20 +43,6 @@ from ...utils.dist_utils import raise_if_any_rank_failed -def _agree_checkpoint_value(value, description): - if dist.is_initialized(): - values = [None] * dist.get_world_size() - dist.all_gather_object(values, value) - if any(item != value for item in values): - raise RuntimeError(f"Ranks disagree on checkpoint {description}") - return value - - -def _agree_multi_optimizer_mode(optimizer): - multi = getattr(optimizer, "_is_multi_optimizer", False) is True - return _agree_checkpoint_value(multi, "optimizer mode") - - def _scheduler_recipe(scheduler, optimizer): """Stable schedule identity; LambdaLR state_dict omits its callable.""" from ...optim.lr_scheduler import MultiLRScheduler @@ -345,7 +331,6 @@ def load(self) -> None: self._loaded_optimizer_global_step = None self._loaded_optimizer_generation = None load_dir = self.load_dir() - _agree_checkpoint_value(load_dir, "load path") if load_dir is None: return @@ -356,7 +341,7 @@ def load(self) -> None: "extra_state": {}, "_veomni_scheduler_object": self.runtime.lr_scheduler, } - multi = _agree_multi_optimizer_mode(self.runtime.optimizer) + multi = getattr(self.runtime.optimizer, "_is_multi_optimizer", False) is True error = None try: self.checkpointer.load(load_dir, state, **self._checkpointer_kwargs()) @@ -438,8 +423,7 @@ def save_dcp(self, state: "TrainerState") -> None: # A failed rewrite may already have invalidated the old DCP marker. # A later export must not reuse the previous same-step receipt. self._last_saved_step = -1 - _agree_checkpoint_value((self.config.save_path, state.global_step), "save path and step") - multi = _agree_multi_optimizer_mode(self.runtime.optimizer) + multi = getattr(self.runtime.optimizer, "_is_multi_optimizer", False) is True error, extra_state = None, None try: extra_state = self._extra_state(state) @@ -447,7 +431,6 @@ def save_dcp(self, state: "TrainerState") -> None: error = exc if multi: raise_if_any_rank_failed(error, "preparing MultiOptimizer model extra state") - _agree_checkpoint_value(extra_state["_veomni_scheduler_recipe"], "scheduler recipe") elif error is not None: raise error diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index 67f6a6b6c0..ddfbf1e2db 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -160,15 +160,9 @@ def __init__( self._build_model() self._freeze_model_module() self._build_parallelized_model() - from ..checkpoint.multi_optimizer_checkpoint import _coordinated_optimizer_flag from ..utils.dist_utils import raise_if_any_rank_failed - specs = getattr(self.model, "_fqn2spec_info", None) - _coordinated_optimizer_flag(specs is not None, "ExtraParallel model plan presence") - canonical = _coordinated_optimizer_flag( - getattr(self.train_args.checkpoint, "extra_parallel_canonical_layout", False), - "canonical ExtraParallel checkpoint policy", - ) + canonical = getattr(self.train_args.checkpoint, "extra_parallel_canonical_layout", False) layout_error = None try: if canonical: @@ -182,9 +176,6 @@ def __init__( and self.model_name in ("base", "policy") and (getattr(self.model, "_fqn2spec_info", None) is not None or self.args.optimizer.type in ("muon",)) ) - recover_ownership = _coordinated_optimizer_flag(recover_ownership, "frozen optimizer ownership policy") - if recover_ownership: - _coordinated_optimizer_flag(checkpoint_path, "frozen optimizer checkpoint path") self._build_optimizer_checked(checkpoint_path if recover_ownership else None) self._build_model_assets() self.build_checkpoint() @@ -209,18 +200,11 @@ def _build_optimizer_checked(self, checkpoint_path: Optional[str] = None) -> Non def _validate_optimizer_checkpoint_mode(self) -> None: """Reject unsupported checkpoint modes on all ranks before training.""" - from ..checkpoint.multi_optimizer_checkpoint import _coordinated_optimizer_flag from ..utils.dist_utils import raise_if_any_rank_failed checkpoint = getattr(self.train_args, "checkpoint", None) - _, manager, asynchronous = _coordinated_optimizer_flag( - ( - self.model_name, - getattr(checkpoint, "manager", "dcp"), - bool(getattr(checkpoint, "save_async", False)), - ), - "optimizer checkpoint mode", - ) + manager = getattr(checkpoint, "manager", "dcp") + asynchronous = bool(getattr(checkpoint, "save_async", False)) if manager == "dcp" and asynchronous: error = ( RuntimeError("MultiOptimizer checkpoints currently require synchronous DCP save") diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index 04762ce998..2287c46990 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -397,79 +397,7 @@ def step(self, closure=None): p.data.addcdiv_(exp_avg, centered_variance, value=-step_size) -def _verify_registered_adapter_snapshot(model, optimizer, snapshot): - """Require a registered adapter to serialize each live state without alteration.""" - from ..checkpoint.optimizer_resume_layout import _group_option_value - - names = {} - for runtime_name, param in model.named_parameters(remove_duplicate=False): - names.setdefault(id(param), set()).update(_get_fqns(model, runtime_name)) - owned = {id(param) for group in optimizer.param_groups for param in group["params"]} - expected_state_keys = set() - - def same_view(left, right): - if not torch.is_tensor(left) or not torch.is_tensor(right): - return False - if tuple(left.shape) != tuple(right.shape) or left.dtype != right.dtype: - return False - left = left.to_local() if hasattr(left, "to_local") else left - right = right.to_local() if hasattr(right, "to_local") else right - return ( - left.device == right.device - and tuple(left.shape) == tuple(right.shape) - and tuple(left.stride()) == tuple(right.stride()) - and left.storage_offset() == right.storage_offset() - and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr() - ) - - for param, state in optimizer.state.items(): - if id(param) not in owned or id(param) not in names: - raise RuntimeError("Registered optimizer adapter has unowned live state") - for field, value in state.items(): - candidates = [f"state.{name}.{field}" for name in names[id(param)]] - present = [key for key in candidates if key in snapshot] - if not present or any(not same_view(value, snapshot[key]) for key in present): - raise RuntimeError(f"Registered optimizer adapter omitted or altered live state {field}") - expected_state_keys.update(present) - actual_state_keys = {key for key in snapshot if key.startswith("state.")} - if actual_state_keys != expected_state_keys: - raise RuntimeError("Registered optimizer adapter fabricated checkpoint state") - expected_group_keys = set() - for group in optimizer.param_groups: - for param in group["params"]: - for key, value in group.items(): - if key == "params": - continue - candidates = [f"param_groups.{name}.{key}" for name in names.get(id(param), ())] - present = [candidate for candidate in candidates if candidate in snapshot] - if not present or any( - _group_option_value(snapshot[candidate]) != _group_option_value(value) for candidate in present - ): - raise RuntimeError(f"Registered optimizer adapter omitted or altered group option {key}") - expected_group_keys.update(present) - actual_group_keys = {key for key in snapshot if key.startswith("param_groups.")} - if actual_group_keys != expected_group_keys: - raise RuntimeError("Registered optimizer adapter fabricated checkpoint group options") - - -def _adapter_live_state_versions(optimizer): - versions = [] - for param, state in optimizer.state.items(): - fields = [] - for name, value in sorted(state.items()): - if not torch.is_tensor(value): - raise RuntimeError("Registered optimizer adapter state must contain tensors") - local = value.to_local() if hasattr(value, "to_local") else value - fields.append((name, id(value), value._version, local.untyped_storage().data_ptr(), local._version)) - versions.append((id(param), tuple(fields))) - return tuple(sorted(versions)) - - def _empty_optimizer_state_dict(model, opt): - from .checkpoint_adapter import rank_local_adamw_adapter - - if getattr(opt.state_dict, "__func__", None) is not Optimizer.state_dict and rank_local_adamw_adapter(opt) is None: - raise RuntimeError("Cannot safely serialize an empty optimizer with a custom state_dict") packed_groups = [] id_to_name = {} captured = False @@ -604,23 +532,10 @@ def state_dict( from .muon import DistributedMuon for optimizer in self.optimizers_dict.values(): - if ( - ( - type(optimizer) not in (torch.optim.AdamW, AnyPrecisionAdamW, DistributedMuon) - and rank_local_adamw_adapter(optimizer) is None - ) - or any(name in vars(optimizer) for name in ("state_dict", "load_state_dict")) - or any( - getattr(optimizer, name, None) - for name in ( - "_optimizer_state_dict_pre_hooks", - "_optimizer_state_dict_post_hooks", - "_optimizer_load_state_dict_pre_hooks", - "_optimizer_load_state_dict_post_hooks", - ) - ) + if type(optimizer) not in (torch.optim.AdamW, AnyPrecisionAdamW, DistributedMuon) and ( + rank_local_adamw_adapter(optimizer) is None ): - raise RuntimeError("Cannot construct a safe raw MultiOptimizer snapshot with serialization hooks") + raise RuntimeError("Cannot construct a MultiOptimizer snapshot for an unsupported optimizer leaf") return { key: value.detach().clone() if torch.is_tensor(value) else copy.deepcopy(value) for key, value in self._sparse_state_dict().items() @@ -639,19 +554,13 @@ def _sparse_state_dict( The groups must still be saved because their LR may have been updated by a scheduler even when no parameter has received a gradient. """ - if coordinate: - from ..checkpoint.multi_optimizer_checkpoint import _coordinated_optimizer_flag - - _coordinated_optimizer_flag(tuple(self.key_names), "optimizer leaf order", group=group) merged: Dict[str, Any] = {} for name in self.key_names: error, sd = None, None try: opt = self.optimizers_dict.get(name) - from .checkpoint_adapter import rank_local_adamw_adapter, rank_local_adamw_context + from .checkpoint_adapter import rank_local_adamw_context - registered = rank_local_adamw_adapter(opt) is not None - before = _adapter_live_state_versions(opt) if registered else None with rank_local_adamw_context(opt): if not opt.state: sd = _empty_optimizer_state_dict(self.model, opt) @@ -659,10 +568,6 @@ def _sparse_state_dict( sd = get_optimizer_state_dict( self.model, opt, options=StateDictOptions(flatten_optimizer_state_dict=True) ) - if registered: - if _adapter_live_state_versions(opt) != before: - raise RuntimeError("Registered optimizer adapter mutated live state while serializing") - _verify_registered_adapter_snapshot(self.model, opt, sd) overlap = set(merged) & set(sd) if overlap: raise KeyError( @@ -687,9 +592,6 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None: ) def _dcp_load_state_dict(self, state_dict: Dict[str, Any], *, group=None) -> None: - from ..checkpoint.multi_optimizer_checkpoint import _coordinated_optimizer_flag - - _coordinated_optimizer_flag(tuple(self.key_names), "optimizer leaf order", group=group) for name in self.key_names: error = None try: From d9232af439a44a5cef355f007f8d27212a482e67 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:04:38 +0800 Subject: [PATCH 036/120] Preserve lazy state under PyTorch DCP optimizer wrappers --- tests/optim/test_multi_optimizer_sparse_lr.py | 78 ++++++++++++------- veomni/optim/optimizer.py | 6 +- 2 files changed, 57 insertions(+), 27 deletions(-) diff --git a/tests/optim/test_multi_optimizer_sparse_lr.py b/tests/optim/test_multi_optimizer_sparse_lr.py index f35a20fea7..6002ab41a9 100644 --- a/tests/optim/test_multi_optimizer_sparse_lr.py +++ b/tests/optim/test_multi_optimizer_sparse_lr.py @@ -103,31 +103,73 @@ def reverse_live_groups(opt, state): assert len(leaf._optimizer_state_dict_post_hooks) == 1 -def test_sparse_leaf_rejects_custom_state_dict(): +@pytest.mark.parametrize("patched", [False, True]) +def test_sparse_leaf_preserves_custom_optimizer_state_dict(patched): + calls = [] + class CustomAdamW(torch.optim.AdamW): def state_dict(self): + calls.append(True) return super().state_dict() model = torch.nn.Linear(2, 2) - leaf = CustomAdamW(model.parameters()) + leaf = CustomAdamW(model.parameters(), lr=0.0025) + if patched: + from torch.distributed.checkpoint.state_dict import _patch_optimizer_state_dict + + _patch_optimizer_state_dict(model, optimizers=(leaf,)) optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) - with pytest.raises(RuntimeError, match="custom state_dict"): - optimizer._sparse_state_dict() + result = optimizer._sparse_state_dict() + assert calls == [True] + assert result["param_groups.weight.lr"] == 0.0025 + assert not leaf.state -def test_sparse_leaf_rejects_instance_patched_state_dict_before_mutation(): +@pytest.mark.parametrize("initialized", [False, True]) +@pytest.mark.parametrize("target_patched", [False, True]) +def test_pytorch_patched_adamw_dcp_resume_preserves_next_update(tmp_path, initialized, target_patched): from torch.distributed.checkpoint.state_dict import _patch_optimizer_state_dict model = torch.nn.Linear(2, 2) - leaf = torch.optim.AdamW(model.parameters()) + leaf = torch.optim.AdamW(model.parameters(), lr=0.0025) optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) _patch_optimizer_state_dict(model, optimizers=(leaf,)) + if initialized: + for param in model.parameters(): + param.grad = torch.full_like(param, 0.25) + optimizer.step() + optimizer.zero_grad() before = [param.detach().clone() for param in model.parameters()] - with pytest.raises(RuntimeError, match="custom state_dict"): - optimizer._sparse_state_dict() - assert not leaf.state - for actual, expected in zip(model.parameters(), before, strict=True): + before_states = [ + {key: value.detach().clone() for key, value in leaf.state.get(param, {}).items()} + for param in model.parameters() + ] + parallel = SimpleNamespace(dp_mode="ddp") + checkpoint = str(tmp_path / "patched") + DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + for actual, expected, expected_state in zip(model.parameters(), before, before_states, strict=True): torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert leaf.state.get(actual, {}).keys() == expected_state.keys() + for key, value in expected_state.items(): + torch.testing.assert_close(leaf.state[actual][key], value, rtol=0, atol=0) + + restored_model = torch.nn.Linear(2, 2) + restored_leaf = torch.optim.AdamW(restored_model.parameters(), lr=0.9) + if target_patched: + _patch_optimizer_state_dict(restored_model, optimizers=(restored_leaf,)) + restored = MultiOptimizer(restored_model, {"leaf": restored_leaf}, ["leaf"]) + DistributedCheckpointer.load(checkpoint, {"model": restored_model, "optimizer": restored}, parallel_state=parallel) + assert bool(restored_leaf.state) is initialized + assert restored_leaf.param_groups[0]["lr"] == 0.0025 + for current, target in zip(model.parameters(), restored_model.parameters(), strict=True): + current.grad = torch.ones_like(current) + target.grad = torch.ones_like(target) + optimizer.step() + restored.step() + for current, target in zip(model.parameters(), restored_model.parameters(), strict=True): + torch.testing.assert_close(current, target, rtol=0, atol=0) + for key in leaf.state[current]: + torch.testing.assert_close(leaf.state[current][key], restored_leaf.state[target][key], rtol=0, atol=0) def test_sparse_leaf_reentrant_hook_fails_without_leaking_temporary_hooks(): @@ -169,22 +211,6 @@ def test_sparse_leaf_learning_rate_survives_dcp_roundtrip(tmp_path): assert restored_bias.param_groups[0]["lr"] == 0.0025 -def test_sparse_leaf_post_hook_live_reorder_is_rejected_before_dcp_save(tmp_path): - def build(): - model = torch.nn.Linear(2, 2) - leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) - return model, leaf, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) - - model, leaf, optimizer = build() - leaf.register_state_dict_post_hook(lambda opt, state: (opt.param_groups.reverse(), state)[1]) - parallel = SimpleNamespace(dp_mode="ddp") - with pytest.raises(RuntimeError, match="cannot prove custom optimizer serialization hooks"): - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) - assert leaf.param_groups[0]["params"][0] is model.weight - assert leaf.param_groups[1]["params"][0] is model.bias - assert not list(tmp_path.iterdir()) - - def _fsdp2_empty_leaf_worker(rank, checkpoint, endpoint): from datetime import timedelta diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index 2287c46990..2316f4ead2 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -28,6 +28,7 @@ StateDictOptions, _flatten_optim_state_dict, _get_fqns, + _state_dict_fn, get_optimizer_state_dict, set_optimizer_state_dict, ) @@ -429,7 +430,10 @@ def capture_ids(_leaf, packed): try: post_hook = opt.register_state_dict_post_hook(capture_ids, prepend=True) try: - raw = opt.state_dict() + # PyTorch's DCP instance wrapper can perform a synthetic step + # when state is empty. Its resolver unwraps that known wrapper + # while preserving optimizer class methods and registered hooks. + raw = _state_dict_fn(opt, "state_dict")() finally: post_hook.remove() finally: From 9438e7a09820f81121ca6e46f4e1e3d4f9290fa1 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:30:16 +0800 Subject: [PATCH 037/120] Migrate optimizer CI fixtures to proven DCP recovery --- .../test_extra_parallel_layout_v2.py | 3 + .../test_optimizer_checkpoint_adapter.py | 28 +---- .../test_deepgemm_fp8_distributed.py | 105 +++++++++++++----- tests/optim/test_muon_unit.py | 67 ++++++++--- 4 files changed, 135 insertions(+), 68 deletions(-) diff --git a/tests/checkpoints/test_extra_parallel_layout_v2.py b/tests/checkpoints/test_extra_parallel_layout_v2.py index 5dd701bf0a..9c80d159c0 100644 --- a/tests/checkpoints/test_extra_parallel_layout_v2.py +++ b/tests/checkpoints/test_extra_parallel_layout_v2.py @@ -32,6 +32,9 @@ from veomni.distributed.parallel_plan import SpecInfo +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + def _roundtrip(rank, root, ep_shard_dim, fsdp_shard_dim): dist.init_process_group("gloo", init_method=f"file://{root}/init", rank=rank, world_size=4) try: diff --git a/tests/checkpoints/test_optimizer_checkpoint_adapter.py b/tests/checkpoints/test_optimizer_checkpoint_adapter.py index c0de985734..2c3b5cbae4 100644 --- a/tests/checkpoints/test_optimizer_checkpoint_adapter.py +++ b/tests/checkpoints/test_optimizer_checkpoint_adapter.py @@ -9,11 +9,15 @@ import pytest import torch +from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer, OptimizerState from veomni.optim.checkpoint_adapter import register_rank_local_adamw from veomni.optim.optimizer import MultiOptimizer +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + class ToyAdamW(torch.optim.Optimizer): def __init__(self, params, scale=1.0): super().__init__(params, {"lr": 0.01}) @@ -174,27 +178,3 @@ def test_registered_single_optimizer_is_rejected_before_any_save_or_model_load(t str(model_dir), {"model": target, "optimizer": ToyAdamW(target.parameters())}, parallel_state=parallel ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) - - -def test_registered_adapter_cannot_fabricate_optimizer_group_options(tmp_path): - class FabricatingToy(ToyAdamW): - def state_dict(self): - state = super().state_dict() - state["param_groups"][0]["factor"] = 2 - return state - - register_rank_local_adamw( - FabricatingToy, - lambda _optimizer: nullcontext(), - lambda _optimizer: ("toy-v1",), - lambda _optimizer, _param, _field: torch.float32, - lambda _optimizer, _param: False, - ) - model = torch.nn.Linear(2, 2) - leaf = FabricatingToy(model.parameters()) - optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) - with pytest.raises(RuntimeError, match="fabricated checkpoint group options"): - DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") - ) - assert not list(tmp_path.iterdir()) diff --git a/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py b/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py index 791b630b3f..0e91ed2a43 100644 --- a/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py +++ b/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py @@ -248,16 +248,22 @@ def forward(self, x, routing, indices): def _builder_worker(rank, world, shard_dim, directory, rendezvous): + from types import SimpleNamespace + import torch.distributed.checkpoint as dcp from veomni.checkpoint.dcp_checkpointer import ( _EXPERT_LAYOUTS_KEY, + DistributedCheckpointer, ModelState, - OptimizerState, _ExtraParallelSavePlanner, _ModelStrictLoadPlanner, ) + from veomni.checkpoint.layout import optimizer_dir + from veomni.checkpoint.optimizer_resume_layout import MULTI_LAYOUT_MARKER + from veomni.checkpoint.rank_local_optimizer import rank_inventory_key from veomni.distributed.torch_parallelize import build_parallelize_model + from veomni.models.model_runtime import _configure_extra_parallel_checkpoint_layout from veomni.ops.kernels.moe_experts.shared.fp8_weight import FP8ExpertWeight from veomni.optim import build_optimizer @@ -278,6 +284,7 @@ def _builder_worker(rank, world, shard_dim, directory, rendezvous): expert_fsdp_shard_dim=shard_dim, enable_gradient_checkpointing=False, ) + _configure_extra_parallel_checkpoint_layout(model, SimpleNamespace(extra_parallel_canonical_layout=True)) assert all(p.placements[0].dim == shard_dim for p in model.parameters()) assert all(spec.fsdp_shard_dim == shard_dim for spec in model._fqn2spec_info.values()) opt = build_optimizer(model, lr=1e-3) @@ -306,34 +313,43 @@ def _builder_worker(rank, world, shard_dim, directory, rendezvous): for param, state in sub.state.items() for key, value in state.items() } - save_kwargs = { - "state_dict": {"model": ModelState(model), "optimizer": OptimizerState(model, opt)}, - "storage_writer": dcp.FileSystemWriter(directory), + opt.zero_grad(set_to_none=True) + # The proven MultiOptimizer format has a fixed owner layout; + # exercise it separately from model-only resharding. + DistributedCheckpointer.save(directory, {"model": model, "optimizer": opt}) + optim_metadata = dcp.FileSystemReader(optimizer_dir(directory)).read_metadata().state_dict_metadata + assert f"optimizer.{MULTI_LAYOUT_MARKER}" in optim_metadata + for optimizer_rank in range(world): + assert f"optimizer.{rank_inventory_key(optimizer_rank)}" in optim_metadata + + # Model-only export can also be loaded with another FSDP axis. + # Keep async model staging coverage without claiming that the + # exact MultiOptimizer layout itself supports async recovery. + model_save = { + "state_dict": {"model": ModelState(model)}, + "storage_writer": dcp.FileSystemWriter(directory + "-model"), "planner": _ExtraParallelSavePlanner(), } if shard_dim == 0: - # Metadata must also survive the production GPU -> CPU - # async staging path, including optimizer tensors. - dcp.async_save(**save_kwargs, process_group=save_group).result() + dcp.async_save(**model_save, process_group=save_group).result() else: - dcp.save(**save_kwargs) - layouts = dcp.FileSystemReader(directory).read_metadata().planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] - assert any(key.startswith("optimizer.") for key in layouts) - assert all(layout["fsdp_shard_dim"] == shard_dim for layout in layouts.values()) + dcp.save(**model_save) + model_layouts = ( + dcp.FileSystemReader(directory + "-model") + .read_metadata() + .planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + ) + assert any(key.startswith("model.") for key in model_layouts) + assert all(layout["fsdp_shard_dim"] == shard_dim for layout in model_layouts.values()) elif step == 1: expected = y.detach().clone() with torch.no_grad(): for p in model.parameters(): p.zero_() - for sub in opt.optimizers_dict.values(): - for values in sub.state.values(): - for value in values.values(): - value.add_(1) - dcp.load( - {"model": ModelState(model), "optimizer": OptimizerState(model, opt, load=True)}, - checkpoint_id=directory, - planner=_ModelStrictLoadPlanner(strict_model=True), - ) + opt.zero_grad(set_to_none=True) + # The target must be fresh; proven recovery rejects warm state. + opt = build_optimizer(model, lr=1e-3) + DistributedCheckpointer.load(directory, {"model": model, "optimizer": opt}) for (group, param, key), expected_value in saved_optim.items(): actual_value = opt.optimizers_dict[group].state[param][key] torch.testing.assert_close(actual_value, expected_value, rtol=0, atol=0) @@ -347,18 +363,12 @@ def _builder_worker(rank, world, shard_dim, directory, rendezvous): expert_fsdp_shard_dim=1 - shard_dim, enable_gradient_checkpointing=False, ) - restored_opt = build_optimizer(restored, lr=1e-3) + _configure_extra_parallel_checkpoint_layout(restored, SimpleNamespace(extra_parallel_canonical_layout=True)) dcp.load( - {"model": ModelState(restored), "optimizer": OptimizerState(restored, restored_opt, load=True)}, - checkpoint_id=directory, + {"model": ModelState(restored)}, + checkpoint_id=directory + "-model", planner=_ModelStrictLoadPlanner(strict_model=True), ) - parameter_map = dict(zip(model.parameters(), restored.parameters())) - for (group, param, key), expected_value in saved_optim.items(): - actual_value = restored_opt.optimizers_dict[group].state[parameter_map[param]][key] - if isinstance(actual_value, DTensor): - actual_value, expected_value = actual_value.full_tensor(), expected_value.full_tensor() - torch.testing.assert_close(actual_value, expected_value, rtol=0, atol=0) for original, loaded in zip(model.parameters(), restored.parameters()): torch.testing.assert_close(original.full_tensor(), loaded.full_tensor(), rtol=0, atol=0) restored_y = restored(x, routing, indices) @@ -369,6 +379,43 @@ def _builder_worker(rank, world, shard_dim, directory, rendezvous): torch.testing.assert_close( original.grad.full_tensor() / 2, loaded.grad.full_tensor(), rtol=1e-5, atol=1e-7 ) + # A second fresh target at the original layout validates the next + # real optimizer update and its state independently of model resharding. + with torch.device("meta"): + same_layout = _MetaModel() + same_layout = build_parallelize_model( + same_layout, + init_device="meta", + expert_weight_fp8_all_gather=True, + expert_fsdp_shard_dim=shard_dim, + enable_gradient_checkpointing=False, + ) + _configure_extra_parallel_checkpoint_layout(same_layout, SimpleNamespace(extra_parallel_canonical_layout=True)) + same_opt = build_optimizer(same_layout, lr=1e-3) + DistributedCheckpointer.load(directory, {"model": same_layout, "optimizer": same_opt}) + same_layout.set_requires_gradient_sync(False) + same_layout(x, routing, indices).float().square().mean().backward() + same_layout.set_requires_gradient_sync(True) + same_layout(x, routing, indices).float().square().mean().backward() + for original, loaded in zip(model.parameters(), same_layout.parameters(), strict=True): + torch.testing.assert_close(original.grad.full_tensor(), loaded.grad.full_tensor(), rtol=1e-5, atol=1e-7) + # Attribute the exact optimizer comparison to checkpoint state, + # not nondeterminism across independent FP8 backward passes. + loaded.grad.copy_(original.grad) + opt.step() + same_opt.step() + parameter_map = dict(zip(model.parameters(), same_layout.parameters(), strict=True)) + for original, loaded in parameter_map.items(): + torch.testing.assert_close(original.full_tensor(), loaded.full_tensor(), rtol=0, atol=0) + for group, sub in opt.optimizers_dict.items(): + for param, state in sub.state.items(): + restored_state = same_opt.optimizers_dict[group].state[parameter_map[param]] + assert state.keys() == restored_state.keys() + for key, value in state.items(): + actual = restored_state[key] + if isinstance(actual, DTensor): + value, actual = value.full_tensor(), actual.full_tensor() + torch.testing.assert_close(value, actual, rtol=0, atol=0) dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/optim/test_muon_unit.py b/tests/optim/test_muon_unit.py index fdc6a3ca19..ba1db4addb 100644 --- a/tests/optim/test_muon_unit.py +++ b/tests/optim/test_muon_unit.py @@ -25,6 +25,7 @@ import torch.nn as nn from torch.distributed.tensor import Partial, Replicate, Shard +from scripts.ci.accelerator.launch_plan import LaunchSpec from tests.tools.training_utils import make_eager_ops_config from veomni.models.configuration_base import as_veomni_config from veomni.optim import muon as muon_impl @@ -61,6 +62,45 @@ def _toy_model() -> nn.Module: ) +def _assert_nested_state_equal(before, after): + if isinstance(before, torch.Tensor): + torch.testing.assert_close(before, after, atol=0.0, rtol=0.0) + elif isinstance(before, dict): + assert before.keys() == after.keys() + for key in before: + _assert_nested_state_equal(before[key], after[key]) + elif isinstance(before, (list, tuple)): + assert type(before) is type(after) and len(before) == len(after) + for left, right in zip(before, after, strict=True): + _assert_nested_state_equal(left, right) + else: + assert before == after + + +def _dcp_restore_optimizer_and_model(tmp_path, source, source_optimizer, target, target_optimizer): + """Exercise the supported identity-checked MultiOptimizer checkpoint path.""" + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer + + checkpoint = str(tmp_path / "checkpoint") + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(checkpoint, {"model": source, "optimizer": source_optimizer}, parallel_state=parallel) + DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": target_optimizer}, parallel_state=parallel) + for before, after in zip(source.parameters(), target.parameters(), strict=True): + torch.testing.assert_close(before, after, atol=0.0, rtol=0.0) + + +def _assert_next_optimizer_update_matches(source, source_optimizer, target, target_optimizer): + for before, after in zip(source.parameters(), target.parameters(), strict=True): + if before.requires_grad: + before.grad = torch.full_like(before, 0.125) + after.grad = torch.full_like(after, 0.125) + source_optimizer.step() + target_optimizer.step() + for before, after in zip(source.parameters(), target.parameters(), strict=True): + torch.testing.assert_close(before, after, atol=0.0, rtol=0.0) + _assert_nested_state_equal(source_optimizer.state_dict(), target_optimizer.state_dict()) + + def _sample_2d(M: int, K: int, dtype: torch.dtype = torch.float32, seed: int = 0) -> torch.Tensor: g = torch.Generator(device="cpu").manual_seed(seed) return torch.randn(M, K, generator=g, dtype=dtype) @@ -338,7 +378,8 @@ def test_single_step_matches_upstream(self): torch.testing.assert_close(a.detach(), b.detach(), atol=0.0, rtol=0.0) - def test_state_dict_roundtrip(self): + @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=1)) + def test_state_dict_roundtrip(self, tmp_path): """Step + serialize + restore preserves both Muon and AdamW state.""" from veomni.optim import build_optimizer @@ -354,6 +395,7 @@ def test_state_dict_roundtrip(self): if p.requires_grad: p.grad = torch.randn_like(p) opt_a.step() + opt_a.zero_grad(set_to_none=True) sd_before = opt_a.state_dict() flat_keys = set(sd_before.keys()) @@ -365,7 +407,6 @@ def test_state_dict_roundtrip(self): ) model_b = _toy_model() - model_b.load_state_dict(model_a.state_dict()) opt_b = build_optimizer( model_b, lr=1e-4, @@ -373,14 +414,11 @@ def test_state_dict_roundtrip(self): optimizer_type="muon", muon_kwargs={"lr": 1e-2, "adjust_lr_fn": "match_rms_adamw"}, ) - opt_b.load_state_dict(sd_before) + _dcp_restore_optimizer_and_model(tmp_path, model_a, opt_a, model_b, opt_b) sd_after = opt_b.state_dict() - assert set(sd_before.keys()) == set(sd_after.keys()) - for k, v_before in sd_before.items(): - v_after = sd_after[k] - if isinstance(v_before, torch.Tensor): - torch.testing.assert_close(v_before, v_after, atol=0.0, rtol=0.0, msg=f"state_dict mismatch at {k}") + _assert_nested_state_equal(sd_before, sd_after) + _assert_next_optimizer_update_matches(model_a, opt_a, model_b, opt_b) class TestBuildOptimizer: @@ -1090,7 +1128,8 @@ def test_head_blocks_only_reaches_the_state_dict_when_splitting(self, head_group tagged = sorted(k[len("param_groups.") : -len(".head_blocks")] for k in opt.state_dict() if "head_blocks" in k) assert tagged == expected - def test_split_run_survives_a_state_dict_roundtrip(self): + @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=1)) + def test_split_run_survives_a_state_dict_roundtrip(self, tmp_path): """Multi-group Muon must save and reload through the MultiOptimizer path.""" from veomni.optim import build_optimizer @@ -1113,19 +1152,17 @@ def _build(model): if p.requires_grad: p.grad = torch.randn_like(p) opt_a.step() + opt_a.zero_grad(set_to_none=True) saved = opt_a.state_dict() model_b = _attention_model() - model_b.load_state_dict(model_a.state_dict()) opt_b = _build(model_b) - opt_b.load_state_dict(saved) + _dcp_restore_optimizer_and_model(tmp_path, model_a, opt_a, model_b, opt_b) assert [g.get("head_blocks", 1) for g in opt_b.optimizers_dict["muon"].param_groups] == [1, 4] reloaded = opt_b.state_dict() - assert set(saved) == set(reloaded) - for key, before in saved.items(): - if isinstance(before, torch.Tensor): - torch.testing.assert_close(before, reloaded[key], atol=0.0, rtol=0.0, msg=f"mismatch at {key}") + _assert_nested_state_equal(saved, reloaded) + _assert_next_optimizer_update_matches(model_a, opt_a, model_b, opt_b) class TestGramQuackFallback: From c4050a00dfa78c4097aac3dc806dd861ad96e151 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 02:37:58 +0800 Subject: [PATCH 038/120] [checkpoint] docs: classify multi-optimizer resume changes --- docs/usage/checkpoint.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 2dcd659168..79ad37600f 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -283,6 +283,25 @@ supported migration API: it can return success without reading any tensor when the namespaces differ. Replicated values must already agree across ranks. ## MultiOptimizer recovery identity +**Breaking change relative to earlier VeOmni checkpoint behavior:** + +- MultiOptimizer DCP save now requires `save_async: false`. Existing async + MultiOptimizer jobs must switch to synchronous save before adopting this + codec; model-only and ordinary-optimizer async paths remain separate. +- A MultiOptimizer checkpoint can only resume with the recorded owner, mesh, + world size and optimizer layout. It cannot be restored across an EP/FSDP + shard-axis or world-size change, even when model weights can be resharded. +- Older MultiOptimizer checkpoints whose parameter-group and scheduler identity + cannot be proven are rejected. The repository does not yet provide a + validated general migration for these checkpoints. Continue with the + original compatible runtime or use an explicitly accepted weights-only + restart with new optimizer/scheduler state; that is not an exact continuation. + +Supported CPU optimizer offload adapters are independent of asynchronous DCP +save and remain available with synchronous checkpointing. A future migration must preserve +owner-specific master weights, moments and step counters as well as optimizer +groups and scheduler state, and verify the next update against the source. + Use `DistributedCheckpointer` for training checkpoints containing a `MultiOptimizer`. The save is synchronous and requires a completed optimizer step with no pending gradients. Setting `save_async: true` for a runtime whose From 4e40d31709d3d3f44408e51d0950f5803b4928bf Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 03:00:00 +0800 Subject: [PATCH 039/120] [checkpoint] fix: declare scheduler for proven resume --- docs/usage/checkpoint.md | 28 +- tests/checkpoints/test_dcp_checkpointer.py | 22 +- .../test_extra_parallel_layout_v2.py | 6 +- .../test_multi_optimizer_proven_resume.py | 574 +++++++++++++----- .../test_multi_optimizer_rope_cache.py | 72 ++- .../test_optimizer_checkpoint_adapter.py | 44 +- .../test_qwen_omni_nested_checkpoint_hooks.py | 8 +- .../checkpoints/test_rank_local_optimizer.py | 154 ++--- tests/checkpoints/test_rank_local_ple.py | 8 +- .../test_deepgemm_fp8_distributed.py | 12 +- tests/optim/test_multi_optimizer_sparse_lr.py | 20 +- tests/optim/test_muon_unit.py | 12 +- veomni/checkpoint/dcp_checkpointer.py | 15 +- veomni/checkpoint/optimizer_resume_layout.py | 81 +-- veomni/checkpoint/rank_local_optimizer.py | 82 --- veomni/optim/lr_scheduler.py | 7 - 16 files changed, 658 insertions(+), 487 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 79ad37600f..6ea8f2e589 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -302,6 +302,23 @@ save and remain available with synchronous checkpointing. A future migration mus owner-specific master weights, moments and step counters as well as optimizer groups and scheduler state, and verify the next update against the source. +Direct `DistributedCheckpointer` calls containing a `MultiOptimizer` must include +`"_veomni_scheduler_object"` in the state, set to the supported live scheduler +or explicitly to `None` for runs without a scheduler. This declaration is a +trusted contract: a caller must not declare `None` when a scheduler is in use. +The framework checkpoint manager supplies the declaration and restores the +scheduler automatically. A direct `DistributedCheckpointer.load` caller must +apply the returned scheduler state to the live scheduler before continuing: + +```python +state = {"model": model, "optimizer": optimizer, "extra_state": {}, + "_veomni_scheduler_object": scheduler} +DistributedCheckpointer.load(path, state) +scheduler.load_state_dict(state["extra_state"]["lr_scheduler"]) +``` + +Unregistered schedulers are not supported by the proven resume codec. + Use `DistributedCheckpointer` for training checkpoints containing a `MultiOptimizer`. The save is synchronous and requires a completed optimizer step with no pending gradients. Setting `save_async: true` for a runtime whose @@ -322,12 +339,13 @@ optimizer shards and does not support elastic resharding. All ranks must use the same trusted model and configuration. Optimizer leaf ordering is checked as part of the existing checkpoint-layout exchange; owner-specific freezing that removes a whole leaf from only some ranks is rejected before per-leaf -checkpoint collectives, rather than supported. Replicated values are assumed synchronized, -while checkpoint shard dtype, shape and geometry are validated without hashing -tensor bytes. Hooks and adapters are trusted. MultiOptimizer asynchronous saves +checkpoint collectives, rather than supported. Replica values, trainability, dtype, shape and geometry +are trusted to be consistent across ranks. Source and target checkpoint layout +metadata are checked during restore; no tensor bytes are hashed. Hooks and +adapters are trusted. MultiOptimizer asynchronous saves and unsupported layouts fail closed; synchronous staged saves are supported. -Partially overlapping shards with different geometries are rejected, even when -their shared values agree; exact replicas and disjoint shards remain supported. +The runtime assumes ranks use a consistent model configuration and compatible +replica geometries; saved source layout is validated on restore. Group/state inventory and generation mismatches are rejected before model copying. If an I/O or optimizer-application error occurs after model copying, the load fails and clears its receipt but does not roll back in-memory objects; diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py index dbb7194568..0c2fd08cc3 100644 --- a/tests/checkpoints/test_dcp_checkpointer.py +++ b/tests/checkpoints/test_dcp_checkpointer.py @@ -79,11 +79,17 @@ def build_model_and_optimizer(): expected = source_optimizer._sparse_state_dict() checkpoint_dir = tmp_path / "ckpt" - DistributedCheckpointer.save(str(checkpoint_dir), {"model": source_model, "optimizer": source_optimizer}) + DistributedCheckpointer.save( + str(checkpoint_dir), + {"model": source_model, "optimizer": source_optimizer, "_veomni_scheduler_object": None}, + ) dist.barrier() target_model, target_optimizer = build_model_and_optimizer() - DistributedCheckpointer.load(str(checkpoint_dir), {"model": target_model, "optimizer": target_optimizer}) + DistributedCheckpointer.load( + str(checkpoint_dir), + {"model": target_model, "optimizer": target_optimizer, "_veomni_scheduler_object": None}, + ) actual = target_optimizer._sparse_state_dict() assert expected.keys() == actual.keys() @@ -183,10 +189,14 @@ def build_model_and_optimizer(): target_model, target_optimizer = build_model_and_optimizer() DistributedCheckpointer.save( - str(tmp_path), {"model": source_model, "optimizer": source_optimizer}, parallel_state=parallel_state + str(tmp_path), + {"model": source_model, "optimizer": source_optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel_state, ) DistributedCheckpointer.load( - str(tmp_path), {"model": target_model, "optimizer": target_optimizer}, parallel_state=parallel_state + str(tmp_path), + {"model": target_model, "optimizer": target_optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel_state, ) actual = target_optimizer._sparse_state_dict() @@ -231,7 +241,9 @@ def test_multi_optimizer_sparse_state_excludes_synthetic(self, tmp_path): parallel_state = SimpleNamespace(dp_mode="fsdp2") DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel_state + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel_state, ) reader = FileSystemReader(tmp_path / "model" / "optimizer") diff --git a/tests/checkpoints/test_extra_parallel_layout_v2.py b/tests/checkpoints/test_extra_parallel_layout_v2.py index 9c80d159c0..11e90c4bc4 100644 --- a/tests/checkpoints/test_extra_parallel_layout_v2.py +++ b/tests/checkpoints/test_extra_parallel_layout_v2.py @@ -190,7 +190,11 @@ def test_v2_optimizer_fails_before_model_or_checkpoint_mutation(tmp_path): with pytest.raises(RuntimeError, match="owner-aware checkpoint codec"): adapter.load_state_dict({}) with pytest.raises(RuntimeError, match="owner-aware checkpoint codec"): - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert not list(tmp_path.iterdir()) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 9c96cb9bbc..9d3f56bdb6 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -45,7 +45,11 @@ def _load_preflight_failure_worker(rank, checkpoint, failure, endpoint): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(checkpoint, {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + checkpoint, + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _model_and_optimizer() original_model = checkpoint_module._ProvenModelState original_optimizer = checkpoint_module._ProvenMultiOptimizerState @@ -86,7 +90,11 @@ def reader(path): patch.object(DistributedCheckpointer, "_create_storage_reader", side_effect=reader), pytest.raises((OSError, RuntimeError), match="injected|failed on another rank"), ): - DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) if rank == 1: assert injected, f"the {failure} fault did not reach its intended load stage" completed = torch.tensor(1) @@ -109,7 +117,7 @@ def test_multi_optimizer_async_checkpoint_rejected_before_writes(tmp_path): with pytest.raises(RuntimeError, match="require synchronous DCP save"): DistributedCheckpointer.save( str(tmp_path), - {"model": model, "optimizer": optimizer}, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, save_async=True, parallel_state=SimpleNamespace(dp_mode="ddp"), ) @@ -165,7 +173,9 @@ def test_shared_mutable_group_options_rejected_before_save(tmp_path, mode): optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) with pytest.raises(RuntimeError, match="Shared mutable optimizer option"): DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not list(tmp_path.iterdir()) @@ -176,7 +186,9 @@ def test_tensor_lr_aliased_to_optimizer_defaults_rejected_before_save(tmp_path): optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) with pytest.raises(RuntimeError, match="Shared mutable optimizer tensor option"): DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not list(tmp_path.iterdir()) @@ -195,9 +207,17 @@ def build(): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build() - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) source.weight.grad = torch.full_like(source.weight, 0.125) target.weight.grad = source.weight.grad.clone() optimizer.step() @@ -258,7 +278,7 @@ def fail_to_remove(path, *args, **kwargs): with pytest.raises(FileExistsError): DistributedCheckpointer.save( str(checkpoint), - {"model": model, "optimizer": optimizer}, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, stage_dir=str(stage), parallel_state=SimpleNamespace(dp_mode="ddp"), ) @@ -282,10 +302,17 @@ def _proven_staged_resume_worker(rank, checkpoint, stage, endpoint): optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") DistributedCheckpointer.save( - checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel, stage_dir=stage + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + stage_dir=stage, ) target, restored = _model_and_optimizer() - DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) for source, destination in zip(model.parameters(), target.parameters(), strict=True): torch.testing.assert_close(source, destination, rtol=0, atol=0) source.grad = torch.full_like(source, 0.5) @@ -326,7 +353,11 @@ def make_writer(cls, path): DistributedCheckpointer._create_storage_writer = classmethod(make_writer) with pytest.raises(RuntimeError, match="writer"): - DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) completed = torch.tensor(1) dist.all_reduce(completed) assert completed.item() == 2 @@ -372,7 +403,11 @@ def build(): return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"]), parallel source, optimizer, parallel = build() - DistributedCheckpointer.save(checkpoint, {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + checkpoint, + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored, parallel = build() attempts = 0 @@ -386,7 +421,11 @@ def reader(cls, path): DistributedCheckpointer._create_storage_reader = classmethod(reader) with pytest.raises((RuntimeError, OSError), match="inventory reader"): - DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert attempts == 2 completed = torch.tensor(1) dist.all_reduce(completed) @@ -491,34 +530,6 @@ def test_inventory_hash_failure_propagates_before_collective(): spawn_with_rendezvous(_inventory_hash_failure_worker, args=(), nprocs=2) -def _trainability_mismatch_worker(rank, checkpoint, endpoint): - from datetime import timedelta - - import torch.distributed as dist - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - model, optimizer = _model_and_optimizer() - if rank == 1: - model.weight.requires_grad_(False) - with pytest.raises(RuntimeError, match="trainability differs"): - DistributedCheckpointer.save( - checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") - ) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -def test_rank_asymmetric_trainability_is_rejected_before_save(tmp_path): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_trainability_mismatch_worker, args=(str(tmp_path),), nprocs=2) - - def _tied_parameter_order_worker(rank, endpoint): from datetime import timedelta @@ -567,7 +578,9 @@ def synchronize(): checkpoint_module.synchronize = synchronize with pytest.raises(RuntimeError, match="synchronize|finalizing"): DistributedCheckpointer.save( - checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) completed = torch.tensor(1) dist.all_reduce(completed) @@ -648,19 +661,28 @@ def make_model(base): optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") DistributedCheckpointer.save( - str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel, trainable_only=True + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + trainable_only=True, ) target, restored = make_model(1.0 if same_base else 2.0) before = target.bias.detach().clone() if same_base: DistributedCheckpointer.load( - str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel, trainable_only=True + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + trainable_only=True, ) torch.testing.assert_close(target.bias, source.bias, rtol=0, atol=0) else: with pytest.raises(RuntimeError, match="identity|layout|checkpoint"): DistributedCheckpointer.load( - str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel, trainable_only=True + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + trainable_only=True, ) torch.testing.assert_close(target.bias, before, rtol=0, atol=0) @@ -672,11 +694,19 @@ def test_multi_optimizer_rejects_changed_group_identity(tmp_path): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, changed = _model_and_optimizer(split=False) with pytest.raises(RuntimeError, match="(?i)optimizer.*(group|layout|identity)"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": changed}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": changed, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) def test_unversioned_multi_optimizer_checkpoint_fails_closed(tmp_path): @@ -693,7 +723,9 @@ def test_unversioned_multi_optimizer_checkpoint_fails_closed(tmp_path): target, restored = _model_and_optimizer() with pytest.raises(RuntimeError, match="(?i)group layout is unproven"): DistributedCheckpointer.load( - str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=SimpleNamespace(dp_mode="ddp") + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) @@ -703,7 +735,11 @@ def test_nonfresh_target_is_rejected_before_any_checkpoint_load(tmp_path): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _model_and_optimizer() restored.optimizers_dict["weight"].state[target.weight] = { @@ -713,7 +749,11 @@ def test_nonfresh_target_is_rejected_before_any_checkpoint_load(tmp_path): } original = target.weight.detach().clone() with pytest.raises(RuntimeError, match="fresh optimizer"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, original, rtol=0, atol=0) @@ -737,7 +777,11 @@ def build(): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build() target_leaf = restored.optimizers_dict["weight"] original_load = AnyPrecisionAdamW.load_state_dict @@ -753,7 +797,11 @@ def tracked_load(self, state_dict): return result monkeypatch.setattr(AnyPrecisionAdamW, "load_state_dict", tracked_load) - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert len(owned_pointers) == 1 before, after = owned_pointers[0] assert before == after @@ -920,9 +968,17 @@ def build(): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build() - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert target.weight not in restored.optimizers_dict["weight"].state model.weight.grad = torch.ones_like(model.weight) target.weight.grad = torch.ones_like(target.weight) @@ -947,7 +1003,7 @@ def test_trainable_only_rejects_trained_then_frozen_optimizer_parameter(tmp_path with pytest.raises(RuntimeError, match="trainable-only save would omit a parameter with optimizer history"): DistributedCheckpointer.save( str(tmp_path), - {"model": model, "optimizer": optimizer}, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=SimpleNamespace(dp_mode="ddp"), trainable_only=True, ) @@ -959,14 +1015,18 @@ def test_trainable_only_resume_rejects_frozen_parameter_with_saved_history(tmp_p optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _model_and_optimizer() target.weight.requires_grad_(False) original = target.weight.detach().clone() with pytest.raises(RuntimeError, match="trainable-only resume cannot prove"): DistributedCheckpointer.load( str(tmp_path), - {"model": target, "optimizer": restored}, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, parallel_state=parallel, trainable_only=True, ) @@ -980,12 +1040,20 @@ def test_multi_optimizer_uninitialized_leaf_uses_its_first_real_step(tmp_path): optimizer.zero_grad() bias_before = model.bias.detach().clone() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(model.bias, bias_before, rtol=0, atol=0) assert model.bias not in optimizer.optimizers_dict["bias"].state target, restored = _model_and_optimizer() - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert target.bias not in restored.optimizers_dict["bias"].state model.bias.grad = torch.full_like(model.bias, 0.4) target.bias.grad = torch.full_like(target.bias, 0.4) @@ -1045,7 +1113,11 @@ def test_multi_model_generation_keeps_standard_torch_export_strictly_loadable(tm model, optimizer = _model_and_optimizer() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) exported = tmp_path / "export.pt" dcp_to_torch_save(str(tmp_path / "model" / "ckpt"), str(exported)) state = torch.load(exported, weights_only=False) @@ -1097,39 +1169,6 @@ def test_multi_lr_scheduler_rejects_missing_leaf(): scheduler.load_state_dict(state) -@pytest.mark.parametrize("override_container", [False, True]) -@pytest.mark.parametrize("method", ["load_state_dict", "step", "get_last_lr"]) -def test_multi_optimizer_rejects_overridden_scheduler_load(tmp_path, override_container, method): - model, optimizer = _model_and_optimizer() - scheduler = build_lr_scheduler(optimizer, train_steps=100, lr_decay_style="linear") - parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save( - str(tmp_path), - { - "model": model, - "optimizer": optimizer, - "extra_state": {"lr_scheduler": scheduler.state_dict()}, - "_veomni_scheduler_object": scheduler, - }, - parallel_state=parallel, - ) - target, restored = _model_and_optimizer() - restored_scheduler = build_lr_scheduler(restored, train_steps=100, lr_decay_style="linear") - owner = restored_scheduler if override_container else restored_scheduler["weight"] - setattr(owner, method, lambda *args, **kwargs: None) - with pytest.raises(RuntimeError, match="overridden scheduler"): - DistributedCheckpointer.load( - str(tmp_path), - { - "model": target, - "optimizer": restored, - "extra_state": {}, - "_veomni_scheduler_object": restored_scheduler, - }, - parallel_state=parallel, - ) - - def test_multi_optimizer_rejects_snapshot_from_an_old_scheduler_binding(tmp_path): model, optimizer = _model_and_optimizer() original = build_lr_scheduler(optimizer, train_steps=100, lr_decay_style="linear") @@ -1146,25 +1185,6 @@ def test_multi_optimizer_rejects_snapshot_from_an_old_scheduler_binding(tmp_path assert newer is not None -def test_multi_optimizer_rejects_replaced_framework_scheduler_recipe(tmp_path): - model, optimizer = _model_and_optimizer() - scheduler = build_lr_scheduler(optimizer, train_steps=100, lr_decay_style="linear") - - def custom(step, *, factor=0.5): - return factor - - scheduler["weight"].lr_lambdas[0] = custom - state = { - "model": model, - "optimizer": optimizer, - "extra_state": {"lr_scheduler": scheduler.state_dict()}, - "_veomni_scheduler_object": scheduler, - } - with pytest.raises(RuntimeError, match="scheduler recipe changed after construction"): - DistributedCheckpointer.save(str(tmp_path), state, parallel_state=SimpleNamespace(dp_mode="ddp")) - assert not (tmp_path / "model").exists() - - def test_multi_optimizer_rejects_modified_scheduler_sidecar_before_loading(tmp_path): model, optimizer = _model_and_optimizer() scheduler = build_lr_scheduler(optimizer, train_steps=100) @@ -1219,7 +1239,12 @@ def test_multi_optimizer_rejects_mixed_checkpoint_generations(tmp_path, swapped) optimizer.zero_grad() DistributedCheckpointer.save( str(path), - {"model": model, "optimizer": optimizer, "extra_state": {"condition_rng": torch.tensor(index)}}, + { + "model": model, + "optimizer": optimizer, + "extra_state": {"condition_rng": torch.tensor(index)}, + "_veomni_scheduler_object": None, + }, parallel_state=parallel, ) source = first / "model" / swapped @@ -1231,7 +1256,9 @@ def test_multi_optimizer_rejects_mixed_checkpoint_generations(tmp_path, swapped) original = target.weight.detach().clone() with pytest.raises(RuntimeError, match="(?i)generation|layout"): DistributedCheckpointer.load( - str(second), {"model": target, "optimizer": restored, "extra_state": {}}, parallel_state=parallel + str(second), + {"model": target, "optimizer": restored, "extra_state": {}, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) torch.testing.assert_close(target.weight, original, rtol=0, atol=0) @@ -1248,7 +1275,11 @@ def test_model_generation_rechecked_by_actual_dcp_load(tmp_path, monkeypatch): model.weight.grad = torch.full_like(model.weight, 0.1 * index) optimizer.step() optimizer.zero_grad() - DistributedCheckpointer.save(str(path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) original_preflight = checkpoint._preflight_multi_optimizer_layout def replace_after_preflight(*args, **kwargs): @@ -1260,7 +1291,11 @@ def replace_after_preflight(*args, **kwargs): monkeypatch.setattr(checkpoint, "_preflight_multi_optimizer_layout", replace_after_preflight) target, restored = _model_and_optimizer() with pytest.raises(BaseException, match="(?i)checkpoint.*(changed|generation)"): - DistributedCheckpointer.load(str(second), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(second), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) def test_checkpoint_generation_marker_is_not_a_hf_parameter(): @@ -1281,10 +1316,18 @@ def build(): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build() - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, model.weight, rtol=0, atol=0) torch.testing.assert_close(target.bias, model.bias, rtol=0, atol=0) assert target.bias not in restored.optimizers_dict["adam"].state @@ -1309,10 +1352,18 @@ def build(head_blocks): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, changed = build(4) with pytest.raises(RuntimeError, match="(?i)optimizer.*layout"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": changed}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": changed, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) def test_multi_optimizer_layout_binds_extra_parallel_mesh_names(): @@ -1342,17 +1393,27 @@ def build(): optimizer.zero_grad() model.weight.requires_grad_(False) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build() target.weight.requires_grad_(not target_frozen) if not target_frozen: with pytest.raises(RuntimeError, match="(?i)optimizer.*layout"): DistributedCheckpointer.load( - str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) return - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert target.weight.requires_grad is not target_frozen model.weight.requires_grad_(True) target.weight.requires_grad_(True) @@ -1374,7 +1435,7 @@ def test_multi_optimizer_rejects_checkpoint_with_pending_gradients(tmp_path): with pytest.raises(RuntimeError, match="cannot omit pending gradients"): DistributedCheckpointer.save( str(tmp_path), - {"model": model, "optimizer": optimizer}, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=SimpleNamespace(dp_mode="ddp"), ) torch.testing.assert_close(model.weight, before, rtol=0, atol=0) @@ -1394,16 +1455,28 @@ def build(): optimizer.zero_grad() source.bias.grad = torch.ones_like(source.bias) with pytest.raises(RuntimeError, match="cannot omit pending gradients"): - DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert not list(tmp_path.iterdir()) source.bias.grad = None - DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build() target.bias.grad = torch.ones_like(target.bias) before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="cannot omit pending gradients"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) @@ -1416,7 +1489,7 @@ def test_incomplete_optimizer_source_state_cannot_be_saved(tmp_path): with pytest.raises(RuntimeError, match="incomplete or unsupported parameter state"): DistributedCheckpointer.save( str(tmp_path), - {"model": model, "optimizer": optimizer}, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not list(tmp_path.iterdir()) @@ -1432,7 +1505,11 @@ def test_incomplete_optimizer_checkpoint_rejected_before_model_copy(tmp_path): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) path = Path(optimizer_dir(str(tmp_path))) / ".metadata" with path.open("rb") as handle: metadata = pickle.load(handle) @@ -1442,7 +1519,11 @@ def test_incomplete_optimizer_checkpoint_rejected_before_model_copy(tmp_path): target, restored = _model_and_optimizer() before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="incomplete or unsupported parameter state|inventory"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) @@ -1458,7 +1539,11 @@ def test_missing_group_or_all_parameter_history_rejected_before_model_copy(tmp_p optimizer.zero_grad() optimizer.optimizers_dict["weight"].param_groups[0]["lr"] = 0.0025 parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) path = Path(optimizer_dir(str(tmp_path))) / ".metadata" with path.open("rb") as handle: metadata = pickle.load(handle) @@ -1478,7 +1563,11 @@ def test_missing_group_or_all_parameter_history_rejected_before_model_copy(tmp_p target, restored = _model_and_optimizer() before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="inventory"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) @@ -1487,11 +1576,15 @@ def test_multi_optimizer_rejects_stale_sidecar_before_overwrite(tmp_path): parallel = SimpleNamespace(dp_mode="ddp") DistributedCheckpointer.save( str(tmp_path), - {"model": model, "optimizer": optimizer, "extra_state": {"cursor": "old"}}, + {"model": model, "optimizer": optimizer, "extra_state": {"cursor": "old"}, "_veomni_scheduler_object": None}, parallel_state=parallel, ) with pytest.raises(RuntimeError, match="sidecar cannot be reused"): - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) def test_multi_optimizer_rejects_duplicate_ownership_before_checkpoint_write(tmp_path): @@ -1500,7 +1593,7 @@ def test_multi_optimizer_rejects_duplicate_ownership_before_checkpoint_write(tmp with pytest.raises(RuntimeError, match="ownership is duplicated"): DistributedCheckpointer.save( str(tmp_path), - {"model": model, "optimizer": optimizer}, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not list(tmp_path.iterdir()) @@ -1512,7 +1605,9 @@ def test_multi_optimizer_rejects_empty_parameter_group_before_checkpoint_write(t optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) with pytest.raises(RuntimeError, match="empty parameter group"): DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not list(tmp_path.iterdir()) @@ -1528,7 +1623,7 @@ def test_optimizer_root_names_must_match_checkpoint_model(tmp_path): with pytest.raises(RuntimeError, match="different canonical parameter identities"): DistributedCheckpointer.save( str(tmp_path), - {"model": model, "optimizer": optimizer}, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not list(tmp_path.iterdir()) @@ -1543,11 +1638,19 @@ def build(dtype): model, optimizer = build(torch.float32) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build(torch.bfloat16) before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) @@ -1562,7 +1665,11 @@ def build(shared): model, optimizer = build(True) parallel = SimpleNamespace(dp_mode="ddp") with pytest.raises(RuntimeError, match="Shared mutable optimizer tensor option"): - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert not list(tmp_path.iterdir()) @@ -1580,10 +1687,12 @@ def test_proven_optimizer_load_sets_identity_receipt(tmp_path): model, optimizer = _model_and_optimizer() parallel = SimpleNamespace(dp_mode="ddp") generation = DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) target, restored = _model_and_optimizer() - state = {"model": target, "optimizer": restored} + state = {"model": target, "optimizer": restored, "_veomni_scheduler_object": None} DistributedCheckpointer.load(str(tmp_path), state, parallel_state=parallel) assert state["_veomni_optimizer_identity_verified"] is True assert isinstance(generation, str) and state["_veomni_optimizer_checkpoint_generation"] == generation @@ -1599,13 +1708,17 @@ def test_failed_second_load_clears_old_identity_receipt(tmp_path, monkeypatch): first.weight.fill_(1) second.weight.fill_(9) DistributedCheckpointer.save( - str(tmp_path / "first"), {"model": first, "optimizer": first_optimizer}, parallel_state=parallel + str(tmp_path / "first"), + {"model": first, "optimizer": first_optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) DistributedCheckpointer.save( - str(tmp_path / "second"), {"model": second, "optimizer": second_optimizer}, parallel_state=parallel + str(tmp_path / "second"), + {"model": second, "optimizer": second_optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) target, restored = _model_and_optimizer() - state = {"model": target, "optimizer": restored} + state = {"model": target, "optimizer": restored, "_veomni_scheduler_object": None} DistributedCheckpointer.load(str(tmp_path / "first"), state, parallel_state=parallel) assert state["_veomni_optimizer_identity_verified"] is True @@ -1624,12 +1737,20 @@ def test_model_buffer_dtype_change_rejected_before_model_load(tmp_path): model, optimizer = _model_and_optimizer() model.register_buffer("sensitive", torch.tensor([0.1], dtype=torch.float32)) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _model_and_optimizer() target.register_buffer("sensitive", torch.tensor([0.2], dtype=torch.bfloat16)) before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="tensor dtype or shape differs"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) @@ -1637,11 +1758,19 @@ def test_source_model_extra_buffer_rejected_before_model_load(tmp_path): model, optimizer = _model_and_optimizer() model.register_buffer("offset", torch.tensor([0.5])) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _model_and_optimizer() before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="model state inventory differs"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) @@ -1660,11 +1789,19 @@ def build(shared): model, optimizer = build(True) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build(False) before = target.p.detach().clone() with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.p, before, rtol=0, atol=0) @@ -1678,7 +1815,11 @@ def test_multi_optimizer_save_preserves_completed_checkpoint(tmp_path, old_model write_manifest(str(tmp_path), global_step=0, world_size=1, model_state=old_model_state) assert checkpoint_is_complete(str(tmp_path)) with pytest.raises(RuntimeError, match="Cannot overwrite a completed checkpoint"): - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert checkpoint_is_complete(str(tmp_path)) @@ -1696,6 +1837,133 @@ def failing_writer(path): monkeypatch.setattr(DistributedCheckpointer, "_create_storage_writer", failing_writer) with pytest.raises(OSError, match="injected optimizer write failure"): DistributedCheckpointer.save( - str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not checkpoint_is_complete(str(tmp_path)) + + +def test_multi_checkpoint_requires_explicit_scheduler_declaration(tmp_path): + model, optimizer = _model_and_optimizer() + # Ordinary PyTorch schedulers can be attached without VeOmni's registry. + # Omission must not silently be interpreted as a claim that none exists. + scheduler = torch.optim.lr_scheduler.StepLR(optimizer.optimizers_dict["weight"], step_size=1) + state = {"model": model, "optimizer": optimizer} + with pytest.raises(RuntimeError, match="explicit scheduler declaration"): + DistributedCheckpointer.save(str(tmp_path), state, parallel_state=SimpleNamespace(dp_mode="ddp")) + assert not (tmp_path / "model").exists() + state["_veomni_scheduler_object"] = scheduler + with pytest.raises(RuntimeError, match="scheduler object has no matching scheduler binding"): + DistributedCheckpointer.save(str(tmp_path), state, parallel_state=SimpleNamespace(dp_mode="ddp")) + assert not (tmp_path / "model").exists() + + +def test_multi_checkpoint_accepts_explicit_no_scheduler(tmp_path): + model, optimizer = _model_and_optimizer() + model.weight.grad = torch.ones_like(model.weight) + model.bias.grad = torch.ones_like(model.bias) + optimizer.step() + optimizer.zero_grad() + parallel = SimpleNamespace(dp_mode="ddp") + source = {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None} + DistributedCheckpointer.save(str(tmp_path), source, parallel_state=parallel) + target_model, target_optimizer = _model_and_optimizer() + target = {"model": target_model, "optimizer": target_optimizer, "_veomni_scheduler_object": None} + DistributedCheckpointer.load(str(tmp_path), target, parallel_state=parallel) + for actual, expected in zip(target_model.parameters(), model.parameters(), strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(target_optimizer.state_dict(), optimizer.state_dict(), rtol=0, atol=0) + for source_param, target_param in zip(model.parameters(), target_model.parameters(), strict=True): + source_param.grad = torch.full_like(source_param, 0.25) + target_param.grad = torch.full_like(target_param, 0.25) + optimizer.step() + target_optimizer.step() + for actual, expected in zip(target_model.parameters(), model.parameters(), strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(target_optimizer.state_dict(), optimizer.state_dict(), rtol=0, atol=0) + + +def test_direct_multi_checkpoint_applies_returned_scheduler_state(tmp_path): + model, optimizer = _model_and_optimizer() + scheduler = build_lr_scheduler(optimizer, train_steps=10, lr_decay_style="linear") + for param in model.parameters(): + param.grad = torch.ones_like(param) + optimizer.step() + optimizer.zero_grad() + scheduler.step() + parallel = SimpleNamespace(dp_mode="ddp") + source = { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": scheduler.state_dict()}, + "_veomni_scheduler_object": scheduler, + } + DistributedCheckpointer.save(str(tmp_path), source, parallel_state=parallel) + target, target_optimizer = _model_and_optimizer() + target_scheduler = build_lr_scheduler(target_optimizer, train_steps=10, lr_decay_style="linear") + restored = { + "model": target, + "optimizer": target_optimizer, + "extra_state": {}, + "_veomni_scheduler_object": target_scheduler, + } + DistributedCheckpointer.load(str(tmp_path), restored, parallel_state=parallel) + target_scheduler.load_state_dict(restored["extra_state"]["lr_scheduler"]) + assert scheduler.state_dict() == target_scheduler.state_dict() + for source_param, target_param in zip(model.parameters(), target.parameters(), strict=True): + source_param.grad = torch.full_like(source_param, 0.25) + target_param.grad = torch.full_like(target_param, 0.25) + optimizer.step() + target_optimizer.step() + scheduler.step() + target_scheduler.step() + for actual, expected in zip(target.parameters(), model.parameters(), strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert scheduler.state_dict() == target_scheduler.state_dict() + + +def _rank_scheduler_declaration_worker(rank, root, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + torch.manual_seed(42) + model, optimizer = _model_and_optimizer() + parallel = SimpleNamespace(dp_mode="ddp") + state = {"model": model, "optimizer": optimizer} + if rank == 1: + state["_veomni_scheduler_object"] = None + message = "explicit scheduler declaration" if rank == 0 else "failed on another rank" + with pytest.raises(RuntimeError, match=message): + DistributedCheckpointer.save(f"{root}/missing", state, parallel_state=parallel) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + assert not (Path(root) / "missing").exists() + state["_veomni_scheduler_object"] = None + DistributedCheckpointer.save(f"{root}/good", state, parallel_state=parallel) + target, restored = _model_and_optimizer() + before = target.weight.detach().clone() + load_state = {"model": target, "optimizer": restored} + if rank == 0: + load_state["_veomni_scheduler_object"] = None + message = "explicit scheduler declaration" if rank == 1 else "failed on another rank" + with pytest.raises(RuntimeError, match=message): + DistributedCheckpointer.load(f"{root}/good", load_state, parallel_state=parallel) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + assert all(not child.state for child in restored.optimizers_dict.values()) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) +def test_missing_scheduler_declaration_fails_all_ranks_before_mutation(tmp_path): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_rank_scheduler_declaration_worker, args=(str(tmp_path),), nprocs=2) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index 6ef8794ef8..1e401c6764 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -111,7 +111,9 @@ def _rope_update(model, optimizer): def _fresh_resume(_rank, checkpoint, expected_path): model, optimizer = _build() DistributedCheckpointer.load( - checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) trace = _rope_update(model, optimizer) expected = torch.load(expected_path, map_location="cpu", weights_only=True) @@ -158,7 +160,9 @@ def test_expanded_dynamic_rope_cache_is_rejected_before_checkpoint_write(tmp_pat checkpoint = tmp_path / "checkpoint" with pytest.raises(RuntimeError, match="Dynamic RoPE cache.*depends on prior sequence lengths"): DistributedCheckpointer.save( - str(checkpoint), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + str(checkpoint), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not checkpoint.exists() @@ -173,22 +177,38 @@ def test_old_unproven_expanded_rope_checkpoint_is_rejected_before_model_load(tmp parallel = SimpleNamespace(dp_mode="ddp") with monkeypatch.context() as old: old.setattr(dynamic_rope, "dynamic_rope_checkpoint_contract", lambda _model: []) - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _build() before = target.projection.weight.detach().clone() with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) def test_changed_dynamic_rope_config_is_rejected_before_model_load(tmp_path): model, optimizer = _build() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _build(theta=5000) before = target.projection.weight.detach().clone() with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) @@ -201,7 +221,9 @@ def _rank_local_rope_error(rank, checkpoint, endpoint): error = "Dynamic RoPE cache.*prior sequence lengths" if rank == 1 else "layout failed on another rank" with pytest.raises(RuntimeError, match=error): DistributedCheckpointer.save( - checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) finally: dist.destroy_process_group() @@ -223,7 +245,9 @@ def test_real_qwen3_dynamic_rope_reset_is_rejected_before_save(tmp_path): checkpoint = tmp_path / "checkpoint" with pytest.raises(RuntimeError, match="Cannot prove dynamic RoPE checkpoint cache"): DistributedCheckpointer.save( - str(checkpoint), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + str(checkpoint), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) assert not checkpoint.exists() @@ -233,7 +257,9 @@ def test_unexpanded_dynamic_rope_fresh_process_matches_training_update(tmp_path) _update(model, optimizer, 1.0) checkpoint = str(tmp_path / "checkpoint") DistributedCheckpointer.save( - checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), ) trace = _rope_update(model, optimizer) expected = tmp_path / "expected.pt" @@ -262,12 +288,20 @@ def build(alias): model.weight.fill_(5) model.persistent.fill_(7) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build(not source_alias) before = target.weight.detach().clone() persistent_before = target.persistent.detach().clone() with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) torch.testing.assert_close(target.persistent, persistent_before, rtol=0, atol=0) @@ -291,11 +325,19 @@ def build(alias): with torch.no_grad(): model.weight.fill_(5) parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build(not source_alias) before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="layout|identity"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) @@ -334,6 +376,8 @@ def build(alias): parallel = SimpleNamespace(dp_mode="ddp") with pytest.raises(RuntimeError, match="Cannot prove dynamic RoPE checkpoint cache"): DistributedCheckpointer.save( - str(tmp_path / "checkpoint"), {"model": model, "optimizer": optimizer}, parallel_state=parallel + str(tmp_path / "checkpoint"), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) assert not (tmp_path / "checkpoint").exists() diff --git a/tests/checkpoints/test_optimizer_checkpoint_adapter.py b/tests/checkpoints/test_optimizer_checkpoint_adapter.py index 2c3b5cbae4..4b825c1a3c 100644 --- a/tests/checkpoints/test_optimizer_checkpoint_adapter.py +++ b/tests/checkpoints/test_optimizer_checkpoint_adapter.py @@ -74,9 +74,17 @@ def test_registered_optimizer_preserves_lazy_state_and_next_update(tmp_path): optimizer.step() optimizer.zero_grad() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _build() - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert target.bias not in restored.optimizers_dict["toy"].state for left, right in zip(source.parameters(), target.parameters(), strict=True): left.grad = torch.full_like(left, 0.125) @@ -106,10 +114,18 @@ def test_registered_optimizer_raw_snapshot_is_detached_but_restore_requires_proo def test_registered_optimizer_rejects_changed_algorithm_before_loading(tmp_path): source, optimizer = _build() parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = _build(scale=2.0) with pytest.raises(RuntimeError, match="layout|identity|checkpoint"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) @pytest.mark.parametrize("changed_field", ["master_weight", "step"]) @@ -135,11 +151,19 @@ def build(dtype): parallel = SimpleNamespace(dp_mode="ddp") source, optimizer = build(torch.float32) assert not optimizer.optimizers_dict["leaf"].state - DistributedCheckpointer.save(str(tmp_path), {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) target, restored = build(torch.bfloat16) before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="layout|identity|checkpoint"): - DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) @@ -157,7 +181,9 @@ def test_registered_single_optimizer_is_rejected_before_any_save_or_model_load(t invalid_dir = tmp_path / "single-optimizer" with pytest.raises(RuntimeError, match="require MultiOptimizer"): DistributedCheckpointer.save( - str(invalid_dir), {"model": source, "optimizer": optimizer}, parallel_state=parallel + str(invalid_dir), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) assert not invalid_dir.exists() @@ -175,6 +201,8 @@ def test_registered_single_optimizer_is_rejected_before_any_save_or_model_load(t before = target.weight.detach().clone() with pytest.raises(RuntimeError, match="require MultiOptimizer"): DistributedCheckpointer.load( - str(model_dir), {"model": target, "optimizer": ToyAdamW(target.parameters())}, parallel_state=parallel + str(model_dir), + {"model": target, "optimizer": ToyAdamW(target.parameters()), "_veomni_scheduler_object": None}, + parallel_state=parallel, ) torch.testing.assert_close(target.weight, before, rtol=0, atol=0) diff --git a/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py b/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py index 590b08d4f2..06e8e9f2c3 100644 --- a/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py +++ b/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py @@ -75,9 +75,13 @@ def build(): expected = {name: value.detach().clone() for name, value in model.state_dict().items()} parallel = SimpleNamespace(dp_mode="ddp") checkpoint = str(tmp_path / "checkpoint") - DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel + ) target, restored = build() - DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + checkpoint, {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, parallel_state=parallel + ) actual = target.state_dict() assert expected.keys() == actual.keys() for name in expected: diff --git a/tests/checkpoints/test_rank_local_optimizer.py b/tests/checkpoints/test_rank_local_optimizer.py index b34067c3ed..8a0c8fb169 100644 --- a/tests/checkpoints/test_rank_local_optimizer.py +++ b/tests/checkpoints/test_rank_local_optimizer.py @@ -18,7 +18,6 @@ local_optimizer_state, recover_frozen_optimizer_ownership, restore_optimizer_state, - verify_replicated_model_shards, ) from veomni.distributed.parallel_plan import SpecInfo from veomni.optim.muon import DistributedMuon @@ -105,11 +104,15 @@ def build(): optimizer.step() optimizer.zero_grad() DistributedCheckpointer.save( - f"{root}/training", {"model": source, "optimizer": optimizer}, parallel_state=parallel + f"{root}/training", + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) target, restored = build() DistributedCheckpointer.load( - f"{root}/training", {"model": target, "optimizer": restored}, parallel_state=parallel + f"{root}/training", + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) torch.testing.assert_close(target.weight.to_local(), source.weight.to_local(), rtol=0, atol=0) source_state = optimizer.optimizers_dict["expert"].state.get(source.weight, {}) @@ -215,7 +218,7 @@ def advance(gradient, *, include_bias=True): optimizer.zero_grad() DistributedCheckpointer.save( f"{root}/cold", - {"model": model, "optimizer": optimizer}, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel, ) before = snapshot() @@ -223,7 +226,9 @@ def advance(gradient, *, include_bias=True): torch.save({"before": before, "after": snapshot()}, f"{root}/reference-{rank}.pt") else: DistributedCheckpointer.load( - f"{root}/cold", {"model": model, "optimizer": optimizer}, parallel_state=parallel + f"{root}/cold", + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) reference = torch.load(f"{root}/reference-{rank}.pt", weights_only=True) torch.testing.assert_close(snapshot(), reference["before"], rtol=0, atol=0) @@ -277,76 +282,6 @@ def test_anyprecision_bf16_state_extra_parallel_fresh_process_resume(tmp_path, s ) -def _global_tensor_metadata_mismatch_worker(rank, root, mismatch): - from torch.distributed.tensor import DTensor - - dist.init_process_group( - "gloo", - init_method=f"file://{root}/global-metadata-init", - rank=rank, - world_size=2, - timeout=timedelta(seconds=30), - ) - try: - mesh = init_device_mesh("cpu", (2,)) - dtype = torch.bfloat16 if mismatch == "dtype" and rank else torch.float32 - global_size = 6 if mismatch == "shape" and rank else 4 - local = torch.ones(global_size // 2, dtype=dtype) - tensor = DTensor.from_local(local, mesh, [Shard(0)], shape=torch.Size([global_size]), stride=(1,)) - with pytest.raises(RuntimeError, match="global dtype or shape"): - verify_replicated_model_shards({"expert": tensor}) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("mismatch", ["dtype", "shape"]) -def test_checkpoint_global_tensor_metadata_mismatch_fails_all_ranks(tmp_path, mismatch): - mp.spawn(_global_tensor_metadata_mismatch_worker, args=(str(tmp_path), mismatch), nprocs=2, join=True) - - -def _overlapping_tensor_shards_worker(rank, root, same_offset): - from torch.distributed.tensor import DTensor - - dist.init_process_group( - "gloo", - init_method=f"file://{root}/overlap-init", - rank=rank, - world_size=2, - timeout=timedelta(seconds=30), - ) - try: - mesh = init_device_mesh("cpu", (2,)) - shard = rank == (0 if same_offset else 1) - placement = Shard(0) if shard else Replicate() - tensor = DTensor.from_local( - torch.ones(2 if shard else 4), mesh, [placement], shape=torch.Size([4]), stride=(1,), run_check=False - ) - with pytest.raises(RuntimeError, match="overlapping nonidentical shards"): - verify_replicated_model_shards({"expert": tensor}) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("same_offset", [False, True]) -def test_checkpoint_overlapping_shard_geometries_fail_all_ranks(tmp_path, same_offset): - mp.spawn(_overlapping_tensor_shards_worker, args=(str(tmp_path), same_offset), nprocs=2, join=True) - - -def test_checkpoint_rejects_empty_shard_with_same_offset_as_nonempty_shard(): - from veomni.checkpoint.rank_local_optimizer import _reject_overlapping_geometries - - with pytest.raises(RuntimeError, match="overlapping nonidentical shards"): - _reject_overlapping_geometries("expert", [((0,), (0,)), ((0,), (2,))]) - - def _owner_frozen_leaf_mismatch_worker(rank, root, mode): from unittest.mock import patch @@ -371,7 +306,9 @@ def _owner_frozen_leaf_mismatch_worker(rank, root, mode): if mode == "load": source = build_extra_parallel_fsdp2_optimizer(model) DistributedCheckpointer.save( - f"{root}/checkpoint", {"model": model, "optimizer": source}, parallel_state=parallel + f"{root}/checkpoint", + {"model": model, "optimizer": source, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) if owner == 1: model.weight.requires_grad_(False) @@ -382,7 +319,11 @@ def _owner_frozen_leaf_mismatch_worker(rank, root, mode): before = model.weight.to_local().clone() with pytest.raises(RuntimeError, match="leaf order differs across ranks"): method = DistributedCheckpointer.save if mode == "save" else DistributedCheckpointer.load - method(f"{root}/checkpoint", {"model": model, "optimizer": optimizer}, parallel_state=parallel) + method( + f"{root}/checkpoint", + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(model.weight.to_local(), before, rtol=0, atol=0) completed = torch.tensor(1) dist.all_reduce(completed) @@ -397,45 +338,6 @@ def test_owner_specific_freeze_missing_leaf_fails_before_checkpoint_collectives( mp.spawn(_owner_frozen_leaf_mismatch_worker, args=(str(tmp_path), mode), nprocs=4, join=True) -def _owner_trainability_worker(rank, root, within_owner_mismatch): - from veomni.checkpoint.optimizer_resume_layout import coordinated_multi_layout - - dist.init_process_group( - "gloo", - init_method=f"file://{root}/owner-trainability-init", - rank=rank, - world_size=4, - timeout=timedelta(seconds=30), - ) - try: - mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ple_fsdp", "ple")) - model = torch.nn.Module() - values = torch.arange(64, dtype=torch.float32).reshape(8, 8) - model.weight = torch.nn.Parameter(distribute_tensor(values, mesh, [Shard(1), Shard(0)])) - model._fqn2spec_info = {"weight": SpecInfo("ple", Shard(0), "weight", mesh, 1, 1, (8, 8))} - optimizer = MultiOptimizer(model, {"leaf": torch.optim.AdamW([model.weight])}, ["leaf"]) - owner = mesh.get_local_rank("ple") - replica = mesh.get_local_rank("ple_fsdp") - if owner == 1 and (not within_owner_mismatch or replica == 1): - model.weight.requires_grad_(False) - if within_owner_mismatch: - with pytest.raises(RuntimeError, match="trainability differs"): - coordinated_multi_layout(model, optimizer, writing=True) - else: - coordinated_multi_layout(model, optimizer, writing=True) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 4 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) -@pytest.mark.parametrize("within_owner_mismatch", [False, True]) -def test_trainability_may_differ_between_owners_but_not_replicas(tmp_path, within_owner_mismatch): - mp.spawn(_owner_trainability_worker, args=(str(tmp_path), within_owner_mismatch), nprocs=4, join=True) - - def _frozen_resume_worker(rank, root, phase): dist.init_process_group( "gloo", @@ -468,7 +370,9 @@ def build(): optimizer.zero_grad() model.weight.requires_grad_(False) DistributedCheckpointer.save( - f"{root}/frozen", {"model": model, "optimizer": optimizer}, parallel_state=parallel + f"{root}/frozen", + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) model.weight.requires_grad_(True) model.weight.grad = torch.full_like(model.weight, 0.2) @@ -481,7 +385,9 @@ def build(): optimizer = build() assert not model.weight.requires_grad DistributedCheckpointer.load( - f"{root}/frozen", {"model": model, "optimizer": optimizer}, parallel_state=parallel + f"{root}/frozen", + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) model.weight.requires_grad_(True) model.weight.grad = torch.full_like(model.weight, 0.2) @@ -529,7 +435,11 @@ def build(): optimizer.step() optimizer.zero_grad() checkpoint = f"{root}/corrupt" - DistributedCheckpointer.save(checkpoint, {"model": source, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + checkpoint, + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) if rank == 0: opt_path = optimizer_dir(checkpoint, "") metadata = dcp.FileSystemReader(opt_path).read_metadata() @@ -543,7 +453,11 @@ def build(): from torch.distributed.checkpoint.api import CheckpointException with pytest.raises((CheckpointException, RuntimeError)): - DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) torch.testing.assert_close(target.weight.to_local(), before, rtol=0, atol=0) completed = torch.tensor(1) dist.all_reduce(completed) diff --git a/tests/checkpoints/test_rank_local_ple.py b/tests/checkpoints/test_rank_local_ple.py index ae97e4ba3b..5478fdf93f 100644 --- a/tests/checkpoints/test_rank_local_ple.py +++ b/tests/checkpoints/test_rank_local_ple.py @@ -53,7 +53,9 @@ def snapshot(): optimizer.step() optimizer.zero_grad() DistributedCheckpointer.save( - f"{root}/checkpoint", {"model": model, "optimizer": optimizer}, parallel_state=parallel + f"{root}/checkpoint", + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) before = snapshot() model.weight.grad = torch.full_like(model.weight, 0.2) @@ -61,7 +63,9 @@ def snapshot(): torch.save({"before": before, "after": snapshot()}, f"{root}/ple-ref-{rank}.pt") else: DistributedCheckpointer.load( - f"{root}/checkpoint", {"model": model, "optimizer": optimizer}, parallel_state=parallel + f"{root}/checkpoint", + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) reference = torch.load(f"{root}/ple-ref-{rank}.pt", weights_only=True) torch.testing.assert_close(snapshot(), reference["before"], rtol=0, atol=0) diff --git a/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py b/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py index 0e91ed2a43..84f4d18a08 100644 --- a/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py +++ b/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py @@ -316,7 +316,9 @@ def _builder_worker(rank, world, shard_dim, directory, rendezvous): opt.zero_grad(set_to_none=True) # The proven MultiOptimizer format has a fixed owner layout; # exercise it separately from model-only resharding. - DistributedCheckpointer.save(directory, {"model": model, "optimizer": opt}) + DistributedCheckpointer.save( + directory, {"model": model, "optimizer": opt, "_veomni_scheduler_object": None} + ) optim_metadata = dcp.FileSystemReader(optimizer_dir(directory)).read_metadata().state_dict_metadata assert f"optimizer.{MULTI_LAYOUT_MARKER}" in optim_metadata for optimizer_rank in range(world): @@ -349,7 +351,9 @@ def _builder_worker(rank, world, shard_dim, directory, rendezvous): opt.zero_grad(set_to_none=True) # The target must be fresh; proven recovery rejects warm state. opt = build_optimizer(model, lr=1e-3) - DistributedCheckpointer.load(directory, {"model": model, "optimizer": opt}) + DistributedCheckpointer.load( + directory, {"model": model, "optimizer": opt, "_veomni_scheduler_object": None} + ) for (group, param, key), expected_value in saved_optim.items(): actual_value = opt.optimizers_dict[group].state[param][key] torch.testing.assert_close(actual_value, expected_value, rtol=0, atol=0) @@ -392,7 +396,9 @@ def _builder_worker(rank, world, shard_dim, directory, rendezvous): ) _configure_extra_parallel_checkpoint_layout(same_layout, SimpleNamespace(extra_parallel_canonical_layout=True)) same_opt = build_optimizer(same_layout, lr=1e-3) - DistributedCheckpointer.load(directory, {"model": same_layout, "optimizer": same_opt}) + DistributedCheckpointer.load( + directory, {"model": same_layout, "optimizer": same_opt, "_veomni_scheduler_object": None} + ) same_layout.set_requires_gradient_sync(False) same_layout(x, routing, indices).float().square().mean().backward() same_layout.set_requires_gradient_sync(True) diff --git a/tests/optim/test_multi_optimizer_sparse_lr.py b/tests/optim/test_multi_optimizer_sparse_lr.py index 6002ab41a9..9a3833aece 100644 --- a/tests/optim/test_multi_optimizer_sparse_lr.py +++ b/tests/optim/test_multi_optimizer_sparse_lr.py @@ -146,7 +146,9 @@ def test_pytorch_patched_adamw_dcp_resume_preserves_next_update(tmp_path, initia ] parallel = SimpleNamespace(dp_mode="ddp") checkpoint = str(tmp_path / "patched") - DistributedCheckpointer.save(checkpoint, {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel + ) for actual, expected, expected_state in zip(model.parameters(), before, before_states, strict=True): torch.testing.assert_close(actual, expected, rtol=0, atol=0) assert leaf.state.get(actual, {}).keys() == expected_state.keys() @@ -158,7 +160,11 @@ def test_pytorch_patched_adamw_dcp_resume_preserves_next_update(tmp_path, initia if target_patched: _patch_optimizer_state_dict(restored_model, optimizers=(restored_leaf,)) restored = MultiOptimizer(restored_model, {"leaf": restored_leaf}, ["leaf"]) - DistributedCheckpointer.load(checkpoint, {"model": restored_model, "optimizer": restored}, parallel_state=parallel) + DistributedCheckpointer.load( + checkpoint, + {"model": restored_model, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert bool(restored_leaf.state) is initialized assert restored_leaf.param_groups[0]["lr"] == 0.0025 for current, target in zip(model.parameters(), restored_model.parameters(), strict=True): @@ -202,11 +208,17 @@ def test_sparse_leaf_learning_rate_survives_dcp_roundtrip(tmp_path): weight.zero_grad() bias.param_groups[0]["lr"] = 0.0025 parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + str(tmp_path), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) assert not bias.state restored_model, _, restored_bias, restored = _build() DistributedCheckpointer.load( - str(tmp_path), {"model": restored_model, "optimizer": restored}, parallel_state=parallel + str(tmp_path), + {"model": restored_model, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, ) assert restored_bias.param_groups[0]["lr"] == 0.0025 diff --git a/tests/optim/test_muon_unit.py b/tests/optim/test_muon_unit.py index ba1db4addb..626862ba5a 100644 --- a/tests/optim/test_muon_unit.py +++ b/tests/optim/test_muon_unit.py @@ -83,8 +83,16 @@ def _dcp_restore_optimizer_and_model(tmp_path, source, source_optimizer, target, checkpoint = str(tmp_path / "checkpoint") parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save(checkpoint, {"model": source, "optimizer": source_optimizer}, parallel_state=parallel) - DistributedCheckpointer.load(checkpoint, {"model": target, "optimizer": target_optimizer}, parallel_state=parallel) + DistributedCheckpointer.save( + checkpoint, + {"model": source, "optimizer": source_optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": target_optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) for before, after in zip(source.parameters(), target.parameters(), strict=True): torch.testing.assert_close(before, after, atol=0.0, rtol=0.0) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 8f725f8038..9e03ef9fc5 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -928,7 +928,6 @@ def __init__(self, *args, optimizer, group, writing=False, **kwargs): super().__init__(*args, **kwargs) self.optimizer = optimizer self.group = group - self.writing = writing def state_dict(self): result, error = None, None @@ -945,10 +944,6 @@ def state_dict(self): except BaseException as exc: error = exc raise_if_any_rank_failed(error, "constructing proven model state", group=self.group) - if self.writing: - from .rank_local_optimizer import verify_replicated_model_shards - - verify_replicated_model_shards(result, group=self.group) return result def load_state_dict(self, state_dict): @@ -1773,6 +1768,11 @@ def save( error = None try: parallel_state = parallel_state if parallel_state is not None else get_parallel_state() + if "_veomni_scheduler_object" not in state: + raise RuntimeError( + "MultiOptimizer checkpoint requires an explicit scheduler declaration; " + "set _veomni_scheduler_object to the scheduler or None" + ) if save_async: raise RuntimeError("Proven MultiOptimizer checkpoints currently require synchronous DCP save") _validate_source_optimizer_states(optimizer, source_model) @@ -2072,6 +2072,11 @@ def _optimizer_state() -> "OptimizerState": error = None try: parallel_state = parallel_state if parallel_state is not None else get_parallel_state() + if "_veomni_scheduler_object" not in state: + raise RuntimeError( + "MultiOptimizer checkpoint requires an explicit scheduler declaration; " + "set _veomni_scheduler_object to the scheduler or None" + ) except BaseException as exc: error = exc raise_if_any_rank_failed(error, "validating MultiOptimizer checkpoint policy", group=process_group) diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index 6ecc50380e..c2aee39d47 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -68,48 +68,15 @@ def _lr_layout(group, key, aliases): def _scheduler_layout(leaf): import math - from ..optim.lr_scheduler import checkpoint_scheduler_binding, checkpoint_scheduler_functions + from ..optim.lr_scheduler import checkpoint_scheduler_binding scheduler = checkpoint_scheduler_binding(leaf) - wrapped = getattr(leaf.step, "_wrapped_by_lr_sched", False) if scheduler is None: - if wrapped or "step" in vars(leaf): - raise RuntimeError("MultiOptimizer scheduler identity is unknown; explicit migration is required") return None - import weakref - - def nested_codes(code): - result = {code} - for value in code.co_consts: - if isinstance(value, types.CodeType): - result.update(nested_codes(value)) - return result - - step = vars(leaf).get("step") - permitted = nested_codes(torch.optim.lr_scheduler.LRScheduler.__init__.__code__) - closure = getattr(step, "__closure__", None) or () - captured = tuple(cell.cell_contents for cell in closure) - if ( - getattr(step, "__code__", None) not in permitted - or not any(item is type(leaf).step for item in captured) - or not any(isinstance(item, weakref.ReferenceType) and item() is leaf for item in captured) - ): - raise RuntimeError("MultiOptimizer cannot prove optimizer step implementation") - if type(scheduler) is not torch.optim.lr_scheduler.LambdaLR or any( - name in vars(scheduler) for name in ("state_dict", "load_state_dict", "step", "get_lr", "get_last_lr") - ): - raise RuntimeError("MultiOptimizer cannot prove an overridden scheduler method") - if scheduler.optimizer is not leaf or not wrapped: + if type(scheduler) is not torch.optim.lr_scheduler.LambdaLR: + raise RuntimeError("MultiOptimizer requires a supported scheduler") + if scheduler.optimizer is not leaf: raise RuntimeError("MultiOptimizer scheduler binding changed") - original_functions = checkpoint_scheduler_functions(scheduler) - if ( - original_functions is None - or len(original_functions) != len(scheduler.lr_lambdas) - or any( - current is not original for current, original in zip(scheduler.lr_lambdas, original_functions, strict=True) - ) - ): - raise RuntimeError("MultiOptimizer scheduler recipe changed after construction") functions = [] for function in scheduler.lr_lambdas: if not isinstance(function, types.FunctionType): @@ -327,9 +294,6 @@ def mesh_identity(mesh): if type(scheduler_object) is not MultiLRScheduler or set(scheduler_object) != set(keys): raise RuntimeError("MultiOptimizer scheduler checkpoint requires its live scheduler object") - if any(name in vars(scheduler_object) for name in ("state_dict", "load_state_dict", "step", "get_last_lr")): - raise RuntimeError("MultiOptimizer cannot prove an overridden scheduler container method") - for name in keys: current_scheduler = scheduler_object[name] if current_scheduler is not checkpoint_scheduler_binding(optimizer.optimizers_dict[name]): @@ -395,7 +359,7 @@ def coordinated_multi_layout( frozen_digest=None, ): """Return the same rank-indexed layout on every participant or fail together.""" - local, trainability = None, None + local = None error = None try: local = _local_layout( @@ -407,31 +371,6 @@ def coordinated_multi_layout( writing=writing, frozen_digest=frozen_digest, ) - from torch.distributed.tensor import Shard - - specs = getattr(model, "_fqn2spec_info", {}) or {} - trainability_by_owner = {} - for runtime_name, param in model.named_parameters(remove_duplicate=False): - for name in _get_fqns(model, runtime_name): - spec = specs.get(runtime_name) or specs.get(name) - owner = None - if spec is not None and isinstance(spec.placement, Shard): - mesh = spec.para_fsdp_mesh - if mesh is None: - raise RuntimeError(f"Cannot prove ExtraParallel owner for trainability of {name}") - owner = (spec.para_name, mesh[spec.para_name].get_local_rank()) - if owner[1] < 0: - raise RuntimeError(f"Current rank is not an ExtraParallel owner of {name}") - values = trainability_by_owner.setdefault(owner, {}) - enabled = bool(param.requires_grad) - previous = values.setdefault(name, enabled) - if previous != enabled: - raise RuntimeError(f"Tied parameter trainability differs for {name}") - trainability = [] - for owner, values in trainability_by_owner.items(): - payload = json.dumps(sorted(values.items()), separators=(",", ":")).encode() - trainability.append((owner, len(values), hashlib.sha256(payload).hexdigest())) - trainability = tuple(sorted(trainability, key=lambda entry: repr(entry[0]))) except BaseException as exc: error = exc raise_if_any_rank_failed(error, "inspecting MultiOptimizer layout", group=group) @@ -442,16 +381,10 @@ def coordinated_multi_layout( # This existing layout exchange also guards the order of subsequent # per-leaf collectives. Owner-specific freezing can otherwise leave # one rank without a leaf and strand peers during checkpoint I/O. - dist.all_gather_object(gathered, (local, trainability, tuple(optimizer.key_names)), group=group) + dist.all_gather_object(gathered, (local, tuple(optimizer.key_names)), group=group) digests = [item[0] for item in gathered] - if any(item[2] != gathered[0][2] for item in gathered): + if any(item[1] != gathered[0][1] for item in gathered): raise RuntimeError("MultiOptimizer checkpoint leaf order differs across ranks") - owners = {} - for _, entries, _ in gathered: - for owner, count, digest in entries: - previous = owners.setdefault(owner, (count, digest)) - if previous != (count, digest): - raise RuntimeError("Model parameter trainability differs across checkpoint replica ranks") ranks = list(range(world)) if group is None else dist.get_process_group_ranks(group) else: ranks, digests = [0], [local] diff --git a/veomni/checkpoint/rank_local_optimizer.py b/veomni/checkpoint/rank_local_optimizer.py index f71c8a9d81..313416d8d4 100644 --- a/veomni/checkpoint/rank_local_optimizer.py +++ b/veomni/checkpoint/rank_local_optimizer.py @@ -9,7 +9,6 @@ mesh identity before use. It does not support elastic resharding. """ -import heapq import io import json from contextlib import contextmanager @@ -162,87 +161,6 @@ def zeros_like_parameter_on_cpu(param, *, dtype): return wrap_local_like_parameter(torch.zeros_like(local, dtype=dtype, device="cpu"), param) -def _reject_overlapping_geometries(name, geometries): - """Reject distinct overlapping boxes before DCP deduplicates by offset.""" - offsets_to_size = {} - nonempty = [] - for offsets, size in geometries: - old_size = offsets_to_size.setdefault(offsets, size) - if old_size != size: - raise RuntimeError(f"Checkpoint tensor {name} has overlapping nonidentical shards") - if all(size): - nonempty.append((offsets, size)) - if len(nonempty) < 2 or not nonempty[0][0]: - return - # Choose a dimension with the most distinct intervals to keep the sweep - # active set small for both row- and column-sharded tensors. - axis = max( - range(len(nonempty[0][0])), - key=lambda dim: len({(offsets[dim], size[dim]) for offsets, size in nonempty}), - ) - ordered = sorted(nonempty, key=lambda item: item[0][axis]) - active, expiration = {}, [] - for index, (offsets, size) in enumerate(ordered): - start = offsets[axis] - while expiration and expiration[0][0] <= start: - _, previous = heapq.heappop(expiration) - del active[previous] - for other_offsets, other_size in active.values(): - if all( - start < other_start + other_length and other_start < start + length - for start, length, other_start, other_length in zip( - offsets, size, other_offsets, other_size, strict=True - ) - ): - raise RuntimeError(f"Checkpoint tensor {name} has overlapping nonidentical shards") - active[index] = (offsets, size) - heapq.heappush(expiration, (start + size[axis], index)) - - -def verify_replicated_model_shards(state: dict, *, group=None): - """Validate shard geometry without reading or hashing tensor values.""" - from ..utils.dist_utils import raise_if_any_rank_failed - - error, records = None, [] - try: - flattened, _ = flatten_state_dict({"model": state}) - for name, value in sorted(flattened.items()): - if not torch.is_tensor(value): - continue - if isinstance(value, DTensor): - chunks = value.__create_chunk_list__() - if len(chunks) != 1: - raise RuntimeError(f"Cannot prove model shard geometry for {name}") - offsets, size = tuple(chunks[0].offsets), tuple(chunks[0].sizes) - local = value.to_local() - else: - local = value - offsets, size = (0,) * value.ndim, tuple(value.shape) - if tuple(local.shape) != size: - raise RuntimeError(f"Cannot prove model shard shape for {name}") - records.append((name, offsets, size, str(local.dtype), tuple(value.shape))) - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "checking model shard geometry", group=group) - if not dist.is_initialized(): - return - gathered = [None] * dist.get_world_size(group) - dist.all_gather_object(gathered, records, group=group) - geometries, signatures, seen = {}, {}, set() - for items in gathered: - for name, offsets, size, dtype, global_shape in items: - signature = (dtype, global_shape) - previous = signatures.setdefault(name, signature) - if previous != signature: - raise RuntimeError(f"Checkpoint tensor {name} has different global dtype or shape across ranks") - identity = (name, offsets, size) - if identity not in seen: - seen.add(identity) - geometries.setdefault(name, []).append((offsets, size)) - for name, shards in geometries.items(): - _reject_overlapping_geometries(name, shards) - - @contextmanager def recover_frozen_optimizer_ownership(model, checkpoint_path, *, module="", group=None): """Include source-recorded frozen parameters while building an optimizer. diff --git a/veomni/optim/lr_scheduler.py b/veomni/optim/lr_scheduler.py index fab230bce7..bd887abd0b 100644 --- a/veomni/optim/lr_scheduler.py +++ b/veomni/optim/lr_scheduler.py @@ -28,7 +28,6 @@ logger = logging.get_logger(__name__) _CHECKPOINT_SCHEDULER_BINDINGS = weakref.WeakKeyDictionary() -_CHECKPOINT_SCHEDULER_FUNCTIONS = weakref.WeakKeyDictionary() def checkpoint_scheduler_binding(optimizer): @@ -37,11 +36,6 @@ def checkpoint_scheduler_binding(optimizer): return None if reference is None else reference() -def checkpoint_scheduler_functions(scheduler): - """Identify the immutable function objects installed by the framework builder.""" - return _CHECKPOINT_SCHEDULER_FUNCTIONS.get(scheduler) - - class MultiLRScheduler(dict): """ Mapping of name -> LR scheduler with convenience methods to mirror a single scheduler. @@ -128,7 +122,6 @@ def build_lr_scheduler( else: raise ValueError(f"Unknown learning rate decay style: {lr_decay_style}.") _CHECKPOINT_SCHEDULER_BINDINGS[optimizer] = weakref.ref(scheduler) - _CHECKPOINT_SCHEDULER_FUNCTIONS[scheduler] = tuple(scheduler.lr_lambdas) return scheduler From 09a428a05fb2c27a6655b23ecf7dd1e68717924f Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 03:57:22 +0800 Subject: [PATCH 040/120] [checkpoint] fix: validate legacy expert sources --- docs/usage/checkpoint.md | 39 +++- scripts/merge_dcp_to_hf.py | 54 ++++- tests/checkpoints/test_expert_shard_layout.py | 220 +++++++++++++++++- veomni/checkpoint/checkpointer.py | 6 +- veomni/checkpoint/dcp_checkpointer.py | 70 ++++-- 5 files changed, 351 insertions(+), 38 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 6ea8f2e589..8a69094eff 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -265,12 +265,45 @@ non-persistent ExtraParallel shard in a checkpoint must use the canonical layout: mixing legacy and canonical shards is rejected. The v2 tensor namespace identifies canonical global values independently of source geometry, so DCP can reshard a model checkpoint across supported FSDP shard axes. Use -`DistributedCheckpointer` for save and resume: it rejects old/unversioned -layouts before copying model tensors. Migrating a legacy checkpoint requires -verified source layout information. Canonical v2 optimizer save and resume use +`DistributedCheckpointer` for save and resume: it rejects ambiguous unversioned +expert layouts before copying model tensors unless the model-only source is +explicitly declared as described below. Migrating a legacy optimizer checkpoint +requires verified source layout information. Canonical v2 optimizer save and resume use the rank-local codec described below. The standard VeOmni HF export strips the namespace. +### Known legacy VeOmni model-only sources + +An unversioned checkpoint does not record its producing EP/FSDP geometry. If +the caller has independently verified that **every expert tensor** came from +the legacy standard VeOmni TP1 implementation with expert axis 0 and FSDP +axis 1, they may explicitly declare `legacy_expert_source_layout="veomni_tp1_ep0_fsdp1"` +on `DistributedCheckpointer.load` or `dcp_to_torch_state_dict`. All ranks must +provide the same declaration. A legacy VeOmni TP1 source that used no expert +parallelism at all can instead be declared as `"veomni_tp1_no_ep"`. A missing +declaration keeps ambiguous sources +fail-closed. This is a trusted assertion about the source, not a conclusion drawn +from target geometry or the DCP chunks. In particular, do not declare this for +same-axis Muon/EP checkpoints or external TP-produced shards. An incorrect +assertion can silently permute experts. The declaration supports model-only +restore/export and is rejected when restoring an optimizer; it does not migrate +optimizer groups, moments, master weights or scheduler state. It does not load +directly into the canonical v2 namespace: export the verified legacy weights +and load them as model weights into the v2 target instead. +This also applies to old non-EP optimizer checkpoints: an unversioned, +FSDP-sharded expert bias or moment without source provenance can be rejected. +Neither source declaration authorizes optimizer-state migration. +When per-tensor source layout metadata is present, a recorded noncanonical +expert is always rejected. Missing metadata for one expert is not inferred +from the presence of metadata for another expert. + +For CLI export, pass `--legacy-expert-source-layout veomni_tp1_ep0_fsdp1` to +`scripts/merge_dcp_to_hf.py` only after verifying the source. Use +`--save-dtype preserve` when the checkpoint contains FP32 tensors that must not +be rounded to BF16; the CLI retains its historical BF16 export default. This +preserves tensor dtypes on disk; consumers must configure their loader to avoid +subsequently casting the FP32 tensors. + All ranks in a checkpoint operation must use the same logical model and checkpoint configuration, including the layout version and corresponding sharding specifications. Normal rank-local shard contents and mesh coordinates diff --git a/scripts/merge_dcp_to_hf.py b/scripts/merge_dcp_to_hf.py index e1595c2389..e6a80a8d09 100644 --- a/scripts/merge_dcp_to_hf.py +++ b/scripts/merge_dcp_to_hf.py @@ -46,6 +46,7 @@ def save_lora_adapter_weights( checkpoint_path: Union[str, os.PathLike], save_dtype: Optional[Union[str, torch.dtype]] = "bfloat16", adapter_config_path: Optional[Union[str, os.PathLike]] = None, + legacy_expert_source_layout: Optional[str] = None, ) -> None: """Convert a DCP checkpoint that contains a PEFT LoRA adapter to ``adapter_model.safetensors``. @@ -67,7 +68,9 @@ def save_lora_adapter_weights( ) logger.info(f"Found {len(lora_keys)} LoRA tensors; loading and re-saving as adapter_model.safetensors") - processed_dict = _process_shard(lora_keys, checkpoint_path, save_dtype) + processed_dict = _process_shard( + lora_keys, checkpoint_path, save_dtype, legacy_expert_source_layout=legacy_expert_source_layout + ) save_path = os.path.join(output_dir, "adapter_model.safetensors") save_file(processed_dict, save_path, metadata={"format": "pt"}) @@ -101,6 +104,7 @@ def save_model_weights( shard_size: int = 2_000_000_000, safe_serialization: bool = True, model_assets: Optional[Sequence["ModelAssets"]] = None, + legacy_expert_source_layout: Optional[str] = None, ) -> None: """Convert DCP checkpoint to HuggingFace format with shard-by-shard processing (memory-efficient).""" os.makedirs(output_dir, exist_ok=True) @@ -135,7 +139,9 @@ def save_model_weights( save_path = os.path.join(output_dir, filename) logger.info(f"Processing shard {shard_idx + 1}/{num_shards}: {filename} ({len(shard_keys)} tensors)") - processed_dict = _process_shard(shard_keys, checkpoint_path, save_dtype) + processed_dict = _process_shard( + shard_keys, checkpoint_path, save_dtype, legacy_expert_source_layout=legacy_expert_source_layout + ) # Save shard if safe_serialization: @@ -176,7 +182,12 @@ def save_model_weights( def merge_to_hf_pt( - load_dir: str, save_path: str, model_assets_dir: Optional[str] = None, shard_size: int = 2_000_000_000 + load_dir: str, + save_path: str, + model_assets_dir: Optional[str] = None, + shard_size: int = 2_000_000_000, + save_dtype: Optional[Union[str, torch.dtype]] = "bfloat16", + legacy_expert_source_layout: Optional[str] = None, ) -> None: """Main conversion function: load DCP from load_dir and save HF format to save_path.""" model_assets = None @@ -198,7 +209,14 @@ def merge_to_hf_pt( if not model_assets: model_assets = None - save_model_weights(save_path, load_dir, shard_size=shard_size, model_assets=model_assets) + save_model_weights( + save_path, + load_dir, + save_dtype=save_dtype, + shard_size=shard_size, + model_assets=model_assets, + legacy_expert_source_layout=legacy_expert_source_layout, + ) def main(): @@ -244,6 +262,22 @@ def main(): "'full' / 'lora' force the corresponding mode." ), ) + parser.add_argument( + "--legacy-expert-source-layout", + choices=("veomni_tp1_ep0_fsdp1", "veomni_tp1_no_ep"), + default=None, + help=( + "Trusted source declaration for legacy VeOmni TP1 checkpoints: expert-axis-0/FSDP-axis-1 " + "or no expert parallelism. " + "Only use after verifying the source layout; a wrong declaration can permute experts." + ), + ) + parser.add_argument( + "--save-dtype", + choices=("bfloat16", "float32", "preserve"), + default="bfloat16", + help="Export dtype; 'preserve' retains each checkpoint tensor's original dtype (default: bfloat16).", + ) parser.add_argument( "--adapter-config-path", type=str, @@ -256,6 +290,7 @@ def main(): ), ) args = parser.parse_args() + save_dtype = None if args.save_dtype == "preserve" else args.save_dtype # A step directory is not itself a DCP directory: it holds model/ckpt and # model/optimizer. Resolve to the weights, but keep the export next to the @@ -283,12 +318,21 @@ def main(): save_lora_adapter_weights( output_dir=save_dir, checkpoint_path=load_dir, + save_dtype=save_dtype, adapter_config_path=args.adapter_config_path, + legacy_expert_source_layout=args.legacy_expert_source_layout, ) else: if args.adapter_config_path is not None: logger.warning("--adapter-config-path is only used in 'lora' mode; ignoring.") - merge_to_hf_pt(load_dir, save_dir, model_assets_dir, shard_size=shard_size) + merge_to_hf_pt( + load_dir, + save_dir, + model_assets_dir, + shard_size=shard_size, + save_dtype=save_dtype, + legacy_expert_source_layout=args.legacy_expert_source_layout, + ) if __name__ == "__main__": diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 3e4ad86c80..af4e98f5a0 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -4,23 +4,26 @@ """CPU/Gloo regressions for canonical EP-first DCP expert coordinates.""" from datetime import timedelta +from types import SimpleNamespace import pytest import torch import torch.distributed as dist import torch.distributed.checkpoint as dcp import torch.multiprocessing as mp -from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh from torch.distributed.tensor import DTensor, Replicate, Shard from scripts.ci.accelerator.launch_plan import LaunchSpec, from_parameter from tests.tools.launch_utils import spawn_with_rendezvous from veomni.checkpoint.dcp_checkpointer import ( _EXPERT_LAYOUTS_KEY, + DistributedCheckpointer, _ExtraParallelSavePlanner, _get_sharding_plan, _ModelStrictLoadPlanner, _process_shard, + dcp_to_torch_state_dict, drop_extra_parallel_dim, restore_extra_parallel_dim, ) @@ -102,6 +105,91 @@ def wrap(dim, *, empty=False): with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=True)) assert target["model"]["weight"].to_local().count_nonzero() == 0 + dcp.load( + target, + checkpoint_id=path, + planner=_ModelStrictLoadPlanner(strict_model=True, legacy_expert_source_layout="veomni_tp1_ep0_fsdp1"), + ) + torch.testing.assert_close(target["model"]["weight"].full_tensor(), global_weight, rtol=0, atol=0) + + # The public model-only load/export paths accept the same independently + # verified declaration; unknown source and optimizer restore stay closed. + path = directory + "known_legacy_model_only" + bf16_bias = torch.tensor([0.25, -0.5], dtype=torch.bfloat16) + dcp.save({"model": {"layer.experts.weight": wrap(1), "layer.bias": bf16_bias}}, checkpoint_id=path) + model = torch.nn.Module() + model.layer = torch.nn.Module() + model.layer.experts = torch.nn.Module() + model.layer.experts.weight = torch.nn.Parameter(torch.zeros_like(global_weight)) + parallel = SimpleNamespace(dp_mode="ddp") + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + DistributedCheckpointer.load( + path, + {"model": model}, + parallel_state=parallel, + legacy_expert_source_layout="veomni_tp1_ep0_fsdp1" if rank == 0 else None, + ) + assert model.layer.experts.weight.count_nonzero() == 0 + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + DistributedCheckpointer.load(path, {"model": model}, parallel_state=parallel) + assert model.layer.experts.weight.count_nonzero() == 0 + with pytest.raises(ValueError, match="only supported for model-only loads"): + DistributedCheckpointer.load( + path, + {"model": model, "optimizer": torch.optim.AdamW(model.parameters())}, + parallel_state=parallel, + legacy_expert_source_layout="veomni_tp1_ep0_fsdp1", + ) + assert model.layer.experts.weight.count_nonzero() == 0 + DistributedCheckpointer.load( + path, + {"model": model}, + parallel_state=parallel, + legacy_expert_source_layout="veomni_tp1_ep0_fsdp1", + ) + torch.testing.assert_close(model.layer.experts.weight, global_weight, rtol=0, atol=0) + if rank == 0: + import subprocess + import sys + + from safetensors.torch import load_file + + from veomni.checkpoint import ckpt_to_state_dict + + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + dcp_to_torch_state_dict(path) + exported = dcp_to_torch_state_dict(path, legacy_expert_source_layout="veomni_tp1_ep0_fsdp1") + torch.testing.assert_close(exported["layer.experts.weight"], global_weight, rtol=0, atol=0) + exported = ckpt_to_state_dict(path, legacy_expert_source_layout="veomni_tp1_ep0_fsdp1") + torch.testing.assert_close(exported["layer.experts.weight"], global_weight, rtol=0, atol=0) + torch.testing.assert_close(exported["layer.bias"], bf16_bias, rtol=0, atol=0) + assert exported["layer.experts.weight"].dtype == torch.float32 + assert exported["layer.bias"].dtype == torch.bfloat16 + cli_dir = directory + "known_legacy_hf" + subprocess.run( + [ + sys.executable, + "scripts/merge_dcp_to_hf.py", + "--load-dir", + path, + "--save-dir", + cli_dir, + "--mode", + "full", + "--legacy-expert-source-layout", + "veomni_tp1_ep0_fsdp1", + "--save-dtype", + "preserve", + ], + check=True, + timeout=60, + ) + exported = load_file(cli_dir + "/model.safetensors") + torch.testing.assert_close(exported["layer.experts.weight"], global_weight, rtol=0, atol=0) + torch.testing.assert_close(exported["layer.bias"], bf16_bias, rtol=0, atol=0) + assert exported["layer.experts.weight"].dtype == torch.float32 + assert exported["layer.bias"].dtype == torch.bfloat16 + dist.barrier() # Legacy metadata used two ordinary Shards, recording the wrong order. local = ep_weight.chunk(2, 0)[fsdp_rank].contiguous() @@ -126,19 +214,33 @@ def wrap(dim, *, empty=False): _process_shard(plan, path) dist.barrier() + # A one-dimensional expert bias may also be EP-sharded. Exporting it + # as a plain tensor must not bypass the same unversioned-source guard. + bias = torch.arange(8, dtype=torch.float32) + local_bias = bias.chunk(2, 0)[ep_rank].chunk(2, 0)[fsdp_rank].contiguous() + placements = [Shard(0), Shard(0)] if world == 4 else [Replicate(), Shard(0), Shard(0)] + legacy_bias = DTensor.from_local(local_bias, mesh, placements) + path = directory + "legacy_ambiguous_bias_export" + dcp.save({"model": {"layer.experts.bias": legacy_bias}}, checkpoint_id=path) + if rank == 0: + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + dcp_to_torch_state_dict(path) + dist.barrier() + # When the source layout is known, the guard also works for a plain # tensor target. Never bless legacy ordinary Shard(0)/Shard(0) saves. path = directory + "marked_legacy" dcp.save({"model": {"weight": legacy}}, checkpoint_id=path, planner=_ExtraParallelSavePlanner()) for destination in (wrap(0, empty=True), wrap(1, empty=True), torch.zeros_like(global_weight)): - with pytest.raises(dcp.CheckpointException, match="noncanonical expert order"): - dcp.load( - {"model": {"weight": destination}}, - checkpoint_id=path, - planner=_ModelStrictLoadPlanner(strict_model=True), - ) - local = destination.to_local() if isinstance(destination, DTensor) else destination - assert local.count_nonzero() == 0 + for declaration in (None, "veomni_tp1_ep0_fsdp1", "veomni_tp1_no_ep"): + with pytest.raises(dcp.CheckpointException, match="noncanonical expert order"): + dcp.load( + {"model": {"weight": destination}}, + checkpoint_id=path, + planner=_ModelStrictLoadPlanner(strict_model=True, legacy_expert_source_layout=declaration), + ) + local = destination.to_local() if isinstance(destination, DTensor) else destination + assert local.count_nonzero() == 0 finally: dist.destroy_process_group() @@ -151,6 +253,106 @@ def test_expert_shard_layout_roundtrip_and_legacy_rejection(world, tmp_path): ) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +def test_unsupported_legacy_expert_source_layout_is_rejected(): + with pytest.raises(ValueError, match="Unsupported legacy expert source layout"): + _ModelStrictLoadPlanner(strict_model=True, legacy_expert_source_layout="unverified") + + +def _non_ep_legacy_expert_bias_worker(rank, root, endpoint): + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + expected = torch.arange(8, dtype=torch.float32) + mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("dp_shard",)) + source = DTensor.from_local(expected.chunk(2)[rank].contiguous(), mesh, [Shard(0)]) + dcp.save({"model": {"layer.experts.bias": source}}, checkpoint_id=root, planner=_ExtraParallelSavePlanner()) + target = {"model": {"layer.experts.bias": torch.zeros_like(expected)}} + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + dcp.load(target, checkpoint_id=root, planner=_ModelStrictLoadPlanner(strict_model=True)) + assert target["model"]["layer.experts.bias"].count_nonzero() == 0 + dcp.load( + target, + checkpoint_id=root, + planner=_ModelStrictLoadPlanner(strict_model=True, legacy_expert_source_layout="veomni_tp1_no_ep"), + ) + torch.testing.assert_close(target["model"]["layer.experts.bias"], expected, rtol=0, atol=0) + if rank == 0: + exported = dcp_to_torch_state_dict(root, legacy_expert_source_layout="veomni_tp1_no_ep") + torch.testing.assert_close(exported["layer.experts.bias"], expected, rtol=0, atol=0) + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_legacy_non_ep_fsdp_expert_bias_requires_source_declaration(tmp_path): + spawn_with_rendezvous(_non_ep_legacy_expert_bias_worker, args=(str(tmp_path / "dcp"),), nprocs=2) + + +def _unversioned_wrong_target_expert_axis_worker(rank, root, endpoint): + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=4, timeout=timedelta(seconds=60)) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ep_fsdp", "ep")) + fsdp, ep = mesh.get_coordinate() + expected = torch.arange(64, dtype=torch.float32).reshape(8, 8) + local = expected.chunk(2, 0)[ep].chunk(2, 0)[fsdp].contiguous() + source = DTensor.from_local(local, mesh, [Shard(0), Shard(0)]) + dcp.save({"model": {"layer.experts.weight": source}}, checkpoint_id=root) + target = DTensor.from_local(torch.zeros(4, 4), mesh, [Shard(0), Shard(1)]) + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + dcp.load( + {"model": {"layer.experts.weight": target}}, + checkpoint_id=root, + planner=_ModelStrictLoadPlanner(strict_model=True), + ) + assert target.to_local().count_nonzero() == 0 + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +def test_unversioned_source_axis_is_not_inferred_from_target(tmp_path): + spawn_with_rendezvous(_unversioned_wrong_target_expert_axis_worker, args=(str(tmp_path / "dcp"),), nprocs=4) + + +def _partial_expert_layout_metadata_worker(rank, root, endpoint): + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=4, timeout=timedelta(seconds=60)) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ep_fsdp", "ep")) + unnamed_mesh = DeviceMesh("cpu", mesh.mesh, mesh_dim_names=("axis0", "axis1")) + fsdp, ep = mesh.get_coordinate() + expected = torch.arange(64, dtype=torch.float32).reshape(8, 8) + good_local = expected.chunk(2, 0)[ep].chunk(2, 1)[fsdp].contiguous() + bad_local = expected.chunk(2, 0)[ep].chunk(2, 0)[fsdp].contiguous() + good = DTensor.from_local(good_local, mesh, [Shard(1), Shard(0)]) + bad = DTensor.from_local(bad_local, unnamed_mesh, [Shard(0), Shard(0)]) + source = {"model": {"layer.experts.good": good, "layer.experts.bad": bad}} + dcp.save(source, checkpoint_id=root, planner=_ExtraParallelSavePlanner()) + recorded = dcp.FileSystemReader(root).read_metadata().planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + assert "model.layer.experts.good" in recorded and "model.layer.experts.bad" not in recorded + target = { + "model": { + f"layer.experts.{name}": DTensor.from_local(torch.zeros(4, 4), mesh, [Shard(1), Shard(0)]) + for name in ("good", "bad") + } + } + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + dcp.load(target, checkpoint_id=root, planner=_ModelStrictLoadPlanner(strict_model=True)) + assert all(value.to_local().count_nonzero() == 0 for value in target["model"].values()) + if rank == 0: + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + dcp_to_torch_state_dict(root) + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +def test_partial_source_layout_metadata_does_not_authorize_other_experts(tmp_path): + spawn_with_rendezvous(_partial_expert_layout_metadata_worker, args=(str(tmp_path / "dcp"),), nprocs=4) + + def _unversioned_composite_expert_worker(rank, root, endpoint): dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=8) try: diff --git a/veomni/checkpoint/checkpointer.py b/veomni/checkpoint/checkpointer.py index ec08272381..928385afc8 100644 --- a/veomni/checkpoint/checkpointer.py +++ b/veomni/checkpoint/checkpointer.py @@ -118,7 +118,9 @@ def dcp_checkpointer(dist_backend: str): @CHECKPOINT_TO_STATE_DICT_REGISTRY.register("dcp") -def dcp_ckpt_to_state_dict(save_checkpoint_path: Union[str, os.PathLike], **kwargs): +def dcp_ckpt_to_state_dict( + save_checkpoint_path: Union[str, os.PathLike], *, legacy_expert_source_layout=None, **kwargs +): from .dcp_checkpointer import dcp_to_torch_state_dict - return dcp_to_torch_state_dict(save_checkpoint_path) + return dcp_to_torch_state_dict(save_checkpoint_path, legacy_expert_source_layout=legacy_expert_source_layout) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 9e03ef9fc5..45a627b679 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -85,6 +85,13 @@ _EXPERT_LAYOUTS_KEY = "__veomni_expert_layouts__" +_LEGACY_VEOMNI_TP1_EP0_FSDP1 = "veomni_tp1_ep0_fsdp1" +_LEGACY_VEOMNI_TP1_NO_EP = "veomni_tp1_no_ep" + + +def _validate_legacy_expert_source_layout(value): + if value is not None and value not in (_LEGACY_VEOMNI_TP1_EP0_FSDP1, _LEGACY_VEOMNI_TP1_NO_EP): + raise ValueError(f"Unsupported legacy expert source layout: {value!r}") def _expert_layout(tensor): @@ -140,10 +147,12 @@ def create_global_plan(self, all_plans): class _ModelStrictLoadPlanner(DefaultLoadPlanner): """Allow partial optimizer state while requiring a complete full-model DCP.""" - def __init__(self, strict_model: bool, expected_save_id=None): + def __init__(self, strict_model: bool, expected_save_id=None, legacy_expert_source_layout=None): super().__init__(allow_partial_load=True) + _validate_legacy_expert_source_layout(legacy_expert_source_layout) self.strict_model = strict_model self.expected_save_id = expected_save_id + self.legacy_expert_source_layout = legacy_expert_source_layout def create_local_plan(self): assert self.metadata is not None @@ -204,16 +213,17 @@ def create_local_plan(self): def _validate_expert_layouts(self): planner_data = self.metadata.planner_data or {} source = planner_data.get(_EXPERT_LAYOUTS_KEY) + source_tensors = {} if source is not None: if source["version"] != 1: raise RuntimeError(f"Unsupported checkpoint expert layout version: {source['version']}") - for key in self.state_dict.keys() & source["tensors"].keys(): - if source["tensors"][key]["expert_order"] != "canonical": + source_tensors = source["tensors"] + for key in self.state_dict.keys() & source_tensors.keys(): + if source_tensors[key]["expert_order"] != "canonical": raise RuntimeError( - f"Checkpoint tensor {key!r} has noncanonical expert order: {source['tensors'][key]}. " + f"Checkpoint tensor {key!r} has noncanonical expert order: {source_tensors[key]}. " "Migrate model and optimizer state using the original training layout before loading." ) - return # Checkpoints produced by the earlier version of this feature already # used canonical coordinates, but did not record individual layouts. @@ -227,22 +237,25 @@ def _validate_expert_layouts(self): return for key, value in self.state_dict.items(): + if key in source_tensors: + continue layout = _expert_layout(value) saved = self.metadata.state_dict_metadata.get(key) - if saved is None: + if saved is None or not torch.is_tensor(value) or not hasattr(saved, "size"): continue - # Unversioned chunks cannot distinguish a safe EP+row-FSDP layout - # from the same-axis legacy bug combined with another parallel - # axis. Only a source chunk spanning all experts is unambiguous. + # Unversioned chunks cannot identify the source EP axis. The + # target axis is not source evidence: a source sharded along one + # axis can appear whole along a different target expert axis. if layout is None: # Offline export uses a plain CPU tensor as its target. Old # stacked expert weights still require the same source check. - if "experts" not in key.split(".") or len(saved.size) < 2: + if "experts" not in key.split("."): continue - ep_dim = 0 - else: - ep_dim = layout["ep_shard_dim"] - if all(chunk.sizes[ep_dim] == saved.size[ep_dim] for chunk in saved.chunks): + if all(tuple(chunk.sizes) == tuple(saved.size) for chunk in saved.chunks): + continue + # This declaration is a trusted statement about the source, not + # an inference from target topology or ambiguous chunk geometry. + if self.legacy_expert_source_layout is not None: continue raise RuntimeError( f"Cannot determine source expert order for {key!r} in this unversioned checkpoint. " @@ -1993,6 +2006,7 @@ def load( module: str = "", trainable_only: bool = False, parallel_state=None, + legacy_expert_source_layout: Optional[str] = None, ) -> Dict[str, Any]: """ load training state from distributed checkpoint @@ -2009,6 +2023,9 @@ def load( 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`. + legacy_expert_source_layout: Optional, trusted declaration of an + unversioned model-only checkpoint's source layout. It must be + identical on all ranks and cannot be used for optimizer restore. trainable_only: when True, ``set_model_state_dict`` runs in non-strict mode (``StateDictOptions(strict=False)``). Use this for LoRA / PEFT resumes where the DCP only contains trainable adapter weights and the @@ -2049,6 +2066,9 @@ def _optimizer_state() -> "OptimizerState": fused_dir, error = None, None try: + _validate_legacy_expert_source_layout(legacy_expert_source_layout) + if legacy_expert_source_layout is not None and optimizer is not None: + raise ValueError("Legacy expert source declarations are only supported for model-only loads") fused_dir = cls._legacy_fused_dir(path, module) except BaseException as exc: error = exc @@ -2205,7 +2225,9 @@ def _optimizer_state() -> "OptimizerState": state_dict=load_state, storage_reader=cls._create_storage_reader(fused_dir), process_group=process_group, - planner=_ModelStrictLoadPlanner(strict_model=not trainable_only), + planner=_ModelStrictLoadPlanner( + strict_model=not trainable_only, legacy_expert_source_layout=legacy_expert_source_layout + ), ) cls._load_extra_state(checkpoint_dir=fused_dir, state=state) logger.info_rank0(f"Loaded pre-split checkpoint from {fused_dir}") @@ -2216,7 +2238,9 @@ def _optimizer_state() -> "OptimizerState": state_dict={"model": _model_state()}, storage_reader=cls._create_storage_reader(weights_dir(path, module)), process_group=process_group, - planner=_ModelStrictLoadPlanner(strict_model=not trainable_only), + planner=_ModelStrictLoadPlanner( + strict_model=not trainable_only, legacy_expert_source_layout=legacy_expert_source_layout + ), ) if wants_optimizer: dcp.load( @@ -2597,7 +2621,9 @@ def _process_shard( shard_keys: Dict[str, str], checkpoint_path: str, save_dtype: Optional[Union[str, torch.dtype]] = None, + legacy_expert_source_layout: Optional[str] = None, ) -> str: + _validate_legacy_expert_source_layout(legacy_expert_source_layout) reader = FileSystemReader(checkpoint_path) metadata = reader.read_metadata() @@ -2627,7 +2653,7 @@ def _process_shard( state_dict, checkpoint_id=checkpoint_path, storage_reader=FileSystemReader(checkpoint_path), - planner=_ModelStrictLoadPlanner(strict_model=True), + planner=_ModelStrictLoadPlanner(strict_model=True, legacy_expert_source_layout=legacy_expert_source_layout), no_dist=True, ) @@ -2660,19 +2686,25 @@ def _process_shard( return processed_dict -def dcp_to_torch_state_dict(save_checkpoint_path: Union[str, os.PathLike]) -> STATE_DICT_TYPE: +def dcp_to_torch_state_dict( + save_checkpoint_path: Union[str, os.PathLike], *, legacy_expert_source_layout: Optional[str] = None +) -> STATE_DICT_TYPE: """ Given a directory containing a DCP checkpoint, this function will convert it into a Torch state_dict. Args: save_checkpoint_path: Directory containing the DCP checkpoint. + legacy_expert_source_layout: Trusted declaration of the source's legacy + expert layout. See ``docs/usage/checkpoint.md`` for its contract. .. warning:: To avoid OOM, it's recommended to only run this function on a single rank. """ shard, _, _ = _get_sharding_plan(save_checkpoint_path) - processed_dict = _process_shard(shard, save_checkpoint_path) + processed_dict = _process_shard( + shard, save_checkpoint_path, legacy_expert_source_layout=legacy_expert_source_layout + ) return processed_dict From 02f0753ac6bb43b7ea9504cf4128cdf74ca9aa64 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 04:03:13 +0800 Subject: [PATCH 041/120] [checkpoint] fix: keep export options keyword-only --- scripts/merge_dcp_to_hf.py | 3 +++ veomni/checkpoint/dcp_checkpointer.py | 1 + 2 files changed, 4 insertions(+) diff --git a/scripts/merge_dcp_to_hf.py b/scripts/merge_dcp_to_hf.py index e6a80a8d09..f7408caa62 100644 --- a/scripts/merge_dcp_to_hf.py +++ b/scripts/merge_dcp_to_hf.py @@ -46,6 +46,7 @@ def save_lora_adapter_weights( checkpoint_path: Union[str, os.PathLike], save_dtype: Optional[Union[str, torch.dtype]] = "bfloat16", adapter_config_path: Optional[Union[str, os.PathLike]] = None, + *, legacy_expert_source_layout: Optional[str] = None, ) -> None: """Convert a DCP checkpoint that contains a PEFT LoRA adapter to ``adapter_model.safetensors``. @@ -104,6 +105,7 @@ def save_model_weights( shard_size: int = 2_000_000_000, safe_serialization: bool = True, model_assets: Optional[Sequence["ModelAssets"]] = None, + *, legacy_expert_source_layout: Optional[str] = None, ) -> None: """Convert DCP checkpoint to HuggingFace format with shard-by-shard processing (memory-efficient).""" @@ -186,6 +188,7 @@ def merge_to_hf_pt( save_path: str, model_assets_dir: Optional[str] = None, shard_size: int = 2_000_000_000, + *, save_dtype: Optional[Union[str, torch.dtype]] = "bfloat16", legacy_expert_source_layout: Optional[str] = None, ) -> None: diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 45a627b679..437e0376ee 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -2621,6 +2621,7 @@ def _process_shard( shard_keys: Dict[str, str], checkpoint_path: str, save_dtype: Optional[Union[str, torch.dtype]] = None, + *, legacy_expert_source_layout: Optional[str] = None, ) -> str: _validate_legacy_expert_source_layout(legacy_expert_source_layout) From f73d4161c486b85d316990762325f0cd03710c56 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 04:29:50 +0800 Subject: [PATCH 042/120] [checkpoint] allow trusted scheduler instrumentation --- docs/usage/checkpoint.md | 11 +++++-- .../test_multi_optimizer_trainer_step.py | 29 +++++++++++++++---- veomni/models/checkpoint/manager.py | 2 -- 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 8a69094eff..136a0e9500 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -374,8 +374,15 @@ as part of the existing checkpoint-layout exchange; owner-specific freezing that removes a whole leaf from only some ranks is rejected before per-leaf checkpoint collectives, rather than supported. Replica values, trainability, dtype, shape and geometry are trusted to be consistent across ranks. Source and target checkpoint layout -metadata are checked during restore; no tensor bytes are hashed. Hooks and -adapters are trusted. MultiOptimizer asynchronous saves +metadata are checked during restore; tensor bytes are not hashed for replica +consistency. Trainable-only MultiOptimizer checkpoints do hash each rank's local +frozen base parameters on save and load to verify the base identity. That may copy +device tensors to CPU and can add substantial time for a large frozen model. +Hooks, adapters and scheduler method wrappers are trusted: a wrapper must +preserve serializable scheduler state and implement the same learning-rate +behavior on both sides of a restore. The persisted recipe checks scheduler +configuration and optimizer binding; it does not inspect arbitrary wrapper +code. MultiOptimizer asynchronous saves and unsupported layouts fail closed; synchronous staged saves are supported. The runtime assumes ranks use a consistent model configuration and compatible replica geometries; saved source layout is validated on restore. diff --git a/tests/checkpoints/test_multi_optimizer_trainer_step.py b/tests/checkpoints/test_multi_optimizer_trainer_step.py index 7eb3cf0e74..597d7b0287 100644 --- a/tests/checkpoints/test_multi_optimizer_trainer_step.py +++ b/tests/checkpoints/test_multi_optimizer_trainer_step.py @@ -99,7 +99,8 @@ def test_trainer_state_cannot_silently_diverge_from_restored_optimizer(tmp_path, @pytest.mark.parametrize("mode", ["missing", "stale_generation", "matching"]) -def test_real_dcp_with_scheduler_requires_matching_trainer_generation(tmp_path, mode): +@pytest.mark.parametrize("instrumented", ["none", "source", "target", "both"]) +def test_real_dcp_with_scheduler_requires_matching_trainer_generation(tmp_path, mode, instrumented): import torch from tests.checkpoints.test_checkpoint_callback import _make_mock_runtime @@ -110,12 +111,22 @@ def test_real_dcp_with_scheduler_requires_matching_trainer_generation(tmp_path, from veomni.optim.optimizer import MultiOptimizer from veomni.trainer.callbacks.base import TrainerState - def build(): + def build(with_instrumentation): runtime = _make_mock_runtime(save_path=str(tmp_path)) runtime.model = torch.nn.Linear(2, 2) leaf = torch.optim.AdamW(runtime.model.parameters(), lr=0.01) runtime.optimizer = MultiOptimizer(runtime.model, {"model": leaf}, ["model"]) runtime.lr_scheduler = build_lr_scheduler(runtime.optimizer, train_steps=100, lr=0.01) + if with_instrumentation: + scheduler = runtime.lr_scheduler["model"] + original_state_dict = scheduler.state_dict + + def observed_state_dict(): + # A profiling wrapper delegates to the scheduler and keeps the + # callable itself out of the serializable scheduler snapshot. + return {key: value for key, value in original_state_dict().items() if key != "state_dict"} + + scheduler.state_dict = observed_state_dict runtime.parallel_state = SimpleNamespace(global_rank=0, dp_mode="ddp", extra_parallel_names=[]) return runtime @@ -123,7 +134,7 @@ def build(): patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer), patch("torch.distributed.barrier"), ): - source = build() + source = build(instrumented in ("source", "both")) manager = ModelCheckpointManager(source) for _ in range(3): source.model(torch.ones(1, 2)).sum().backward() @@ -147,7 +158,7 @@ def build(): os.makedirs(os.path.dirname(sidecar), exist_ok=True) torch.save({"global_step": 3, "_veomni_optimizer_generations": {"model": "stale"}}, sidecar) - target = build() + target = build(instrumented in ("target", "both")) target.train_args.checkpoint.load_path = str(tmp_path / "global_step_3") target.checkpoint = ModelCheckpointManager(target) target.checkpoint.load() @@ -161,11 +172,19 @@ def build(): if mode == "matching": GlobalStateCallback(trainer).load_global_state() assert trainer.state.global_step == 3 + for runtime in (source, target): + runtime.optimizer.zero_grad() + runtime.model(torch.ones(1, 2)).sum().backward() + runtime.optimizer.step() + runtime.lr_scheduler.step() + for actual, expected in zip(target.model.parameters(), source.model.parameters(), strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert target.lr_scheduler["model"].get_last_lr() == source.lr_scheduler["model"].get_last_lr() else: with pytest.raises(RuntimeError, match="Trainer state is missing|generations differ"): GlobalStateCallback(trainer).load_global_state() assert trainer.state.global_step == 0 - assert target.lr_scheduler["model"].last_epoch == 3 + assert target.lr_scheduler["model"].last_epoch == (4 if mode == "matching" else 3) def test_train_end_sidecar_keeps_generation_after_optimizer_release(tmp_path): diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 9938f2d576..9695fee77a 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -59,8 +59,6 @@ def _scheduler_recipe(scheduler, optimizer): raise RuntimeError("MultiOptimizer requires a scheduler for each optimizer leaf") if type(scheduler) not in (LambdaLR, StepLR) or scheduler.optimizer is not optimizer: raise RuntimeError("MultiOptimizer scheduler implementation or binding is unsupported") - if any(name in vars(scheduler) for name in ("state_dict", "load_state_dict", "step", "get_lr", "get_last_lr")): - raise RuntimeError("MultiOptimizer scheduler methods are overridden") if type(scheduler) is StepLR: return {"type": "StepLR", "step_size": scheduler.step_size, "gamma": scheduler.gamma} From 478ae4c3f34dd5e44db39b871bf84e1218fd1e29 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 10:33:43 +0800 Subject: [PATCH 043/120] [checkpoint] prove EP1 source and simplify layout validation --- docs/usage/checkpoint.md | 9 +- .../test_expert_noep_checkpoint.py | 206 +++++++++++++++ tests/checkpoints/test_expert_shard_layout.py | 8 - .../test_multi_optimizer_proven_resume.py | 93 +++++++ veomni/checkpoint/dcp_checkpointer.py | 238 ++++++++---------- .../checkpoint/multi_optimizer_checkpoint.py | 5 - veomni/checkpoint/optimizer_resume_layout.py | 52 +++- 7 files changed, 463 insertions(+), 148 deletions(-) create mode 100644 tests/checkpoints/test_expert_noep_checkpoint.py diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 136a0e9500..59948e5aff 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -272,6 +272,12 @@ requires verified source layout information. Canonical v2 optimizer save and res the rank-local codec described below. The standard VeOmni HF export strips the namespace. +New TP1/EP1 checkpoints also record no-EP provenance for expert tensors on the +standard VeOmni FSDP mesh. Model and optimizer restoration can use this source +record without a legacy declaration. The record is not inferred for custom +meshes, tensor parallelism, active EP, or older checkpoints missing provenance; +those cases retain the existing source-layout checks. + ### Known legacy VeOmni model-only sources An unversioned checkpoint does not record its producing EP/FSDP geometry. If @@ -380,7 +386,8 @@ frozen base parameters on save and load to verify the base identity. That may co device tensors to CPU and can add substantial time for a large frozen model. Hooks, adapters and scheduler method wrappers are trusted: a wrapper must preserve serializable scheduler state and implement the same learning-rate -behavior on both sides of a restore. The persisted recipe checks scheduler +behavior on both sides of a restore, and hooks must not mutate checkpoint +model, optimizer, or layout during a save/load call. The persisted recipe checks scheduler configuration and optimizer binding; it does not inspect arbitrary wrapper code. MultiOptimizer asynchronous saves and unsupported layouts fail closed; synchronous staged saves are supported. diff --git a/tests/checkpoints/test_expert_noep_checkpoint.py b/tests/checkpoints/test_expert_noep_checkpoint.py new file mode 100644 index 0000000000..dde0285150 --- /dev/null +++ b/tests/checkpoints/test_expert_noep_checkpoint.py @@ -0,0 +1,206 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 + +"""New EP1 checkpoints preserve the source layout across a fresh-process resume.""" + +from datetime import timedelta +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from tests.tools.launch_utils import spawn_with_rendezvous +from veomni.checkpoint.dcp_checkpointer import ( + _EXPERT_LAYOUTS_KEY, + DistributedCheckpointer, + _ExtraParallelSavePlanner, + _ModelStrictLoadPlanner, + _no_ep_expert_layout, + dcp_to_torch_state_dict, +) +from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state +from veomni.optim.optimizer import MultiOptimizer + + +def _build_noep_model(parallel_state, optimizer_kind, real_fsdp=False): + model = torch.nn.Module() + model.layer = torch.nn.Module() + model.layer.experts = torch.nn.Module() + mesh = parallel_state.fsdp_mesh + placements = [Replicate()] * (mesh.ndim - 1) + [Shard(0)] + for name, shape in (("bias", (8,)), ("matrix", (8, 4)), ("stack", (8, 4, 2))): + value = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.float32).reshape(shape) / 100 + if not real_fsdp: + value = distribute_tensor(value, mesh, placements) + setattr(model.layer.experts, name, torch.nn.Parameter(value)) + if real_fsdp: + from torch.distributed.fsdp import fully_shard + + fully_shard(model, mesh=mesh) + leaf = torch.optim.AdamW(model.parameters(), lr=0.01, betas=(0.9, 0.95), foreach=False) + optimizer = MultiOptimizer(model, {"expert": leaf}, ["expert"]) if optimizer_kind == "multi" else leaf + return model, optimizer, leaf + + +def _local(tensor): + if isinstance(tensor, DTensor): + tensor = tensor.to_local() + return tensor.detach().cpu().clone() if torch.is_tensor(tensor) else tensor + + +def _snapshot(model, leaf): + result = {} + for name, parameter in model.named_parameters(): + result[name] = { + "weight": _local(parameter), + "state": {key: _local(value) for key, value in leaf.state[parameter].items()}, + } + return result + + +def _step(model, optimizer, value): + for parameter in model.parameters(): + parameter.grad = torch.full_like(parameter, value) + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + +def _noep_worker(rank, directory, geometry, optimizer_kind, phase, endpoint): + world = 4 if geometry == "hsdp" else 2 + dist.init_process_group( + "gloo", store=endpoint.connect(), rank=rank, world_size=world, timeout=timedelta(seconds=60) + ) + try: + geometry_args = { + "fsdp": dict(dp_size=2, dp_shard_size=2), + "fsdp2": dict(dp_size=2, dp_shard_size=2), + "sp": dict(dp_size=1, dp_shard_size=1, ulysses_size=2), + "hsdp": dict(dp_size=4, dp_replicate_size=2, dp_shard_size=2), + }[geometry] + parallel_state = _init_parallel_state(device_type="cpu", **geometry_args) + model, optimizer, leaf = _build_noep_model(parallel_state, optimizer_kind, real_fsdp=geometry == "fsdp2") + checkpoint = str(Path(directory) / "checkpoint") + state = {"model": model, "optimizer": optimizer} + if optimizer_kind == "multi": + state["_veomni_scheduler_object"] = None + if phase == "save": + _step(model, optimizer, 0.25) + DistributedCheckpointer.save(checkpoint, state, parallel_state=parallel_state) + reference_before = _snapshot(model, leaf) + full_model = {name: parameter.detach().full_tensor().cpu() for name, parameter in model.named_parameters()} + _step(model, optimizer, -0.125) + torch.save( + {"before": reference_before, "after": _snapshot(model, leaf)}, + Path(directory) / f"reference-{rank}.pt", + ) + if rank == 0: + torch.save(full_model, Path(directory) / "full-model.pt") + else: + reference = torch.load(Path(directory) / f"reference-{rank}.pt", weights_only=True) + DistributedCheckpointer.load(checkpoint, state, parallel_state=parallel_state) + torch.testing.assert_close(_snapshot(model, leaf), reference["before"], rtol=0, atol=0) + _step(model, optimizer, -0.125) + torch.testing.assert_close(_snapshot(model, leaf), reference["after"], rtol=0, atol=0) + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.parametrize("geometry", ["fsdp", "fsdp2", "sp", "hsdp"]) +@pytest.mark.parametrize("optimizer_kind", ["adamw", "multi"]) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +def test_new_noep_expert_checkpoint_fresh_process_roundtrip(tmp_path, geometry, optimizer_kind): + world = 4 if geometry == "hsdp" else 2 + spawn_with_rendezvous(_noep_worker, args=(str(tmp_path), geometry, optimizer_kind, "save"), nprocs=world) + model_dir = tmp_path / "checkpoint" / "model" / "ckpt" + optimizer_dir = tmp_path / "checkpoint" / "model" / "optimizer" + for root, prefix in ((model_dir, "model.layer.experts."), (optimizer_dir, "optimizer.state.layer.experts.")): + layouts = dcp.FileSystemReader(root).read_metadata().planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + for name in ("bias", "matrix", "stack"): + key = prefix + name + (".exp_avg" if root == optimizer_dir else "") + assert layouts[key]["source_kind"] == "veomni_tp1_no_ep" + assert layouts[key]["expert_order"] == "canonical" + exported = dcp_to_torch_state_dict(model_dir) + expected = torch.load(tmp_path / "full-model.pt", weights_only=True) + torch.testing.assert_close(exported, expected, rtol=0, atol=0) + spawn_with_rendezvous(_noep_worker, args=(str(tmp_path), geometry, optimizer_kind, "resume"), nprocs=world) + + +def _unproven_source_worker(rank, directory, mode, endpoint): + world = 4 if mode == "tp" else 2 + dist.init_process_group( + "gloo", store=endpoint.connect(), rank=rank, world_size=world, timeout=timedelta(seconds=60) + ) + try: + if mode == "ep": + parallel_state = _init_parallel_state( + dp_size=2, dp_shard_size=2, extra_parallel_sizes=(2,), device_type="cpu" + ) + elif mode == "tp": + parallel_state = _init_parallel_state(dp_size=2, dp_shard_size=2, tp_size=2, device_type="cpu") + else: + parallel_state = _init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cpu") + mesh = parallel_state.fsdp_mesh + if mode == "custom": + mesh = DeviceMesh("cpu", list(range(world)), mesh_dim_names=("custom_axis",)) + value = distribute_tensor(torch.arange(32, dtype=torch.float32).reshape(8, 4), mesh, [Shard(0)]) + key = "model.layer.experts.weight" + assert _no_ep_expert_layout(key, value, parallel_state) is None + path = str(Path(directory) / "unproven") + dcp.save( + {"model": {"layer.experts.weight": value}}, + checkpoint_id=path, + planner=_ExtraParallelSavePlanner(parallel_state=parallel_state), + ) + layouts = dcp.FileSystemReader(path).read_metadata().planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + assert key not in layouts + target = {"model": {"layer.experts.weight": torch.zeros(8, 4)}} + with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): + dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=True)) + assert target["model"]["layer.experts.weight"].count_nonzero() == 0 + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.parametrize("mode", ["ep", "tp", "custom"]) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +def test_ep_tp_and_custom_mesh_are_not_inferred_as_noep(tmp_path, mode): + spawn_with_rendezvous(_unproven_source_worker, args=(str(tmp_path), mode), nprocs=4 if mode == "tp" else 2) + + +def _uninitialized_parallel_state_worker(rank, directory, endpoint): + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=1, timeout=timedelta(seconds=30)) + try: + mesh = DeviceMesh("cpu", [0], mesh_dim_names=("dp_shard",)) + + def build(): + model = torch.nn.Module() + model.layer = torch.nn.Module() + model.layer.experts = torch.nn.Module() + model.layer.experts.bias = torch.nn.Parameter( + distribute_tensor(torch.arange(8, dtype=torch.float32), mesh, [Shard(0)]) + ) + return model + + model = build() + checkpoint = str(Path(directory) / "uninitialized") + DistributedCheckpointer.save(checkpoint, {"model": model}) + target = build() + DistributedCheckpointer.load(checkpoint, {"model": target}) + torch.testing.assert_close( + target.layer.experts.bias.to_local(), model.layer.experts.bias.to_local(), rtol=0, atol=0 + ) + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=1)) +def test_uninitialized_parallel_state_preserves_single_rank_save(tmp_path): + spawn_with_rendezvous(_uninitialized_parallel_state_worker, args=(str(tmp_path),), nprocs=1) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index af4e98f5a0..ea0408415c 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -122,14 +122,6 @@ def wrap(dim, *, empty=False): model.layer.experts = torch.nn.Module() model.layer.experts.weight = torch.nn.Parameter(torch.zeros_like(global_weight)) parallel = SimpleNamespace(dp_mode="ddp") - with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): - DistributedCheckpointer.load( - path, - {"model": model}, - parallel_state=parallel, - legacy_expert_source_layout="veomni_tp1_ep0_fsdp1" if rank == 0 else None, - ) - assert model.layer.experts.weight.count_nonzero() == 0 with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"): DistributedCheckpointer.load(path, {"model": model}, parallel_state=parallel) assert model.layer.experts.weight.count_nonzero() == 0 diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 9d3f56bdb6..29cc54a9ae 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -28,6 +28,76 @@ def _model_and_optimizer(*, split=True): return model, MultiOptimizer(model, {"merged": merged}, ["merged"]) +def test_captured_layout_binds_sidecar_without_rechecking_runtime(): + from veomni.checkpoint.optimizer_resume_layout import ( + capture_multi_layout, + coordinated_multi_layout, + finish_multi_layout, + ) + + model, optimizer = _model_and_optimizer() + captured = capture_multi_layout(model, optimizer, None, None) + expected = coordinated_multi_layout(model, optimizer, None, None, "sidecar-digest", writing=True) + assert finish_multi_layout(captured, "sidecar-digest") == expected + assert finish_multi_layout(captured, "other-sidecar-digest") != expected + + +def _corrupt_inventory_worker(rank, endpoint): + import hashlib + import json + from datetime import timedelta + + import torch.distributed as dist + + from veomni.checkpoint.dcp_checkpointer import _ProvenMultiOptimizerState + from veomni.checkpoint.multi_optimizer_checkpoint import _MULTI_INVENTORY_MARKER + from veomni.checkpoint.optimizer_resume_layout import MULTI_LAYOUT_MARKER + + class OptimizerProbe: + def __init__(self): + self.entered = False + + def _dcp_load_state_dict(self, state_dict, *, group=None): + self.entered = True + dist.all_reduce(torch.ones(1), group=group) + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + optimizer = OptimizerProbe() + state = _ProvenMultiOptimizerState( + model=torch.nn.Linear(1, 1), + optimizer=optimizer, + parallel_state=SimpleNamespace(dp_mode="ddp"), + generation="generation", + expected_layout="layout", + group=None, + load=True, + ) + inventory = ["optimizer.weight"] + digest = hashlib.sha256(json.dumps(inventory, separators=(",", ":")).encode()).hexdigest() + marker = {"generation": "generation", "layout": "layout", "inventory_digests": [digest, digest]} + if rank == 0: + marker["inventory_digests"][0] = "corrupt" + payload = { + "weight": torch.ones(1), + _MULTI_INVENTORY_MARKER: json.dumps(inventory), + MULTI_LAYOUT_MARKER: json.dumps(marker), + } + with pytest.raises(RuntimeError, match="inventory identity"): + state.load_state_dict(payload) + assert not optimizer.entered + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_corrupt_optimizer_inventory_fails_on_all_ranks(): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_corrupt_inventory_worker, nprocs=2) + + def _load_preflight_failure_worker(rank, checkpoint, failure, endpoint): from datetime import timedelta from unittest.mock import patch @@ -286,6 +356,29 @@ def fail_to_remove(path, *args, **kwargs): assert not checkpoint.exists() +def test_staged_layout_finalization_failure_cleans_temporary_sidecar(tmp_path, monkeypatch): + import veomni.checkpoint.dcp_checkpointer as checkpoint_module + + checkpoint = tmp_path / "final" + stage = tmp_path / "stage" + staged = stage / checkpoint_module._STAGE_ROOT / checkpoint_module._stage_key(str(checkpoint)) + model, optimizer = _model_and_optimizer() + + def fail(*args, **kwargs): + raise RuntimeError("injected layout finalization failure") + + monkeypatch.setattr(checkpoint_module, "finish_multi_layout", fail) + with pytest.raises(RuntimeError, match="injected layout finalization failure"): + DistributedCheckpointer.save( + str(checkpoint), + {"model": model, "optimizer": optimizer, "extra_state": {}, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + stage_dir=str(stage), + ) + assert not staged.exists() + assert not (checkpoint / "model").exists() + + def _proven_staged_resume_worker(rank, checkpoint, stage, endpoint): import os from datetime import timedelta diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 437e0376ee..135f121971 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -47,7 +47,7 @@ from torch.distributed.checkpoint.stateful import Stateful from torch.distributed.tensor.placement_types import _StridedShard -from ..distributed.parallel_state import get_parallel_state +from ..distributed.parallel_state import ParallelState, get_parallel_state from ..utils import logging from ..utils.device import empty_cache, synchronize from ..utils.dist_utils import any_rank_failed, raise_if_any_rank_failed @@ -74,7 +74,12 @@ _prepare_multi_optimizer_saved_states, _validate_source_optimizer_states, ) -from .optimizer_resume_layout import MULTI_LAYOUT_MARKER, coordinated_multi_layout, frozen_parameter_digest +from .optimizer_resume_layout import ( + MULTI_LAYOUT_MARKER, + capture_multi_layout, + finish_multi_layout, + frozen_parameter_digest, +) logger = logging.get_logger(__name__) @@ -85,6 +90,7 @@ _EXPERT_LAYOUTS_KEY = "__veomni_expert_layouts__" +_UNKNOWN_EXPERT_LAYOUTS_KEY = "__veomni_unknown_expert_layouts__" _LEGACY_VEOMNI_TP1_EP0_FSDP1 = "veomni_tp1_ep0_fsdp1" _LEGACY_VEOMNI_TP1_NO_EP = "veomni_tp1_no_ep" @@ -122,23 +128,91 @@ def _expert_layout(tensor): } +def _no_ep_expert_layout(key, tensor, parallel_state): + """Record no-EP provenance only from a known TP1/EP1 VeOmni source.""" + if "experts" not in key.split(".") or not isinstance(tensor, DTensor) or not tensor.ndim: + return None + if not isinstance(parallel_state, ParallelState): + return None + if ( + parallel_state.tp_size != 1 + or parallel_state.extra_parallel_sizes.get("ep") != 1 + or parallel_state.device_mesh is None + ): + return None + mesh = tensor.device_mesh + names = mesh.mesh_dim_names or () + if names not in ( + ("dp_shard",), + ("dp_shard_sp",), + ("dp_replicate", "dp_shard"), + ("dp_replicate", "dp_shard_sp"), + ): + return None + if mesh != parallel_state.fsdp_mesh or len(tensor.placements) != len(names): + return None + if len(names) == 2 and not isinstance(tensor.placements[0], Replicate): + return None + shard = tensor.placements[-1] + if type(shard) not in (Shard, Replicate): + return None + return { + "ep_shard_dim": None, + "fsdp_shard_dim": shard.dim % tensor.ndim if isinstance(shard, Shard) else None, + "ep_size": 1, + "fsdp_size": mesh.size(-1), + "expert_order": "canonical", + "source_kind": "veomni_tp1_no_ep", + } + + class _ExtraParallelSavePlanner(DefaultSavePlanner): """Record per-tensor source layouts in DCP metadata, including optimizer state.""" + def __init__(self, *args, parallel_state=None, **kwargs): + super().__init__(*args, **kwargs) + self.parallel_state = parallel_state + def create_local_plan(self): plan = super().create_local_plan() - layouts = {key: layout for key, value in self.state_dict.items() if (layout := _expert_layout(value))} - return replace(plan, planner_data={**(plan.planner_data or {}), _EXPERT_LAYOUTS_KEY: layouts}) + layouts = { + key: layout + for key, value in self.state_dict.items() + if (layout := _expert_layout(value) or _no_ep_expert_layout(key, value, self.parallel_state)) + } + unknown = [ + key + for key, value in self.state_dict.items() + if "experts" in key.split(".") and isinstance(value, DTensor) and key not in layouts + ] + return replace( + plan, + planner_data={ + **(plan.planner_data or {}), + _EXPERT_LAYOUTS_KEY: layouts, + _UNKNOWN_EXPERT_LAYOUTS_KEY: unknown, + }, + ) def create_global_plan(self, all_plans): plans, metadata = super().create_global_plan(all_plans) # Reuse DCP's plan exchange; sparse optimizer states may appear on only # some ranks. Keep the ordinary FQN -> state-dict path mappings intact. - layouts = {} + layouts, unknown = {}, set() for plan in all_plans: - layouts.update(plan.planner_data[_EXPERT_LAYOUTS_KEY]) + unknown.update(plan.planner_data[_UNKNOWN_EXPERT_LAYOUTS_KEY]) + for key, value in plan.planner_data[_EXPERT_LAYOUTS_KEY].items(): + if key in layouts and layouts[key] != value: + raise RuntimeError(f"Checkpoint tensor {key!r} has inconsistent source layouts across ranks") + layouts[key] = value + if ambiguous := sorted(layouts.keys() & unknown): + raise RuntimeError(f"Checkpoint tensors have both known and unknown source layouts: {ambiguous}") metadata.planner_data = { - **(metadata.planner_data or {}), + **{ + key: value + for key, value in (metadata.planner_data or {}).items() + if key != _UNKNOWN_EXPERT_LAYOUTS_KEY + }, _EXPERT_LAYOUTS_KEY: {"version": 1, "tensors": layouts}, } return plans, metadata @@ -352,12 +426,12 @@ def _check_canonical_spec_coverage(model, specs, state_keys, *, trainable_only): def _reject_unproven_extra_parallel_optimizer( state, parallel_state, *, group=None, trainable_only=False, allow_proven_optimizer=False ): - error, signature = None, None + error, has_canonical_layout = None, False try: specs = getattr(state.get("model"), "_fqn2spec_info", None) if _needs_extra_parallel_layout_marker(specs): _extra_parallel_layout_key(specs) - signature = _runtime_extra_parallel_layout_signature(specs) + has_canonical_layout = True if state.get("optimizer") is not None: ps = parallel_state if parallel_state is not None else get_parallel_state() else: @@ -370,7 +444,7 @@ def _reject_unproven_extra_parallel_optimizer( raise_if_any_rank_failed(error, "checking canonical ExtraParallel optimizer state", group=group) error = None try: - if signature is not None: + if has_canonical_layout: model = state["model"] options = StateDictOptions(ignore_frozen_params=True) if trainable_only else None canonical_state = get_model_state_dict(model, options=options) @@ -409,25 +483,6 @@ def _extra_parallel_layout_key(specs): return f"{_EXTRA_PARALLEL_LAYOUT_MARKER}{digest}" -def _runtime_extra_parallel_layout_signature(specs): - """Detect mutation while a single Stateful adapter is in use.""" - if _needs_extra_parallel_layout_marker(specs): - _extra_parallel_layout_key(specs) - return tuple( - ( - name, - spec.para_name, - repr(spec.placement), - spec.fsdp_shard_dim, - getattr(spec, "persistent_fsdp_shard_dim", None), - getattr(spec, "global_shape", None), - getattr(spec, "checkpoint_layout_v2", False), - tuple(spec.para_fsdp_mesh.shape) if spec.para_fsdp_mesh else None, - ) - for name, spec in sorted((specs or {}).items()) - ) - - def _strip_extra_parallel_layout_marker(state_dict, required, marker_key): if not required: return state_dict @@ -721,17 +776,12 @@ def __init__(self, model, trainable_only: bool = False, parallel_state=None): if _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info) else None ) - self._layout_signature = _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) self.should_extra_parallel_aware = ( self.extra_parallel_fqn2spec_info is not None and self.parallel_state.dp_mode == "fsdp2" ) @torch.no_grad() def state_dict(self): - if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( - self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) - ): - raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") options = StateDictOptions(ignore_frozen_params=True) if self.trainable_only else None model_state_dict = get_model_state_dict(model=self.model, options=options) if self.should_extra_parallel_aware and _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info): @@ -759,10 +809,6 @@ def load_state_dict(self, state_dict): need to drop ExtraParallel-dim when loading from DCP checkpoints so that ExtraParallel-FSDP would not be confused """ - if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( - self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) - ): - raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") model_state_dict = _strip_extra_parallel_layout_marker( state_dict, self.should_extra_parallel_aware @@ -838,7 +884,6 @@ def __init__(self, model, optimizer, parallel_state=None, *, load: bool = False) if _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info) else None ) - self._layout_signature = _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) self.should_extra_parallel_aware = ( self.extra_parallel_fqn2spec_info is not None and self.parallel_state.dp_mode == "fsdp2" ) @@ -847,10 +892,6 @@ def __init__(self, model, optimizer, parallel_state=None, *, load: bool = False) def state_dict(self): if _requires_rank_local_optimizer_checkpoint(self.optimizer): raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") - if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( - self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) - ): - raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") if self.should_extra_parallel_aware and self._layout_marker_key is not None: raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") if getattr(self.optimizer, "_is_multi_optimizer", False): @@ -888,10 +929,6 @@ def load_state_dict(self, state_dict): raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") if self.should_extra_parallel_aware and self._layout_marker_key is not None: raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") - if self.extra_parallel_fqn2spec_info is not getattr(self.model, "_fqn2spec_info", None) or ( - self._layout_signature != _runtime_extra_parallel_layout_signature(self.extra_parallel_fqn2spec_info) - ): - raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active") optim_state_from_dcp_load = _strip_extra_parallel_layout_marker( state_dict, self.should_extra_parallel_aware @@ -975,22 +1012,14 @@ def __init__( self, *args, generation, - scheduler_snapshot, - scheduler_object, - extra_digest, expected_layout, group, - frozen_digest=None, **kwargs, ): super().__init__(*args, **kwargs) self.generation = generation - self.scheduler_snapshot = scheduler_snapshot - self.scheduler_object = scheduler_object - self.extra_digest = extra_digest self.expected_layout = expected_layout self.group = group - self.frozen_digest = frozen_digest def state_dict(self): result, error = None, None @@ -1019,31 +1048,10 @@ def state_dict(self): if self._load: result[MULTI_LAYOUT_MARKER] = "" return result - initial = coordinated_multi_layout( - self.model, - self.optimizer, - self.scheduler_snapshot, - self.scheduler_object, - writing=True, - group=self.group, - frozen_digest=self.frozen_digest, - ) - if initial != self.expected_layout: - raise RuntimeError("MultiOptimizer layout changed while saving checkpoint") - current = coordinated_multi_layout( - self.model, - self.optimizer, - self.scheduler_snapshot, - self.scheduler_object, - self.extra_digest, - writing=True, - group=self.group, - frozen_digest=self.frozen_digest, - ) result[MULTI_LAYOUT_MARKER] = json.dumps( { "generation": self.generation, - "layout": current, + "layout": self.expected_layout, "inventory_digests": _inventory_digests(inventory, group=self.group), }, sort_keys=True, @@ -1085,19 +1093,13 @@ def load_state_dict(self, state_dict): except BaseException as exc: error = exc raise_if_any_rank_failed(error, "validating loaded MultiOptimizer state", group=self.group) - if marker["inventory_digests"] != _inventory_digests(inventory, group=self.group): - raise RuntimeError("MultiOptimizer inventory identity changed during tensor loading") - current = coordinated_multi_layout( - self.model, - self.optimizer, - self.scheduler_snapshot, - self.scheduler_object, - self.extra_digest, - group=self.group, - frozen_digest=self.frozen_digest, + digests = _inventory_digests(inventory, group=self.group) + error = ( + RuntimeError("MultiOptimizer inventory identity changed during tensor loading") + if marker["inventory_digests"] != digests + else None ) - if current != self.expected_layout: - raise RuntimeError("MultiOptimizer layout changed while loading optimizer state") + raise_if_any_rank_failed(error, "validating MultiOptimizer inventory identity", group=self.group) error = None try: self.optimizer._dcp_load_state_dict(state_dict, group=self.group) @@ -1795,8 +1797,8 @@ def save( error = exc raise_if_any_rank_failed(error, "validating MultiOptimizer checkpoint policy") frozen_digest = frozen_parameter_digest(state["model"], enabled=trainable_only) - initial_layout = coordinated_multi_layout( - state["model"], optimizer, snapshot, scheduler_object, writing=True, frozen_digest=frozen_digest + captured_layout = capture_multi_layout( + state["model"], optimizer, snapshot, scheduler_object, frozen_digest=frozen_digest ) error, generation_value = None, None if not dist.is_initialized() or dist.get_rank() == 0: @@ -1867,6 +1869,12 @@ def save( except BaseException as exc: error = exc raise_if_any_rank_failed(error, "hashing MultiOptimizer extra state") + try: + initial_layout = finish_multi_layout(captured_layout, extra_digest) + except BaseException: + if stage_path is not None and _local_rank() == 0: + shutil.rmtree(stage_path, ignore_errors=True) + raise def prepare_writer_state(factory, what): result, error = None, None @@ -1907,19 +1915,9 @@ def prepare_writer_state(factory, what): save_to_lowest_rank=save_to_lowest_rank, slot=WEIGHTS_DIRNAME, timeout_seconds=save_timeout_seconds, + parallel_state=parallel_state if parallel_state is not None else get_parallel_state(), ) - if multi: - error = None - if state.get("model") is not source_model or state.get("optimizer") is not optimizer: - error = RuntimeError("Model or optimizer object changed while saving model state") - elif ( - state.get("extra_state") is not extra - or state.get("_veomni_scheduler_object") is not scheduler_object - ): - error = RuntimeError("Scheduler or extra-state object changed while saving model state") - raise_if_any_rank_failed(error, "rechecking MultiOptimizer checkpoint ownership") - if "optimizer" in state and state["optimizer"] is not None: optim_state, optim_writer = prepare_writer_state( lambda: ( @@ -1928,12 +1926,8 @@ def prepare_writer_state(factory, what): optimizer=optimizer, parallel_state=parallel_state, generation=generation, - scheduler_snapshot=snapshot, - scheduler_object=scheduler_object, - extra_digest=extra_digest, expected_layout=initial_layout, group=None, - frozen_digest=frozen_digest, ) if multi else OptimizerState(model=state["model"], optimizer=optimizer, parallel_state=parallel_state), @@ -1948,6 +1942,7 @@ def prepare_writer_state(factory, what): save_to_lowest_rank=save_to_lowest_rank, slot=OPTIMIZER_DIRNAME, timeout_seconds=save_timeout_seconds, + parallel_state=parallel_state if parallel_state is not None else get_parallel_state(), ) except BaseException: if stage_path is not None and _local_rank() == 0: @@ -2155,28 +2150,6 @@ def _optimizer_state() -> "OptimizerState": process_group=process_group, planner=planner, ) - error = ( - RuntimeError("Optimizer object changed while loading model state") - if state.get("optimizer") is not optimizer - else None - ) - raise_if_any_rank_failed(error, "rechecking MultiOptimizer identity", group=process_group) - snapshot = extra.get("lr_scheduler") if extra is not None else None - current = coordinated_multi_layout( - state["model"], - optimizer, - snapshot, - scheduler_object, - extra_digest, - group=process_group, - frozen_digest=frozen_digest, - ) - error = ( - RuntimeError("MultiOptimizer layout changed while loading model state") - if current != expected_layout - else None - ) - raise_if_any_rank_failed(error, "rechecking MultiOptimizer layout", group=process_group) _prepare_multi_optimizer_saved_states( state["model"], optimizer, @@ -2194,12 +2167,8 @@ def _optimizer_state() -> "OptimizerState": parallel_state=parallel_state, load=True, generation=generation, - scheduler_snapshot=snapshot, - scheduler_object=scheduler_object, - extra_digest=extra_digest, expected_layout=expected_layout, group=process_group, - frozen_digest=frozen_digest, ) reader = cls._create_storage_reader(optimizer_path) planner = _ModelStrictLoadPlanner(strict_model=False, expected_save_id=optimizer_id) @@ -2374,6 +2343,7 @@ def execute_save( save_to_lowest_rank: bool = False, slot: str = WEIGHTS_DIRNAME, timeout_seconds: Optional[int] = None, + parallel_state=None, ) -> None: """Execute DCP save with optional async support. @@ -2386,7 +2356,9 @@ def execute_save( outlive its own next write. ``timeout_seconds`` is that slot's group timeout, applied when the group is created. """ - planner = _ExtraParallelSavePlanner(dedup_save_to_lowest_rank=save_to_lowest_rank) + planner = _ExtraParallelSavePlanner( + dedup_save_to_lowest_rank=save_to_lowest_rank, parallel_state=parallel_state + ) if save_async: # Lazily create this slot's dedicated Gloo process group. Creating a # group is itself collective, and every rank runs the same save diff --git a/veomni/checkpoint/multi_optimizer_checkpoint.py b/veomni/checkpoint/multi_optimizer_checkpoint.py index 4c79a7b481..805389a58f 100644 --- a/veomni/checkpoint/multi_optimizer_checkpoint.py +++ b/veomni/checkpoint/multi_optimizer_checkpoint.py @@ -425,11 +425,6 @@ def _prepare_multi_optimizer_saved_states( from ..optim.checkpoint_adapter import rank_local_adamw_adapter from ..optim.optimizer import AnyPrecisionAdamW - for leaf in _iter_optimizer_leaves(optimizer): - if leaf.state or any( - param.grad is not None for param_group in leaf.param_groups for param in param_group["params"] - ): - raise RuntimeError("MultiOptimizer state or gradients changed while loading model state") metadata = reader_factory(checkpoint_path).read_metadata() save_id = getattr(getattr(metadata, "storage_meta", None), "save_id", None) if expected_save_id is not None and save_id != expected_save_id: diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index c2aee39d47..59d3869802 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -103,7 +103,7 @@ def normalize(code): def _local_layout( - model, optimizer, scheduler_snapshot, scheduler_object, extra_digest, *, writing, frozen_digest=None + model, optimizer, scheduler_snapshot, scheduler_object, extra_digest, *, writing, frozen_digest=None, capture=False ): from ..optim.checkpoint_adapter import rank_local_adamw_adapter from ..optim.muon import DistributedMuon @@ -320,6 +320,10 @@ def mesh_identity(mesh): # Earlier checkpoints did not prove that transient cache history was # reconstructible by a fresh process. layout["dynamic_rope_cache_v1"] = rope_contract + return layout if capture else _layout_digest(layout) + + +def _layout_digest(layout): return hashlib.sha256( json.dumps(layout, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() ).hexdigest() @@ -347,6 +351,52 @@ def frozen_parameter_digest(model, *, enabled, group=None, chunk_bytes=4 << 20): return digest.hexdigest() +def capture_multi_layout(model, optimizer, scheduler_snapshot, scheduler_object, *, group=None, frozen_digest=None): + """Validate the live save layout once, before creating checkpoint files.""" + encoded, names, error = None, None, None + try: + payload = _local_layout( + model, + optimizer, + scheduler_snapshot, + scheduler_object, + None, + writing=True, + frozen_digest=frozen_digest, + capture=True, + ) + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), allow_nan=False) + names = tuple(optimizer.key_names) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "capturing MultiOptimizer checkpoint layout", group=group) + if dist.is_initialized(): + gathered = [None] * dist.get_world_size(group) + dist.all_gather_object(gathered, names, group=group) + if any(name != gathered[0] for name in gathered): + raise RuntimeError("MultiOptimizer checkpoint leaf order differs across ranks") + return encoded + + +def finish_multi_layout(encoded, extra_digest, *, group=None): + """Bind saved sidecar bytes without inspecting live model or optimizer state again.""" + local, error = None, None + try: + payload = json.loads(encoded) + payload["extra_digest"] = extra_digest + local = _layout_digest(payload) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "finalizing MultiOptimizer checkpoint layout", group=group) + if dist.is_initialized(): + digests = [None] * dist.get_world_size(group) + dist.all_gather_object(digests, local, group=group) + ranks = list(range(len(digests))) if group is None else dist.get_process_group_ranks(group) + else: + ranks, digests = [0], [local] + return json.dumps({"version": 1, "ranks": ranks, "layouts": digests}, sort_keys=True, separators=(",", ":")) + + def coordinated_multi_layout( model, optimizer, From 830d6254a9d44ddfdd865749e94415b5e729fa5b Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 11:58:44 +0800 Subject: [PATCH 044/120] [checkpoint] remove redundant leaf consensus checks --- docs/usage/checkpoint.md | 13 +++-- .../checkpoints/test_rank_local_optimizer.py | 56 ------------------- veomni/checkpoint/optimizer_resume_layout.py | 17 +----- 3 files changed, 9 insertions(+), 77 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 59948e5aff..1d074ede8a 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -375,10 +375,10 @@ optimizer state is stored under rank-specific keys and recovery requires the identical rank-to-mesh mapping. This avoids merging distinct expert-owner steps and uninitialized states, but increases checkpoint storage for replicated optimizer shards and does not support elastic resharding. All ranks must use -the same trusted model and configuration. Optimizer leaf ordering is checked -as part of the existing checkpoint-layout exchange; owner-specific freezing -that removes a whole leaf from only some ranks is rejected before per-leaf -checkpoint collectives, rather than supported. Replica values, trainability, dtype, shape and geometry +the same trusted model and configuration, including optimizer leaf ordering. +Owner-specific freezing that removes a whole leaf from only some ranks is not +supported or checked and can cause checkpoint collectives to hang. Replica +values, trainability, dtype, shape and geometry are trusted to be consistent across ranks. Source and target checkpoint layout metadata are checked during restore; tensor bytes are not hashed for replica consistency. Trainable-only MultiOptimizer checkpoints do hash each rank's local @@ -389,8 +389,9 @@ preserve serializable scheduler state and implement the same learning-rate behavior on both sides of a restore, and hooks must not mutate checkpoint model, optimizer, or layout during a save/load call. The persisted recipe checks scheduler configuration and optimizer binding; it does not inspect arbitrary wrapper -code. MultiOptimizer asynchronous saves -and unsupported layouts fail closed; synchronous staged saves are supported. +code. MultiOptimizer asynchronous saves and detected unsupported layouts fail +closed; the cross-rank leaf mismatch above is not detected. Synchronous staged +saves are supported. The runtime assumes ranks use a consistent model configuration and compatible replica geometries; saved source layout is validated on restore. Group/state inventory and generation mismatches are rejected before model diff --git a/tests/checkpoints/test_rank_local_optimizer.py b/tests/checkpoints/test_rank_local_optimizer.py index 8a0c8fb169..a7767f1132 100644 --- a/tests/checkpoints/test_rank_local_optimizer.py +++ b/tests/checkpoints/test_rank_local_optimizer.py @@ -282,62 +282,6 @@ def test_anyprecision_bf16_state_extra_parallel_fresh_process_resume(tmp_path, s ) -def _owner_frozen_leaf_mismatch_worker(rank, root, mode): - from unittest.mock import patch - - from veomni.optim.optimizer import build_extra_parallel_fsdp2_optimizer - - dist.init_process_group( - "gloo", init_method=f"file://{root}/leaf-init", rank=rank, world_size=4, timeout=timedelta(seconds=30) - ) - try: - mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ple_fsdp", "ple")) - model = torch.nn.Module() - model.weight = torch.nn.Parameter( - distribute_tensor(torch.arange(64, dtype=torch.float32).reshape(8, 8), mesh, [Shard(1), Shard(0)]) - ) - model.bias = torch.nn.Parameter(torch.ones(2)) - model._fqn2spec_info = {"weight": SpecInfo("ple", Shard(0), "weight", mesh, 1, 1, (8, 8))} - owner = mesh.get_local_rank("ple") - parallel = SimpleNamespace( - dp_mode="fsdp2", extra_parallel_names=("ple",), extra_parallel_fsdp_device_mesh={"ple": mesh} - ) - with patch("veomni.optim.optimizer.get_parallel_state", return_value=parallel): - if mode == "load": - source = build_extra_parallel_fsdp2_optimizer(model) - DistributedCheckpointer.save( - f"{root}/checkpoint", - {"model": model, "optimizer": source, "_veomni_scheduler_object": None}, - parallel_state=parallel, - ) - if owner == 1: - model.weight.requires_grad_(False) - optimizer = build_extra_parallel_fsdp2_optimizer(model) - assert list(optimizer.optimizers_dict) == ( - ["ple", "non_extra_parallel"] if owner == 0 else ["non_extra_parallel"] - ) - before = model.weight.to_local().clone() - with pytest.raises(RuntimeError, match="leaf order differs across ranks"): - method = DistributedCheckpointer.save if mode == "save" else DistributedCheckpointer.load - method( - f"{root}/checkpoint", - {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, - parallel_state=parallel, - ) - torch.testing.assert_close(model.weight.to_local(), before, rtol=0, atol=0) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 4 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) -@pytest.mark.parametrize("mode", ["save", "load"]) -def test_owner_specific_freeze_missing_leaf_fails_before_checkpoint_collectives(tmp_path, mode): - mp.spawn(_owner_frozen_leaf_mismatch_worker, args=(str(tmp_path), mode), nprocs=4, join=True) - - def _frozen_resume_worker(rank, root, phase): dist.init_process_group( "gloo", diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index 59d3869802..87317b3f3f 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -353,7 +353,7 @@ def frozen_parameter_digest(model, *, enabled, group=None, chunk_bytes=4 << 20): def capture_multi_layout(model, optimizer, scheduler_snapshot, scheduler_object, *, group=None, frozen_digest=None): """Validate the live save layout once, before creating checkpoint files.""" - encoded, names, error = None, None, None + encoded, error = None, None try: payload = _local_layout( model, @@ -366,15 +366,9 @@ def capture_multi_layout(model, optimizer, scheduler_snapshot, scheduler_object, capture=True, ) encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), allow_nan=False) - names = tuple(optimizer.key_names) except BaseException as exc: error = exc raise_if_any_rank_failed(error, "capturing MultiOptimizer checkpoint layout", group=group) - if dist.is_initialized(): - gathered = [None] * dist.get_world_size(group) - dist.all_gather_object(gathered, names, group=group) - if any(name != gathered[0] for name in gathered): - raise RuntimeError("MultiOptimizer checkpoint leaf order differs across ranks") return encoded @@ -427,14 +421,7 @@ def coordinated_multi_layout( if dist.is_initialized(): world = dist.get_world_size(group) digests = [None] * world - gathered = [None] * world - # This existing layout exchange also guards the order of subsequent - # per-leaf collectives. Owner-specific freezing can otherwise leave - # one rank without a leaf and strand peers during checkpoint I/O. - dist.all_gather_object(gathered, (local, tuple(optimizer.key_names)), group=group) - digests = [item[0] for item in gathered] - if any(item[1] != gathered[0][1] for item in gathered): - raise RuntimeError("MultiOptimizer checkpoint leaf order differs across ranks") + dist.all_gather_object(digests, local, group=group) ranks = list(range(world)) if group is None else dist.get_process_group_ranks(group) else: ranks, digests = [0], [local] From 2ae3848de2a14b3dff3ae8c027184e9bbb425e14 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 13:22:54 +0800 Subject: [PATCH 045/120] [checkpoint] reuse proven MultiOptimizer export save --- .../checkpoints/test_multi_export_receipt.py | 172 ++++++++++++++++++ .../test_multi_optimizer_rope_cache.py | 28 --- .../test_optimizer_checkpoint_preflight.py | 15 +- veomni/models/checkpoint/manager.py | 18 +- .../trainer/callbacks/checkpoint_callback.py | 82 +++++++++ 5 files changed, 278 insertions(+), 37 deletions(-) create mode 100644 tests/checkpoints/test_multi_export_receipt.py diff --git a/tests/checkpoints/test_multi_export_receipt.py b/tests/checkpoints/test_multi_export_receipt.py new file mode 100644 index 0000000000..3d4b91685c --- /dev/null +++ b/tests/checkpoints/test_multi_export_receipt.py @@ -0,0 +1,172 @@ +"""Same-step MultiOptimizer DCP and HF export must share a proven receipt.""" + +import gc +import hashlib +import weakref +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from tests.checkpoints.test_checkpoint_callback import _make_mock_runtime, _make_mock_trainer +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer +from veomni.models.checkpoint.manager import ModelCheckpointManager +from veomni.optim.optimizer import MultiOptimizer +from veomni.trainer.callbacks.base import TrainerState +from veomni.trainer.callbacks.checkpoint_callback import CheckpointCallback +from veomni.trainer.callbacks.global_state_callback import GlobalStateCallback + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +def _multi_optimizer(model): + return MultiOptimizer(model, {"model": torch.optim.AdamW(model.parameters())}, ["model"]) + + +def _runtime(tmp_path): + runtime = _make_mock_runtime(save_path=str(tmp_path)) + runtime.model = torch.nn.Linear(2, 2) + runtime.optimizer = _multi_optimizer(runtime.model) + runtime.lr_scheduler = None + runtime.parallel_state = SimpleNamespace(global_rank=0, dp_mode="ddp", extra_parallel_names=[]) + return runtime + + +@pytest.mark.parametrize("replacement", ["multi", "ordinary", "none"]) +def test_replaced_optimizer_cannot_reuse_multi_receipt(tmp_path, replacement): + runtime = _runtime(tmp_path) + trainer = _make_mock_trainer(save_path=str(tmp_path)) + state = TrainerState(global_step=1, epoch=0) + with ( + patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer), + patch("torch.distributed.barrier"), + ): + manager = ModelCheckpointManager(runtime) + runtime.checkpoint = manager + trainer.model = runtime + trainer.save_dcp.side_effect = manager.save_dcp + trainer.save_hf_or_lora.side_effect = lambda state, stage="step_end": manager.prepare_export(state, stage) + callback = CheckpointCallback(trainer) + callback._save_hf(state) + runtime.optimizer = { + "multi": lambda: _multi_optimizer(runtime.model), + "ordinary": lambda: torch.optim.AdamW(runtime.model.parameters()), + "none": lambda: None, + }[replacement]() + with pytest.raises(RuntimeError, match="receipt changed"): + callback.on_epoch_end(state) + trainer.save_dcp.assert_called_once_with(state) + + +def test_train_end_does_not_retain_optimizer_during_hf_export(tmp_path): + runtime = _runtime(tmp_path) + optimizer_ref = weakref.ref(runtime.optimizer) + trainer = _make_mock_trainer(save_path=str(tmp_path)) + state = TrainerState(global_step=1, epoch=0) + with ( + patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer), + patch("torch.distributed.barrier"), + ): + manager = ModelCheckpointManager(runtime) + runtime.checkpoint = manager + trainer.model = runtime + trainer.save_dcp.side_effect = manager.save_dcp + + def export(state, stage="step_end"): + manager.prepare_export(state, stage) + gc.collect() + assert optimizer_ref() is None + + trainer.save_hf_or_lora.side_effect = export + CheckpointCallback(trainer)._save_hf(state, stage="train_end") + + +def test_unproven_same_step_receipt_is_not_overwritten(tmp_path): + runtime = _runtime(tmp_path) + trainer = _make_mock_trainer(save_path=str(tmp_path)) + with patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer): + manager = ModelCheckpointManager(runtime) + runtime.checkpoint = manager + manager._last_saved_step = 1 + manager._last_saved_optimizer_generation = "unproven" + trainer.model = runtime + callback = CheckpointCallback(trainer) + with pytest.raises(RuntimeError, match="no matching same-step checkpoint receipt"): + callback._save_hf(TrainerState(global_step=1)) + trainer.save_dcp.assert_not_called() + trainer.save_hf_or_lora.assert_not_called() + + +def _real_multi_save_worker(rank, directory, endpoint): + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + _init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cpu") + torch.manual_seed(123) + runtime = _runtime(directory) + runtime.parallel_state = SimpleNamespace(global_rank=rank, dp_mode="ddp", extra_parallel_names=[]) + runtime.model(torch.ones(1, 2)).sum().backward() + runtime.optimizer.step() + runtime.optimizer.zero_grad() + trainer = _make_mock_trainer(save_path=directory) + trainer.args.train.global_rank = rank + trainer.args.train.world_size = 2 + trainer.args.train.checkpoint.save_steps = 0 + trainer.args.train.checkpoint.save_epochs = 1 + trainer.args.train.checkpoint.hf_save_steps = 1 + trainer.args.train.checkpoint.hf_save_epochs = 0 + trainer.train_dataloader = None + trainer.data_iterator = None + trainer.environ_meter.state_dict.return_value = {} + trainer.channel_loss_callback = None + trainer.offline_cache_writer = None + with patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer): + manager = ModelCheckpointManager(runtime) + runtime.checkpoint = manager + trainer.model = runtime + trainer.save_dcp.side_effect = manager.save_dcp + trainer.save_hf_or_lora.side_effect = lambda state, stage="step_end": manager.prepare_export(state, stage) + callback = CheckpointCallback(trainer) + global_callback = GlobalStateCallback(trainer) + state = TrainerState(global_step=1, epoch=0) + callback.on_step_end(state) + global_callback.on_step_end(state) + generation = manager._last_saved_optimizer_generation + assert generation and callback._last_multi_receipt == generation + snapshot = None + if rank == 0: + root = Path(directory) + snapshot = { + p.relative_to(root): hashlib.sha256(p.read_bytes()).digest() for p in root.rglob("*") if p.is_file() + } + assert snapshot + callback.on_epoch_end(state) + global_callback.on_epoch_end(state) + trainer.save_dcp.assert_called_once_with(state) + assert manager._last_saved_optimizer_generation == generation + if rank == 0: + assert snapshot == { + p.relative_to(root): hashlib.sha256(p.read_bytes()).digest() for p in root.rglob("*") if p.is_file() + } + if rank == 1: + manager._last_saved_optimizer_generation = None + with pytest.raises(RuntimeError, match="receipt changed"): + callback.on_epoch_end(state) + dist.barrier() + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_real_multi_dcp_hf_epoch_on_two_ranks(tmp_path): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_real_multi_save_worker, args=(str(tmp_path),), nprocs=2) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index 1e401c6764..db4d5b8762 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -39,34 +39,6 @@ def test_runtime_preflight_rejects_unsupported_dynamic_rope_before_training(): VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) -def _runtime_rope_preflight_failure(rank, endpoint): - from veomni.models.model_runtime import VeOmniModelRuntime - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - model, optimizer = _build() - if rank == 0: - model.rope.rope_type = "default" - runtime = SimpleNamespace( - model=model, - optimizer=optimizer, - model_name="policy", - train_args=SimpleNamespace(checkpoint=SimpleNamespace(manager="dcp", save_async=False)), - ) - with pytest.raises(RuntimeError, match="dynamic RoPE|another rank"): - VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) - done = torch.tensor(1) - dist.all_reduce(done) - assert done.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_dynamic_rope_preflight_failure_reaches_all_ranks(): - spawn_with_rendezvous(_runtime_rope_preflight_failure, nprocs=2) - - class _TinyQwenRope(torch.nn.Module): def __init__(self, theta=10000): super().__init__() diff --git a/tests/models/test_optimizer_checkpoint_preflight.py b/tests/models/test_optimizer_checkpoint_preflight.py index a30feb03d7..6dcf580d3f 100644 --- a/tests/models/test_optimizer_checkpoint_preflight.py +++ b/tests/models/test_optimizer_checkpoint_preflight.py @@ -114,14 +114,13 @@ def test_omni_module_scope_checks_async_before_checkpoint(monkeypatch, multi): assert reached_checkpoint == [True] -def _mismatch_worker(rank, endpoint): +def _preflight_failure_worker(rank, endpoint): dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) try: model_runtime.use_parallel_state = lambda *_: nullcontext() + _runtime(True, False) with pytest.raises(RuntimeError, match="synchronous|checking MultiOptimizer checkpoint mode"): - _runtime(True if rank == 0 else None, True) - # A frozen local shard must not prevent synchronous checkpoint setup. - _runtime(True if rank == 0 else None, False) + _runtime(True, True) with pytest.raises(RuntimeError, match="injected optimizer build failure|building model optimizer"): _runtime("fail" if rank == 0 else False, False) done = torch.tensor(1) @@ -132,10 +131,10 @@ def _mismatch_worker(rank, endpoint): @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -def test_mismatched_optimizer_checkpoint_modes_fail_on_all_ranks(): +def test_optimizer_checkpoint_failures_reach_all_ranks(): from tests.tools.launch_utils import spawn_with_rendezvous - spawn_with_rendezvous(_mismatch_worker, args=(), nprocs=2) + spawn_with_rendezvous(_preflight_failure_worker, args=(), nprocs=2) def _accelerator_worker(rank, endpoint): @@ -148,9 +147,9 @@ def _accelerator_worker(rank, endpoint): ) try: model_runtime.use_parallel_state = lambda *_: nullcontext() - _runtime(True if rank == 0 else None, False) + _runtime(True, False) with pytest.raises(RuntimeError, match="synchronous|checking MultiOptimizer checkpoint mode"): - _runtime(True if rank == 0 else None, True) + _runtime(True, True) with pytest.raises(RuntimeError, match="injected optimizer build failure|building model optimizer"): _runtime("fail" if rank == 0 else False, False) done = torch.ones(1, device=device) diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 9695fee77a..2947d4839c 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -33,6 +33,7 @@ import marshal import math import types +import weakref from typing import TYPE_CHECKING, Any, Dict, Optional import torch.distributed as dist @@ -203,6 +204,7 @@ def __init__(self, runtime: "VeOmniModelRuntime"): self._loaded_optimizer_global_step = None self._loaded_optimizer_generation = None self._last_saved_optimizer_generation = None + self._last_saved_multi_optimizer_ref = None self.checkpointer: CheckpointerBase = build_checkpointer( ckpt_manager=self.config.manager, dist_backend=runtime.args.accelerator.fsdp_config.fsdp_mode, @@ -326,6 +328,7 @@ def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: def load(self) -> None: """Restore model, optimizer and model-bound extra state from ``load_path``.""" + self._last_saved_multi_optimizer_ref = None self._loaded_optimizer_global_step = None self._loaded_optimizer_generation = None load_dir = self.load_dir() @@ -417,11 +420,13 @@ def save_dcp(self, state: "TrainerState") -> None: dataloader is, the rng — has its own writer. """ self._last_saved_optimizer_generation = None + self._last_saved_multi_optimizer_ref = None if self._last_saved_step == state.global_step: # A failed rewrite may already have invalidated the old DCP marker. # A later export must not reuse the previous same-step receipt. self._last_saved_step = -1 - multi = getattr(self.runtime.optimizer, "_is_multi_optimizer", False) is True + optimizer = self.runtime.optimizer + multi = getattr(optimizer, "_is_multi_optimizer", False) is True error, extra_state = None, None try: extra_state = self._extra_state(state) @@ -469,6 +474,17 @@ def clear_cache(stage): clear_cache("clearing cache after MultiOptimizer checkpoint save") dist.barrier() if multi: + error = ( + RuntimeError("MultiOptimizer changed while saving its checkpoint") + if self.runtime.optimizer is not optimizer + else None + ) + group = dist.group.WORLD if dist.is_initialized() else None + raise_if_any_rank_failed(error, "checking saved MultiOptimizer identity", group=group) + try: + self._last_saved_multi_optimizer_ref = weakref.ref(optimizer) + except TypeError: + self._last_saved_multi_optimizer_ref = None self._last_saved_optimizer_generation = generation self._last_saved_step = state.global_step logger.info_rank0(f"Distributed checkpoint saved at {self.save_dir(state)} successfully!") diff --git a/veomni/trainer/callbacks/checkpoint_callback.py b/veomni/trainer/callbacks/checkpoint_callback.py index bb05980d89..184d1284c3 100644 --- a/veomni/trainer/callbacks/checkpoint_callback.py +++ b/veomni/trainer/callbacks/checkpoint_callback.py @@ -30,8 +30,11 @@ :mod:`~veomni.trainer.callbacks.global_state_callback`. """ +from collections.abc import Mapping from typing import TYPE_CHECKING +import torch.distributed as dist + from ...utils import helper from .base import Callback, TrainerState @@ -43,6 +46,47 @@ logger = helper.create_logger(__name__) +def _single_multi_checkpoint(trainer): + """Limit the same-step receipt to the synchronous single-model route.""" + runtime = getattr(trainer, "model", None) + if isinstance(getattr(runtime, "module_runtimes", None), Mapping): + return None + optimizer = getattr(runtime, "optimizer", None) + if getattr(optimizer, "_is_multi_optimizer", False) is not True: + return None + return runtime, optimizer, getattr(runtime, "checkpoint", None) + + +def _agree_multi_value(value): + if not dist.is_initialized(): + return [value] + values = [None] * dist.get_world_size() + dist.all_gather_object(values, value) + return values + + +def _current_multi_receipt(runtime, optimizer, checkpoint, step): + from ...models.checkpoint.manager import ModelCheckpointManager + + if ( + not isinstance(checkpoint, ModelCheckpointManager) + or getattr(checkpoint, "runtime", None) is not runtime + or getattr(runtime, "optimizer", None) is not optimizer + ): + return None + ref = getattr(checkpoint, "_last_saved_multi_optimizer_ref", None) + generation = getattr(checkpoint, "_last_saved_optimizer_generation", None) + if ( + getattr(checkpoint, "_last_saved_step", None) != step + or not callable(ref) + or ref() is not optimizer + or not isinstance(generation, str) + or not generation + ): + return None + return generation + + class CheckpointCallback(Callback): """Schedule DCP / HF / LoRA I/O and the one-shot tokenizer/config export.""" @@ -57,6 +101,8 @@ def __init__(self, trainer: "BaseTrainer"): self.hf_every_n_epochs = ckpt.hf_save_epochs self._last_dcp_step: int = -1 self._last_hf_step: int = -1 + self._last_multi_receipt = None + self._uses_single_multi_checkpoint = _single_multi_checkpoint(trainer) is not None def on_train_begin(self, state: TrainerState, **kwargs) -> None: self.trainer.save_model_assets() @@ -84,6 +130,15 @@ def on_step_end(self, state: TrainerState, **kwargs): def on_epoch_end(self, state: TrainerState, **kwargs): if self.dcp_every_n_epochs and (state.epoch + 1) % self.dcp_every_n_epochs == 0: + if self._uses_single_multi_checkpoint: + multi = _single_multi_checkpoint(self.trainer) + recorded = self._last_multi_receipt if self._last_dcp_step == state.global_step else None + current = _current_multi_receipt(*multi, state.global_step) if multi is not None else None + values = _agree_multi_value((state.global_step, recorded, current)) + if any(value[1] is not None for value in values) and ( + any(value != values[0] for value in values) or recorded is None or recorded != current + ): + raise RuntimeError("Same-step MultiOptimizer checkpoint receipt changed across ranks") if state.global_step != self._last_dcp_step: self._save_dcp(state) else: @@ -103,8 +158,35 @@ def on_epoch_end(self, state: TrainerState, **kwargs): def _save_dcp(self, state: TrainerState): self.trainer.save_dcp(state) self._last_dcp_step = state.global_step + self._last_multi_receipt = None def _save_hf(self, state: TrainerState, stage: str = "step_end"): + if self._uses_single_multi_checkpoint: + from ...models.checkpoint.manager import ModelCheckpointManager + + multi = _single_multi_checkpoint(self.trainer) + runtime, optimizer, checkpoint = multi if multi is not None else (None, None, None) + known_owner = ( + isinstance(checkpoint, ModelCheckpointManager) + and checkpoint.runtime is runtime + and not getattr(checkpoint.config, "save_async", False) + ) + needs_save = known_owner and checkpoint._last_saved_step != state.global_step + decisions = _agree_multi_value((state.global_step, known_owner, needs_save)) + if not known_owner or any(decision != decisions[0] for decision in decisions): + raise RuntimeError("MultiOptimizer export save decision differs across ranks") + if needs_save: + self.trainer.save_dcp(state) + checkpoint = getattr(runtime, "checkpoint", None) + generation = _current_multi_receipt(runtime, optimizer, checkpoint, state.global_step) + receipts = _agree_multi_value((state.global_step, generation)) + if generation is None or any(receipt != receipts[0] for receipt in receipts): + raise RuntimeError("MultiOptimizer export has no matching same-step checkpoint receipt") + self._last_dcp_step = state.global_step + self._last_multi_receipt = generation + # At train end the exporter may release the optimizer before the HF + # gather. Do not keep it alive through local strong references. + del optimizer, multi self.trainer.save_hf_or_lora(state, stage=stage) self._last_hf_step = state.global_step From 462cdac43322f895a4f81a444ce71a08baaa85e2 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 13:30:53 +0800 Subject: [PATCH 046/120] [test] preserve same-config RoPE failure coverage --- .../test_multi_optimizer_rope_cache.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index db4d5b8762..c0d7f91fc2 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -39,6 +39,45 @@ def test_runtime_preflight_rejects_unsupported_dynamic_rope_before_training(): VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) +def _runtime_rope_guard_io_failure(rank, endpoint): + from unittest.mock import patch + + from veomni.checkpoint import dynamic_rope + from veomni.models.model_runtime import VeOmniModelRuntime + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + torch.manual_seed(7) + model, optimizer = _build() + model.rope.rope_type = "default" # Same supported model configuration on every rank. + runtime = SimpleNamespace( + model=model, + optimizer=optimizer, + model_name="policy", + train_args=SimpleNamespace(checkpoint=SimpleNamespace(manager="dcp", save_async=False)), + ) + original = dynamic_rope.reject_dynamic_rope_multi_optimizer_training + failure = OSError("injected local RoPE validator failure") if rank == 1 else None + with patch.object( + dynamic_rope, "reject_dynamic_rope_multi_optimizer_training", wraps=original, side_effect=failure + ) as guard: + error = OSError if rank == 1 else RuntimeError + message = ( + "injected local RoPE validator failure" if rank == 1 else "dynamic RoPE checkpoint mode.*another rank" + ) + with pytest.raises(error, match=message): + VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) + guard.assert_called_once_with(model) + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_dynamic_rope_validator_local_failure_reaches_all_ranks(): + spawn_with_rendezvous(_runtime_rope_guard_io_failure, nprocs=2) + + class _TinyQwenRope(torch.nn.Module): def __init__(self, theta=10000): super().__init__() From d18fd1175b4386d5ae4d4e1e457ccb2df412dc63 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:31:22 +0800 Subject: [PATCH 047/120] Reuse same-step exports across module checkpoint owners --- tests/checkpoints/test_checkpoint_callback.py | 22 +++ .../checkpoints/test_multi_export_receipt.py | 137 +++++++++--------- veomni/models/checkpoint/manager.py | 15 -- .../trainer/callbacks/checkpoint_callback.py | 133 ++++++----------- 4 files changed, 138 insertions(+), 169 deletions(-) diff --git a/tests/checkpoints/test_checkpoint_callback.py b/tests/checkpoints/test_checkpoint_callback.py index 2da910ddf0..cc790aa487 100644 --- a/tests/checkpoints/test_checkpoint_callback.py +++ b/tests/checkpoints/test_checkpoint_callback.py @@ -269,6 +269,28 @@ def test_train_end_propagates_async_save_failure(self, mock_helper): with pytest.raises(RuntimeError, match="HDFS write failed"): cb.on_train_end(TrainerState(global_step=60)) + def test_failed_async_epoch_wait_cannot_be_retried_as_success(self, mock_helper, monkeypatch): + from concurrent.futures import Future + + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer + + failed = Future() + failed.set_exception(RuntimeError("HDFS write failed")) + monkeypatch.setattr(DistributedCheckpointer, "_save_futures", {"weights": failed}) + monkeypatch.setattr(DistributedCheckpointer, "_async_process_groups", {}) + trainer = _make_mock_trainer(save_async=True) + trainer.wait_for_pending_save.side_effect = DistributedCheckpointer.wait_for_pending_save + callback = CheckpointCallback(trainer) + callback._last_dcp_step = 5 + state = TrainerState(global_step=5, epoch=0) + + with pytest.raises(RuntimeError, match="HDFS write failed"): + callback.on_epoch_end(state) + assert not DistributedCheckpointer._save_futures + with pytest.raises(RuntimeError, match="restart from a complete checkpoint"): + callback.on_epoch_end(state) + trainer.save_dcp.assert_not_called() + def test_train_end_waits_even_without_async(self, mock_helper): """The call is unconditional; wait_for_pending_save is a no-op when nothing is pending.""" trainer = _make_mock_trainer(save_async=False) diff --git a/tests/checkpoints/test_multi_export_receipt.py b/tests/checkpoints/test_multi_export_receipt.py index 3d4b91685c..eb88f6d7f7 100644 --- a/tests/checkpoints/test_multi_export_receipt.py +++ b/tests/checkpoints/test_multi_export_receipt.py @@ -1,4 +1,4 @@ -"""Same-step MultiOptimizer DCP and HF export must share a proven receipt.""" +"""Same-step HF and epoch saves reuse the model owners' DCP transaction.""" import gc import hashlib @@ -6,7 +6,7 @@ from datetime import timedelta from pathlib import Path from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch @@ -14,8 +14,10 @@ from scripts.ci.accelerator.launch_plan import LaunchSpec from tests.checkpoints.test_checkpoint_callback import _make_mock_runtime, _make_mock_trainer +from veomni.checkpoint import layout from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.models.checkpoint.manager import ModelCheckpointManager +from veomni.models.seed_omni.utils.checkpoint import OmniModuleCheckpointManager from veomni.optim.optimizer import MultiOptimizer from veomni.trainer.callbacks.base import TrainerState from veomni.trainer.callbacks.checkpoint_callback import CheckpointCallback @@ -38,32 +40,6 @@ def _runtime(tmp_path): return runtime -@pytest.mark.parametrize("replacement", ["multi", "ordinary", "none"]) -def test_replaced_optimizer_cannot_reuse_multi_receipt(tmp_path, replacement): - runtime = _runtime(tmp_path) - trainer = _make_mock_trainer(save_path=str(tmp_path)) - state = TrainerState(global_step=1, epoch=0) - with ( - patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer), - patch("torch.distributed.barrier"), - ): - manager = ModelCheckpointManager(runtime) - runtime.checkpoint = manager - trainer.model = runtime - trainer.save_dcp.side_effect = manager.save_dcp - trainer.save_hf_or_lora.side_effect = lambda state, stage="step_end": manager.prepare_export(state, stage) - callback = CheckpointCallback(trainer) - callback._save_hf(state) - runtime.optimizer = { - "multi": lambda: _multi_optimizer(runtime.model), - "ordinary": lambda: torch.optim.AdamW(runtime.model.parameters()), - "none": lambda: None, - }[replacement]() - with pytest.raises(RuntimeError, match="receipt changed"): - callback.on_epoch_end(state) - trainer.save_dcp.assert_called_once_with(state) - - def test_train_end_does_not_retain_optimizer_during_hf_export(tmp_path): runtime = _runtime(tmp_path) optimizer_ref = weakref.ref(runtime.optimizer) @@ -76,7 +52,6 @@ def test_train_end_does_not_retain_optimizer_during_hf_export(tmp_path): manager = ModelCheckpointManager(runtime) runtime.checkpoint = manager trainer.model = runtime - trainer.save_dcp.side_effect = manager.save_dcp def export(state, stage="step_end"): manager.prepare_export(state, stage) @@ -87,34 +62,32 @@ def export(state, stage="step_end"): CheckpointCallback(trainer)._save_hf(state, stage="train_end") -def test_unproven_same_step_receipt_is_not_overwritten(tmp_path): - runtime = _runtime(tmp_path) - trainer = _make_mock_trainer(save_path=str(tmp_path)) - with patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer): - manager = ModelCheckpointManager(runtime) - runtime.checkpoint = manager - manager._last_saved_step = 1 - manager._last_saved_optimizer_generation = "unproven" - trainer.model = runtime - callback = CheckpointCallback(trainer) - with pytest.raises(RuntimeError, match="no matching same-step checkpoint receipt"): - callback._save_hf(TrainerState(global_step=1)) - trainer.save_dcp.assert_not_called() - trainer.save_hf_or_lora.assert_not_called() - - -def _real_multi_save_worker(rank, directory, endpoint): +def _real_multi_save_worker(rank, directory, scope, endpoint): from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state + from veomni.models.seed_omni.accelerated.omni_model.omni_model_runtime import OmniModelRuntime dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) try: _init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cpu") torch.manual_seed(123) - runtime = _runtime(directory) - runtime.parallel_state = SimpleNamespace(global_rank=rank, dp_mode="ddp", extra_parallel_names=[]) - runtime.model(torch.ones(1, 2)).sum().backward() - runtime.optimizer.step() - runtime.optimizer.zero_grad() + + def make_runtime(name=None): + runtime = _runtime(directory) + runtime.module_name = name + runtime.parallel_state = SimpleNamespace(global_rank=rank, dp_mode="ddp", extra_parallel_names=[]) + with patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer): + manager = OmniModuleCheckpointManager(runtime) if name else ModelCheckpointManager(runtime) + runtime.checkpoint = manager + runtime.save_dcp = lambda state: runtime.checkpoint.save_dcp(state) + runtime.save_hf_or_lora = lambda state, stage="step_end": runtime.checkpoint.prepare_export(state, stage) + return runtime + + names = ("language", "vision") if scope == "omni" else (None,) + runtimes = [make_runtime(name) for name in names] + for runtime in runtimes: + runtime.model(torch.ones(1, 2)).sum().backward() + runtime.optimizer.step() + runtime.optimizer.zero_grad() trainer = _make_mock_trainer(save_path=directory) trainer.args.train.global_rank = rank trainer.args.train.world_size = 2 @@ -127,19 +100,31 @@ def _real_multi_save_worker(rank, directory, endpoint): trainer.environ_meter.state_dict.return_value = {} trainer.channel_loss_callback = None trainer.offline_cache_writer = None - with patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer): - manager = ModelCheckpointManager(runtime) - runtime.checkpoint = manager - trainer.model = runtime - trainer.save_dcp.side_effect = manager.save_dcp - trainer.save_hf_or_lora.side_effect = lambda state, stage="step_end": manager.prepare_export(state, stage) + if scope == "omni": + modules = {runtime.module_name: runtime for runtime in runtimes} + modules["frozen"] = SimpleNamespace( + optimizer=None, + checkpoint=None, + has_trainable_parameters=False, + save_dcp=lambda state: None, + save_hf_or_lora=lambda state, stage: None, + ) + trainer.model = OmniModelRuntime(SimpleNamespace(), module_runtimes=modules) + trainer.save_dcp.side_effect = trainer.model.save_dcp + trainer.save_hf_or_lora.side_effect = trainer.model.save_hf_or_lora + else: + trainer.model = runtimes[0] + trainer.save_dcp.side_effect = runtimes[0].save_dcp + trainer.save_hf_or_lora.side_effect = runtimes[0].save_hf_or_lora + for runtime in runtimes: + runtime.checkpoint.save_dcp = Mock(wraps=runtime.checkpoint.save_dcp) callback = CheckpointCallback(trainer) global_callback = GlobalStateCallback(trainer) state = TrainerState(global_step=1, epoch=0) callback.on_step_end(state) global_callback.on_step_end(state) - generation = manager._last_saved_optimizer_generation - assert generation and callback._last_multi_receipt == generation + generations = [runtime.checkpoint._last_saved_optimizer_generation for runtime in runtimes] + assert all(generations) snapshot = None if rank == 0: root = Path(directory) @@ -150,23 +135,43 @@ def _real_multi_save_worker(rank, directory, endpoint): callback.on_epoch_end(state) global_callback.on_epoch_end(state) trainer.save_dcp.assert_called_once_with(state) - assert manager._last_saved_optimizer_generation == generation + for runtime in runtimes: + runtime.checkpoint.save_dcp.assert_called_once_with(state) + assert [runtime.checkpoint._last_saved_optimizer_generation for runtime in runtimes] == generations if rank == 0: assert snapshot == { p.relative_to(root): hashlib.sha256(p.read_bytes()).digest() for p in root.rglob("*") if p.is_file() } - if rank == 1: - manager._last_saved_optimizer_generation = None - with pytest.raises(RuntimeError, match="receipt changed"): - callback.on_epoch_end(state) + + restored = [make_runtime(name) for name in names] + for before, after, generation in zip(runtimes, restored, generations): + after.train_args.checkpoint.load_path = layout.step_dir(directory, 1) + after.checkpoint.load() + assert after.checkpoint._loaded_optimizer_generation == generation + for original_param, loaded_param in zip(before.model.parameters(), after.model.parameters()): + torch.testing.assert_close(original_param, loaded_param, rtol=0, atol=0) + original_state = before.optimizer.optimizers_dict["model"].state[original_param] + loaded_state = after.optimizer.optimizers_dict["model"].state[loaded_param] + for key in ("step", "exp_avg", "exp_avg_sq"): + torch.testing.assert_close(original_state[key], loaded_state[key], rtol=0, atol=0) + if scope == "omni": + modules = {runtime.module_name: runtime for runtime in restored} + modules["frozen"] = SimpleNamespace(optimizer=None, checkpoint=None, has_trainable_parameters=False) + trainer.model = OmniModelRuntime(SimpleNamespace(), module_runtimes=modules) + else: + trainer.model = restored[0] + trainer.args.train.checkpoint.load_path = layout.step_dir(directory, 1) + trainer.state = TrainerState() + assert GlobalStateCallback(trainer).load_global_state()["global_step"] == 1 dist.barrier() finally: clear_parallel_state() dist.destroy_process_group() +@pytest.mark.parametrize("scope", ["single", "omni"]) @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_real_multi_dcp_hf_epoch_on_two_ranks(tmp_path): +def test_real_multi_dcp_hf_epoch_on_two_ranks(tmp_path, scope): from tests.tools.launch_utils import spawn_with_rendezvous - spawn_with_rendezvous(_real_multi_save_worker, args=(str(tmp_path),), nprocs=2) + spawn_with_rendezvous(_real_multi_save_worker, args=(str(tmp_path), scope), nprocs=2) diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 2947d4839c..b01ff0b108 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -33,7 +33,6 @@ import marshal import math import types -import weakref from typing import TYPE_CHECKING, Any, Dict, Optional import torch.distributed as dist @@ -204,7 +203,6 @@ def __init__(self, runtime: "VeOmniModelRuntime"): self._loaded_optimizer_global_step = None self._loaded_optimizer_generation = None self._last_saved_optimizer_generation = None - self._last_saved_multi_optimizer_ref = None self.checkpointer: CheckpointerBase = build_checkpointer( ckpt_manager=self.config.manager, dist_backend=runtime.args.accelerator.fsdp_config.fsdp_mode, @@ -328,7 +326,6 @@ def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: def load(self) -> None: """Restore model, optimizer and model-bound extra state from ``load_path``.""" - self._last_saved_multi_optimizer_ref = None self._loaded_optimizer_global_step = None self._loaded_optimizer_generation = None load_dir = self.load_dir() @@ -420,7 +417,6 @@ def save_dcp(self, state: "TrainerState") -> None: dataloader is, the rng — has its own writer. """ self._last_saved_optimizer_generation = None - self._last_saved_multi_optimizer_ref = None if self._last_saved_step == state.global_step: # A failed rewrite may already have invalidated the old DCP marker. # A later export must not reuse the previous same-step receipt. @@ -474,17 +470,6 @@ def clear_cache(stage): clear_cache("clearing cache after MultiOptimizer checkpoint save") dist.barrier() if multi: - error = ( - RuntimeError("MultiOptimizer changed while saving its checkpoint") - if self.runtime.optimizer is not optimizer - else None - ) - group = dist.group.WORLD if dist.is_initialized() else None - raise_if_any_rank_failed(error, "checking saved MultiOptimizer identity", group=group) - try: - self._last_saved_multi_optimizer_ref = weakref.ref(optimizer) - except TypeError: - self._last_saved_multi_optimizer_ref = None self._last_saved_optimizer_generation = generation self._last_saved_step = state.global_step logger.info_rank0(f"Distributed checkpoint saved at {self.save_dir(state)} successfully!") diff --git a/veomni/trainer/callbacks/checkpoint_callback.py b/veomni/trainer/callbacks/checkpoint_callback.py index 184d1284c3..fcad79cde0 100644 --- a/veomni/trainer/callbacks/checkpoint_callback.py +++ b/veomni/trainer/callbacks/checkpoint_callback.py @@ -30,11 +30,8 @@ :mod:`~veomni.trainer.callbacks.global_state_callback`. """ -from collections.abc import Mapping from typing import TYPE_CHECKING -import torch.distributed as dist - from ...utils import helper from .base import Callback, TrainerState @@ -46,47 +43,6 @@ logger = helper.create_logger(__name__) -def _single_multi_checkpoint(trainer): - """Limit the same-step receipt to the synchronous single-model route.""" - runtime = getattr(trainer, "model", None) - if isinstance(getattr(runtime, "module_runtimes", None), Mapping): - return None - optimizer = getattr(runtime, "optimizer", None) - if getattr(optimizer, "_is_multi_optimizer", False) is not True: - return None - return runtime, optimizer, getattr(runtime, "checkpoint", None) - - -def _agree_multi_value(value): - if not dist.is_initialized(): - return [value] - values = [None] * dist.get_world_size() - dist.all_gather_object(values, value) - return values - - -def _current_multi_receipt(runtime, optimizer, checkpoint, step): - from ...models.checkpoint.manager import ModelCheckpointManager - - if ( - not isinstance(checkpoint, ModelCheckpointManager) - or getattr(checkpoint, "runtime", None) is not runtime - or getattr(runtime, "optimizer", None) is not optimizer - ): - return None - ref = getattr(checkpoint, "_last_saved_multi_optimizer_ref", None) - generation = getattr(checkpoint, "_last_saved_optimizer_generation", None) - if ( - getattr(checkpoint, "_last_saved_step", None) != step - or not callable(ref) - or ref() is not optimizer - or not isinstance(generation, str) - or not generation - ): - return None - return generation - - class CheckpointCallback(Callback): """Schedule DCP / HF / LoRA I/O and the one-shot tokenizer/config export.""" @@ -101,8 +57,20 @@ def __init__(self, trainer: "BaseTrainer"): self.hf_every_n_epochs = ckpt.hf_save_epochs self._last_dcp_step: int = -1 self._last_hf_step: int = -1 - self._last_multi_receipt = None - self._uses_single_multi_checkpoint = _single_multi_checkpoint(trainer) is not None + self._async_save_failed = False + + def _check_async_save(self): + if self._async_save_failed: + raise RuntimeError("A previous asynchronous checkpoint save failed; restart from a complete checkpoint") + + def _wait_for_pending_save(self): + try: + self.trainer.wait_for_pending_save() + except BaseException: + if getattr(self.trainer.args.train.checkpoint, "save_async", False): + self._async_save_failed = True + self._last_dcp_step = -1 + raise def on_train_begin(self, state: TrainerState, **kwargs) -> None: self.trainer.save_model_assets() @@ -110,6 +78,7 @@ def on_train_begin(self, state: TrainerState, **kwargs) -> None: helper.empty_cache() def on_train_end(self, state: TrainerState, **kwargs) -> None: + self._check_async_save() if self.save_hf_weights: if state.global_step != self._last_hf_step: self._save_hf(state, stage="train_end") @@ -118,30 +87,33 @@ def on_train_end(self, state: TrainerState, **kwargs) -> None: f"Skipping duplicate HF checkpoint save at train_end (global_step {state.global_step} " f"already saved)." ) - self.trainer.wait_for_pending_save() + self._wait_for_pending_save() else: - self.trainer.wait_for_pending_save() + self._wait_for_pending_save() def on_step_end(self, state: TrainerState, **kwargs): - if self.dcp_every_n_steps and state.global_step % self.dcp_every_n_steps == 0: + self._check_async_save() + dcp_due = self.dcp_every_n_steps and state.global_step % self.dcp_every_n_steps == 0 + hf_due = self.save_hf_weights and self.hf_every_n_steps and state.global_step % self.hf_every_n_steps == 0 + if dcp_due and state.global_step != self._last_dcp_step: self._save_dcp(state) - if self.save_hf_weights and self.hf_every_n_steps and state.global_step % self.hf_every_n_steps == 0: + if hf_due: + # Export needs a same-step DCP. Schedule it through the trainer so + # all model owners share the receipt used by the epoch cadence. + if self.dcp_every_n_epochs and state.global_step != self._last_dcp_step: + if getattr(self.trainer.args.train.checkpoint, "save_async", False): + self._wait_for_pending_save() + self._save_dcp(state) self._save_hf(state) def on_epoch_end(self, state: TrainerState, **kwargs): + self._check_async_save() if self.dcp_every_n_epochs and (state.epoch + 1) % self.dcp_every_n_epochs == 0: - if self._uses_single_multi_checkpoint: - multi = _single_multi_checkpoint(self.trainer) - recorded = self._last_multi_receipt if self._last_dcp_step == state.global_step else None - current = _current_multi_receipt(*multi, state.global_step) if multi is not None else None - values = _agree_multi_value((state.global_step, recorded, current)) - if any(value[1] is not None for value in values) and ( - any(value != values[0] for value in values) or recorded is None or recorded != current - ): - raise RuntimeError("Same-step MultiOptimizer checkpoint receipt changed across ranks") if state.global_step != self._last_dcp_step: self._save_dcp(state) else: + if getattr(self.trainer.args.train.checkpoint, "save_async", False): + self._wait_for_pending_save() logger.info_rank0( f"Skipping duplicate checkpoint save at epoch_end (global_step {state.global_step} " f"already saved at step_end)." @@ -156,38 +128,23 @@ def on_epoch_end(self, state: TrainerState, **kwargs): ) def _save_dcp(self, state: TrainerState): - self.trainer.save_dcp(state) + try: + self.trainer.save_dcp(state) + except BaseException: + if getattr(self.trainer.args.train.checkpoint, "save_async", False): + self._async_save_failed = True + self._last_dcp_step = -1 + raise self._last_dcp_step = state.global_step - self._last_multi_receipt = None def _save_hf(self, state: TrainerState, stage: str = "step_end"): - if self._uses_single_multi_checkpoint: - from ...models.checkpoint.manager import ModelCheckpointManager - - multi = _single_multi_checkpoint(self.trainer) - runtime, optimizer, checkpoint = multi if multi is not None else (None, None, None) - known_owner = ( - isinstance(checkpoint, ModelCheckpointManager) - and checkpoint.runtime is runtime - and not getattr(checkpoint.config, "save_async", False) - ) - needs_save = known_owner and checkpoint._last_saved_step != state.global_step - decisions = _agree_multi_value((state.global_step, known_owner, needs_save)) - if not known_owner or any(decision != decisions[0] for decision in decisions): - raise RuntimeError("MultiOptimizer export save decision differs across ranks") - if needs_save: - self.trainer.save_dcp(state) - checkpoint = getattr(runtime, "checkpoint", None) - generation = _current_multi_receipt(runtime, optimizer, checkpoint, state.global_step) - receipts = _agree_multi_value((state.global_step, generation)) - if generation is None or any(receipt != receipts[0] for receipt in receipts): - raise RuntimeError("MultiOptimizer export has no matching same-step checkpoint receipt") - self._last_dcp_step = state.global_step - self._last_multi_receipt = generation - # At train end the exporter may release the optimizer before the HF - # gather. Do not keep it alive through local strong references. - del optimizer, multi - self.trainer.save_hf_or_lora(state, stage=stage) + try: + self.trainer.save_hf_or_lora(state, stage=stage) + except BaseException: + if getattr(self.trainer.args.train.checkpoint, "save_async", False): + self._async_save_failed = True + self._last_dcp_step = -1 + raise self._last_hf_step = state.global_step From c917429cf995de30755196ab3951a34ae2e1c4e5 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:57:52 +0800 Subject: [PATCH 048/120] Reuse completed exports across empty epochs --- .../checkpoints/test_multi_export_receipt.py | 43 +++++++++++++------ .../trainer/callbacks/checkpoint_callback.py | 19 ++++++++ 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/tests/checkpoints/test_multi_export_receipt.py b/tests/checkpoints/test_multi_export_receipt.py index eb88f6d7f7..62eabf3e98 100644 --- a/tests/checkpoints/test_multi_export_receipt.py +++ b/tests/checkpoints/test_multi_export_receipt.py @@ -20,7 +20,7 @@ from veomni.models.seed_omni.utils.checkpoint import OmniModuleCheckpointManager from veomni.optim.optimizer import MultiOptimizer from veomni.trainer.callbacks.base import TrainerState -from veomni.trainer.callbacks.checkpoint_callback import CheckpointCallback +from veomni.trainer.callbacks.checkpoint_callback import CheckpointCallback, _export_wrote_dcp from veomni.trainer.callbacks.global_state_callback import GlobalStateCallback @@ -40,6 +40,16 @@ def _runtime(tmp_path): return runtime +def test_export_does_not_reuse_an_unknown_owner(): + saved = SimpleNamespace(checkpoint=SimpleNamespace(_last_saved_step=1)) + frozen = SimpleNamespace(checkpoint=None, has_trainable_parameters=False) + unknown = SimpleNamespace(checkpoint=None) + trainer = SimpleNamespace(model=SimpleNamespace(module_runtimes={"saved": saved, "frozen": frozen})) + assert _export_wrote_dcp(trainer, 1) + trainer.model.module_runtimes["unknown"] = unknown + assert not _export_wrote_dcp(trainer, 1) + + def test_train_end_does_not_retain_optimizer_during_hf_export(tmp_path): runtime = _runtime(tmp_path) optimizer_ref = weakref.ref(runtime.optimizer) @@ -62,7 +72,7 @@ def export(state, stage="step_end"): CheckpointCallback(trainer)._save_hf(state, stage="train_end") -def _real_multi_save_worker(rank, directory, scope, endpoint): +def _real_multi_save_worker(rank, directory, scope, cadence, endpoint): from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state from veomni.models.seed_omni.accelerated.omni_model.omni_model_runtime import OmniModelRuntime @@ -92,9 +102,9 @@ def make_runtime(name=None): trainer.args.train.global_rank = rank trainer.args.train.world_size = 2 trainer.args.train.checkpoint.save_steps = 0 - trainer.args.train.checkpoint.save_epochs = 1 - trainer.args.train.checkpoint.hf_save_steps = 1 - trainer.args.train.checkpoint.hf_save_epochs = 0 + trainer.args.train.checkpoint.save_epochs = 1 if cadence == "step" else 2 + trainer.args.train.checkpoint.hf_save_steps = 1 if cadence == "step" else 0 + trainer.args.train.checkpoint.hf_save_epochs = 0 if cadence == "step" else 1 trainer.train_dataloader = None trainer.data_iterator = None trainer.environ_meter.state_dict.return_value = {} @@ -121,8 +131,12 @@ def make_runtime(name=None): callback = CheckpointCallback(trainer) global_callback = GlobalStateCallback(trainer) state = TrainerState(global_step=1, epoch=0) - callback.on_step_end(state) - global_callback.on_step_end(state) + if cadence == "step": + callback.on_step_end(state) + global_callback.on_step_end(state) + else: + callback.on_epoch_end(state) + global_callback.on_epoch_end(state) generations = [runtime.checkpoint._last_saved_optimizer_generation for runtime in runtimes] assert all(generations) snapshot = None @@ -132,9 +146,13 @@ def make_runtime(name=None): p.relative_to(root): hashlib.sha256(p.read_bytes()).digest() for p in root.rglob("*") if p.is_file() } assert snapshot - callback.on_epoch_end(state) - global_callback.on_epoch_end(state) - trainer.save_dcp.assert_called_once_with(state) + next_state = state if cadence == "step" else TrainerState(global_step=1, epoch=1) + callback.on_epoch_end(next_state) + global_callback.on_epoch_end(next_state) + if cadence == "step": + trainer.save_dcp.assert_called_once_with(state) + else: + trainer.save_dcp.assert_not_called() for runtime in runtimes: runtime.checkpoint.save_dcp.assert_called_once_with(state) assert [runtime.checkpoint._last_saved_optimizer_generation for runtime in runtimes] == generations @@ -170,8 +188,9 @@ def make_runtime(name=None): @pytest.mark.parametrize("scope", ["single", "omni"]) +@pytest.mark.parametrize("cadence", ["step", "empty_epoch"]) @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_real_multi_dcp_hf_epoch_on_two_ranks(tmp_path, scope): +def test_real_multi_dcp_hf_epoch_on_two_ranks(tmp_path, scope, cadence): from tests.tools.launch_utils import spawn_with_rendezvous - spawn_with_rendezvous(_real_multi_save_worker, args=(str(tmp_path), scope), nprocs=2) + spawn_with_rendezvous(_real_multi_save_worker, args=(str(tmp_path), scope, cadence), nprocs=2) diff --git a/veomni/trainer/callbacks/checkpoint_callback.py b/veomni/trainer/callbacks/checkpoint_callback.py index fcad79cde0..dcdcfaefe7 100644 --- a/veomni/trainer/callbacks/checkpoint_callback.py +++ b/veomni/trainer/callbacks/checkpoint_callback.py @@ -30,6 +30,7 @@ :mod:`~veomni.trainer.callbacks.global_state_callback`. """ +from collections.abc import Mapping from typing import TYPE_CHECKING from ...utils import helper @@ -43,6 +44,22 @@ logger = helper.create_logger(__name__) +def _export_wrote_dcp(trainer, step): + """Reuse a completed export only when each trainable owner saved this step.""" + root = getattr(trainer, "model", None) + modules = getattr(root, "module_runtimes", None) + runtimes = list(modules.values()) if isinstance(modules, Mapping) and modules else [root] + owners = [] + for runtime in runtimes: + checkpoint = getattr(runtime, "checkpoint", None) + if checkpoint is None and getattr(runtime, "has_trainable_parameters", None) is False: + continue + if checkpoint is None or getattr(checkpoint, "_last_saved_step", None) != step: + return False + owners.append(checkpoint) + return bool(owners) + + class CheckpointCallback(Callback): """Schedule DCP / HF / LoRA I/O and the one-shot tokenizer/config export.""" @@ -146,6 +163,8 @@ def _save_hf(self, state: TrainerState, stage: str = "step_end"): self._last_dcp_step = -1 raise self._last_hf_step = state.global_step + if _export_wrote_dcp(self.trainer, state.global_step): + self._last_dcp_step = state.global_step __all__ = ["CheckpointCallback"] From 228c0af95b7c7f151a2c30dbaccb2bf5c8f2d067 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 15:29:36 +0800 Subject: [PATCH 049/120] Declare expert-only fixture without GDN layers --- tests/ci/_npu_moe_fp32_cold_worker.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/ci/_npu_moe_fp32_cold_worker.py b/tests/ci/_npu_moe_fp32_cold_worker.py index ee1e67993f..dd8f4c073f 100644 --- a/tests/ci/_npu_moe_fp32_cold_worker.py +++ b/tests/ci/_npu_moe_fp32_cold_worker.py @@ -54,7 +54,8 @@ class ToyQwen4(nn.Module): def __init__(self): super().__init__() - self.config = SimpleNamespace(model_type="qwen4_exp") + # This expert-only fixture deliberately has no GDN layers. + self.config = SimpleNamespace(model_type="qwen4_exp", layer_types=["full_attention"]) self.layer = ToyDecoder() def forward(self, x, indices, weights): From a267fae8874db9ad2770a242887e4866bf779180 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 16:30:19 +0800 Subject: [PATCH 050/120] Preserve proven checkpoint transactions after cold resume --- .../test_extra_parallel_layout_v2.py | 25 -- .../checkpoints/test_multi_export_receipt.py | 323 +++++++++++++++++- .../test_multi_optimizer_proven_resume.py | 46 --- veomni/checkpoint/layout.py | 27 ++ veomni/models/checkpoint/manager.py | 14 +- .../trainer/callbacks/checkpoint_callback.py | 33 ++ .../callbacks/global_state_callback.py | 161 +++++++-- 7 files changed, 513 insertions(+), 116 deletions(-) diff --git a/tests/checkpoints/test_extra_parallel_layout_v2.py b/tests/checkpoints/test_extra_parallel_layout_v2.py index 11e90c4bc4..036d80818e 100644 --- a/tests/checkpoints/test_extra_parallel_layout_v2.py +++ b/tests/checkpoints/test_extra_parallel_layout_v2.py @@ -198,31 +198,6 @@ def test_v2_optimizer_fails_before_model_or_checkpoint_mutation(tmp_path): assert not list(tmp_path.iterdir()) -def _rank_local_invalid_layout_worker(rank, root): - from datetime import timedelta - - dist.init_process_group( - "gloo", init_method=f"file://{root}/init", rank=rank, world_size=2, timeout=timedelta(seconds=20) - ) - try: - model = torch.nn.Linear(2, 2, bias=False) - mesh = SimpleNamespace(shape=(2, 2)) - shape = None if rank == 0 else (2, 2) - model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(0), "weight", mesh, 0, None, shape, True)} - with pytest.raises(RuntimeError): - DistributedCheckpointer.save( - f"{root}/checkpoint", {"model": model}, parallel_state=SimpleNamespace(dp_mode="fsdp2") - ) - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_rank_local_invalid_layout_stops_all_ranks_before_write(tmp_path): - torch.multiprocessing.spawn(_rank_local_invalid_layout_worker, args=(str(tmp_path),), nprocs=2, join=True) - assert not (tmp_path / "checkpoint").exists() - - def test_same_v2_wrapper_can_save_and_load(tmp_path, monkeypatch): model = torch.nn.Linear(2, 2, bias=False) mesh = SimpleNamespace(shape=(2, 2)) diff --git a/tests/checkpoints/test_multi_export_receipt.py b/tests/checkpoints/test_multi_export_receipt.py index 62eabf3e98..fc99a1e0c2 100644 --- a/tests/checkpoints/test_multi_export_receipt.py +++ b/tests/checkpoints/test_multi_export_receipt.py @@ -62,6 +62,7 @@ def test_train_end_does_not_retain_optimizer_during_hf_export(tmp_path): manager = ModelCheckpointManager(runtime) runtime.checkpoint = manager trainer.model = runtime + trainer.save_dcp.side_effect = manager.save_dcp def export(state, stage="step_end"): manager.prepare_export(state, stage) @@ -72,7 +73,7 @@ def export(state, stage="step_end"): CheckpointCallback(trainer)._save_hf(state, stage="train_end") -def _real_multi_save_worker(rank, directory, scope, cadence, endpoint): +def _real_multi_save_worker(rank, directory, scope, cadence, same_output, target_exists, endpoint): from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state from veomni.models.seed_omni.accelerated.omni_model.omni_model_runtime import OmniModelRuntime @@ -83,6 +84,8 @@ def _real_multi_save_worker(rank, directory, scope, cadence, endpoint): def make_runtime(name=None): runtime = _runtime(directory) + if scope == "mixed" and name == "vision": + runtime.optimizer = torch.optim.AdamW(runtime.model.parameters()) runtime.module_name = name runtime.parallel_state = SimpleNamespace(global_rank=rank, dp_mode="ddp", extra_parallel_names=[]) with patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer): @@ -92,7 +95,7 @@ def make_runtime(name=None): runtime.save_hf_or_lora = lambda state, stage="step_end": runtime.checkpoint.prepare_export(state, stage) return runtime - names = ("language", "vision") if scope == "omni" else (None,) + names = ("vision", "language") if scope == "mixed" else ("language", "vision") if scope == "omni" else (None,) runtimes = [make_runtime(name) for name in names] for runtime in runtimes: runtime.model(torch.ones(1, 2)).sum().backward() @@ -110,7 +113,7 @@ def make_runtime(name=None): trainer.environ_meter.state_dict.return_value = {} trainer.channel_loss_callback = None trainer.offline_cache_writer = None - if scope == "omni": + if scope in ("omni", "mixed"): modules = {runtime.module_name: runtime for runtime in runtimes} modules["frozen"] = SimpleNamespace( optimizer=None, @@ -138,7 +141,7 @@ def make_runtime(name=None): callback.on_epoch_end(state) global_callback.on_epoch_end(state) generations = [runtime.checkpoint._last_saved_optimizer_generation for runtime in runtimes] - assert all(generations) + assert any(generations) and (scope != "omni" or all(generations)) snapshot = None if rank == 0: root = Path(directory) @@ -149,10 +152,7 @@ def make_runtime(name=None): next_state = state if cadence == "step" else TrainerState(global_step=1, epoch=1) callback.on_epoch_end(next_state) global_callback.on_epoch_end(next_state) - if cadence == "step": - trainer.save_dcp.assert_called_once_with(state) - else: - trainer.save_dcp.assert_not_called() + trainer.save_dcp.assert_called_once_with(state) for runtime in runtimes: runtime.checkpoint.save_dcp.assert_called_once_with(state) assert [runtime.checkpoint._last_saved_optimizer_generation for runtime in runtimes] == generations @@ -161,6 +161,43 @@ def make_runtime(name=None): p.relative_to(root): hashlib.sha256(p.read_bytes()).digest() for p in root.rglob("*") if p.is_file() } + target_snapshot = None + if target_exists: + target_path = Path(directory) / "new-output" + if target_exists == "stat_error": + if rank == 0: + target_path.mkdir(parents=True, exist_ok=True) + (target_path / "sentinel").write_bytes(b"keep") + dist.barrier() + else: + torch.manual_seed(456) + target_runtimes = [make_runtime(name) for name in names] + for runtime in target_runtimes: + runtime.train_args.checkpoint.save_path = str(target_path) + runtime.model(torch.ones(1, 2)).sum().backward() + runtime.optimizer.step() + runtime.optimizer.zero_grad() + runtime.checkpoint.save_dcp(state) + target_trainer = _make_mock_trainer(save_path=str(target_path)) + target_trainer.args.train.global_rank = rank + target_trainer.args.train.world_size = 2 + target_trainer.train_dataloader = None + target_trainer.data_iterator = None + target_trainer.environ_meter.state_dict.return_value = {} + target_trainer.channel_loss_callback = None + target_trainer.offline_cache_writer = None + target_trainer.model = OmniModelRuntime( + SimpleNamespace(), module_runtimes={runtime.module_name: runtime for runtime in target_runtimes} + ) + GlobalStateCallback(target_trainer).save_global_state(state) + if rank == 0: + target_snapshot = { + p.relative_to(target_path): hashlib.sha256(p.read_bytes()).digest() + for p in target_path.rglob("*") + if p.is_file() + } + assert target_snapshot + restored = [make_runtime(name) for name in names] for before, after, generation in zip(runtimes, restored, generations): after.train_args.checkpoint.load_path = layout.step_dir(directory, 1) @@ -168,29 +205,283 @@ def make_runtime(name=None): assert after.checkpoint._loaded_optimizer_generation == generation for original_param, loaded_param in zip(before.model.parameters(), after.model.parameters()): torch.testing.assert_close(original_param, loaded_param, rtol=0, atol=0) - original_state = before.optimizer.optimizers_dict["model"].state[original_param] - loaded_state = after.optimizer.optimizers_dict["model"].state[loaded_param] + original_optimizer = getattr(before.optimizer, "optimizers_dict", {}).get("model", before.optimizer) + loaded_optimizer = getattr(after.optimizer, "optimizers_dict", {}).get("model", after.optimizer) + original_state = original_optimizer.state[original_param] + loaded_state = loaded_optimizer.state[loaded_param] for key in ("step", "exp_avg", "exp_avg_sq"): torch.testing.assert_close(original_state[key], loaded_state[key], rtol=0, atol=0) - if scope == "omni": + if scope in ("omni", "mixed"): modules = {runtime.module_name: runtime for runtime in restored} - modules["frozen"] = SimpleNamespace(optimizer=None, checkpoint=None, has_trainable_parameters=False) + modules["frozen"] = SimpleNamespace( + optimizer=None, + checkpoint=None, + has_trainable_parameters=False, + save_dcp=lambda state: None, + save_hf_or_lora=lambda state, stage: None, + ) trainer.model = OmniModelRuntime(SimpleNamespace(), module_runtimes=modules) else: trainer.model = restored[0] trainer.args.train.checkpoint.load_path = layout.step_dir(directory, 1) + if not same_output: + new_output = str(Path(directory) / "new-output") + trainer.args.train.checkpoint.save_path = new_output + for runtime in restored: + runtime.train_args.checkpoint.save_path = new_output trainer.state = TrainerState() - assert GlobalStateCallback(trainer).load_global_state()["global_step"] == 1 + trainer.checkpoint_callback = CheckpointCallback(trainer) + resume_state = GlobalStateCallback(trainer) + assert resume_state.load_global_state()["global_step"] == 1 + for runtime in restored: + runtime.checkpoint.save_dcp = Mock(wraps=runtime.checkpoint.save_dcp) + trainer.save_dcp.side_effect = trainer.model.save_dcp + trainer.save_hf_or_lora.side_effect = trainer.model.save_hf_or_lora + resume_step = TrainerState(global_step=1, epoch=1) + if target_exists: + import errno + import os + + original_stat = os.stat + manifest = layout.manifest_path(layout.step_dir(str(target_path), 1)) + + def stat_with_error(path, *args, **kwargs): + if target_exists == "stat_error" and rank == 1 and str(path) == manifest: + raise OSError(errno.EIO, "injected manifest stat error") + return original_stat(path, *args, **kwargs) + + with ( + patch("veomni.trainer.callbacks.checkpoint_callback.os.stat", side_effect=stat_with_error), + pytest.raises( + (OSError, RuntimeError), + match="Cannot overwrite a completed checkpoint|failed on another rank|manifest stat error", + ) as caught, + ): + if target_exists == "dcp": + trainer.checkpoint_callback.dcp_every_n_steps = 1 + trainer.checkpoint_callback.on_step_end(resume_step) + else: + trainer.checkpoint_callback.on_train_end(resume_step) + if target_exists == "stat_error": + expected_error = "manifest stat error" if rank == 1 else "failed on another rank" + assert expected_error in str(caught.value) + for runtime in restored: + runtime.checkpoint.save_dcp.assert_not_called() + if rank == 0: + assert target_snapshot == { + p.relative_to(target_path): hashlib.sha256(p.read_bytes()).digest() + for p in target_path.rglob("*") + if p.is_file() + } + dist.barrier() + return + trainer.checkpoint_callback.on_train_end(resume_step) + resume_state.on_train_end(resume_step) + for runtime in restored: + if same_output: + runtime.checkpoint.save_dcp.assert_not_called() + else: + runtime.checkpoint.save_dcp.assert_called_once_with(resume_step) + if same_output and rank == 0: + assert snapshot == { + p.relative_to(root): hashlib.sha256(p.read_bytes()).digest() for p in root.rglob("*") if p.is_file() + } dist.barrier() finally: clear_parallel_state() dist.destroy_process_group() -@pytest.mark.parametrize("scope", ["single", "omni"]) +@pytest.mark.parametrize("scope", ["single", "omni", "mixed"]) @pytest.mark.parametrize("cadence", ["step", "empty_epoch"]) +@pytest.mark.parametrize("same_output", [True, False]) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_real_multi_dcp_hf_epoch_on_two_ranks(tmp_path, scope, cadence, same_output): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_real_multi_save_worker, args=(str(tmp_path), scope, cadence, same_output, False), nprocs=2) + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +@pytest.mark.parametrize("entry", ["dcp", "hf", "stat_error"]) +def test_mixed_omni_existing_target_is_rejected_before_any_owner_writes(tmp_path, entry): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_real_multi_save_worker, args=(str(tmp_path), "mixed", "step", False, entry), nprocs=2) + + +def _fresh_process_export_worker(rank, root, phase, endpoint): + import json + + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + + class StatefulLoader: + def __init__(self, position): + self.position = position + + def state_dict(self): + return {"position": self.position} + + def load_state_dict(self, state): + self.position = state["position"] + + _init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cpu") + torch.manual_seed(123) + directory = Path(root) / "checkpoint-output" + runtime = _runtime(directory) + runtime.parallel_state = SimpleNamespace(global_rank=rank, dp_mode="ddp", extra_parallel_names=[]) + with patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer): + manager = ModelCheckpointManager(runtime) + runtime.checkpoint = manager + trainer = _make_mock_trainer(save_path=str(directory)) + trainer.args.train.global_rank = rank + trainer.args.train.world_size = 2 + trainer.args.train.checkpoint.save_steps = 0 + trainer.args.train.checkpoint.save_epochs = 1 + trainer.args.train.checkpoint.hf_save_steps = 1 + trainer.args.train.checkpoint.hf_save_epochs = 0 + trainer.model = runtime + trainer.train_dataloader = StatefulLoader(7 + rank if phase == "save" else -1) + trainer.data_iterator = None + trainer.environ_meter.state_dict.return_value = {} + trainer.channel_loss_callback = None + trainer.offline_cache_writer = None + trainer.save_dcp.side_effect = manager.save_dcp + trainer.save_hf_or_lora.side_effect = lambda state, stage="step_end": manager.prepare_export(state, stage) + trainer.checkpoint_callback = CheckpointCallback(trainer) + global_callback = GlobalStateCallback(trainer) + state = TrainerState(global_step=1, epoch=0) + + def snapshot_optimizer(): + leaf = runtime.optimizer.optimizers_dict["model"] + return { + name: { + "parameter": parameter.detach().clone(), + **{key: leaf.state[parameter][key].detach().clone() for key in ("step", "exp_avg", "exp_avg_sq")}, + } + for name, parameter in runtime.model.named_parameters() + } + + def assert_snapshot(expected): + current = snapshot_optimizer() + assert current.keys() == expected.keys() + for name in current: + for key in current[name]: + torch.testing.assert_close(current[name][key], expected[name][key], rtol=0, atol=0) + + def next_update(): + runtime.optimizer.zero_grad() + runtime.model(torch.full((1, 2), 2.0)).sum().backward() + runtime.optimizer.step() + + if phase == "save": + runtime.model(torch.ones(1, 2)).sum().backward() + runtime.optimizer.step() + runtime.optimizer.zero_grad() + trainer.checkpoint_callback.on_step_end(state) + global_callback.on_step_end(state) + if rank == 0: + hashes = { + str(p.relative_to(directory)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in directory.rglob("*") + if p.is_file() + } + (Path(root) / "saved-hashes.json").write_text(json.dumps(hashes, sort_keys=True)) + before = snapshot_optimizer() + next_update() + torch.save({"before": before, "after": snapshot_optimizer()}, Path(root) / "expected.pt") + else: + runtime.train_args.checkpoint.load_path = layout.step_dir(str(directory), 1) + trainer.args.train.checkpoint.load_path = runtime.train_args.checkpoint.load_path + manager.load() + trainer.state = TrainerState() + saved_sidecar = torch.load( + layout.extra_state_path(runtime.train_args.checkpoint.load_path, rank), weights_only=False + ) + torch.manual_seed(777 + rank) + assert not torch.equal(torch.get_rng_state(), saved_sidecar["torch_rng_state"]) + if phase in ("resume_bad_rng", "resume_missing_cursor"): + with pytest.raises( + RuntimeError, match="conflicting CPU RNG|proven data cursor|failed on another rank" + ): + global_callback.load_global_state() + assert trainer.checkpoint_callback._last_dcp_step == -1 + assert manager._last_saved_step == -1 + dist.barrier() + return + assert global_callback.load_global_state()["global_step"] == 1 + assert torch.equal(torch.get_rng_state(), saved_sidecar["torch_rng_state"]) + assert trainer.train_dataloader.position == 7 + rank + expected = torch.load(Path(root) / "expected.pt", weights_only=True) + assert_snapshot(expected["before"]) + if phase == "resume_update": + next_update() + assert_snapshot(expected["after"]) + else: + manager.save_dcp = Mock(wraps=manager.save_dcp) + trainer.checkpoint_callback.on_train_end(state) + global_callback.on_train_end(state) + manager.save_dcp.assert_not_called() + if rank == 0: + before = json.loads((Path(root) / "saved-hashes.json").read_text()) + after = { + str(p.relative_to(directory)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in directory.rglob("*") + if p.is_file() + } + assert before == after + dist.barrier() + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_fresh_process_same_output_final_export_reuses_proven_dcp(tmp_path): + from tests.tools.launch_utils import spawn_with_rendezvous + + for phase in ("save", "resume_export", "resume_update"): + spawn_with_rendezvous(_fresh_process_export_worker, args=(str(tmp_path), phase), nprocs=2) + sidecar = Path(layout.extra_state_path(layout.step_dir(str(tmp_path / "checkpoint-output"), 1), 1)) + payload = torch.load(sidecar, weights_only=False) + payload["device_rng_state"] = payload["torch_rng_state"].clone() + payload["device_rng_state"][0] ^= 1 + torch.save(payload, sidecar) + spawn_with_rendezvous(_fresh_process_export_worker, args=(str(tmp_path), "resume_bad_rng"), nprocs=2) + payload["device_rng_state"] = payload["torch_rng_state"].clone() + torch.save(payload, sidecar) + loader_file = Path(layout.loader_path(layout.step_dir(str(tmp_path / "checkpoint-output"), 1), 1)) + torch.save({}, loader_file) + spawn_with_rendezvous(_fresh_process_export_worker, args=(str(tmp_path), "resume_missing_cursor"), nprocs=2) + + +def _ordinary_restore_failure_worker(rank, root, endpoint): + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + _init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cpu") + runtime = _runtime(root) + runtime.optimizer = torch.optim.AdamW(runtime.model.parameters()) + runtime.train_args.checkpoint.load_path = str(root) + if rank == 1: + runtime.load_extra_state.side_effect = OSError("local scheduler read failed") + with patch("veomni.models.checkpoint.manager.build_checkpointer"): + manager = ModelCheckpointManager(runtime) + with pytest.raises((OSError, RuntimeError), match="local scheduler read failed|failed on another rank"): + manager.load() + assert manager._loaded_checkpoint_path is None + dist.barrier() + finally: + clear_parallel_state() + dist.destroy_process_group() + + @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_real_multi_dcp_hf_epoch_on_two_ranks(tmp_path, scope, cadence): +def test_ordinary_owner_restore_failure_reaches_every_rank(tmp_path): from tests.tools.launch_utils import spawn_with_rendezvous - spawn_with_rendezvous(_real_multi_save_worker, args=(str(tmp_path), scope, cadence), nprocs=2) + spawn_with_rendezvous(_ordinary_restore_failure_worker, args=(str(tmp_path),), nprocs=2) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 29cc54a9ae..35cda79036 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -2014,49 +2014,3 @@ def test_direct_multi_checkpoint_applies_returned_scheduler_state(tmp_path): for actual, expected in zip(target.parameters(), model.parameters(), strict=True): torch.testing.assert_close(actual, expected, rtol=0, atol=0) assert scheduler.state_dict() == target_scheduler.state_dict() - - -def _rank_scheduler_declaration_worker(rank, root, endpoint): - from datetime import timedelta - - import torch.distributed as dist - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - torch.manual_seed(42) - model, optimizer = _model_and_optimizer() - parallel = SimpleNamespace(dp_mode="ddp") - state = {"model": model, "optimizer": optimizer} - if rank == 1: - state["_veomni_scheduler_object"] = None - message = "explicit scheduler declaration" if rank == 0 else "failed on another rank" - with pytest.raises(RuntimeError, match=message): - DistributedCheckpointer.save(f"{root}/missing", state, parallel_state=parallel) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - assert not (Path(root) / "missing").exists() - state["_veomni_scheduler_object"] = None - DistributedCheckpointer.save(f"{root}/good", state, parallel_state=parallel) - target, restored = _model_and_optimizer() - before = target.weight.detach().clone() - load_state = {"model": target, "optimizer": restored} - if rank == 0: - load_state["_veomni_scheduler_object"] = None - message = "explicit scheduler declaration" if rank == 1 else "failed on another rank" - with pytest.raises(RuntimeError, match=message): - DistributedCheckpointer.load(f"{root}/good", load_state, parallel_state=parallel) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - torch.testing.assert_close(target.weight, before, rtol=0, atol=0) - assert all(not child.state for child in restored.optimizers_dict.values()) - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -def test_missing_scheduler_declaration_fails_all_ranks_before_mutation(tmp_path): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_rank_scheduler_declaration_worker, args=(str(tmp_path),), nprocs=2) diff --git a/veomni/checkpoint/layout.py b/veomni/checkpoint/layout.py index f0032aeb8a..000107a5bf 100644 --- a/veomni/checkpoint/layout.py +++ b/veomni/checkpoint/layout.py @@ -33,11 +33,38 @@ import json import os +import posixpath from typing import Any, Sequence +from urllib.parse import urlsplit, urlunsplit from ..utils.checkpoint_utils import _GLOBAL_STEP_PREFIX +def same_checkpoint_path(left, right) -> bool: + """Compare logical URIs or existing local directories, including symlinks.""" + + def normalize(path): + raw = os.fspath(path) + if not isinstance(raw, str): + raise TypeError("Checkpoint paths must be strings") + parsed = urlsplit(raw) + if parsed.scheme and parsed.netloc: + return urlunsplit( + (parsed.scheme.lower(), parsed.netloc, posixpath.normpath(parsed.path), parsed.query, parsed.fragment) + ) + return os.path.realpath(raw) + + left_normalized, right_normalized = normalize(left), normalize(right) + if left_normalized == right_normalized: + return True + left_uri, right_uri = urlsplit(os.fspath(left)), urlsplit(os.fspath(right)) + if (left_uri.scheme and left_uri.netloc) or (right_uri.scheme and right_uri.netloc): + return False + if os.path.exists(left_normalized) and os.path.exists(right_normalized): + return os.path.samefile(left_normalized, right_normalized) + return False + + # Resume tree. Everything needed to continue training, and nothing else. MODEL_DIRNAME = "model" WEIGHTS_DIRNAME = "ckpt" diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index b01ff0b108..ea10e31246 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -203,6 +203,7 @@ def __init__(self, runtime: "VeOmniModelRuntime"): self._loaded_optimizer_global_step = None self._loaded_optimizer_generation = None self._last_saved_optimizer_generation = None + self._loaded_checkpoint_path = None self.checkpointer: CheckpointerBase = build_checkpointer( ckpt_manager=self.config.manager, dist_backend=runtime.args.accelerator.fsdp_config.fsdp_mode, @@ -326,6 +327,9 @@ def _load_extra_state(self, extra_state: Dict[str, Any]) -> None: def load(self) -> None: """Restore model, optimizer and model-bound extra state from ``load_path``.""" + self._last_saved_step = -1 + self._last_saved_optimizer_generation = None + self._loaded_checkpoint_path = None self._loaded_optimizer_global_step = None self._loaded_optimizer_generation = None load_dir = self.load_dir() @@ -345,10 +349,7 @@ def load(self) -> None: self.checkpointer.load(load_dir, state, **self._checkpointer_kwargs()) except BaseException as exc: error = exc - if multi: - raise_if_any_rank_failed(error, "loading MultiOptimizer checkpoint") - elif error is not None: - raise error + raise_if_any_rank_failed(error, "loading checkpoint") error, step = None, None try: if multi: @@ -380,8 +381,8 @@ def load(self) -> None: self._load_extra_state(state["extra_state"]) except BaseException as exc: error = exc + raise_if_any_rank_failed(error, "restoring trainer step and scheduler") if multi: - raise_if_any_rank_failed(error, "restoring MultiOptimizer trainer step and scheduler") generation = state.get("_veomni_optimizer_checkpoint_generation") error = ( RuntimeError("MultiOptimizer checkpoint generation is missing") @@ -389,8 +390,6 @@ def load(self) -> None: else None ) raise_if_any_rank_failed(error, "checking restored MultiOptimizer generation") - elif error is not None: - raise error try: dist.barrier() except BaseException: @@ -408,6 +407,7 @@ def load(self) -> None: if multi: self._loaded_optimizer_global_step = step self._loaded_optimizer_generation = generation + self._loaded_checkpoint_path = load_dir logger.info_rank0(f"Load distributed checkpoint from {load_dir} successfully!") def save_dcp(self, state: "TrainerState") -> None: diff --git a/veomni/trainer/callbacks/checkpoint_callback.py b/veomni/trainer/callbacks/checkpoint_callback.py index dcdcfaefe7..bd33c4bd82 100644 --- a/veomni/trainer/callbacks/checkpoint_callback.py +++ b/veomni/trainer/callbacks/checkpoint_callback.py @@ -30,10 +30,13 @@ :mod:`~veomni.trainer.callbacks.global_state_callback`. """ +import os from collections.abc import Mapping from typing import TYPE_CHECKING +from ...checkpoint import layout from ...utils import helper +from ...utils.dist_utils import raise_if_any_rank_failed from .base import Callback, TrainerState @@ -60,6 +63,33 @@ def _export_wrote_dcp(trainer, step): return bool(owners) +def _has_multi_optimizer(trainer): + root = getattr(trainer, "model", None) + modules = getattr(root, "module_runtimes", None) + runtimes = list(modules.values()) if isinstance(modules, Mapping) and modules else [root] + return any( + getattr(getattr(runtime, "optimizer", None), "_is_multi_optimizer", False) is True for runtime in runtimes + ) + + +def _check_multi_destination(trainer, step): + """Reject an existing transaction before any mixed-owner write begins.""" + if not _has_multi_optimizer(trainer): + return + error = None + try: + step_root = layout.step_dir(trainer.args.train.checkpoint.save_path, step) + try: + os.stat(layout.manifest_path(step_root)) + except FileNotFoundError: + pass + else: + raise RuntimeError("Cannot overwrite a completed checkpoint transaction with MultiOptimizer state") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking MultiOptimizer checkpoint destination") + + class CheckpointCallback(Callback): """Schedule DCP / HF / LoRA I/O and the one-shot tokenizer/config export.""" @@ -146,6 +176,7 @@ def on_epoch_end(self, state: TrainerState, **kwargs): def _save_dcp(self, state: TrainerState): try: + _check_multi_destination(self.trainer, state.global_step) self.trainer.save_dcp(state) except BaseException: if getattr(self.trainer.args.train.checkpoint, "save_async", False): @@ -156,6 +187,8 @@ def _save_dcp(self, state: TrainerState): def _save_hf(self, state: TrainerState, stage: str = "step_end"): try: + if _has_multi_optimizer(self.trainer) and self._last_dcp_step != state.global_step: + self._save_dcp(state) self.trainer.save_hf_or_lora(state, stage=stage) except BaseException: if getattr(self.trainer.args.train.checkpoint, "save_async", False): diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index daf3e5e1a4..4e4bfa171e 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -45,7 +45,7 @@ from ...checkpoint import layout from ...utils import helper -from ...utils.device import get_device_rng_state, get_device_type, set_device_rng_state +from ...utils.device import get_device_rng_state, get_device_type, get_torch_device, set_device_rng_state from ...utils.dist_utils import raise_if_any_rank_failed from .base import Callback, TrainerState @@ -117,6 +117,59 @@ def _saved_multi_optimizer_generations(trainer, step): return generations +def _same_output_resume_owners(trainer, load_path, step): + """Return verified owners only when this completed step is the output step.""" + config = trainer.args.train.checkpoint + same_output, error = False, None + try: + same_output = layout.same_checkpoint_path(load_path, layout.step_dir(config.save_path, step)) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking restored checkpoint output path") + if dist.is_initialized(): + decisions = [None] * dist.get_world_size() + dist.all_gather_object(decisions, same_output) + if any(decision != same_output for decision in decisions): + raise RuntimeError("Restored checkpoint output path differs across ranks") + if not same_output: + return [] + + owners, error = [], None + try: + from ...models.checkpoint.manager import ModelCheckpointManager + + manifest = layout.read_manifest(load_path) + if ( + not isinstance(manifest, dict) + or not layout.checkpoint_is_complete(load_path) + or manifest.get("model_state") is not True + or manifest.get("global_step") != step + or manifest.get("world_size") != trainer.args.train.world_size + ): + raise RuntimeError("Cannot reuse an incomplete restored optimizer checkpoint") + root = getattr(trainer, "model", None) + modules = getattr(root, "module_runtimes", None) + runtimes = list(modules.values()) if isinstance(modules, Mapping) and modules else [root] + for runtime in runtimes: + checkpoint = getattr(runtime, "checkpoint", None) + if checkpoint is None and getattr(runtime, "has_trainable_parameters", None) is False: + continue + loaded_path = getattr(checkpoint, "_loaded_checkpoint_path", None) + if ( + not isinstance(checkpoint, ModelCheckpointManager) + or loaded_path is None + or not layout.same_checkpoint_path(loaded_path, load_path) + ): + raise RuntimeError("Cannot reuse an unverified model checkpoint owner") + owners.append(checkpoint) + if not owners: + raise RuntimeError("Cannot reuse a checkpoint without a model owner") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "validating the restored checkpoint for reuse") + return owners + + class GlobalStateCallback(Callback): """Save and resume the state that belongs to the job rather than to a model. @@ -375,33 +428,97 @@ def load_global_state(self) -> Optional[Dict[str, Any]]: elif global_state.get("_veomni_optimizer_generations") != generations: error = RuntimeError("Trainer and MultiOptimizer checkpoint generations differ or are missing") raise_if_any_rank_failed(error, "checking the trainer checkpoint step") - self.trainer.state.global_step = global_state["global_step"] - self._restore_position(global_state) + reuse_owners = ( + _same_output_resume_owners(self.trainer, load_path, restored_step) if restored_step is not None else [] + ) + error = None + if reuse_owners: + cpu_rng = global_state.get("torch_rng_state") + device_rng = global_state.get("device_rng_state") + cursor = global_state.get("train_dataloader") + if not torch.is_tensor(cpu_rng): + error = RuntimeError("Cannot reuse a checkpoint without the trainer RNG state") + elif ( + get_device_type() == "cpu" + and device_rng is not None + and (not torch.is_tensor(device_rng) or not torch.equal(cpu_rng, device_rng)) + ): + error = RuntimeError("Cannot reuse a checkpoint with conflicting CPU RNG states") + elif get_device_type() != "cpu" and ( + device_rng is None or not callable(getattr(get_torch_device(), "set_rng_state", None)) + ): + error = RuntimeError("Cannot reuse a checkpoint without the device RNG state") + elif self.trainer.train_dataloader is not None and ( + not isinstance(cursor, dict) or (restored_step > 0 and not cursor) + ): + error = RuntimeError("Cannot reuse a checkpoint without a proven data cursor") + raise_if_any_rank_failed(error, "checking restored checkpoint RNG and cursor") + + def restore_local(action, description): + if restored_step is None: + action() + return + error = None + try: + action() + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, description) + + def restore_position(): + self.trainer.state.global_step = global_state["global_step"] + self._restore_position(global_state) + restore_local(restore_position, "restoring trainer position") channel_loss_state = global_state.get("channel_loss_callback") channel_loss_callback = getattr(self.trainer, "channel_loss_callback", None) - if channel_loss_state is not None and channel_loss_callback is not None: - channel_loss_callback.load_state_dict(channel_loss_state) - + restore_local( + lambda: channel_loss_callback.load_state_dict(channel_loss_state) + if channel_loss_state is not None and channel_loss_callback is not None + else None, + "restoring channel-loss state", + ) offline_cache_writer_state = global_state.get("offline_cache_writer") offline_cache_writer = getattr(self.trainer, "offline_cache_writer", None) - if offline_cache_writer_state is not None and offline_cache_writer is not None: - offline_cache_writer.load_state_dict(offline_cache_writer_state) - - if self.trainer.train_dataloader is not None and global_state.get("train_dataloader") is not None: - self.trainer.train_dataloader.load_state_dict(global_state["train_dataloader"]) - - # OmniTrainer checkpoints written before it had a meter store None here. - if global_state.get("environ_meter") is not None: - self.trainer.environ_meter.load_state_dict(global_state["environ_meter"]) - rng_state = global_state.get("torch_rng_state") - if rng_state is not None: - torch.set_rng_state(rng_state) - # ``.get`` keeps checkpoints written before these fields existed loadable; - # the device namespace may also not support RNG state at all. - set_device_rng_state(global_state.get("device_rng_state")) + restore_local( + lambda: offline_cache_writer.load_state_dict(offline_cache_writer_state) + if offline_cache_writer_state is not None and offline_cache_writer is not None + else None, + "restoring offline-cache state", + ) + restore_local( + lambda: self.trainer.train_dataloader.load_state_dict(global_state["train_dataloader"]) + if self.trainer.train_dataloader is not None and global_state.get("train_dataloader") is not None + else None, + "restoring the training data cursor", + ) + restore_local( + lambda: self.trainer.environ_meter.load_state_dict(global_state["environ_meter"]) + if global_state.get("environ_meter") is not None + else None, + "restoring training meter", + ) + + def restore_rng(): + rng_state = global_state.get("torch_rng_state") + if rng_state is not None: + torch.set_rng_state(rng_state) + set_device_rng_state(global_state.get("device_rng_state")) + + restore_local(restore_rng, "restoring training RNG") if self.trainer.start_step == 0 and self.trainer.train_dataloader is not None: - iter(self.trainer.train_dataloader) + restore_local(lambda: iter(self.trainer.train_dataloader), "initializing the restored data iterator") + + if reuse_owners: + for checkpoint in reuse_owners: + checkpoint._last_saved_step = restored_step + checkpoint._last_saved_optimizer_generation = getattr(checkpoint, "_loaded_optimizer_generation", None) + self._last_saved_step = restored_step + from .checkpoint_callback import CheckpointCallback + + checkpoint_callback = getattr(self.trainer, "checkpoint_callback", None) + if isinstance(checkpoint_callback, CheckpointCallback): + checkpoint_callback._last_dcp_step = restored_step logger.info_rank0( f"Restored trainer state from {state_path} (global_step={self.trainer.state.global_step}, " From 20219eb20f969052df165801be5f83ccaffdc971 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:41:45 +0800 Subject: [PATCH 051/120] Fail closed on checkpoint destination lookup errors --- docs/usage/checkpoint.md | 10 ++- .../test_multi_optimizer_proven_resume.py | 61 +++++++++++++++++++ tests/optim/test_multi_optimizer_sparse_lr.py | 53 ---------------- veomni/checkpoint/dcp_checkpointer.py | 28 ++++++--- veomni/optim/optimizer.py | 5 -- 5 files changed, 86 insertions(+), 71 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 1d074ede8a..a1284a8695 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -387,9 +387,13 @@ device tensors to CPU and can add substantial time for a large frozen model. Hooks, adapters and scheduler method wrappers are trusted: a wrapper must preserve serializable scheduler state and implement the same learning-rate behavior on both sides of a restore, and hooks must not mutate checkpoint -model, optimizer, or layout during a save/load call. The persisted recipe checks scheduler -configuration and optimizer binding; it does not inspect arbitrary wrapper -code. MultiOptimizer asynchronous saves and detected unsupported layouts fail +model, optimizer, or layout during a save/load call. A hook may transform the +serialized representation within the codec's supported inventory, dtypes, and +shapes, but exact continuation requires the restored optimizer state, +parameter-group bindings, and scheduler behavior to match; when necessary, a +matching load hook must undo the transformation. The persisted recipe does not +verify arbitrary hook or +wrapper code. MultiOptimizer asynchronous saves and detected unsupported layouts fail closed; the cross-rank leaf mismatch above is not detected. Synchronous staged saves are supported. The runtime assumes ranks use a consistent model configuration and compatible diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 35cda79036..51464d05b2 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -1898,6 +1898,67 @@ def build(shared): torch.testing.assert_close(target.p, before, rtol=0, atol=0) +def _save_destination_stat_error_worker(rank, checkpoint, location, endpoint): + import errno + import os + from datetime import timedelta + from unittest.mock import patch + + import torch.distributed as dist + + from veomni.checkpoint.layout import manifest_path, model_dir + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + root = Path(checkpoint) + if rank == 0: + root.mkdir(parents=True) + (root / "sentinel").write_bytes(b"preserve-existing-data") + dist.barrier() + target = ( + manifest_path(checkpoint) + if location == "manifest" + else os.path.join(model_dir(checkpoint, ""), "extra_state", "extra_state_rank_1.pt") + ) + original_stat = os.stat + injected = False + + def stat(path, *args, **kwargs): + nonlocal injected + if rank == 1 and isinstance(path, (str, os.PathLike)) and os.fspath(path) == target: + injected = True + raise OSError(errno.EIO, "injected checkpoint stat failure") + return original_stat(path, *args, **kwargs) + + model, optimizer = _model_and_optimizer() + with ( + patch("veomni.checkpoint.dcp_checkpointer.os.stat", side_effect=stat), + pytest.raises((OSError, RuntimeError), match="stat failure|failed on another rank") as captured, + ): + DistributedCheckpointer.save( + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + if rank == 1: + assert injected and "injected checkpoint stat failure" in str(captured.value) + else: + assert "failed on another rank" in str(captured.value) + assert (root / "sentinel").read_bytes() == b"preserve-existing-data" + assert not Path(model_dir(checkpoint, "")).exists() + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) +@pytest.mark.parametrize("location", ["manifest", "sidecar"]) +def test_multi_optimizer_save_coordinates_destination_stat_errors(tmp_path, location): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous(_save_destination_stat_error_worker, args=(str(tmp_path / "ckpt"), location), nprocs=2) + + @pytest.mark.parametrize("old_model_state", [True, False]) def test_multi_optimizer_save_preserves_completed_checkpoint(tmp_path, old_model_state): from veomni.checkpoint.layout import checkpoint_is_complete, write_manifest diff --git a/tests/optim/test_multi_optimizer_sparse_lr.py b/tests/optim/test_multi_optimizer_sparse_lr.py index 9a3833aece..747afd15ed 100644 --- a/tests/optim/test_multi_optimizer_sparse_lr.py +++ b/tests/optim/test_multi_optimizer_sparse_lr.py @@ -56,19 +56,6 @@ def test_sparse_leaf_preserves_tensor_learning_rate(): assert not leaf.state -def test_sparse_leaf_serializes_state_created_by_pre_hook(): - model = torch.nn.Linear(2, 2) - leaf = torch.optim.AdamW([model.bias], lr=0.01) - optimizer = MultiOptimizer(model, {"bias": leaf}, ["bias"]) - - def add_state(opt): - opt.state[model.bias]["custom"] = torch.ones_like(model.bias) - - leaf.register_state_dict_pre_hook(add_state) - result = optimizer._sparse_state_dict() - torch.testing.assert_close(result["state.bias.custom"], torch.ones_like(model.bias), rtol=0, atol=0) - - def test_sparse_leaf_respects_post_hook_parameter_reordering(): model = torch.nn.Linear(2, 2) leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) @@ -86,23 +73,6 @@ def swap_serialized_groups(_optimizer, state): assert not leaf.state -def test_sparse_leaf_respects_live_group_reordering_in_post_hook(): - model = torch.nn.Linear(2, 2) - leaf = torch.optim.AdamW([{"params": [model.weight], "lr": 0.01}, {"params": [model.bias], "lr": 0.02}]) - optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) - - def reverse_live_groups(opt, state): - opt.param_groups.reverse() - return state - - leaf.register_state_dict_post_hook(reverse_live_groups) - result = optimizer._sparse_state_dict() - assert result["param_groups.weight.lr"] == 0.01 - assert result["param_groups.bias.lr"] == 0.02 - assert len(leaf._optimizer_state_dict_pre_hooks) == 0 - assert len(leaf._optimizer_state_dict_post_hooks) == 1 - - @pytest.mark.parametrize("patched", [False, True]) def test_sparse_leaf_preserves_custom_optimizer_state_dict(patched): calls = [] @@ -178,29 +148,6 @@ def test_pytorch_patched_adamw_dcp_resume_preserves_next_update(tmp_path, initia torch.testing.assert_close(leaf.state[current][key], restored_leaf.state[target][key], rtol=0, atol=0) -def test_sparse_leaf_reentrant_hook_fails_without_leaking_temporary_hooks(): - model = torch.nn.Linear(2, 2) - leaf = torch.optim.AdamW(model.parameters()) - optimizer = MultiOptimizer(model, {"leaf": leaf}, ["leaf"]) - entered = False - - def reenter(opt): - nonlocal entered - if not entered: - entered = True - try: - opt.state_dict() - finally: - entered = False - - leaf.register_state_dict_pre_hook(reenter) - with pytest.raises(RuntimeError, match="Cannot reenter"): - optimizer._sparse_state_dict() - assert not leaf.state - assert len(leaf._optimizer_state_dict_pre_hooks) == 1 - assert not leaf._optimizer_state_dict_post_hooks - - def test_sparse_leaf_learning_rate_survives_dcp_roundtrip(tmp_path): model, weight, bias, optimizer = _build() model.weight.grad = torch.ones_like(model.weight) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 135f121971..bef3e31b69 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -1817,20 +1817,28 @@ def save( checkpoint_dir = step_dir(path, global_steps) if global_steps is not None else path model_root = model_dir(checkpoint_dir, module) if multi: - error = ( - RuntimeError("Cannot overwrite a completed checkpoint transaction with MultiOptimizer state") - if os.path.exists(manifest_path(checkpoint_dir)) - else None - ) + error = None + try: + os.stat(manifest_path(checkpoint_dir)) + except FileNotFoundError: + pass + except BaseException as exc: + error = exc + else: + error = RuntimeError("Cannot overwrite a completed checkpoint transaction with MultiOptimizer state") raise_if_any_rank_failed(error, "checking MultiOptimizer checkpoint destination") if multi and extra is None: rank = dist.get_rank() if dist.is_initialized() else 0 extra_path = os.path.join(model_root, _EXTRA_STATE_DIRNAME, _EXTRA_STATE_FORMAT.format(rank)) - error = ( - RuntimeError("Existing checkpoint sidecar cannot be reused when extra_state is omitted") - if os.path.exists(extra_path) - else None - ) + error = None + try: + os.stat(extra_path) + except FileNotFoundError: + pass + except BaseException as exc: + error = exc + else: + error = RuntimeError("Existing checkpoint sidecar cannot be reused when extra_state is omitted") raise_if_any_rank_failed(error, "checking for stale MultiOptimizer extra state") # Keyed on the module's own run-level path, not just ``path``: a diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index 2316f4ead2..d44b4b756d 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -401,13 +401,8 @@ def step(self, closure=None): def _empty_optimizer_state_dict(model, opt): packed_groups = [] id_to_name = {} - captured = False def capture_params(leaf): - nonlocal captured - if captured: - raise RuntimeError("Cannot reenter optimizer state_dict while saving an empty optimizer") - captured = True packed_groups.extend([list(group["params"]) for group in leaf.param_groups]) def capture_ids(_leaf, packed): From f3b43282e73f0d30bc743092f1383d5b0495c884 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 18:53:16 +0800 Subject: [PATCH 052/120] [checkpoint] remove unsupported rank-dependent fixture --- .../test_multi_optimizer_proven_resume.py | 44 +++---------------- 1 file changed, 5 insertions(+), 39 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 51464d05b2..dcaa9b54ca 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -623,35 +623,6 @@ def test_inventory_hash_failure_propagates_before_collective(): spawn_with_rendezvous(_inventory_hash_failure_worker, args=(), nprocs=2) -def _tied_parameter_order_worker(rank, endpoint): - from datetime import timedelta - - import torch.distributed as dist - - from veomni.checkpoint.optimizer_resume_layout import coordinated_multi_layout - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - model = torch.nn.Module() - shared = torch.nn.Parameter(torch.ones(2, 2)) - for name in ("a", "b") if rank == 0 else ("b", "a"): - model.register_parameter(name, shared) - optimizer = MultiOptimizer(model, {"leaf": torch.optim.AdamW([shared])}, ["leaf"]) - coordinated_multi_layout(model, optimizer, writing=True) - completed = torch.tensor(1) - dist.all_reduce(completed) - assert completed.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -def test_tied_parameter_registration_order_does_not_change_trainability(): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_tied_parameter_order_worker, args=(), nprocs=2) - - def _finalize_error_worker(rank, checkpoint, endpoint): from datetime import timedelta @@ -689,7 +660,7 @@ def test_save_finalize_error_propagates_to_all_ranks(tmp_path): spawn_with_rendezvous(_finalize_error_worker, args=(str(tmp_path),), nprocs=2) -def _runtime_plan_mismatch_worker(rank, mode, endpoint): +def _runtime_plan_mismatch_worker(rank, endpoint): import contextlib from datetime import timedelta @@ -706,8 +677,6 @@ def setup(self): def _build_model(self): self.model = torch.nn.Linear(2, 2) - if mode != "missing" and (mode != "plan" or rank == 0): - self.model._fqn2spec_info = {} def _freeze_model_module(self): pass @@ -716,11 +685,9 @@ def _build_parallelized_model(self): pass runtime_module.use_parallel_state = lambda *_args, **_kwargs: contextlib.nullcontext() - checkpoint = SimpleNamespace( - extra_parallel_canonical_layout=(mode == "missing" or (mode == "flag" and rank == 1)) - ) + checkpoint = SimpleNamespace(extra_parallel_canonical_layout=True) args = SimpleNamespace(optimizer=SimpleNamespace(type="adamw")) - with pytest.raises(RuntimeError, match="differs across ranks|requires an ExtraParallel model plan"): + with pytest.raises(RuntimeError, match="requires an ExtraParallel model plan"): TestRuntime(args, train=SimpleNamespace(checkpoint=checkpoint)) completed = torch.tensor(1) dist.all_reduce(completed) @@ -730,11 +697,10 @@ def _build_parallelized_model(self): @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -@pytest.mark.parametrize("mode", ["missing"]) -def test_runtime_checkpoint_plan_mismatch_fails_all_ranks(mode): +def test_runtime_checkpoint_plan_mismatch_fails_all_ranks(): from tests.tools.launch_utils import spawn_with_rendezvous - spawn_with_rendezvous(_runtime_plan_mismatch_worker, args=(mode,), nprocs=2) + spawn_with_rendezvous(_runtime_plan_mismatch_worker, args=(), nprocs=2) @pytest.mark.parametrize("same_base", [True, False]) From 8a4fbf0a8da93772aa410cebb5be5bca901c027a Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 19:15:45 +0800 Subject: [PATCH 053/120] [checkpoint] prove custom expert layouts and safe exports --- docs/usage/checkpoint.md | 19 + tests/checkpoints/test_expert_shard_layout.py | 632 +++++++++++++++++- veomni/checkpoint/dcp_checkpointer.py | 327 ++++++++- veomni/models/checkpoint/manager.py | 1 + veomni/utils/save_safetensor_utils.py | 100 ++- 5 files changed, 1015 insertions(+), 64 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 8d30a15b38..23cc6941f5 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -266,6 +266,25 @@ under an `experts` module when their source order cannot be established. Other legacy tensor naming schemes may still need explicit source-layout verification or migration before export. +New planner metadata records supported named ExtraParallel/FSDP meshes rather +than only `ep`. Ordinary FSDP expert shards are marked canonical only when the +actual TP1 runtime proves that every ExtraParallel dimension has size one, or +the model's applied plan and parameter aliases prove that the expert is not +sharded by any active ExtraParallel dimension. +An expert on the EP-complementary FSDP mesh is accepted only for trainable +shared-expert LoRA A/B parameters whose wrapper installs cross-EP gradient +synchronization; other parameters on that mesh remain unproven. A normal +single-process save assumes the caller supplied the complete model, not a +manually extracted expert shard. The singleton case is not a source-layout +proof for external, manually sharded tensors. +Older version-1 metadata may omit non-`ep` shards: an unrecorded expert tensor +is rejected unless its source layout can be verified and migrated. This also +rejects an indistinguishable old ordinary-FSDP expert tensor. Historical non-expert tensors +with unrecorded ExtraParallel axes require separate source-layout verification +before loading into a different mesh or offline export; their provenance cannot +be inferred from chunks alone, and not all such cases can be rejected +automatically when the target has no ExtraParallel mesh. + All ranks in a checkpoint operation must use the same logical model and checkpoint configuration, including the layout version and corresponding sharding specifications. Normal rank-local shard contents and mesh coordinates diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 151b205bb9..05c95f9666 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -4,19 +4,25 @@ """CPU/Gloo regressions for canonical EP-first DCP expert coordinates.""" from datetime import timedelta +from types import SimpleNamespace import pytest import torch import torch.distributed as dist import torch.distributed.checkpoint as dcp import torch.multiprocessing as mp -from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.checkpoint.default_planner import DefaultSavePlanner +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh from torch.distributed.tensor import DTensor, Replicate, Shard from scripts.ci.accelerator.launch_plan import from_parameter from tests.tools.launch_utils import spawn_with_rendezvous from veomni.checkpoint.dcp_checkpointer import ( _EXPERT_LAYOUTS_KEY, + ModelState, + _expert_layout, + _extra_parallel_first_mesh, + _extra_parallel_layout_key, _ExtraParallelSavePlanner, _ModelStrictLoadPlanner, dcp_to_torch_state_dict, @@ -215,19 +221,28 @@ def test_lora_adapter_export_preserves_expert_shard_layout(world, shard_dim, neg def _ordinary_fsdp_expert_export_worker(rank, root, ndim, endpoint): from torch.distributed.tensor import distribute_tensor + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state, get_parallel_state + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) try: - mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("fsdp",)) + _init_parallel_state(dp_size=2, dp_shard_size=2, device_type="cpu") + parallel = get_parallel_state() + mesh = parallel.fsdp_mesh shape = (8,) + (4,) * (ndim - 1) expected = torch.arange(8 * 4 ** (ndim - 1), dtype=torch.float32).reshape(shape) weight = distribute_tensor(expected, mesh, [Shard(0)]) - dcp.save({"model": {"experts": {"weight": weight}}}, checkpoint_id=root, planner=_ExtraParallelSavePlanner()) + dcp.save( + {"model": {"experts": {"weight": weight}}}, + checkpoint_id=root, + planner=_ExtraParallelSavePlanner(parallel_state=parallel), + ) dist.barrier() result = [None] if rank == 0: try: metadata = dcp.FileSystemReader(root).read_metadata() - assert metadata.planner_data[_EXPERT_LAYOUTS_KEY] == {"version": 1, "tensors": {}} + record = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"]["model.experts.weight"] + assert record["source_kind"] == "veomni_tp1_no_ep" exported = dcp_to_torch_state_dict(root) torch.testing.assert_close(exported["experts.weight"], expected, rtol=0, atol=0) except Exception as exc: @@ -235,6 +250,83 @@ def _ordinary_fsdp_expert_export_worker(rank, root, ndim, endpoint): dist.broadcast_object_list(result, src=0) assert result[0] is None, result[0] + # A full plain tensor is safe only when the actual runtime proves no + # ExtraParallel source; without that proof it could be EP-local. + for proven in (True, False): + path = f"{root}-plain-{proven}" + planner = _ExtraParallelSavePlanner(parallel_state=parallel if proven else None) + dcp.save({"model": {"experts": {"weight": expected}}}, checkpoint_id=path, planner=planner) + result = [None] + if rank == 0: + try: + metadata = dcp.FileSystemReader(path).read_metadata() + record = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"]["model.experts.weight"] + assert record["expert_order"] == ("canonical" if proven else "unproven") + if proven: + torch.testing.assert_close( + dcp_to_torch_state_dict(path)["experts.weight"], expected, rtol=0, atol=0 + ) + else: + with pytest.raises(RuntimeError, match="noncanonical expert order"): + dcp_to_torch_state_dict(path) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + + # The same FSDP-looking tensor is not proof that an earlier EP axis + # was never stripped; an unknown source must not be certified. + unproven = root + "-unproven" + dcp.save( + {"model": {"experts": {"weight": weight}}}, checkpoint_id=unproven, planner=_ExtraParallelSavePlanner() + ) + result = [None] + if rank == 0: + try: + metadata = dcp.FileSystemReader(unproven).read_metadata() + record = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"]["model.experts.weight"] + assert record["expert_order"] == "unproven" + with pytest.raises(RuntimeError, match="noncanonical expert order"): + dcp_to_torch_state_dict(unproven) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + + from unittest.mock import patch + + from veomni.utils.save_safetensor_utils import save_hf_safetensor + + with ( + patch("veomni.utils.save_safetensor_utils.should_use_distributed_hf_safetensor_save", return_value=False), + pytest.raises(RuntimeError, match="noncanonical expert order|failed on another rank"), + ): + save_hf_safetensor( + save_hf_safetensor_path=f"{root}-hf", + save_checkpoint_path=unproven, + ckpt_manager="dcp", + is_rank_0=rank == 0, + ) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + + # Earlier planner metadata did not prove that non-EP dimensions were + # absent, even when this particular tensor used ordinary FSDP. + old_v1 = root + "-old-v1" + dcp.save( + {"model": {"experts": {"weight": weight}}}, checkpoint_id=old_v1, planner=_LegacyNamedExtraSavePlanner() + ) + result = [None] + if rank == 0: + try: + with pytest.raises(RuntimeError, match="Cannot determine source expert order"): + dcp_to_torch_state_dict(old_v1) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + # An old expert-only chunk has no provenance, even when the target is # ordinary FSDP. It must not be loaded and re-saved as canonical. legacy = root + "-legacy" @@ -244,12 +336,542 @@ def _ordinary_fsdp_expert_export_worker(rank, root, ndim, endpoint): dcp.load(target, checkpoint_id=legacy, planner=_ModelStrictLoadPlanner(strict_model=True)) assert target["model"]["experts"]["weight"].to_local().count_nonzero() == 0 finally: + clear_parallel_state() dist.destroy_process_group() @pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) @pytest.mark.parametrize("world", [2]) -@pytest.mark.parametrize("ndim", [2, 3]) +@pytest.mark.parametrize("ndim", [1, 2, 3]) def test_recorded_non_ep_fsdp_expert_export(world, ndim, tmp_path): assert world == 2 spawn_with_rendezvous(_ordinary_fsdp_expert_export_worker, args=(str(tmp_path / "dcp"), ndim), nprocs=world) + + +@pytest.mark.parametrize("outer", [Shard(1), Replicate()]) +def test_unexplained_mesh_axis_overlap_is_not_certified(outer): + mesh = DeviceMesh( + "cpu", + [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], + mesh_dim_names=("x", "ep_fsdp", "ep"), + _init_backend=False, + _rank=0, + ) + value = DTensor.from_local(torch.zeros(2, 2), mesh, [outer, Shard(1), Shard(0)], run_check=False) + assert _expert_layout(value)["expert_order"] == ("unproven" if isinstance(outer, Shard) else "canonical") + + +def _ple_and_ordinary_expert_worker(rank, root, endpoint): + from torch.distributed.tensor import distribute_tensor + + from veomni.distributed.parallel_plan import ParallelPlan + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state, get_parallel_state + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + _init_parallel_state( + dp_size=2, + dp_shard_size=2, + device_type="cpu", + extra_parallel_names=("ep", "ple"), + extra_parallel_sizes=(1, 2), + extra_parallel_placement_innermost=(False, False), + ) + parallel = get_parallel_state() + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.experts = torch.nn.Module() + self.experts.weight = torch.nn.Parameter(torch.arange(32, dtype=torch.float32).reshape(8, 4)) + self.experts.alias = self.experts.weight + self.ple = torch.nn.Embedding(8, 4) + self.shard_expert_alias = False + + def get_parallel_plan(self): + entries = {"ple.weight": Shard(0)} + if self.shard_expert_alias: + entries["experts.alias"] = Shard(0) + return ParallelPlan({"ple": entries}) + + model = Model() + model._fqn2spec_info = model.get_parallel_plan().apply(model, parallel.extra_parallel_fsdp_device_mesh) + expected = model.experts.weight.detach().clone() + sharded = torch.nn.Parameter(distribute_tensor(expected, parallel.fsdp_mesh, [Shard(0)])) + model.experts.weight = sharded + model.experts.alias = sharded + for conflicting_plan in (False, True): + model.shard_expert_alias = conflicting_plan + path = f"{root}-{conflicting_plan}" + planner = _ExtraParallelSavePlanner(parallel_state=parallel, source_model=model) + dcp.save({"model": {"experts": {"weight": sharded}}}, checkpoint_id=path, planner=planner) + result = [None] + if rank == 0: + try: + metadata = dcp.FileSystemReader(path).read_metadata() + record = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"]["model.experts.weight"] + assert record["expert_order"] == ("unproven" if conflicting_plan else "canonical") + if conflicting_plan: + with pytest.raises(RuntimeError, match="noncanonical expert order"): + dcp_to_torch_state_dict(path) + else: + torch.testing.assert_close( + dcp_to_torch_state_dict(path)["experts.weight"], expected, rtol=0, atol=0 + ) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.parametrize( + ("key", "expected"), + [ + ("model.layers.0.experts.weight", True), + ("optimizer.state.layers.0.experts.weight.exp_avg", True), + ("optimizer.state.layers.0.experts.weight.exp_avg.extra", True), + ("optimizer.state.layers.0.experts.weight_suffix.exp_avg", False), + ("model.layers.0.experts.weight_suffix", False), + ], +) +def test_expert_source_key_matches_parameter_boundaries(key, expected): + from veomni.checkpoint.dcp_checkpointer import _key_belongs_to_fqns + + assert _key_belongs_to_fqns(key, {"layers.0.experts.weight"}) is expected + + +def _single_rank_uninitialized_expert_worker(rank, root, endpoint): + from torch.distributed.tensor import distribute_tensor + + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer + from veomni.distributed.parallel_state import clear_parallel_state + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=1, timeout=timedelta(seconds=30)) + try: + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("dp_shard",)) + + def build(): + model = torch.nn.Module() + model.experts = torch.nn.Module() + model.experts.weight = torch.nn.Parameter( + distribute_tensor(torch.arange(8, dtype=torch.float32), mesh, [Shard(0)]) + ) + return model + + model = build() + DistributedCheckpointer.save(root, {"model": model}) + target = build() + target.experts.weight.data.zero_() + DistributedCheckpointer.load(root, {"model": target}) + torch.testing.assert_close(target.experts.weight.to_local(), model.experts.weight.to_local(), rtol=0, atol=0) + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [1]) +def test_single_rank_without_parallel_state_retains_expert_checkpoint(world, tmp_path): + spawn_with_rendezvous(_single_rank_uninitialized_expert_worker, args=(str(tmp_path / "dcp"),), nprocs=world) + + +def test_no_distributed_initialization_retains_expert_checkpoint(tmp_path, monkeypatch): + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer + from veomni.distributed.parallel_state import clear_parallel_state + + monkeypatch.setenv("WORLD_SIZE", "1") + clear_parallel_state() + assert not dist.is_initialized() + model = torch.nn.Module() + model.experts = torch.nn.Linear(2, 2) + expected = model.experts.weight.detach().clone() + DistributedCheckpointer.save(str(tmp_path), {"model": model}) + with torch.no_grad(): + model.experts.weight.zero_() + DistributedCheckpointer.load(str(tmp_path), {"model": model}) + torch.testing.assert_close(model.experts.weight, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("shared", [False, True]) +def test_single_process_lora_export_without_parallel_state(tmp_path, monkeypatch, shared): + from veomni.distributed.parallel_state import clear_parallel_state + from veomni.lora import VeOmniLoraConfig + from veomni.lora.state_dict import get_lora_state_dict, load_adapter_state_dict + from veomni.utils.save_safetensor_utils import save_lora_adapter_with_dcp + + monkeypatch.setenv("WORLD_SIZE", "1") + clear_parallel_state() + assert not dist.is_initialized() + config = VeOmniLoraConfig.from_yaml( + { + "rank": 8, + "alpha": 16, + "target_parameters": ["experts.gate_up_proj", "experts.down_proj"], + "share_expert_lora": shared, + } + ) + model = _lora_model(config) + expected = { + key: value.detach().to(torch.bfloat16).cpu() + for key, value in get_lora_state_dict(model, config=model.get_lora_config()).items() + } + save_lora_adapter_with_dcp(model, str(tmp_path)) + actual = load_adapter_state_dict(str(tmp_path)) + assert actual.keys() == expected.keys() + for key, value in expected.items(): + torch.testing.assert_close(actual[key], value, rtol=0, atol=0) + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [2]) +def test_ple_with_ep1_proves_ordinary_expert_from_plan(world, tmp_path): + spawn_with_rendezvous(_ple_and_ordinary_expert_worker, args=(str(tmp_path / "dcp"),), nprocs=world) + + +def _ple_ep1_lora_worker(rank, root, unproven, endpoint): + from torch.distributed.tensor import distribute_tensor + + from veomni.distributed.parallel_plan import ParallelPlan, get_runtime_parallel_plan + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state, get_parallel_state + from veomni.lora import VeOmniLoraConfig + from veomni.lora.state_dict import get_lora_state_dict, load_adapter_state_dict, strip_adapter_name + from veomni.utils.save_safetensor_utils import save_lora_adapter_with_dcp + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + _init_parallel_state( + dp_size=2, + dp_shard_size=2, + device_type="cpu", + extra_parallel_names=("ep", "ple"), + extra_parallel_sizes=(1, 2), + extra_parallel_placement_innermost=(False, False), + ) + parallel = get_parallel_state() + torch.manual_seed(42) + config = VeOmniLoraConfig.from_yaml( + {"rank": 8, "alpha": 16, "target_parameters": ["experts.gate_up_proj", "experts.down_proj"], "bias": "all"} + ) + model = _lora_model(config) + base = model.base_model.model + base.experts.bias = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + base.ple = torch.nn.Embedding(8, 4) + base.get_parallel_plan = lambda: ParallelPlan({"ple": {"ple.weight": Shard(0)}}) + plan = get_runtime_parallel_plan(model) + model._fqn2spec_info = plan.apply(model, parallel.extra_parallel_fsdp_device_mesh) + selected = get_lora_state_dict(model, config=model.get_lora_config()) + assert any("experts.bias" in name for name in selected) + for name, param in list(model.named_parameters()): + if strip_adapter_name(name, "default") not in selected: + continue + parent, attr = name.rsplit(".", 1) + setattr( + model.get_submodule(parent), + attr, + torch.nn.Parameter(distribute_tensor(param.detach(), parallel.fsdp_mesh, [Shard(0)])), + ) + selected = get_lora_state_dict(model, config=model.get_lora_config()) + expected = { + name: (value.full_tensor() if isinstance(value, DTensor) else value).detach().to(torch.bfloat16).cpu() + for name, value in selected.items() + } + if unproven: + base.get_parallel_plan = lambda: None + with pytest.raises(RuntimeError, match="noncanonical expert order|failed on another rank"): + save_lora_adapter_with_dcp(model, root, parallel_state=parallel) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + return + save_lora_adapter_with_dcp(model, root, parallel_state=parallel) + exported = load_adapter_state_dict(root) + assert exported.keys() == expected.keys() + for name, value in expected.items(): + torch.testing.assert_close(exported[name], value, rtol=0, atol=0) + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [2]) +@pytest.mark.parametrize("unproven", [False, True]) +def test_ple_ep1_lora_bias_all_checkpoint_export(world, unproven, tmp_path): + spawn_with_rendezvous(_ple_ep1_lora_worker, args=(str(tmp_path / "adapter"), unproven), nprocs=world) + + +def _ep_shared_lora_worker(rank, root, mode, endpoint): + from torch.distributed.tensor import distribute_tensor + + from veomni.distributed.parallel_plan import ParallelPlan, get_runtime_parallel_plan + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state, get_parallel_state + from veomni.lora import VeOmniLoraConfig + from veomni.lora.moe_layers import LoraSharedExperts + from veomni.lora.state_dict import get_lora_state_dict, load_adapter_state_dict, strip_adapter_name + from veomni.utils.save_safetensor_utils import save_lora_adapter_with_dcp + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + _init_parallel_state( + dp_size=2, + dp_shard_size=2, + device_type="cpu", + extra_parallel_names=("ep",), + extra_parallel_sizes=(2,), + extra_parallel_placement_innermost=(False,), + ) + parallel = get_parallel_state() + torch.manual_seed(42) + config = VeOmniLoraConfig.from_yaml( + { + "rank": 8, + "alpha": 16, + "target_parameters": ["experts.gate_up_proj", "experts.down_proj"], + "share_expert_lora": True, + "bias": "all", + } + ) + model = _lora_model(config) + base = model.base_model.model + assert isinstance(base.experts, LoraSharedExperts) + if mode == "rank_divergent_bias": + base.experts.bias = torch.nn.Parameter(torch.full((8, 8), float(rank))) + base.get_parallel_plan = lambda: ParallelPlan( + {"ep": {"experts.gate_up_proj": Shard(0), "experts.down_proj": Shard(0)}} + ) + plan = get_runtime_parallel_plan(model) + model._fqn2spec_info = plan.apply(model, parallel.extra_parallel_fsdp_device_mesh) + selected = get_lora_state_dict(model, config=model.get_lora_config()) + ep_mesh = parallel.extra_parallel_fsdp_device_mesh["ep"] + fsdp_mesh = ep_mesh[ep_mesh.mesh_dim_names[:-1]] + for name, param in list(model.named_parameters()): + if strip_adapter_name(name, "default") not in selected: + continue + parent, attr = name.rsplit(".", 1) + setattr( + model.get_submodule(parent), + attr, + torch.nn.Parameter(distribute_tensor(param.detach(), fsdp_mesh, [Shard(1)])), + ) + selected = get_lora_state_dict(model, config=model.get_lora_config()) + expected = {name: value.full_tensor().detach().to(torch.bfloat16).cpu() for name, value in selected.items()} + if mode == "plan_conflict": + base.get_parallel_plan = lambda: ParallelPlan( + {"ep": {"experts.gate_proj.lora_A.default.weight": Shard(0)}} + ) + if mode != "valid": + with pytest.raises(RuntimeError, match="noncanonical expert order|failed on another rank"): + save_lora_adapter_with_dcp(model, root, parallel_state=parallel) + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + return + save_lora_adapter_with_dcp(model, root, parallel_state=parallel) + exported = load_adapter_state_dict(root) + assert exported.keys() == expected.keys() + for name, value in expected.items(): + torch.testing.assert_close(exported[name], value, rtol=0, atol=0) + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [2]) +@pytest.mark.parametrize("mode", ["valid", "plan_conflict", "rank_divergent_bias"]) +def test_ep_replicated_shared_expert_lora_export(world, mode, tmp_path): + spawn_with_rendezvous(_ep_shared_lora_worker, args=(str(tmp_path / "adapter"), mode), nprocs=world) + + +def _pure_extra_parallel_worker(rank, root, name, endpoint): + from veomni.distributed.parallel_plan import SpecInfo + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + mesh = init_device_mesh("cpu", (1, 2), mesh_dim_names=(f"{name}_fsdp", name)) + fsdp_mesh = mesh[f"{name}_fsdp"] + expected = torch.arange(16, dtype=torch.float32).reshape(8, 2) + local = expected.chunk(2, dim=0)[rank].contiguous() + restored_axis = restore_extra_parallel_dim(local, mesh, fsdp_mesh, ep_shard_dim=0) + torch.testing.assert_close(restored_axis.full_tensor(), expected, rtol=0, atol=0) + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=(name,), extra_parallel_fsdp_device_mesh={name: mesh} + ) + + def model_with(local_value): + model = torch.nn.Module() + model.experts = torch.nn.Module() + model.embedding = torch.nn.Module() + for module_name in ("experts", "embedding"): + getattr(model, module_name).weight = torch.nn.Parameter(local_value.clone()) + model._fqn2spec_info = { + f"{module_name}.weight": SpecInfo(name, Shard(0), f"{module_name}.weight", mesh, 1) + for module_name in ("experts", "embedding") + } + return model + + source = model_with(local) + planner = _ExtraParallelSavePlanner(parallel_state=parallel) + dcp.save({"model": ModelState(source, parallel_state=parallel)}, checkpoint_id=root, planner=planner) + result = [None] + if rank == 0: + try: + exported = dcp_to_torch_state_dict(root) + for key in ("experts.weight", "embedding.weight"): + torch.testing.assert_close(exported[key], expected, rtol=0, atol=0) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + target = model_with(torch.zeros_like(local)) + dcp.load( + {"model": ModelState(target, parallel_state=parallel)}, + checkpoint_id=root, + planner=_ModelStrictLoadPlanner(strict_model=True), + ) + for module_name in ("experts", "embedding"): + torch.testing.assert_close(getattr(target, module_name).weight, local, rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [2]) +@pytest.mark.parametrize("name", ["ep", "ple"]) +def test_plain_pure_extra_parallel_preserves_global_values(world, name, tmp_path): + spawn_with_rendezvous(_pure_extra_parallel_worker, args=(str(tmp_path / "dcp"), name), nprocs=world) + + +class _LegacyNamedExtraSavePlanner(DefaultSavePlanner): + def create_global_plan(self, all_plans): + plans, metadata = super().create_global_plan(all_plans) + metadata.planner_data = {**(metadata.planner_data or {}), _EXPERT_LAYOUTS_KEY: {"version": 1, "tensors": {}}} + return plans, metadata + + +def _named_extra_expert_export_worker(rank, root, endpoint): + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=4, timeout=timedelta(seconds=60)) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("custom_fsdp", "custom")) + fsdp_mesh = mesh["custom_fsdp"] + fsdp_rank, extra_rank = mesh.get_coordinate() + expected = torch.arange(8 * 16 * 4, dtype=torch.float32).reshape(8, 16, 4) + local = expected.chunk(2, 0)[extra_rank].chunk(2, 0)[fsdp_rank].contiguous() + local_dt = DTensor.from_local(local, fsdp_mesh, [Shard(0)]) + canonical = restore_extra_parallel_dim(local_dt, mesh, fsdp_mesh, 0, 0) + legacy = DTensor.from_local(local, mesh, [Shard(0), Shard(0)]) + key = "model.layers.0.mlp.experts.gate_up_proj" + for order, source, planner in ( + ("canonical", canonical, _ExtraParallelSavePlanner()), + ("legacy_fsdp_first", legacy, _ExtraParallelSavePlanner()), + ("legacy_unrecorded", legacy, _LegacyNamedExtraSavePlanner()), + ): + path = f"{root}-{order}" + dcp.save( + { + "model": {key: source}, + "optimizer": {"experts": {"moment": source.clone(), "step": torch.tensor(1.0)}}, + }, + checkpoint_id=path, + planner=planner, + ) + result = [None] + if rank == 0: + try: + metadata = dcp.FileSystemReader(path).read_metadata() + records = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + if order != "legacy_unrecorded": + assert records[f"model.{key}"]["expert_order"] == order + assert records["optimizer.experts.moment"]["expert_order"] == order + if order == "canonical": + torch.testing.assert_close(dcp_to_torch_state_dict(path)[key], expected, rtol=0, atol=0) + else: + with pytest.raises(RuntimeError, match="noncanonical expert order|Cannot determine source"): + dcp_to_torch_state_dict(path) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + if order == "legacy_unrecorded": + target = {"model": {key: torch.zeros_like(canonical)}} + with pytest.raises(dcp.CheckpointException, match="Cannot determine source"): + dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=True)) + target = {"optimizer": {"experts": {"moment": torch.zeros_like(canonical), "step": torch.tensor(0.0)}}} + if order == "canonical": + dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=False)) + torch.testing.assert_close( + target["optimizer"]["experts"]["moment"].full_tensor(), expected, rtol=0, atol=0 + ) + torch.testing.assert_close(target["optimizer"]["experts"]["step"], torch.tensor(1.0), rtol=0, atol=0) + else: + with pytest.raises(dcp.CheckpointException, match="noncanonical expert order|Cannot determine source"): + dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=False)) + + scalar_path = f"{root}-legacy-step" + dcp.save( + {"optimizer": {"experts": {"step": torch.tensor(1.0)}}}, + checkpoint_id=scalar_path, + planner=_LegacyNamedExtraSavePlanner(), + ) + scalar_target = {"optimizer": {"experts": {"step": torch.tensor(0.0)}}} + dcp.load(scalar_target, checkpoint_id=scalar_path, planner=_ModelStrictLoadPlanner(strict_model=False)) + torch.testing.assert_close(scalar_target["optimizer"]["experts"]["step"], torch.tensor(1.0), rtol=0, atol=0) + + # The old per-tensor planner did not recognize ``custom``, but the + # canonical v2 namespace itself proves this checkpoint's source order. + marker = _extra_parallel_layout_key({}) + canonical_mesh = _extra_parallel_first_mesh(mesh) + namespaced = DTensor.from_local( + local, + canonical_mesh, + [Shard(0), Shard(0)], + run_check=False, + shape=expected.shape, + stride=expected.stride(), + ) + path = f"{root}-legacy-v2" + dcp.save( + {"model": {marker: "ep-outer-v1", f"{marker}__{key}": namespaced}}, + checkpoint_id=path, + planner=_LegacyNamedExtraSavePlanner(), + ) + target = {"model": {marker: "", f"{marker}__{key}": torch.zeros_like(namespaced)}} + dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=True)) + torch.testing.assert_close(target["model"][f"{marker}__{key}"].full_tensor(), expected, rtol=0, atol=0) + result = [None] + if rank == 0: + try: + torch.testing.assert_close(dcp_to_torch_state_dict(path)[key], expected, rtol=0, atol=0) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + + for names in (("axis_a", "axis_b"), ("fsdp", "ep")): + unknown_mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=names) + unknown = DTensor.from_local(local, unknown_mesh, [Shard(0), Shard(0)]) + path = f"{root}-unknown-{names[0]}" + dcp.save({"model": {key: unknown}}, checkpoint_id=path, planner=_ExtraParallelSavePlanner()) + result = [None] + if rank == 0: + try: + metadata = dcp.FileSystemReader(path).read_metadata() + record = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"][f"model.{key}"] + assert record["expert_order"] == "unproven" + with pytest.raises(RuntimeError, match="noncanonical expert order"): + dcp_to_torch_state_dict(path) + except Exception as exc: + result[0] = repr(exc) + dist.broadcast_object_list(result, src=0) + assert result[0] is None, result[0] + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [4]) +def test_named_extra_parallel_expert_export_checks_source_order(world, tmp_path): + spawn_with_rendezvous(_named_extra_expert_export_worker, args=(str(tmp_path / "dcp"),), nprocs=world) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 818ebc17e9..f9349fe3c4 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -16,6 +16,7 @@ import gc import hashlib import os +import re import shutil import weakref from collections import OrderedDict @@ -45,7 +46,7 @@ from torch.distributed.checkpoint.stateful import Stateful from torch.distributed.tensor.placement_types import _StridedShard -from ..distributed.parallel_state import get_parallel_state +from ..distributed.parallel_state import ParallelState, get_parallel_state, is_parallel_state_initialized from ..utils import logging from ..utils.device import empty_cache, synchronize from ..utils.dist_utils import any_rank_failed, raise_if_any_rank_failed @@ -72,39 +73,280 @@ _EXPERT_LAYOUTS_KEY = "__veomni_expert_layouts__" -def _expert_layout(tensor): - """Describe the source EP/FSDP layout, independently of the load target.""" +def _expert_layout(tensor, extra_parallel_names=()): + """Describe a source ExtraParallel/FSDP layout independently of the target.""" if not isinstance(tensor, DTensor): return None mesh = tensor.device_mesh names = mesh.mesh_dim_names or () - if "ep" not in names: + extra_names = [name for name in names if name == "ep" or name in extra_parallel_names or f"{name}_fsdp" in names] + if not extra_names: return None - ep_axis = names.index("ep") + if len(extra_names) != 1: + raise RuntimeError(f"Cannot determine the ExtraParallel axis from mesh names {names}") + extra_name = extra_names[0] + ep_axis = names.index(extra_name) ep = tensor.placements[ep_axis] if not isinstance(ep, Shard): return None - fsdp_axis = names.index("ep_fsdp") if "ep_fsdp" in names else None + fsdp_name = f"{extra_name}_fsdp" + fsdp_axis = names.index(fsdp_name) if fsdp_name in names else None fsdp = tensor.placements[fsdp_axis] if fsdp_axis is not None else Replicate() - fsdp_dim = fsdp.dim if isinstance(fsdp, (Shard, _StridedShard)) else None + fsdp_dim = fsdp.dim % tensor.ndim if isinstance(fsdp, (Shard, _StridedShard)) else None + ep_dim = ep.dim % tensor.ndim fsdp_size = mesh.size(fsdp_axis) if fsdp_axis is not None else 1 - same_axis = ep.dim == fsdp_dim and mesh.size(ep_axis) > 1 and fsdp_size > 1 - canonical = not same_axis or isinstance(fsdp, _StridedShard) or ep_axis < fsdp_axis + same_axis = ep_dim == fsdp_dim and mesh.size(ep_axis) > 1 and fsdp_size > 1 + shard_axes = [ + (axis, placement.dim % tensor.ndim) + for axis, placement in enumerate(tensor.placements) + if isinstance(placement, Shard) and mesh.size(axis) > 1 + ] + unknown_overlap = any( + axis not in (ep_axis, fsdp_axis) + and any(other_axis != axis and other_dim == dim for other_axis, other_dim in shard_axes) + for axis, dim in shard_axes + ) + canonical = not same_axis or isinstance(fsdp, _StridedShard) or (fsdp_axis is not None and ep_axis < fsdp_axis) + order = "canonical" if canonical else "legacy_fsdp_first" + if unknown_overlap: + order = "unproven" return { - "ep_shard_dim": ep.dim, + "ep_shard_dim": ep_dim, "fsdp_shard_dim": fsdp_dim, "ep_size": mesh.size(ep_axis), "fsdp_size": fsdp_size, - "expert_order": "canonical" if canonical else "legacy_fsdp_first", + "expert_order": order, + } + + +def _ordinary_expert_fqns(model, parallel_state): + """Prove which parameters are outside every active ExtraParallel plan.""" + if not isinstance(model, torch.nn.Module) or not isinstance(getattr(model, "_fqn2spec_info", None), dict): + return set() + from torch.distributed.checkpoint.state_dict import _get_fqns + + from ..distributed.parallel_plan import ParallelPlan, SpecInfo, get_runtime_parallel_plan + + plan = get_runtime_parallel_plan(model) + if not isinstance(plan, ParallelPlan): + return set() + patterns = [ + pattern + for name, entries in plan.extra_parallel_plan.items() + if parallel_state.extra_parallel_sizes.get(name, 1) > 1 + for pattern in entries + ] + exact_patterns = {pattern for pattern in patterns if "*" not in pattern} + wildcard_patterns = [ + re.compile("^" + re.escape(pattern).replace(r"\*", ".*") + "$") for pattern in patterns if "*" in pattern + ] + aliases = {} + for name, param in model.named_parameters(remove_duplicate=False): + aliases.setdefault(id(param), []).append(name) + result, owners, ambiguous = set(), {}, set() + for owner, raw_names in aliases.items(): + if not any("experts" in name.split(".") for name in raw_names): + continue + names = [(name, _get_fqns(model, name)) for name in raw_names] + specs = [model._fqn2spec_info.get(name) for name in raw_names] + if any( + not isinstance(spec, SpecInfo) + or not isinstance(spec.placement, Replicate) + or parallel_state.extra_parallel_sizes.get(spec.para_name, 1) <= 1 + or spec.para_fsdp_mesh != parallel_state.extra_parallel_fsdp_device_mesh.get(spec.para_name) + for spec in specs + ): + continue + if any( + name in exact_patterns or any(pattern.match(name) for pattern in wildcard_patterns) for name in raw_names + ): + continue + if any(len(canonical) != 1 for _, canonical in names): + continue + for _, canonical in names: + fqn = next(iter(canonical)) + if fqn in owners and owners[fqn] != owner: + ambiguous.add(fqn) + owners[fqn] = owner + result.add(fqn) + return result - ambiguous + + +def _key_belongs_to_fqns(key, fqns): + if key.startswith("model.") and (key in fqns or key.removeprefix("model.") in fqns): + return True + if key.startswith("optimizer.state."): + rest = key.removeprefix("optimizer.state.") + while "." in rest: + rest = rest.rsplit(".", 1)[0] + if rest in fqns: + return True + return False + + +def _shared_lora_fqns(model, ordinary_fqns): + """Identify the shared LoRA parameters whose EP gradients are synchronized.""" + from torch.distributed.checkpoint.state_dict import _get_fqns + + from ..lora.moe_layers import LoraSharedExperts, _is_lora_param_name + + shared_ids = { + id(param) + for module in model.modules() + if isinstance(module, LoraSharedExperts) + for name, param in module.named_parameters(remove_duplicate=False) + if _is_lora_param_name(name) and param.requires_grad + } + return { + fqn + for name, param in model.named_parameters(remove_duplicate=False) + if id(param) in shared_ids + for fqn in _get_fqns(model, name) + if fqn in ordinary_fqns + } + + +def _no_ep_expert_layout(key, tensor, parallel_state, ordinary_fqns=(), shared_lora_fqns=()): + """Prove ordinary FSDP from the actual runtime and, if needed, model plan.""" + if "experts" not in key.split(".") or not torch.is_tensor(tensor) or not tensor.ndim: + return None + if not isinstance(parallel_state, ParallelState): + return None + if parallel_state.tp_size != 1 or not parallel_state.extra_parallel_sizes: + return None + single_process = dist.get_world_size() == 1 if dist.is_initialized() else os.environ.get("WORLD_SIZE", "1") == "1" + if ( + parallel_state.device_mesh is None + and single_process + and all(size == 1 for size in parallel_state.extra_parallel_sizes.values()) + and (not isinstance(tensor, DTensor) or tensor.device_mesh.size() == 1) + ): + return { + "ep_shard_dim": None, + "fsdp_shard_dim": None, + "ep_size": 1, + "fsdp_size": 1, + "expert_order": "canonical", + "source_kind": "veomni_tp1_no_ep", + } + ordinary = _key_belongs_to_fqns(key, ordinary_fqns) + if any(size != 1 for size in parallel_state.extra_parallel_sizes.values()) and not ordinary: + return None + if parallel_state.device_mesh is None and ( + not is_parallel_state_initialized() + or parallel_state is not get_parallel_state() + or (dist.is_initialized() and dist.get_world_size() != 1) + ): + return None + if not isinstance(tensor, DTensor): + if any(size != 1 for size in parallel_state.extra_parallel_sizes.values()): + return None + return { + "ep_shard_dim": None, + "fsdp_shard_dim": None, + "ep_size": 1, + "fsdp_size": 1, + "expert_order": "canonical", + "source_kind": "veomni_tp1_no_ep", + } + if parallel_state.device_mesh is None: + return None + mesh = tensor.device_mesh + names = mesh.mesh_dim_names or () + normal_fsdp = ( + names + in ( + ("dp_shard",), + ("dp_shard_sp",), + ("dp_replicate", "dp_shard"), + ("dp_replicate", "dp_shard_sp"), + ) + and mesh == parallel_state.fsdp_mesh + ) + # Only the shared-LoRA wrapper installs the cross-EP gradient hooks needed + # to keep replicas on complementary FSDP meshes consistent. + shared_lora = _key_belongs_to_fqns(key, shared_lora_fqns) + ep_mesh = parallel_state.extra_parallel_fsdp_device_mesh.get("ep") + replicated_extra_fsdp = ( + ordinary and shared_lora and ep_mesh is not None and mesh == ep_mesh[ep_mesh.mesh_dim_names[:-1]] + ) + if not (normal_fsdp or replicated_extra_fsdp) or len(tensor.placements) != len(names): + return None + if len(names) == 2 and not isinstance(tensor.placements[0], Replicate): + return None + shard = tensor.placements[-1] + if type(shard) not in (Shard, Replicate): + return None + return { + "ep_shard_dim": None, + "fsdp_shard_dim": shard.dim % tensor.ndim if isinstance(shard, Shard) else None, + "ep_size": 1, + "fsdp_size": mesh.size(-1), + "expert_order": "canonical", + "source_kind": "veomni_tp1_no_ep", } +def _unidentified_same_axis_shards(tensor): + """Do not certify a multi-axis layout whose mesh names hide its origin.""" + if not isinstance(tensor, DTensor) or not tensor.ndim: + return False + dims = [ + placement.dim % tensor.ndim + for axis, placement in enumerate(tensor.placements) + if isinstance(placement, Shard) and tensor.device_mesh.size(axis) > 1 + ] + return len(dims) != len(set(dims)) + + class _ExtraParallelSavePlanner(DefaultSavePlanner): """Record per-tensor source layouts in DCP metadata, including optimizer state.""" + def __init__(self, *args, parallel_state=None, source_model=None, adapter_name=None, **kwargs): + super().__init__(*args, **kwargs) + self.parallel_state = parallel_state + self.ordinary_fqns = ( + _ordinary_expert_fqns(source_model, parallel_state) + if isinstance(parallel_state, ParallelState) + and any(size != 1 for size in parallel_state.extra_parallel_sizes.values()) + else set() + ) + self.shared_lora_fqns = _shared_lora_fqns(source_model, self.ordinary_fqns) if self.ordinary_fqns else set() + if adapter_name is not None and self.ordinary_fqns: + from torch.distributed.checkpoint.state_dict import _get_fqns + + from ..lora.state_dict import get_lora_state_dict, strip_adapter_name + + exported = get_lora_state_dict( + source_model, adapter_name=adapter_name, config=source_model.get_lora_config() + ) + owners, candidates, shared_candidates = {}, set(), set() + for name, param in source_model.named_parameters(remove_duplicate=False): + exported_name = strip_adapter_name(name, adapter_name) + if exported_name not in exported: + continue + owners.setdefault(exported_name, set()).add(id(param)) + if name in self.ordinary_fqns or _get_fqns(source_model, name) & self.ordinary_fqns: + candidates.add(exported_name) + if name in self.shared_lora_fqns or _get_fqns(source_model, name) & self.shared_lora_fqns: + shared_candidates.add(exported_name) + self.ordinary_fqns = {name for name in candidates if len(owners[name]) == 1} + self.shared_lora_fqns = {name for name in shared_candidates if len(owners[name]) == 1} + def create_local_plan(self): plan = super().create_local_plan() - layouts = {key: layout for key, value in self.state_dict.items() if (layout := _expert_layout(value))} + layouts = {} + for key, value in self.state_dict.items(): + extra_names = getattr(self.parallel_state, "extra_parallel_names", ()) + layout = _expert_layout(value, extra_names) or _no_ep_expert_layout( + key, value, self.parallel_state, self.ordinary_fqns, self.shared_lora_fqns + ) + if layout is not None: + layouts[key] = layout + elif _unidentified_same_axis_shards(value) or ( + "experts" in key.split(".") and torch.is_tensor(value) and value.ndim > 0 + ): + layouts[key] = {"expert_order": "unproven"} return replace(plan, planner_data={**(plan.planner_data or {}), _EXPERT_LAYOUTS_KEY: layouts}) def create_global_plan(self, all_plans): @@ -113,7 +355,10 @@ def create_global_plan(self, all_plans): # some ranks. Keep the ordinary FQN -> state-dict path mappings intact. layouts = {} for plan in all_plans: - layouts.update(plan.planner_data[_EXPERT_LAYOUTS_KEY]) + for key, value in plan.planner_data[_EXPERT_LAYOUTS_KEY].items(): + if key in layouts and layouts[key] != value: + raise RuntimeError(f"Checkpoint tensor {key!r} has inconsistent source layouts across ranks") + layouts[key] = value metadata.planner_data = { **(metadata.planner_data or {}), _EXPERT_LAYOUTS_KEY: {"version": 1, "tensors": layouts}, @@ -183,16 +428,17 @@ def create_local_plan(self): def _validate_expert_layouts(self): planner_data = self.metadata.planner_data or {} source = planner_data.get(_EXPERT_LAYOUTS_KEY) + recorded = {} if source is not None: if source["version"] != 1: raise RuntimeError(f"Unsupported checkpoint expert layout version: {source['version']}") - for key in self.state_dict.keys() & source["tensors"].keys(): - if source["tensors"][key]["expert_order"] != "canonical": + recorded = source["tensors"] + for key in self.state_dict.keys() & recorded.keys(): + if recorded[key]["expert_order"] != "canonical": raise RuntimeError( - f"Checkpoint tensor {key!r} has noncanonical expert order: {source['tensors'][key]}. " + f"Checkpoint tensor {key!r} has noncanonical expert order: {recorded[key]}. " "Migrate model and optimizer state using the original training layout before loading." ) - return # Checkpoints produced by the earlier version of this feature already # used canonical coordinates, but did not record individual layouts. @@ -207,10 +453,17 @@ def _validate_expert_layouts(self): return for key, value in self.state_dict.items(): + if key in recorded: + continue layout = _expert_layout(value) saved = self.metadata.state_dict_metadata.get(key) if saved is None: continue + if source is not None and "experts" in key.split(".") and bool(getattr(saved, "size", ())): + raise RuntimeError( + f"Cannot determine source expert order for {key!r}: version-1 metadata omits its layout. " + "Verify and migrate the original training layout before loading." + ) # For legacy VeOmni EP weights, the faulty same-axis layout only # split the expert dimension. A source chunk that splits matrix # rows (old dim-1 FSDP), or spans all experts, cannot have that @@ -220,7 +473,7 @@ def _validate_expert_layouts(self): # letting it re-save would incorrectly mark that data as # canonical. The stacked VeOmni expert axis is dimension 0. shape = getattr(saved, "size", ()) - if "experts" not in key.split(".") or len(shape) < 2: + if "experts" not in key.split(".") or not shape: continue ep_dim = 0 else: @@ -979,8 +1232,15 @@ def restore_extra_parallel_dim( placements = [fsdp_placement, Shard(ep_shard_dim)] dtensor = DTensor.from_local(orgin_tensor._local_tensor, device_mesh=fsdp_mesh, placements=placements) elif torch.is_tensor(orgin_tensor): - # If there is no FSDP but only ExtraParallel - dtensor = DTensor.from_local(orgin_tensor, device_mesh=extra_parallel_fsdp_mesh, placements=[Shard(0)]) + # A plain tensor has no FSDP shard. Restore its actual ExtraParallel + # axis, not the FSDP submesh, so DCP sees the complete global tensor. + ndim = orgin_tensor.ndim + if not -ndim <= ep_shard_dim < ndim: + raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") + if fsdp_mesh.size(-2) != 1 or fsdp_mesh.mesh_dim_names is None: + raise RuntimeError("A plain ExtraParallel tensor requires a size-one named FSDP axis") + extra_mesh = fsdp_mesh[fsdp_mesh.mesh_dim_names[-1]] + dtensor = DTensor.from_local(orgin_tensor, device_mesh=extra_mesh, placements=[Shard(ep_shard_dim % ndim)]) else: raise RuntimeError(f"origin_tensor - {orgin_tensor} is not a tensor!") @@ -1588,6 +1848,8 @@ def save( save_to_lowest_rank=save_to_lowest_rank, slot=WEIGHTS_DIRNAME, timeout_seconds=save_timeout_seconds, + parallel_state=parallel_state, + source_model=state["model"], ) if "optimizer" in state and state["optimizer"] is not None: @@ -1605,6 +1867,8 @@ def save( save_to_lowest_rank=save_to_lowest_rank, slot=OPTIMIZER_DIRNAME, timeout_seconds=save_timeout_seconds, + parallel_state=parallel_state, + source_model=state["model"], ) except BaseException: if stage_path is not None and _local_rank() == 0: @@ -1872,6 +2136,8 @@ def execute_save( save_to_lowest_rank: bool = False, slot: str = WEIGHTS_DIRNAME, timeout_seconds: Optional[int] = None, + parallel_state=None, + source_model=None, ) -> None: """Execute DCP save with optional async support. @@ -1884,7 +2150,16 @@ def execute_save( outlive its own next write. ``timeout_seconds`` is that slot's group timeout, applied when the group is created. """ - planner = _ExtraParallelSavePlanner(dedup_save_to_lowest_rank=save_to_lowest_rank) + planner, error = None, None + try: + planner = _ExtraParallelSavePlanner( + dedup_save_to_lowest_rank=save_to_lowest_rank, + parallel_state=parallel_state if parallel_state is not None else get_parallel_state(), + source_model=source_model, + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "preparing the checkpoint save planner") if save_async: # Lazily create this slot's dedicated Gloo process group. Creating a # group is itself collective, and every rank runs the same save @@ -2023,9 +2298,6 @@ def _validate_expert_export_provenance(metadata: Metadata, model_keys) -> None: raise RuntimeError( f"Checkpoint tensor {key!r} has noncanonical expert order; migrate it before export" ) - # The writer records every EP tensor; omissions are ordinary tensors, - # including FSDP-only expert shards, whose coordinates are canonical. - return marker_key = "model." + _extra_parallel_layout_key({}) if marker_key in metadata.state_dict_metadata or planner_data.get("__veomni_ep_first_layout__") == 1: @@ -2036,8 +2308,13 @@ def _validate_expert_export_provenance(metadata: Metadata, model_keys) -> None: for key in model_keys: tensor_metadata = metadata.state_dict_metadata[key] shape = getattr(tensor_metadata, "size", ()) - if "experts" not in key.split(".") or len(shape) < 2 or key in layouts: + if "experts" not in key.split(".") or not shape or key in layouts: continue + if recorded is not None: + raise RuntimeError( + f"Cannot determine source expert order for {key!r}: version-1 metadata omits its layout. " + "Verify and migrate the original training layout before exporting." + ) if any( chunk.sizes[0] < shape[0] and all(chunk.sizes[dim] == shape[dim] for dim in range(1, len(shape))) for chunk in tensor_metadata.chunks diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 8644cf2c9f..cb9b248968 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -293,6 +293,7 @@ def save_lora(self, state: "TrainerState", stage: str = "step_end", adapter_name model=self.runtime.model, save_path=self.lora_export_dir(state), adapter_name=adapter_name, + parallel_state=self.parallel_state, ) helper.empty_cache() dist.barrier() diff --git a/veomni/utils/save_safetensor_utils.py b/veomni/utils/save_safetensor_utils.py index 049064ef71..4c95178c35 100644 --- a/veomni/utils/save_safetensor_utils.py +++ b/veomni/utils/save_safetensor_utils.py @@ -292,9 +292,9 @@ def save_hf_safetensor( """Save model weights in HuggingFace safetensors format. This function is self-contained w.r.t. synchronization: it calls ``synchronize()`` at - entry to flush pending GPU operations before reading tensor data, and calls - ``dist.barrier()`` before returning to ensure all ranks complete the save. Callers - do not need to add external synchronization around this function. + entry to flush pending GPU operations before reading tensor data. All ranks + coordinate completion and errors before returning; only rank 0 writes in legacy + mode. Callers do not need to add external synchronization around this function. Supports two modes: - Distributed mode (PyTorch >= 2.9, ckpt_manager="dcp", non-LoRA): Uses HuggingFaceStorageWriter @@ -335,15 +335,23 @@ def save_hf_safetensor( model, save_hf_safetensor_path, fqn_to_index_mapping, model_assets, parallel_state=parallel_state ) else: - # The single-writer path is rank-0 only; non-rank-0 waits at the barrier below - if is_rank_0: - _save_hf_safetensor_legacy( - save_checkpoint_path, - save_hf_safetensor_path, - model_assets, - ckpt_manager, - output_dir, - ) + from veomni.utils.dist_utils import raise_if_any_rank_failed + + error = None + try: + if is_rank_0: + _save_hf_safetensor_legacy( + save_checkpoint_path, + save_hf_safetensor_path, + model_assets, + ckpt_manager, + output_dir, + ) + except BaseException as exc: + error = exc + group = dist.group.WORLD if dist.is_initialized() else None + raise_if_any_rank_failed(error, "rank-0 HuggingFace export", group=group) + return # Ensure all ranks finish saving before anyone proceeds if dist.is_initialized(): @@ -356,6 +364,7 @@ def save_lora_adapter_with_dcp( save_path: str, adapter_name: str = "default", dcp_subdir: str = ".lora_dcp_tmp", + parallel_state=None, ) -> None: """Save LoRA adapter with DCP parallel write and rank-0 consolidation. @@ -372,6 +381,7 @@ def save_lora_adapter_with_dcp( """ from torch.distributed.tensor import Shard + from veomni.distributed.parallel_state import get_parallel_state, is_parallel_state_initialized from veomni.lora import is_veomni_lora_model from veomni.lora.state_dict import get_lora_state_dict @@ -493,37 +503,59 @@ def save_lora_adapter_with_dcp( single_file_per_rank=True, sync_files=False, ) + planner, error = None, None + try: + single_process = ( + dist.get_world_size() == 1 if dist.is_initialized() else os.environ.get("WORLD_SIZE", "1") == "1" + ) + planner = _ExtraParallelSavePlanner( + parallel_state=( + parallel_state + if parallel_state is not None + else get_parallel_state() + if is_parallel_state_initialized() or single_process + else None + ), + source_model=model, + adapter_name=adapter_name, + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "preparing the LoRA checkpoint save planner") dcp.save( state_dict=dcp_lora_state, storage_writer=storage_writer, - planner=_ExtraParallelSavePlanner(), + planner=planner, ) if dist.is_initialized(): dist.barrier() is_rank_0 = not dist.is_initialized() or dist.get_rank() == 0 - if is_rank_0: - consolidated_state = ckpt_to_state_dict( - save_checkpoint_path=dcp_save_path, - ckpt_manager="dcp", - ) - # safetensors needs contiguous, non-aliased tensors; consolidation returns - # fresh tensors but ``.contiguous()`` is a cheap no-op when already so. - consolidated_state = {k: v.contiguous() for k, v in consolidated_state.items()} - adapter_model_file = os.path.join(save_path, "adapter_model.safetensors") - _save_state_dict(consolidated_state, adapter_model_file, safe_serialization=True) - - # VeOmniLoraModel writes a single PEFT-loadable adapter_config.json - # (MoE metadata, if any, lives in its ``veomni_lora`` block -- no - # separate sidecar). - model.get_lora_config().save_pretrained(save_path) - - shutil.rmtree(dcp_save_path, ignore_errors=True) - logger.info_rank0(f"LoRA adapter saved at {save_path} successfully!") - - if dist.is_initialized(): - dist.barrier() + error = None + try: + if is_rank_0: + consolidated_state = ckpt_to_state_dict( + save_checkpoint_path=dcp_save_path, + ckpt_manager="dcp", + ) + # safetensors needs contiguous, non-aliased tensors; consolidation returns + # fresh tensors but ``.contiguous()`` is a cheap no-op when already so. + consolidated_state = {k: v.contiguous() for k, v in consolidated_state.items()} + adapter_model_file = os.path.join(save_path, "adapter_model.safetensors") + _save_state_dict(consolidated_state, adapter_model_file, safe_serialization=True) + + # VeOmniLoraModel writes a single PEFT-loadable adapter_config.json + # (MoE metadata, if any, lives in its ``veomni_lora`` block -- no + # separate sidecar). + model.get_lora_config().save_pretrained(save_path) + + shutil.rmtree(dcp_save_path, ignore_errors=True) + logger.info_rank0(f"LoRA adapter saved at {save_path} successfully!") + except BaseException as exc: + error = exc + group = dist.group.WORLD if dist.is_initialized() else None + raise_if_any_rank_failed(error, "rank-0 LoRA adapter export", group=group) gc.collect() helper.empty_cache() From 32871dd72a04ff9c407999573bee6680281d73f7 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 19:54:40 +0800 Subject: [PATCH 054/120] [checkpoint] guard sparse and legacy extra-parallel states --- docs/usage/checkpoint.md | 4 + tests/checkpoints/test_expert_shard_layout.py | 174 +++++++++++++++++- veomni/checkpoint/dcp_checkpointer.py | 73 ++++++-- 3 files changed, 227 insertions(+), 24 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 23cc6941f5..d0c0dd1815 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -277,6 +277,10 @@ synchronization; other parameters on that mesh remain unproven. A normal single-process save assumes the caller supplied the complete model, not a manually extracted expert shard. The singleton case is not a source-layout proof for external, manually sharded tensors. +For a plain model tensor sharded on an ExtraParallel axis and replicated across +its FSDP/HSDP axes, all replicas at the same ExtraParallel coordinate must have +identical values; checkpoint serialization does not verify or synchronize them. +Plain optimizer shards on multiple such ranks require a rank-local state codec. Older version-1 metadata may omit non-`ep` shards: an unrecorded expert tensor is rejected unless its source layout can be verified and migrated. This also rejects an indistinguishable old ordinary-FSDP expert tensor. Historical non-expert tensors diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 05c95f9666..e3fd28486f 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -13,7 +13,8 @@ import torch.multiprocessing as mp from torch.distributed.checkpoint.default_planner import DefaultSavePlanner from torch.distributed.device_mesh import DeviceMesh, init_device_mesh -from torch.distributed.tensor import DTensor, Replicate, Shard +from torch.distributed.tensor import DTensor, Partial, Replicate, Shard +from torch.distributed.tensor.placement_types import _StridedShard from scripts.ci.accelerator.launch_plan import from_parameter from tests.tools.launch_utils import spawn_with_rendezvous @@ -361,6 +362,62 @@ def test_unexplained_mesh_axis_overlap_is_not_certified(outer): assert _expert_layout(value)["expert_order"] == ("unproven" if isinstance(outer, Shard) else "canonical") +@pytest.mark.parametrize( + "placements", + [ + [Shard(1), _StridedShard(1, split_factor=2), Shard(0)], + [Replicate(), _StridedShard(0, split_factor=3), Shard(0)], + [Replicate(), Shard(1), _StridedShard(0, split_factor=2)], + ], +) +def test_unproven_strided_extra_layout_is_not_certified(placements): + mesh = DeviceMesh( + "cpu", + [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], + mesh_dim_names=("x", "ep_fsdp", "ep"), + _init_backend=False, + _rank=0, + ) + value = DTensor.from_local(torch.zeros(3, 4), mesh, placements, run_check=False) + assert _expert_layout(value)["expert_order"] == "unproven" + + +def test_strided_unidentified_mesh_is_not_certified(): + from veomni.checkpoint.dcp_checkpointer import _unidentified_same_axis_shards + + mesh = DeviceMesh("cpu", [[0, 1], [2, 3]], mesh_dim_names=("x", "y"), _init_backend=False, _rank=0) + value = DTensor.from_local( + torch.zeros(2, 2), mesh, [_StridedShard(0, split_factor=2), Replicate()], run_check=False + ) + assert _unidentified_same_axis_shards(value) + + +def test_unexplained_strided_axis_does_not_short_circuit_provenance(): + mesh = DeviceMesh( + "cpu", + [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], + mesh_dim_names=("x", "ep_fsdp", "ep"), + _init_backend=False, + _rank=0, + ) + value = DTensor.from_local( + torch.zeros(3, 4, 5), mesh, [_StridedShard(1, split_factor=2), Shard(2), Shard(0)], run_check=False + ) + assert _expert_layout(value)["expert_order"] == "unproven" + + +def test_partial_placement_does_not_certify_source_values(): + mesh = DeviceMesh( + "cpu", + [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], + mesh_dim_names=("x", "ep_fsdp", "ep"), + _init_backend=False, + _rank=0, + ) + value = DTensor.from_local(torch.zeros(3, 4), mesh, [Partial(), Shard(1), Shard(0)], run_check=False) + assert _expert_layout(value)["expert_order"] == "unproven" + + def _ple_and_ordinary_expert_worker(rank, root, endpoint): from torch.distributed.tensor import distribute_tensor @@ -686,15 +743,19 @@ def test_ep_replicated_shared_expert_lora_export(world, mode, tmp_path): spawn_with_rendezvous(_ep_shared_lora_worker, args=(str(tmp_path / "adapter"), mode), nprocs=world) -def _pure_extra_parallel_worker(rank, root, name, endpoint): +def _pure_extra_parallel_worker(rank, root, name, world, endpoint): from veomni.distributed.parallel_plan import SpecInfo - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + dist.init_process_group( + "gloo", store=endpoint.connect(), rank=rank, world_size=world, timeout=timedelta(seconds=30) + ) try: - mesh = init_device_mesh("cpu", (1, 2), mesh_dim_names=(f"{name}_fsdp", name)) - fsdp_mesh = mesh[f"{name}_fsdp"] + shape = (2, 2, 2) if world == 8 else (world // 2, 2) + names = (f"{name}_replicate", f"{name}_fsdp", name) if world == 8 else (f"{name}_fsdp", name) + mesh = init_device_mesh("cpu", shape, mesh_dim_names=names) + fsdp_mesh = mesh[names[:-1]] expected = torch.arange(16, dtype=torch.float32).reshape(8, 2) - local = expected.chunk(2, dim=0)[rank].contiguous() + local = expected.chunk(2, dim=0)[mesh.get_local_rank(name)].contiguous() restored_axis = restore_extra_parallel_dim(local, mesh, fsdp_mesh, ep_shard_dim=0) torch.testing.assert_close(restored_axis.full_tensor(), expected, rtol=0, atol=0) parallel = SimpleNamespace( @@ -739,10 +800,107 @@ def model_with(local_value): @pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) -@pytest.mark.parametrize("world", [2]) +@pytest.mark.parametrize("world", [2, 4, 8]) @pytest.mark.parametrize("name", ["ep", "ple"]) def test_plain_pure_extra_parallel_preserves_global_values(world, name, tmp_path): - spawn_with_rendezvous(_pure_extra_parallel_worker, args=(str(tmp_path / "dcp"), name), nprocs=world) + spawn_with_rendezvous(_pure_extra_parallel_worker, args=(str(tmp_path / "dcp"), name, world), nprocs=world) + + +def _plain_optimizer_unproven_worker(rank, hsdp, endpoint): + from veomni.checkpoint.dcp_checkpointer import _apply_extra_parallel_dim + from veomni.distributed.parallel_plan import SpecInfo + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=4, timeout=timedelta(seconds=30)) + try: + shape = (2, 1, 2) if hsdp else (2, 2) + names = ("ep_replicate", "ep_fsdp", "ep") if hsdp else ("ep_fsdp", "ep") + mesh = init_device_mesh("cpu", shape, mesh_dim_names=names) + parallel = SimpleNamespace(extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh}) + spec = SpecInfo("ep", Shard(0), "experts.weight", mesh, 1) + # Values differ between FSDP ranks and must not be deduplicated as replicas. + state = {"state.experts.weight.exp_avg": torch.full((4, 4), float(rank))} + with pytest.raises(RuntimeError, match="no proven replica layout"): + _apply_extra_parallel_dim(state, {"experts.weight": spec}, parallel, "restore", key_match="substring") + torch.testing.assert_close(state["state.experts.weight.exp_avg"], torch.full((4, 4), float(rank))) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [4]) +@pytest.mark.parametrize("hsdp", [False, True]) +def test_plain_optimizer_shards_are_not_treated_as_replicas(world, hsdp): + spawn_with_rendezvous(_plain_optimizer_unproven_worker, args=(hsdp,), nprocs=world) + + +def _sparse_plain_optimizer_failure_worker(rank, endpoint): + from veomni.checkpoint.dcp_checkpointer import OptimizerState + from veomni.distributed.parallel_plan import SpecInfo + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=4, timeout=timedelta(seconds=30)) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ep_fsdp", "ep")) + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + model = torch.nn.Module() + model._fqn2spec_info = {"experts.weight": SpecInfo("ep", Shard(0), "experts.weight", mesh, 1)} + + class SparseOptimizer: + _is_multi_optimizer = True + + def state_dict(self): + return {"state.experts.weight.exp_avg": torch.ones(4, 4)} if rank == 0 else {} + + state = OptimizerState(model, SparseOptimizer(), parallel_state=parallel) + with pytest.raises(RuntimeError, match="no proven replica layout|failed on another rank"): + state.state_dict() + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 4 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [4]) +def test_sparse_plain_optimizer_failure_exits_all_ranks(world): + spawn_with_rendezvous(_sparse_plain_optimizer_failure_worker, args=(), nprocs=world) + + +def _subgroup_sparse_optimizer_failure_worker(rank, endpoint): + from veomni.checkpoint.dcp_checkpointer import OptimizerState + from veomni.distributed.parallel_plan import SpecInfo + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=4, timeout=timedelta(seconds=30)) + try: + group = dist.new_group([0, 1], timeout=timedelta(seconds=20)) + mesh = DeviceMesh("cpu", [[0], [1]], mesh_dim_names=("ep_fsdp", "ep")) + if rank < 2: + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + model = torch.nn.Module() + model._fqn2spec_info = {"experts.weight": SpecInfo("ep", Shard(0), "experts.weight", mesh, 1)} + + class SparseOptimizer: + _is_multi_optimizer = True + + def state_dict(self): + return {"state.experts.weight.exp_avg": torch.ones(4, 4)} if rank == 0 else {} + + state = OptimizerState(model, SparseOptimizer(), parallel_state=parallel, load=True, process_group=group) + with pytest.raises(RuntimeError, match="no proven replica layout|failed on another rank"): + state.state_dict() + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [4]) +def test_sparse_optimizer_failure_coordinates_only_the_load_group(world): + spawn_with_rendezvous(_subgroup_sparse_optimizer_failure_worker, args=(), nprocs=world) class _LegacyNamedExtraSavePlanner(DefaultSavePlanner): diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index f9349fe3c4..0863a1d9c3 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -28,7 +28,7 @@ import torch import torch.distributed as dist import torch.distributed.checkpoint as dcp -from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard +from torch.distributed._tensor import DeviceMesh, DTensor, Partial, Replicate, Shard from torch.distributed.checkpoint import ( FileSystemReader, FileSystemWriter, @@ -79,6 +79,8 @@ def _expert_layout(tensor, extra_parallel_names=()): return None mesh = tensor.device_mesh names = mesh.mesh_dim_names or () + if any(isinstance(p, Partial) and mesh.size(axis) > 1 for axis, p in enumerate(tensor.placements)): + return {"expert_order": "unproven"} extra_names = [name for name in names if name == "ep" or name in extra_parallel_names or f"{name}_fsdp" in names] if not extra_names: return None @@ -87,6 +89,8 @@ def _expert_layout(tensor, extra_parallel_names=()): extra_name = extra_names[0] ep_axis = names.index(extra_name) ep = tensor.placements[ep_axis] + if isinstance(ep, _StridedShard): + return {"expert_order": "unproven"} if not isinstance(ep, Shard): return None fsdp_name = f"{extra_name}_fsdp" @@ -99,7 +103,7 @@ def _expert_layout(tensor, extra_parallel_names=()): shard_axes = [ (axis, placement.dim % tensor.ndim) for axis, placement in enumerate(tensor.placements) - if isinstance(placement, Shard) and mesh.size(axis) > 1 + if isinstance(placement, (Shard, _StridedShard)) and mesh.size(axis) > 1 ] unknown_overlap = any( axis not in (ep_axis, fsdp_axis) @@ -108,7 +112,15 @@ def _expert_layout(tensor, extra_parallel_names=()): ) canonical = not same_axis or isinstance(fsdp, _StridedShard) or (fsdp_axis is not None and ep_axis < fsdp_axis) order = "canonical" if canonical else "legacy_fsdp_first" - if unknown_overlap: + unknown_strided = any( + axis not in (ep_axis, fsdp_axis) and isinstance(placement, _StridedShard) and mesh.size(axis) > 1 + for axis, placement in enumerate(tensor.placements) + ) + if ( + unknown_overlap + or unknown_strided + or (isinstance(fsdp, _StridedShard) and (fsdp_dim != ep_dim or fsdp.split_factor != mesh.size(ep_axis))) + ): order = "unproven" return { "ep_shard_dim": ep_dim, @@ -291,10 +303,15 @@ def _unidentified_same_axis_shards(tensor): """Do not certify a multi-axis layout whose mesh names hide its origin.""" if not isinstance(tensor, DTensor) or not tensor.ndim: return False + if any( + isinstance(placement, (_StridedShard, Partial)) and tensor.device_mesh.size(axis) > 1 + for axis, placement in enumerate(tensor.placements) + ): + return True dims = [ placement.dim % tensor.ndim for axis, placement in enumerate(tensor.placements) - if isinstance(placement, Shard) and tensor.device_mesh.size(axis) > 1 + if isinstance(placement, (Shard, _StridedShard)) and tensor.device_mesh.size(axis) > 1 ] return len(dims) != len(set(dims)) @@ -886,6 +903,16 @@ def _apply_extra_parallel_dim( # ``(ep_replicate, ep_fsdp)`` for HSDP. Mirrors the mesh slicing used # when sharding the module in ``torch_parallelize``. fsdp_submesh = spec_info.para_fsdp_mesh[spec_info.para_fsdp_mesh.mesh_dim_names[:-1]] + if ( + action == "restore" + and key_match == "substring" + and not isinstance(tensor, DTensor) + and fsdp_submesh.size() > 1 + ): + raise RuntimeError( + "Plain optimizer state on multiple ExtraParallel FSDP ranks has no proven replica layout; " + "use a rank-local optimizer checkpoint codec." + ) # The ExtraParallel (e.g. ``ep``) dim shards the tensor along # ``spec_info.placement.dim`` (dim 0 for experts), while FSDP shards it # along ``spec_info.fsdp_shard_dim`` (1 by default, 0 for Muon zero-comm). @@ -1042,9 +1069,10 @@ class OptimizerState(Stateful): MoE experts) in the checkpoint. """ - def __init__(self, model, optimizer, parallel_state=None, *, load: bool = False): + def __init__(self, model, optimizer, parallel_state=None, *, load: bool = False, process_group=None): self.model = model self.optimizer = optimizer + self.process_group = process_group self.parallel_state = parallel_state if parallel_state is not None else get_parallel_state() self.extra_parallel_fqn2spec_info = getattr(self.model, "_fqn2spec_info", None) self._layout_marker_key = ( @@ -1073,8 +1101,15 @@ def state_dict(self): "ExtraParallel is enabled but optimizer is not a MultiOptimizer instance" ) vanilla_optim_sd = self.optimizer.state_dict() - optim_sd_with_extra_parallel_dim = self.get_state_dict_with_extra_parallel_dim_preprocess( - vanilla_optim_sd, "restore" + optim_sd_with_extra_parallel_dim, error = None, None + try: + optim_sd_with_extra_parallel_dim = self.get_state_dict_with_extra_parallel_dim_preprocess( + vanilla_optim_sd, "restore" + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed( + error, "preparing ExtraParallel optimizer checkpoint state", group=self.process_group ) if _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info): optim_sd_with_extra_parallel_dim = _namespace_extra_parallel_layout_state( @@ -1151,7 +1186,8 @@ def drop_extra_parallel_dim(loaded_tensor: torch.Tensor, device_mesh: DeviceMesh the Muon zero-comm layout). The number of placements on the loaded DTensor reflects the full saved mesh: - * 1 placement: pure ExtraParallel, no FSDP -> return the plain local tensor. + * A pure ExtraParallel shard, with only replicated FSDP axes -> return the + plain local tensor. * 2 placements: ``(ep_fsdp, ep)`` -> keep FSDP ``Shard(fsdp_shard_dim)`` on the 1D sub-mesh. * 3 placements: ``(ep_replicate, ep_fsdp, ep)`` (HSDP) -> keep @@ -1159,7 +1195,7 @@ def drop_extra_parallel_dim(loaded_tensor: torch.Tensor, device_mesh: DeviceMesh """ num_placements = len(loaded_tensor.placements) - if num_placements == 1: + if num_placements == 1 or all(isinstance(p, Replicate) for p in loaded_tensor.placements[:-1]): tensor_to_put = loaded_tensor.to_local() elif num_placements == 2: tensor_to_put = DTensor.from_local( @@ -1214,8 +1250,10 @@ def restore_extra_parallel_dim( if isinstance(orgin_tensor, DTensor): ndim = orgin_tensor.ndim - if not (-ndim <= ep_shard_dim < ndim and -ndim <= fsdp_shard_dim < ndim): + if not -ndim <= ep_shard_dim < ndim: raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") + if not -ndim <= fsdp_shard_dim < ndim: + raise RuntimeError("FSDP shard dimension is outside the tensor rank") ep_shard_dim %= ndim fsdp_shard_dim %= ndim # ExtraParallel+FSDP2. mesh order is (ep_fsdp, ep) or, for HSDP, @@ -1232,15 +1270,17 @@ def restore_extra_parallel_dim( placements = [fsdp_placement, Shard(ep_shard_dim)] dtensor = DTensor.from_local(orgin_tensor._local_tensor, device_mesh=fsdp_mesh, placements=placements) elif torch.is_tensor(orgin_tensor): - # A plain tensor has no FSDP shard. Restore its actual ExtraParallel - # axis, not the FSDP submesh, so DCP sees the complete global tensor. + # A plain model tensor has no FSDP shard. The caller must ensure every + # FSDP/HSDP replica for the same ExtraParallel coordinate has identical + # values; from_local does not check or synchronize them. Optimizer + # states without this guarantee are rejected by the caller. ndim = orgin_tensor.ndim if not -ndim <= ep_shard_dim < ndim: raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") - if fsdp_mesh.size(-2) != 1 or fsdp_mesh.mesh_dim_names is None: - raise RuntimeError("A plain ExtraParallel tensor requires a size-one named FSDP axis") - extra_mesh = fsdp_mesh[fsdp_mesh.mesh_dim_names[-1]] - dtensor = DTensor.from_local(orgin_tensor, device_mesh=extra_mesh, placements=[Shard(ep_shard_dim % ndim)]) + if fsdp_mesh.mesh_dim_names is None: + raise RuntimeError("A plain ExtraParallel tensor requires named mesh dimensions") + placements = [Replicate()] * (fsdp_mesh.ndim - 1) + [Shard(ep_shard_dim % ndim)] + dtensor = DTensor.from_local(orgin_tensor, device_mesh=fsdp_mesh, placements=placements) else: raise RuntimeError(f"origin_tensor - {orgin_tensor} is not a tensor!") @@ -1969,6 +2009,7 @@ def _optimizer_state() -> "OptimizerState": optimizer=state["optimizer"], parallel_state=parallel_state, load=True, + process_group=process_group, ) wants_optimizer = state.get("optimizer") is not None From 9385940c5230c301ec38e8b60dd065c164b05b75 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 20:56:36 +0800 Subject: [PATCH 055/120] [checkpoint] remove counterfeit adapter fixture --- .../test_ranklocal_provenance_integration.py | 23 ------------------- 1 file changed, 23 deletions(-) diff --git a/tests/checkpoints/test_ranklocal_provenance_integration.py b/tests/checkpoints/test_ranklocal_provenance_integration.py index 1951ea94d0..3ba877599b 100644 --- a/tests/checkpoints/test_ranklocal_provenance_integration.py +++ b/tests/checkpoints/test_ranklocal_provenance_integration.py @@ -58,26 +58,3 @@ def build(seed): @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) def test_ranklocal_expert_optimizer_checkpoint_roundtrip(tmp_path): spawn_with_rendezvous(_worker, args=(str(tmp_path / "ckpt"),), nprocs=2) - - -def test_rank_prefix_or_fake_adapter_does_not_prove_ownership(tmp_path): - import torch.distributed.checkpoint as dcp - - from veomni.checkpoint.dcp_checkpointer import _ExtraParallelSavePlanner, _ModelStrictLoadPlanner - - key = "__veomni_rank_0__state.experts.weight.exp_avg" - full_key = f"optimizer.{key}" - fake = SimpleNamespace( - _verified_rank_local_keys=frozenset({full_key}), _load=False, should_extra_parallel_aware=True - ) - source = {"optimizer": {key: torch.ones(2, 2)}} - dcp.save(source, checkpoint_id=tmp_path, planner=_ExtraParallelSavePlanner(rank_local_adapter=fake)) - fake._load = True - target = {"optimizer": {key: torch.zeros(2, 2)}} - with pytest.raises(dcp.CheckpointException, match="noncanonical expert order"): - dcp.load( - target, - checkpoint_id=tmp_path, - planner=_ModelStrictLoadPlanner(strict_model=False, rank_local_adapter=fake), - ) - assert target["optimizer"][key].count_nonzero() == 0 From 7d7d86db1861a38ea91d1c46f2cb3c77bc5bd722 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 21:03:48 +0800 Subject: [PATCH 056/120] [checkpoint] use consistent replica values in test --- tests/checkpoints/test_expert_shard_layout.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index e3fd28486f..6893f24b9f 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -694,8 +694,8 @@ def _ep_shared_lora_worker(rank, root, mode, endpoint): model = _lora_model(config) base = model.base_model.model assert isinstance(base.experts, LoraSharedExperts) - if mode == "rank_divergent_bias": - base.experts.bias = torch.nn.Parameter(torch.full((8, 8), float(rank))) + if mode == "unproven_bias": + base.experts.bias = torch.nn.Parameter(torch.ones(8, 8)) base.get_parallel_plan = lambda: ParallelPlan( {"ep": {"experts.gate_up_proj": Shard(0), "experts.down_proj": Shard(0)}} ) @@ -738,7 +738,7 @@ def _ep_shared_lora_worker(rank, root, mode, endpoint): @pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) @pytest.mark.parametrize("world", [2]) -@pytest.mark.parametrize("mode", ["valid", "plan_conflict", "rank_divergent_bias"]) +@pytest.mark.parametrize("mode", ["valid", "plan_conflict", "unproven_bias"]) def test_ep_replicated_shared_expert_lora_export(world, mode, tmp_path): spawn_with_rendezvous(_ep_shared_lora_worker, args=(str(tmp_path / "adapter"), mode), nprocs=world) From dbef082d4798d048c6b0da38ff5567be7dfdb1b6 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 23:34:03 +0800 Subject: [PATCH 057/120] [test] declare absent scheduler in expert cold test --- tests/ci/_npu_moe_fp32_cold_worker.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/ci/_npu_moe_fp32_cold_worker.py b/tests/ci/_npu_moe_fp32_cold_worker.py index dd8f4c073f..d61e4bf04d 100644 --- a/tests/ci/_npu_moe_fp32_cold_worker.py +++ b/tests/ci/_npu_moe_fp32_cold_worker.py @@ -176,14 +176,18 @@ def run(rank, phase, output_dir, ep, endpoint=None): _local(state["exp_avg_sq"])[1].fill_((rank + 1) / 2048) optimizer.zero_grad() DistributedCheckpointer.save( - ckpt, {"model": model, "optimizer": optimizer}, parallel_state=get_parallel_state() + ckpt, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=get_parallel_state(), ) else: with torch.no_grad(): for param in model.parameters(): _local(param).fill_(-0.05 - rank * 0.01) DistributedCheckpointer.load( - ckpt, {"model": model, "optimizer": optimizer}, parallel_state=get_parallel_state() + ckpt, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=get_parallel_state(), ) before = _snapshot(model, optimizer) after = _step(model, optimizer, rank, ep, 2) From 3745392a45902c21ce8975ab69b99f1fd60b41e7 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 3 Nov 2026 23:35:11 +0800 Subject: [PATCH 058/120] [checkpoint] keep configuration preflight local --- .../test_multi_optimizer_rope_cache.py | 39 --------------- .../runtime/test_omni_model_runtime.py | 48 ------------------- veomni/models/model_runtime.py | 12 ++--- .../omni_model/omni_model_runtime.py | 18 +++---- 4 files changed, 9 insertions(+), 108 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index c0d7f91fc2..db4d5b8762 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -39,45 +39,6 @@ def test_runtime_preflight_rejects_unsupported_dynamic_rope_before_training(): VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) -def _runtime_rope_guard_io_failure(rank, endpoint): - from unittest.mock import patch - - from veomni.checkpoint import dynamic_rope - from veomni.models.model_runtime import VeOmniModelRuntime - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - torch.manual_seed(7) - model, optimizer = _build() - model.rope.rope_type = "default" # Same supported model configuration on every rank. - runtime = SimpleNamespace( - model=model, - optimizer=optimizer, - model_name="policy", - train_args=SimpleNamespace(checkpoint=SimpleNamespace(manager="dcp", save_async=False)), - ) - original = dynamic_rope.reject_dynamic_rope_multi_optimizer_training - failure = OSError("injected local RoPE validator failure") if rank == 1 else None - with patch.object( - dynamic_rope, "reject_dynamic_rope_multi_optimizer_training", wraps=original, side_effect=failure - ) as guard: - error = OSError if rank == 1 else RuntimeError - message = ( - "injected local RoPE validator failure" if rank == 1 else "dynamic RoPE checkpoint mode.*another rank" - ) - with pytest.raises(error, match=message): - VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) - guard.assert_called_once_with(model) - dist.barrier() - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_dynamic_rope_validator_local_failure_reaches_all_ranks(): - spawn_with_rendezvous(_runtime_rope_guard_io_failure, nprocs=2) - - class _TinyQwenRope(torch.nn.Module): def __init__(self, theta=10000): super().__init__() diff --git a/tests/seed_omni/runtime/test_omni_model_runtime.py b/tests/seed_omni/runtime/test_omni_model_runtime.py index 65d06c60dc..6915a2c02a 100644 --- a/tests/seed_omni/runtime/test_omni_model_runtime.py +++ b/tests/seed_omni/runtime/test_omni_model_runtime.py @@ -437,54 +437,6 @@ def test_composed_canonical_extra_parallel_rejected_before_wrap(monkeypatch): assert tuple(model._modules) == children -def _composed_canonical_error_worker(rank, endpoint): - from datetime import timedelta - from unittest.mock import patch - - import torch.distributed as dist - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - model = _model() - children = tuple(model._modules) - train = SimpleNamespace(checkpoint=SimpleNamespace(extra_parallel_canonical_layout=True)) - runtime = OmniModelRuntime(model, train_args=train) - runtime.omni_model_runtime_args = SimpleNamespace(accelerator=SimpleNamespace()) - - def read_state(): - if rank == 1: - raise RuntimeError("injected state failure") - return SimpleNamespace(any_extra_parallel_enabled=False) - - with ( - patch( - "veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime.composed_model_owns_wrap", - return_value=True, - ), - patch( - "veomni.models.seed_omni.accelerated.omni_model.omni_model_runtime.use_parallel_state", - return_value=nullcontext(), - ), - patch( - "veomni.distributed.parallel_state.get_parallel_state", - side_effect=read_state, - ), - ): - message = "injected state failure" if rank == 1 else "failed on another rank" - with pytest.raises(RuntimeError, match=message): - runtime._parallelize_composed_model() - assert tuple(model._modules) == children - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_composed_canonical_extra_parallel_error_is_rank_coordinated(): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_composed_canonical_error_worker, nprocs=2) - - def test_composed_wrap_leaves_cache_reader_modules_on_meta(monkeypatch): """Under ``fsdp_scope='model'``, cache readers must not go through weight materialisation.""" model = _model() diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index ddfbf1e2db..2cb32ab389 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -212,16 +212,10 @@ def _validate_optimizer_checkpoint_mode(self) -> None: else None ) raise_if_any_rank_failed(error, "checking MultiOptimizer checkpoint mode") - if manager == "dcp": - error = None - try: - if getattr(self.optimizer, "_is_multi_optimizer", False) is True: - from ..checkpoint.dynamic_rope import reject_dynamic_rope_multi_optimizer_training + if manager == "dcp" and getattr(self.optimizer, "_is_multi_optimizer", False) is True: + from ..checkpoint.dynamic_rope import reject_dynamic_rope_multi_optimizer_training - reject_dynamic_rope_multi_optimizer_training(self.model) - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "checking MultiOptimizer dynamic RoPE checkpoint mode") + reject_dynamic_rope_multi_optimizer_training(self.model) def __getattr__(self, name: str) -> Any: """Forward unshadowed :class:`torch.nn.Module` APIs to the wrapped model.""" diff --git a/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py b/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py index f6397aabd6..448b1288e2 100644 --- a/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py +++ b/veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py @@ -263,20 +263,14 @@ def _parallelize_composed_model(self, *, for_inference: bool = False) -> None: checkpoint = getattr(self.train_args, "checkpoint", None) if not for_inference and getattr(checkpoint, "extra_parallel_canonical_layout", False): from .....distributed.parallel_state import get_parallel_state - from .....utils.dist_utils import raise_if_any_rank_failed with use_parallel_state("base"): - error = None - try: - state = get_parallel_state() - if getattr(state, "any_extra_parallel_enabled", False): - raise RuntimeError( - "Canonical ExtraParallel checkpoints are not supported for SeedOmni fsdp_scope='model'; " - "use fsdp_scope='module'." - ) - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "checking composed Omni checkpoint layout") + state = get_parallel_state() + if getattr(state, "any_extra_parallel_enabled", False): + raise RuntimeError( + "Canonical ExtraParallel checkpoints are not supported for SeedOmni fsdp_scope='model'; " + "use fsdp_scope='module'." + ) modules = self.module_runtimes # Cache readers stay on meta and must not go through the composed From b4bc8a191e01da2fa4de9e179fd1f791f3da735a Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:02:52 +0800 Subject: [PATCH 059/120] [docs] link optimizer adapter consumers and tests --- .../rank_local_adamw_checkpoint_adapter.md | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/docs/design/rank_local_adamw_checkpoint_adapter.md b/docs/design/rank_local_adamw_checkpoint_adapter.md index ae2cd3286d..06e3430adb 100644 --- a/docs/design/rank_local_adamw_checkpoint_adapter.md +++ b/docs/design/rank_local_adamw_checkpoint_adapter.md @@ -62,15 +62,27 @@ the saved checkpoint. Change these only with an explicitly verified migration. ### Save and load call sites -`optimizer_resume_layout._local_layout` records `algorithm_identity`, each -parameter's `state_dtype`, and `state_is_offloaded` in the checkpoint identity; -a mismatch rejects resume before applying state. During save, -`multi_optimizer_checkpoint._validate_source_optimizer_states` uses -`state_dtype` to check initialized fields against their declared precision. -During load, `_prepare_multi_optimizer_saved_states` checks stored field dtypes -and uses `state_is_offloaded` to allocate missing destination state on CPU when -the adapter specifies offload. `context_factory` wraps the optimizer's actual -state read and load; adapter implementations own correct placement and values. +[`_local_layout`](../../veomni/checkpoint/optimizer_resume_layout.py#L105) +records `algorithm_identity`, each parameter's `state_dtype`, and +`state_is_offloaded` in the checkpoint identity; a mismatch rejects resume +before applying state. During save, +[`_validate_source_optimizer_states`](../../veomni/checkpoint/multi_optimizer_checkpoint.py#L89) +uses `state_dtype` to check initialized fields against their declared precision. +During load, +[`_prepare_multi_optimizer_saved_states`](../../veomni/checkpoint/multi_optimizer_checkpoint.py#L408) +checks stored field dtypes and uses `state_is_offloaded` to allocate missing +destination state on CPU when the adapter specifies offload. `context_factory` +wraps the optimizer's actual state read and load; adapter implementations own +correct placement and values. + +The adapter contract tests include the +[`test_registered_optimizer_preserves_lazy_state_and_next_update`](../../tests/checkpoints/test_optimizer_checkpoint_adapter.py#L71) +round trip and the policy-mismatch checks +[`test_registered_optimizer_rejects_changed_algorithm_before_loading`](../../tests/checkpoints/test_optimizer_checkpoint_adapter.py#L114) +and +[`test_lazy_adapter_state_dtype_policy_is_part_of_checkpoint_identity`](../../tests/checkpoints/test_optimizer_checkpoint_adapter.py#L132). +The toy adapter in these tests reports no offload; actual device-to-CPU placement +must also be validated for an offloading implementation. ## Framework call sites and restore behavior From 81c8bbea42dfa01860100268c2792219cc982d5d Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 10:05:50 +0800 Subject: [PATCH 060/120] [docs] link optimizer adapter context call sites --- docs/design/rank_local_adamw_checkpoint_adapter.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/design/rank_local_adamw_checkpoint_adapter.md b/docs/design/rank_local_adamw_checkpoint_adapter.md index 06e3430adb..53fb98f9f5 100644 --- a/docs/design/rank_local_adamw_checkpoint_adapter.md +++ b/docs/design/rank_local_adamw_checkpoint_adapter.md @@ -71,9 +71,10 @@ uses `state_dtype` to check initialized fields against their declared precision. During load, [`_prepare_multi_optimizer_saved_states`](../../veomni/checkpoint/multi_optimizer_checkpoint.py#L408) checks stored field dtypes and uses `state_is_offloaded` to allocate missing -destination state on CPU when the adapter specifies offload. `context_factory` -wraps the optimizer's actual state read and load; adapter implementations own -correct placement and values. +destination state on CPU when the adapter specifies offload. +[`rank_local_adamw_context`](../../veomni/optim/checkpoint_adapter.py#L55) +uses `context_factory` around the optimizer's actual state read and load; +adapter implementations own correct placement and values. The adapter contract tests include the [`test_registered_optimizer_preserves_lazy_state_and_next_update`](../../tests/checkpoints/test_optimizer_checkpoint_adapter.py#L71) @@ -86,7 +87,8 @@ must also be validated for an offloading implementation. ## Framework call sites and restore behavior -1. `MultiOptimizer._sparse_state_dict()` enters the context around PyTorch's +1. [`MultiOptimizer._sparse_state_dict()`](../../veomni/optim/optimizer.py#L563) + enters the context around PyTorch's optimizer-state serialization. Registered adapter implementations must preserve live state and serialize all initialized fields and group options faithfully. State-empty leaves retain their parameter groups without @@ -96,7 +98,8 @@ must also be validated for an offloading implementation. load checks the checkpoint's identity, field names, dtypes and shapes before preparing initialized state for the saved owners. Rank-local ownership and the supported parallel geometry are not automatically resharded. -3. `MultiOptimizer._dcp_load_state_dict()` enters the context again while +3. [`MultiOptimizer._dcp_load_state_dict()`](../../veomni/optim/optimizer.py#L601) + enters the context again while PyTorch applies the optimizer state. The trusted adapter must restore every validated field without changing its value or dtype. An optimizer may implement its own precision-preserving loader; the adapter callbacks alone From 50c60ffb7cd3ff1c7b6641a165fe73a56b32c753 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:06:28 +0800 Subject: [PATCH 061/120] [checkpoint] defer dynamic RoPE checks to save and load --- docs/usage/checkpoint.md | 10 +++++++--- tests/checkpoints/test_multi_optimizer_rope_cache.py | 12 +++++++----- veomni/checkpoint/dynamic_rope.py | 12 ------------ veomni/models/model_runtime.py | 4 ---- 4 files changed, 14 insertions(+), 24 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index dbd14a98dc..153813d1c1 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -459,9 +459,13 @@ or longrope target, resume requires a matching cache contract; older checkpoints without that proof need a verified migration. Older checkpoints do not identify the source's RoPE configuration, so the operator must also independently verify that the source and target configurations match. Saving an expanded cache is not -yet supported. DCP training with MultiOptimizer rejects dynamic and longrope -modes before training until their transient history can be restored exactly; -ordinary static/default RoPE is unaffected. +yet supported. Dynamic and longrope training may run without saving, but a +MultiOptimizer DCP save fails if its cache cannot be reconstructed; operators +should account for this when planning checkpoint cadence. Ordinary +static/default RoPE is unaffected. Disabling only `save_steps` does not disable +all saves: `save_epochs` defaults to 1 and `save_hf_weights` defaults to true. +Setting both intervals to 0 and `save_hf_weights` to false disables scheduled +DCP saves and the final HF export. For unversioned expert checkpoints whose chunks divide the expert axis, the saved offsets alone do not prove source order, even if another axis is also diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index db4d5b8762..7d998297c8 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -22,7 +22,7 @@ pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) -def test_runtime_preflight_rejects_unsupported_dynamic_rope_before_training(): +def test_runtime_preflight_allows_dynamic_rope_training_without_checkpoint(): from veomni.models.model_runtime import VeOmniModelRuntime model, optimizer = _build() @@ -32,11 +32,13 @@ def test_runtime_preflight_rejects_unsupported_dynamic_rope_before_training(): model_name="policy", train_args=SimpleNamespace(checkpoint=SimpleNamespace(manager="dcp", save_async=False)), ) - with pytest.raises(RuntimeError, match="cannot yet safely resume dynamic RoPE"): - VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) - # A normal, static RoPE module and the standard optimizer remain allowed. - model.rope.rope_type = "default" VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) + trace = _rope_update(model, optimizer) + assert torch.isfinite(trace["loss"]) + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + assert int(model.rope.max_seq_len_cached) == 32 + VeOmniModelRuntime._validate_optimizer_checkpoint_mode(runtime) + _update(model, optimizer, 0.25) class _TinyQwenRope(torch.nn.Module): diff --git a/veomni/checkpoint/dynamic_rope.py b/veomni/checkpoint/dynamic_rope.py index 10891f0739..9e30a81dff 100644 --- a/veomni/checkpoint/dynamic_rope.py +++ b/veomni/checkpoint/dynamic_rope.py @@ -9,18 +9,6 @@ import torch -def reject_dynamic_rope_multi_optimizer_training(model): - """Reject modes whose transient cache history cannot yet be resumed.""" - for name, module in model.named_modules(): - rope_type = getattr(module, "rope_type", None) - kinds = rope_type.values() if isinstance(rope_type, dict) else (rope_type,) - if any(isinstance(kind, str) and ("dynamic" in kind or kind == "longrope") for kind in kinds): - raise RuntimeError( - f"MultiOptimizer DCP training cannot yet safely resume dynamic RoPE for {name or ''}; " - "its nonpersistent cache can depend on prior sequence lengths" - ) - - def dynamic_rope_checkpoint_contract(model): """Reject dynamic caches that a fresh runtime cannot reconstruct exactly. diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index 2cb32ab389..234150cd2c 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -212,10 +212,6 @@ def _validate_optimizer_checkpoint_mode(self) -> None: else None ) raise_if_any_rank_failed(error, "checking MultiOptimizer checkpoint mode") - if manager == "dcp" and getattr(self.optimizer, "_is_multi_optimizer", False) is True: - from ..checkpoint.dynamic_rope import reject_dynamic_rope_multi_optimizer_training - - reject_dynamic_rope_multi_optimizer_training(self.model) def __getattr__(self, name: str) -> Any: """Forward unshadowed :class:`torch.nn.Module` APIs to the wrapped model.""" From 98836e0153633ac597b24ab029b14659a715ae32 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:17:57 +0800 Subject: [PATCH 062/120] [checkpoint] reject unproven TP expert saves early --- docs/usage/checkpoint.md | 6 + .../test_tp_expert_save_preflight.py | 135 ++++++++++++++++++ veomni/checkpoint/dcp_checkpointer.py | 45 +++++- 3 files changed, 179 insertions(+), 7 deletions(-) create mode 100644 tests/checkpoints/test_tp_expert_save_preflight.py diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 153813d1c1..b44442068b 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -320,6 +320,12 @@ distributed reader rejects ambiguous sources for EP/FSDP as well as ordinary FSDP or unsharded targets. An old expert-only FSDP checkpoint without layout metadata also needs verified source-layout migration, since it cannot be distinguished safely from an old EP checkpoint. +For TP greater than one, `DistributedCheckpointer.save` currently rejects a +model containing non-scalar `experts` parameters, persistent buffers, or tensors +assigned to an EP shard by the model's parallel plan before writing. This is a +conservative restriction: it also covers frozen experts and custom layouts that +might otherwise be canonical. TP1 expert checkpoints and TP models without +expert tensors are unaffected. Training without a checkpoint save is unaffected. Offline export also rejects unversioned expert-only chunks for stacked tensors under an `experts` module when their source order cannot be established. Other legacy tensor naming schemes may still need explicit source-layout verification diff --git a/tests/checkpoints/test_tp_expert_save_preflight.py b/tests/checkpoints/test_tp_expert_save_preflight.py new file mode 100644 index 0000000000..d2447d7d6b --- /dev/null +++ b/tests/checkpoints/test_tp_expert_save_preflight.py @@ -0,0 +1,135 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""TP expert checkpoint errors must precede writes and reach every rank.""" + +import os +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist + +from scripts.ci.accelerator.launch_plan import LaunchSpec, runner_capabilities +from tests.tools.launch_utils import spawn_with_rendezvous +from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer +from veomni.distributed.parallel_plan import SpecInfo +from veomni.distributed.parallel_state import ParallelState +from veomni.utils.device import get_device_type, get_dist_comm_backend, get_torch_device + + +def _model(kind): + model = torch.nn.Module() + model.dense = torch.nn.Linear(2, 2, bias=False) + if kind == "dense": + return model + if kind == "registered_ep": + from torch.distributed.tensor import Shard + + model._fqn2spec_info = {"dense.weight": SpecInfo("ep", Shard(0), "dense.weight", None)} + return model + model.experts = torch.nn.Module() + if kind == "buffer": + model.experts.register_buffer("weight", torch.ones(2)) + elif kind == "alias": + model.experts.weight = model.dense.weight + model.dense.weight.requires_grad_(False) + else: + model.experts.weight = torch.nn.Parameter(torch.ones(2), requires_grad=kind != "frozen") + return model + + +@pytest.mark.parametrize("kind", ["parameter", "frozen", "alias", "buffer", "registered_ep"]) +@pytest.mark.parametrize("mode", ["sync", "async", "stage"]) +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +def test_tp_expert_save_rejected_before_writes(tmp_path, kind, mode): + checkpoint = tmp_path / "checkpoint" + marker = checkpoint / "model" / "ckpt" / ".metadata" + marker.parent.mkdir(parents=True) + marker.write_bytes(b"preserve") + stage = tmp_path / "stage" if mode == "stage" else None + with ( + patch.object(DistributedCheckpointer, "_remove_dcp_markers", side_effect=AssertionError("writer started")), + patch("veomni.checkpoint.dcp_checkpointer._prepare_stage_dir", side_effect=AssertionError("staging started")), + ): + with pytest.raises(RuntimeError, match="tensor.parallel.*expert"): + DistributedCheckpointer.save( + str(checkpoint), + {"model": _model(kind)}, + save_async=mode == "async", + trainable_only=kind in ("frozen", "alias"), + stage_dir=str(stage) if stage else None, + parallel_state=SimpleNamespace(tp_size=2, dp_mode="ddp"), + ) + assert marker.read_bytes() == b"preserve" + assert stage is None or not stage.exists() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +def test_tp_dense_model_still_saves_and_loads(tmp_path): + parallel = SimpleNamespace(tp_size=2, dp_mode="ddp") + model = _model("dense") + path = str(tmp_path / "checkpoint") + DistributedCheckpointer.save(path, {"model": model}, parallel_state=parallel) + restored = _model("dense") + DistributedCheckpointer.load(path, {"model": restored}, parallel_state=parallel) + torch.testing.assert_close(restored.dense.weight, model.dense.weight, rtol=0, atol=0) + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +def test_tp_expert_target_can_load_canonical_checkpoint(tmp_path): + model = _model("parameter") + path = str(tmp_path / "checkpoint") + DistributedCheckpointer.save(path, {"model": model}, parallel_state=ParallelState(dp_mode="ddp")) + restored = _model("parameter") + DistributedCheckpointer.load(path, {"model": restored}, parallel_state=SimpleNamespace(tp_size=2, dp_mode="ddp")) + torch.testing.assert_close(restored.experts.weight, model.experts.weight, rtol=0, atol=0) + + +def _distributed_worker(rank, root, accelerator, endpoint): + if accelerator: + device = get_device_type() + get_torch_device().set_device(f"{device}:{rank}") + backend = get_dist_comm_backend() + else: + device, backend = "cpu", "gloo" + dist.init_process_group(backend, store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + model = _model("parameter" if rank == 0 else "dense") + expected = "tensor.parallel.*expert" if rank == 0 else "checking canonical ExtraParallel optimizer state" + with patch.object( + DistributedCheckpointer, "_remove_dcp_markers", side_effect=AssertionError("writer started") + ): + with pytest.raises(RuntimeError, match=expected): + DistributedCheckpointer.save( + root, {"model": model}, parallel_state=SimpleNamespace(tp_size=2, dp_mode="ddp") + ) + done = torch.ones(1, device=device) + dist.all_reduce(done) + assert done.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) +def test_tp_expert_failure_is_coordinated_before_writes(tmp_path): + path = tmp_path / "checkpoint" + spawn_with_rendezvous(_distributed_worker, args=(str(path), False), nprocs=2) + assert not path.exists() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="accelerator")) +def test_native_accelerator_tp_expert_checkpoint_preflight(tmp_path): + device = get_device_type() + available = device != "cpu" and get_torch_device().is_available() and get_torch_device().device_count() >= 2 + require_npu = os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" or ( + "VEOMNI_CI_CAPABILITIES" in os.environ and runner_capabilities().backend == "npu" + ) + if not available or (require_npu and device != "npu"): + if require_npu: + pytest.fail("The NPU checkpoint preflight requires two NPU devices") + pytest.skip("requires two accelerators") + path = tmp_path / "checkpoint" + spawn_with_rendezvous(_distributed_worker, args=(str(path), True), nprocs=2) + assert not path.exists() diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 46101dd5f3..1f6dd58822 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -649,19 +649,50 @@ def _check_canonical_spec_coverage(model, specs, state_keys, *, trainable_only): raise RuntimeError(f"ExtraParallel checkpoint spec is not present in the canonical state dict: {name}") +def _has_live_expert_checkpoint_tensors(model): + """Conservatively include aliases, frozen parameters, and persistent buffers.""" + if not isinstance(model, torch.nn.Module): + return False + + specs = getattr(model, "_fqn2spec_info", None) or {} + + def is_expert(name, tensor): + spec = specs.get(name) + return tensor.ndim and ( + "experts" in name.split(".") + or (getattr(spec, "para_name", None) == "ep" and isinstance(getattr(spec, "placement", None), Shard)) + ) + + if any(is_expert(name, param) for name, param in model.named_parameters(remove_duplicate=False)): + return True + for prefix, module in model.named_modules(remove_duplicate=False): + for name, buffer in module._buffers.items(): + fqn = f"{prefix}.{name}" if prefix else name + if name not in module._non_persistent_buffers_set and buffer is not None and is_expert(fqn, buffer): + return True + return False + + def _reject_unproven_extra_parallel_optimizer( - state, parallel_state, *, group=None, trainable_only=False, allow_proven_optimizer=False + state, parallel_state, *, group=None, trainable_only=False, allow_proven_optimizer=False, saving=False ): error, has_canonical_layout = None, False try: - specs = getattr(state.get("model"), "_fqn2spec_info", None) + model = state.get("model") + ps = ( + (parallel_state if parallel_state is not None else get_parallel_state()) + if saving or state.get("optimizer") is not None + else None + ) + if saving and getattr(ps, "tp_size", 1) > 1 and _has_live_expert_checkpoint_tensors(model): + raise RuntimeError( + "Cannot save a tensor-parallel model containing expert tensors: " + "the expert checkpoint source layout is not proven for TP > 1" + ) + specs = getattr(model, "_fqn2spec_info", None) if _needs_extra_parallel_layout_marker(specs): _extra_parallel_layout_key(specs) has_canonical_layout = True - if state.get("optimizer") is not None: - ps = parallel_state if parallel_state is not None else get_parallel_state() - else: - ps = None if state.get("optimizer") is not None and getattr(ps, "dp_mode", None) == "fsdp2": if _needs_extra_parallel_layout_marker(specs) and not allow_proven_optimizer: raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") @@ -2022,7 +2053,7 @@ def save( if single_exact: raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") _reject_unproven_extra_parallel_optimizer( - state, parallel_state, trainable_only=trainable_only, allow_proven_optimizer=multi + state, parallel_state, trainable_only=trainable_only, allow_proven_optimizer=multi, saving=True ) scheduler_object = state.get("_veomni_scheduler_object") if multi else None source_model = state.get("model") From bf2a37f9f97c160ada9c93b1050360cae0893117 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:19:06 +0800 Subject: [PATCH 063/120] [docs] clarify verified checkpoint migration wording --- docs/usage/checkpoint.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index b44442068b..6a422bfe2d 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -268,7 +268,7 @@ reshard a model checkpoint across supported FSDP shard axes. Use `DistributedCheckpointer` for save and resume: it rejects ambiguous unversioned expert layouts before copying model tensors unless the model-only source is explicitly declared as described below. Migrating a legacy optimizer checkpoint -requires verified source layout information. Canonical v2 optimizer save and resume use +requires source layout information that has been verified. Canonical v2 optimizer save and resume use the rank-local codec described below. The standard VeOmni HF export strips the namespace. @@ -318,7 +318,7 @@ migration. Conversely, an old unversioned checkpoint with expert-only chunks may not prove its source order even when resumed with the same topology. The distributed reader rejects ambiguous sources for EP/FSDP as well as ordinary FSDP or unsharded targets. An old expert-only FSDP checkpoint without layout -metadata also needs verified source-layout migration, since it cannot be +metadata also needs a verified source-layout migration, since it cannot be distinguished safely from an old EP checkpoint. For TP greater than one, `DistributedCheckpointer.save` currently rejects a model containing non-scalar `experts` parameters, persistent buffers, or tensors From 8feba69a71fba175c532d8d54ef033b6cd40b151 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:20:10 +0800 Subject: [PATCH 064/120] [test] cover replicated EP checkpoint preflight --- tests/checkpoints/test_tp_expert_save_preflight.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/checkpoints/test_tp_expert_save_preflight.py b/tests/checkpoints/test_tp_expert_save_preflight.py index d2447d7d6b..d570153bfc 100644 --- a/tests/checkpoints/test_tp_expert_save_preflight.py +++ b/tests/checkpoints/test_tp_expert_save_preflight.py @@ -66,10 +66,15 @@ def test_tp_expert_save_rejected_before_writes(tmp_path, kind, mode): assert stage is None or not stage.exists() +@pytest.mark.parametrize("replicated_ep", [False, True]) @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) -def test_tp_dense_model_still_saves_and_loads(tmp_path): +def test_tp_dense_model_still_saves_and_loads(tmp_path, replicated_ep): parallel = SimpleNamespace(tp_size=2, dp_mode="ddp") model = _model("dense") + if replicated_ep: + from torch.distributed.tensor import Replicate + + model._fqn2spec_info = {"dense.weight": SpecInfo("ep", Replicate(), "dense.weight", None)} path = str(tmp_path / "checkpoint") DistributedCheckpointer.save(path, {"model": model}, parallel_state=parallel) restored = _model("dense") From 6d5f772282233c82107e12e26f74e60da9a3c902 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:48:03 +0800 Subject: [PATCH 065/120] [docs] fix checkpoint manifest field link --- docs/usage/checkpoint.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 6a422bfe2d..5c6c8d6cca 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -132,7 +132,7 @@ half-loaded. The one step with a single half is a run that saves no model state. SeedOmni's `offline_embedding` freezes every module, so no module has a checkpoint manager and the step has no `model/`. Its manifest records -[`model_state: false`](../../veomni/checkpoint/layout.py#L259), and the manifest +[`model_state: false`](../../veomni/checkpoint/layout.py#L190), and the manifest alone then marks the step complete. Without that field, the same tree would look like a step whose shards never started. A manifest written before the field existed counts as carrying model state. From a50fdb6c4cd04e90109204a7f3158e59419535c8 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 13:37:59 +0800 Subject: [PATCH 066/120] [test] keep TP checkpoint checks within model contract --- .../test_tp_expert_save_preflight.py | 68 ++++--------------- 1 file changed, 12 insertions(+), 56 deletions(-) diff --git a/tests/checkpoints/test_tp_expert_save_preflight.py b/tests/checkpoints/test_tp_expert_save_preflight.py index d570153bfc..1bc4a2c191 100644 --- a/tests/checkpoints/test_tp_expert_save_preflight.py +++ b/tests/checkpoints/test_tp_expert_save_preflight.py @@ -1,22 +1,18 @@ # Copyright 2026 Bytedance Ltd. and/or its affiliates # SPDX-License-Identifier: Apache-2.0 -"""TP expert checkpoint errors must precede writes and reach every rank.""" +"""Unsupported TP expert saves must fail before checkpoint writes.""" -import os -from datetime import timedelta from types import SimpleNamespace from unittest.mock import patch import pytest import torch -import torch.distributed as dist -from scripts.ci.accelerator.launch_plan import LaunchSpec, runner_capabilities -from tests.tools.launch_utils import spawn_with_rendezvous +from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.distributed.parallel_plan import SpecInfo from veomni.distributed.parallel_state import ParallelState -from veomni.utils.device import get_device_type, get_dist_comm_backend, get_torch_device +from veomni.utils.dist_utils import raise_if_any_rank_failed def _model(kind): @@ -52,6 +48,7 @@ def test_tp_expert_save_rejected_before_writes(tmp_path, kind, mode): with ( patch.object(DistributedCheckpointer, "_remove_dcp_markers", side_effect=AssertionError("writer started")), patch("veomni.checkpoint.dcp_checkpointer._prepare_stage_dir", side_effect=AssertionError("staging started")), + patch("veomni.checkpoint.dcp_checkpointer.raise_if_any_rank_failed", wraps=raise_if_any_rank_failed) as agree, ): with pytest.raises(RuntimeError, match="tensor.parallel.*expert"): DistributedCheckpointer.save( @@ -62,6 +59,8 @@ def test_tp_expert_save_rejected_before_writes(tmp_path, kind, mode): stage_dir=str(stage) if stage else None, parallel_state=SimpleNamespace(tp_size=2, dp_mode="ddp"), ) + assert isinstance(agree.call_args.args[0], RuntimeError) + assert agree.call_args.args[1] == "checking canonical ExtraParallel optimizer state" assert marker.read_bytes() == b"preserve" assert stage is None or not stage.exists() @@ -76,7 +75,12 @@ def test_tp_dense_model_still_saves_and_loads(tmp_path, replicated_ep): model._fqn2spec_info = {"dense.weight": SpecInfo("ep", Replicate(), "dense.weight", None)} path = str(tmp_path / "checkpoint") - DistributedCheckpointer.save(path, {"model": model}, parallel_state=parallel) + with patch("veomni.checkpoint.dcp_checkpointer.raise_if_any_rank_failed", wraps=raise_if_any_rank_failed) as agree: + DistributedCheckpointer.save(path, {"model": model}, parallel_state=parallel) + assert any( + call.args[0] is None and call.args[1] == "checking canonical ExtraParallel optimizer state" + for call in agree.call_args_list + ) restored = _model("dense") DistributedCheckpointer.load(path, {"model": restored}, parallel_state=parallel) torch.testing.assert_close(restored.dense.weight, model.dense.weight, rtol=0, atol=0) @@ -90,51 +94,3 @@ def test_tp_expert_target_can_load_canonical_checkpoint(tmp_path): restored = _model("parameter") DistributedCheckpointer.load(path, {"model": restored}, parallel_state=SimpleNamespace(tp_size=2, dp_mode="ddp")) torch.testing.assert_close(restored.experts.weight, model.experts.weight, rtol=0, atol=0) - - -def _distributed_worker(rank, root, accelerator, endpoint): - if accelerator: - device = get_device_type() - get_torch_device().set_device(f"{device}:{rank}") - backend = get_dist_comm_backend() - else: - device, backend = "cpu", "gloo" - dist.init_process_group(backend, store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) - try: - model = _model("parameter" if rank == 0 else "dense") - expected = "tensor.parallel.*expert" if rank == 0 else "checking canonical ExtraParallel optimizer state" - with patch.object( - DistributedCheckpointer, "_remove_dcp_markers", side_effect=AssertionError("writer started") - ): - with pytest.raises(RuntimeError, match=expected): - DistributedCheckpointer.save( - root, {"model": model}, parallel_state=SimpleNamespace(tp_size=2, dp_mode="ddp") - ) - done = torch.ones(1, device=device) - dist.all_reduce(done) - assert done.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) -def test_tp_expert_failure_is_coordinated_before_writes(tmp_path): - path = tmp_path / "checkpoint" - spawn_with_rendezvous(_distributed_worker, args=(str(path), False), nprocs=2) - assert not path.exists() - - -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="accelerator")) -def test_native_accelerator_tp_expert_checkpoint_preflight(tmp_path): - device = get_device_type() - available = device != "cpu" and get_torch_device().is_available() and get_torch_device().device_count() >= 2 - require_npu = os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" or ( - "VEOMNI_CI_CAPABILITIES" in os.environ and runner_capabilities().backend == "npu" - ) - if not available or (require_npu and device != "npu"): - if require_npu: - pytest.fail("The NPU checkpoint preflight requires two NPU devices") - pytest.skip("requires two accelerators") - path = tmp_path / "checkpoint" - spawn_with_rendezvous(_distributed_worker, args=(str(path), True), nprocs=2) - assert not path.exists() From 3eef3973122a32c2232160a671744892f6b462c0 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 15:27:14 +0800 Subject: [PATCH 067/120] [checkpoint] test: enforce native preflight capacity --- .../test_tp_expert_save_preflight.py | 38 +++++++++++++ .../test_optimizer_checkpoint_preflight.py | 57 +++++++++++++++++-- 2 files changed, 90 insertions(+), 5 deletions(-) diff --git a/tests/checkpoints/test_tp_expert_save_preflight.py b/tests/checkpoints/test_tp_expert_save_preflight.py index 1bc4a2c191..6f1bccc634 100644 --- a/tests/checkpoints/test_tp_expert_save_preflight.py +++ b/tests/checkpoints/test_tp_expert_save_preflight.py @@ -7,6 +7,7 @@ import pytest import torch +import torch.distributed as dist from scripts.ci.accelerator.launch_plan import LaunchSpec from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer @@ -94,3 +95,40 @@ def test_tp_expert_target_can_load_canonical_checkpoint(tmp_path): restored = _model("parameter") DistributedCheckpointer.load(path, {"model": restored}, parallel_state=SimpleNamespace(tp_size=2, dp_mode="ddp")) torch.testing.assert_close(restored.experts.weight, model.experts.weight, rtol=0, atol=0) + + +def _uniform_tp_expert_worker(rank, checkpoint, stage, endpoint): + from datetime import timedelta + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + torch.manual_seed(1) + with ( + patch( + "veomni.checkpoint.dcp_checkpointer.raise_if_any_rank_failed", wraps=raise_if_any_rank_failed + ) as agree, + pytest.raises(RuntimeError, match="tensor.parallel.*expert"), + ): + DistributedCheckpointer.save( + checkpoint, + {"model": _model("parameter")}, + stage_dir=stage, + parallel_state=SimpleNamespace(tp_size=2, dp_mode="ddp"), + ) + assert isinstance(agree.call_args.args[0], RuntimeError) + assert agree.call_args.args[1] == "checking canonical ExtraParallel optimizer state" + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) +def test_uniform_tp_expert_rejection_reaches_all_ranks_before_writes(tmp_path): + from tests.tools.launch_utils import spawn_with_rendezvous + + checkpoint, stage = tmp_path / "checkpoint", tmp_path / "stage" + spawn_with_rendezvous(_uniform_tp_expert_worker, args=(str(checkpoint), str(stage)), nprocs=2) + assert not checkpoint.exists() + assert not stage.exists() diff --git a/tests/models/test_optimizer_checkpoint_preflight.py b/tests/models/test_optimizer_checkpoint_preflight.py index 6dcf580d3f..0d917d32c5 100644 --- a/tests/models/test_optimizer_checkpoint_preflight.py +++ b/tests/models/test_optimizer_checkpoint_preflight.py @@ -3,6 +3,7 @@ """MultiOptimizer checkpoint compatibility is checked before any training step.""" +import os from contextlib import nullcontext from datetime import timedelta from types import SimpleNamespace @@ -11,8 +12,9 @@ import torch import torch.distributed as dist -from scripts.ci.accelerator.launch_plan import LaunchSpec +from scripts.ci.accelerator.launch_plan import LaunchSpec, runner_capabilities from veomni.models import model_runtime +from veomni.utils.device import get_device_type pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) @@ -159,12 +161,57 @@ def _accelerator_worker(rank, endpoint): dist.destroy_process_group() -@pytest.mark.skipif( - not torch.accelerator.is_available() or torch.accelerator.device_count() < 2, - reason="requires two accelerators", -) +def _require_two_native_devices(): + """Only optional local runs may skip the two-device preflight.""" + capabilities = runner_capabilities() if os.environ.get("VEOMNI_CI_CAPABILITIES") else None + assigned = capabilities is not None and capabilities.backend != "cpu" + if not torch.accelerator.is_available() or torch.accelerator.device_count() < 2: + if assigned: + pytest.fail("Assigned accelerator runtime must provide two usable devices") + pytest.skip("requires two accelerators") + if assigned and get_device_type() != capabilities.backend: + pytest.fail("Assigned accelerator backend does not match the runtime") + + +@pytest.mark.parametrize("count", [0, 1]) +def test_native_preflight_fails_when_assigned_devices_are_unusable(monkeypatch, count): + monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "declared") + monkeypatch.setitem(globals(), "runner_capabilities", lambda: SimpleNamespace(backend="npu")) + monkeypatch.setattr(torch.accelerator, "is_available", lambda: count > 0) + monkeypatch.setattr(torch.accelerator, "device_count", lambda: count) + with pytest.raises(pytest.fail.Exception, match="two usable devices"): + _require_two_native_devices() + + +def test_native_preflight_skips_without_optional_local_devices(monkeypatch): + monkeypatch.delenv("VEOMNI_CI_CAPABILITIES", raising=False) + monkeypatch.setattr(torch.accelerator, "is_available", lambda: False) + with pytest.raises(pytest.skip.Exception, match="two accelerators"): + _require_two_native_devices() + + +def test_native_preflight_accepts_two_assigned_devices(monkeypatch): + monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "declared") + monkeypatch.setitem(globals(), "runner_capabilities", lambda: SimpleNamespace(backend="npu")) + monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) + monkeypatch.setattr(torch.accelerator, "device_count", lambda: 2) + monkeypatch.setitem(globals(), "get_device_type", lambda: "npu") + _require_two_native_devices() + + +def test_native_preflight_rejects_mismatched_backend(monkeypatch): + monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "declared") + monkeypatch.setitem(globals(), "runner_capabilities", lambda: SimpleNamespace(backend="npu")) + monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) + monkeypatch.setattr(torch.accelerator, "device_count", lambda: 2) + monkeypatch.setitem(globals(), "get_device_type", lambda: "cuda") + with pytest.raises(pytest.fail.Exception, match="backend does not match"): + _require_two_native_devices() + + @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="accelerator")) def test_native_accelerator_optimizer_checkpoint_preflight(): from tests.tools.launch_utils import spawn_with_rendezvous + _require_two_native_devices() spawn_with_rendezvous(_accelerator_worker, args=(), nprocs=2) From 366fa23f9bc3bda063b2ff03a869dfcbde732c26 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 16:14:37 +0800 Subject: [PATCH 068/120] [checkpoint] prove TP FSDP source layouts --- docs/usage/checkpoint.md | 17 ++ tests/checkpoints/test_tp_fsdp_provenance.py | 241 +++++++++++++++++++ veomni/checkpoint/dcp_checkpointer.py | 127 +++++++++- 3 files changed, 382 insertions(+), 3 deletions(-) create mode 100644 tests/checkpoints/test_tp_fsdp_provenance.py diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index d0c0dd1815..2fa4cf5151 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -289,6 +289,23 @@ before loading into a different mesh or offline export; their provenance cannot be inferred from chunks alone, and not all such cases can be rejected automatically when the target has no ExtraParallel mesh. +There is currently no general, verified migration utility for these ambiguous +legacy expert layouts. Preserve the original checkpoint and its training code, +rank-to-expert mapping, mesh, and optimizer configuration. A model-specific +migration needs to reconstruct both the global expert tensors and their matching +optimizer states from that provenance, compare every expert against an +independent reference, and save to a new location with the current writer. A +fresh-process restore and next-update comparison are required before resuming. +Without enough provenance to establish the mapping, do not resume or export the +ambiguous state as if it were canonical; start from independently verified +weights and a fresh optimizer instead. This describes necessary checks, not a +provided or generally validated migration procedure. + +For ordinary TP/FSDP-only models, the writer may certify a strided same-axis +layout when the active runtime mesh and a unique source-model parameter owner +prove its provenance. Parameters with unknown ownership, unexpected meshes, +partial placements, or any active ExtraParallel dimension remain unproven. + All ranks in a checkpoint operation must use the same logical model and checkpoint configuration, including the layout version and corresponding sharding specifications. Normal rank-local shard contents and mesh coordinates diff --git a/tests/checkpoints/test_tp_fsdp_provenance.py b/tests/checkpoints/test_tp_fsdp_provenance.py new file mode 100644 index 0000000000..4f41d25323 --- /dev/null +++ b/tests/checkpoints/test_tp_fsdp_provenance.py @@ -0,0 +1,241 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 + +"""TP/FSDP checkpoint provenance must not weaken unknown expert guards.""" + +import os +from copy import copy, deepcopy +from datetime import timedelta +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed._state_dict_utils import _create_cpu_state_dict +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import DTensor, Partial, Shard +from torch.distributed.tensor.parallel import ColwiseParallel, parallelize_module +from torch.distributed.tensor.placement_types import _StridedShard + +from scripts.ci.accelerator.launch_plan import LaunchSpec, runner_capabilities +from tests.tools.launch_utils import spawn_with_rendezvous +from veomni.checkpoint.dcp_checkpointer import ( + _EXPERT_LAYOUTS_KEY, + DistributedCheckpointer, + _ExtraParallelSavePlanner, + _ModelStrictLoadPlanner, + _tp_only_source_layout, + _tp_only_source_parameters, + dcp_to_torch_state_dict, +) +from veomni.checkpoint.layout import optimizer_dir, weights_dir +from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state, get_parallel_state +from veomni.utils.device import IS_NPU_AVAILABLE + + +def _tp_fsdp_worker(rank, world, directory, device_type, endpoint): + device = torch.device(device_type, rank) if device_type == "npu" else torch.device("cpu") + if device_type == "npu": + torch.npu.set_device(device) + dist.init_process_group( + "hccl" if device_type == "npu" else "gloo", + store=endpoint.connect(), + rank=rank, + world_size=world, + timeout=timedelta(seconds=90), + ) + try: + ambient = _init_parallel_state(dp_size=4, device_type=device_type, name="base") + state = _init_parallel_state(dp_size=2, tp_size=2, device_type=device_type, name="policy") + assert get_parallel_state() is ambient and state is not ambient + torch.manual_seed(7) + model = torch.nn.Sequential(torch.nn.Linear(8, 8, bias=False)) + reference = deepcopy(model) + reference_optimizer = torch.optim.AdamW(reference.parameters(), lr=0.001) + model.to(device) + parallelize_module(model, state.tp_mesh, {"0": ColwiseParallel()}) + fully_shard(model, mesh=state.fsdp_mesh) + optimizer = torch.optim.AdamW(model.parameters(), lr=0.001) + inputs = torch.arange(16, dtype=torch.float32).view(2, 8) / 16 + model(inputs.to(device)).sum().backward() + optimizer.step() + optimizer.zero_grad(set_to_none=True) + reference(inputs).sum().backward() + reference_optimizer.step() + reference_optimizer.zero_grad(set_to_none=True) + param = model[0].weight + assert isinstance(param, DTensor) + assert param.placements == (_StridedShard(0, split_factor=2), Shard(0)) + owners, mesh = _tp_only_source_parameters(model, state) + assert owners["0.weight"] is param + layout = _tp_only_source_layout("model.0.weight", param, owners, mesh) + assert layout and layout["expert_order"] == "canonical" + + # Optimizer-shaped state must map to the same parameter and layout. + avg = optimizer.state[param]["exp_avg"] + variance = optimizer.state[param]["exp_avg_sq"] + torch.testing.assert_close(param.full_tensor().cpu(), reference[0].weight, atol=1e-6, rtol=1e-5) + torch.testing.assert_close( + avg.full_tensor().cpu(), reference_optimizer.state[reference[0].weight]["exp_avg"], atol=1e-6, rtol=1e-5 + ) + assert _tp_only_source_layout("optimizer.state.0.weight.exp_avg", avg, owners, mesh) == layout + assert _tp_only_source_layout("optimizer.state.missing.exp_avg", avg, owners, mesh) is None + wrong_placement = DTensor.from_local( + torch.zeros_like(param.to_local()), mesh, [Partial(), Shard(0)], shape=param.shape, stride=param.stride() + ) + assert _tp_only_source_layout("model.0.weight", wrong_placement, owners, mesh) is None + wrong_split = DTensor.from_local( + torch.zeros_like(param.to_local()), + mesh, + [_StridedShard(0, split_factor=3), Shard(0)], + shape=param.shape, + stride=param.stride(), + ) + assert _tp_only_source_layout("model.0.weight", wrong_split, owners, mesh) is None + other_root = DeviceMesh(device_type, [[0, 1], [2, 3]], mesh_dim_names=mesh.mesh_dim_names, _init_backend=False) + wrong_root = DTensor.from_local( + torch.zeros_like(param.to_local()), other_root, param.placements, shape=param.shape, stride=param.stride() + ) + assert _tp_only_source_layout("model.0.weight", wrong_root, owners, mesh) is None + assert _tp_only_source_layout("model.0.weight", wrong_root, owners, mesh, allow_staged_copies=True) is None + staged = _create_cpu_state_dict({"weight": param})["weight"] + if staged.device_mesh._get_root_mesh() is not mesh._get_root_mesh(): + assert _tp_only_source_layout("model.0.weight", staged, owners, mesh) is None + assert _tp_only_source_layout("model.0.weight", staged, owners, mesh, allow_staged_copies=True) == layout + assert _tp_only_source_parameters(None, state) == ({}, None) + assert _tp_only_source_parameters(model, copy(state)) == ({}, None) + with patch.object(dist, "is_initialized", return_value=False): + assert _tp_only_source_parameters(model, state) == ({}, None) + extra_state = _init_parallel_state( + dp_size=2, tp_size=2, extra_parallel_sizes=(2,), device_type=device_type, name="policy_ep" + ) + assert _tp_only_source_parameters(model, extra_state) == ({}, None) + aliased = torch.nn.Module() + aliased.left = model[0] + aliased.right = model[0] + assert not _tp_only_source_parameters(aliased, state)[0] + + checkpoint = directory + "/dense" + dcp.save( + { + "model": {"0.weight": param}, + "optimizer": {"state": {"0.weight": {"exp_avg": avg, "exp_avg_sq": variance}}}, + }, + checkpoint_id=checkpoint, + planner=_ExtraParallelSavePlanner(source_model=model, parallel_state=state), + ) + metadata = dcp.FileSystemReader(checkpoint).read_metadata().planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + assert metadata["model.0.weight"]["expert_order"] == "canonical" + assert metadata["optimizer.state.0.weight.exp_avg"]["expert_order"] == "canonical" + expected = param.full_tensor().cpu() + target = {"model": {"0.weight": torch.zeros_like(expected)}} + dcp.load(target, checkpoint_id=checkpoint, planner=_ModelStrictLoadPlanner(strict_model=True)) + torch.testing.assert_close(target["model"]["0.weight"], expected, atol=0, rtol=0) + if rank == 0: + exported = dcp_to_torch_state_dict(checkpoint) + torch.testing.assert_close(exported["0.weight"], expected, atol=0, rtol=0) + dist.barrier() + + production = directory + "/production" + DistributedCheckpointer.save(production, {"model": model, "optimizer": optimizer}, parallel_state=state) + for path, key in ( + (weights_dir(production), "model.0.weight"), + (optimizer_dir(production), "optimizer.state.0.weight.exp_avg"), + (optimizer_dir(production), "optimizer.state.0.weight.exp_avg_sq"), + ): + recorded = dcp.FileSystemReader(path).read_metadata().planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + assert recorded[key]["expert_order"] == "canonical" + + torch.manual_seed(9) + restored_model = torch.nn.Sequential(torch.nn.Linear(8, 8, bias=False)) + restored_model.to(device) + parallelize_module(restored_model, state.tp_mesh, {"0": ColwiseParallel()}) + fully_shard(restored_model, mesh=state.fsdp_mesh) + restored_param = restored_model[0].weight + restored_optimizer = torch.optim.AdamW(restored_model.parameters(), lr=0.001) + DistributedCheckpointer.load( + production, {"model": restored_model, "optimizer": restored_optimizer}, parallel_state=state + ) + restored_avg = restored_optimizer.state[restored_param]["exp_avg"] + restored_variance = restored_optimizer.state[restored_param]["exp_avg_sq"] + torch.testing.assert_close(restored_param.to_local(), param.to_local(), atol=0, rtol=0) + torch.testing.assert_close(restored_avg.to_local(), avg.to_local(), atol=0, rtol=0) + torch.testing.assert_close(restored_variance.to_local(), variance.to_local(), atol=0, rtol=0) + torch.testing.assert_close(restored_optimizer.state[restored_param]["step"], optimizer.state[param]["step"]) + + async_path = directory + "/production_async" + DistributedCheckpointer.save( + async_path, {"model": model, "optimizer": optimizer}, parallel_state=state, save_async=True + ) + DistributedCheckpointer.wait_for_pending_save() + for path, key in ( + (weights_dir(async_path), "model.0.weight"), + (optimizer_dir(async_path), "optimizer.state.0.weight.exp_avg"), + (optimizer_dir(async_path), "optimizer.state.0.weight.exp_avg_sq"), + ): + recorded = dcp.FileSystemReader(path).read_metadata().planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + assert recorded[key]["expert_order"] == "canonical" + DistributedCheckpointer.load( + async_path, {"model": restored_model, "optimizer": restored_optimizer}, parallel_state=state + ) + restored_avg = restored_optimizer.state[restored_param]["exp_avg"] + restored_variance = restored_optimizer.state[restored_param]["exp_avg_sq"] + torch.testing.assert_close(restored_param.to_local(), param.to_local(), atol=0, rtol=0) + torch.testing.assert_close(restored_avg.to_local(), avg.to_local(), atol=0, rtol=0) + torch.testing.assert_close(restored_variance.to_local(), variance.to_local(), atol=0, rtol=0) + model(inputs.to(device)).sum().backward() + restored_model(inputs.to(device)).sum().backward() + reference(inputs).sum().backward() + optimizer.step() + restored_optimizer.step() + reference_optimizer.step() + torch.testing.assert_close(param.full_tensor().cpu(), reference[0].weight, atol=1e-6, rtol=1e-5) + torch.testing.assert_close(restored_param.to_local(), param.to_local(), atol=0, rtol=0) + torch.testing.assert_close(restored_avg.to_local(), avg.to_local(), atol=0, rtol=0) + torch.testing.assert_close(restored_variance.to_local(), variance.to_local(), atol=0, rtol=0) + + # Even an innocently named tensor must fail closed without a proven owner. + unknown = directory + "/unknown" + dcp.save( + {"model": {"unknown": param}}, + checkpoint_id=unknown, + planner=_ExtraParallelSavePlanner(source_model=model, parallel_state=state), + ) + target = {"model": {"unknown": torch.zeros_like(expected)}} + with pytest.raises(dcp.CheckpointException, match="noncanonical expert order"): + dcp.load(target, checkpoint_id=unknown, planner=_ModelStrictLoadPlanner(strict_model=True)) + if rank == 0: + with pytest.raises(RuntimeError, match="noncanonical expert order"): + dcp_to_torch_state_dict(unknown) + dist.barrier() + finally: + clear_parallel_state() + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu")) +def test_tp_fsdp_checkpoint_requires_provenance(tmp_path): + spawn_with_rendezvous(_tp_fsdp_worker, args=(4, str(tmp_path), "cpu"), nprocs=4) + + +def _npu_launch_plan(context): + if os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" and context.capabilities.backend != "npu": + # Let the test fail explicitly even under --ci-batch on a CPU runner. + return LaunchSpec(device="cpu") + return LaunchSpec(nproc=4) + + +@pytest.mark.launch_plan(provider=_npu_launch_plan) +def test_tp_fsdp_checkpoint_native_npu(tmp_path): + if not IS_NPU_AVAILABLE: + require_npu = os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" or ( + bool(os.environ.get("VEOMNI_CI_CAPABILITIES")) and runner_capabilities().backend == "npu" + ) + if require_npu: + pytest.fail("NPU CI or an explicit NPU test requires an available NPU") + pytest.skip("native TP/FSDP requires NPU") + if torch.npu.device_count() < 4: + pytest.fail("native TP/FSDP requires four visible NPUs") + spawn_with_rendezvous(_tp_fsdp_worker, args=(4, str(tmp_path), "npu"), nprocs=4) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 0863a1d9c3..463e435e2a 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -316,12 +316,124 @@ def _unidentified_same_axis_shards(tensor): return len(dims) != len(set(dims)) +def _tp_only_source_parameters(model, parallel_state): + """Identify unambiguous model owners on the runtime's TP/FSDP mesh. + + This does not certify custom sharding or any active ExtraParallel plan. + Like the rest of the checkpointer, it assumes the supplied runtime model + was built from a complete logical model rather than pre-sliced parameters. + """ + if not isinstance(model, torch.nn.Module) or not isinstance(parallel_state, ParallelState): + return {}, None + if ( + parallel_state.device_mesh is None + or not dist.is_initialized() + or parallel_state.device_mesh._get_root_mesh().size() != dist.get_world_size() + or parallel_state.tp_size <= 1 + or not parallel_state.extra_parallel_sizes + or any(size != 1 for size in parallel_state.extra_parallel_sizes.values()) + ): + return {}, None + from ..distributed.parallel_state import _PARALLEL_STATE_CACHE + + # A multi-model trainer may save a registered non-ambient state. Require a + # state created by the runtime's topology cache, not just matching names. + if not any(parallel_state is state for state in _PARALLEL_STATE_CACHE.values()): + return {}, None + + from torch.distributed.checkpoint.state_dict import _get_fqns + + try: + mesh = DeviceMesh._concatenate([parallel_state.fsdp_mesh, parallel_state.tp_mesh]) + except (AttributeError, RuntimeError, TypeError, ValueError): + return {}, None + if mesh._get_root_mesh() is not parallel_state.device_mesh._get_root_mesh(): + return {}, None + + aliases = {} + for name, param in model.named_parameters(remove_duplicate=False): + aliases.setdefault(id(param), []).append((name, param)) + candidates, ambiguous = {}, set() + for names in aliases.values(): + if len(names) != 1: + continue + name, param = names[0] + fqns = _get_fqns(model, name) + if len(fqns) != 1: + continue + fqn = next(iter(fqns)) + if fqn in candidates: + ambiguous.add(fqn) + candidates[fqn] = param + for fqn in ambiguous: + candidates.pop(fqn, None) + return candidates, mesh + + +def _tp_only_source_layout(key, tensor, parameters, mesh, *, allow_staged_copies=False): + """Certify only a model parameter or same-layout state of its unique owner.""" + if mesh is None or not isinstance(tensor, DTensor) or not tensor.ndim: + return None + if key.startswith("model."): + fqn = key.removeprefix("model.") + elif key.startswith("optimizer.state."): + fqn = key.removeprefix("optimizer.state.") + while fqn not in parameters and "." in fqn: + fqn = fqn.rsplit(".", 1)[0] + else: + return None + param = parameters.get(fqn) + if not isinstance(param, DTensor) or tuple(tensor.shape) != tuple(param.shape): + return None + if param.device_mesh._get_root_mesh() is not mesh._get_root_mesh(): + return None + if tensor.device_mesh != mesh or param.device_mesh != mesh or tensor.placements != param.placements: + return None + if tensor.device_mesh._get_root_mesh() is not mesh._get_root_mesh(): + if not allow_staged_copies: + return None + # Async CPU staging deep-copies DTensor meshes. Admit a staged copy + # only for the production async writer and only if every axis still + # resolves to the runtime's identical process group. + try: + if any(tensor.device_mesh.get_group(i) is not mesh.get_group(i) for i in range(mesh.ndim)): + return None + except (AssertionError, AttributeError, IndexError, KeyError, RuntimeError, ValueError): + return None + placements = tensor.placements + if len(placements) not in (2, 3) or any(isinstance(placement, Partial) for placement in placements): + return None + if len(placements) == 3 and not isinstance(placements[0], Replicate): + return None + fsdp, tp = placements[-2:] + if type(tp) not in (Shard, Replicate) or type(fsdp) not in (Shard, _StridedShard): + return None + same_axis = isinstance(tp, Shard) and fsdp.dim % tensor.ndim == tp.dim % tensor.ndim + if same_axis: + if type(fsdp) is not _StridedShard or fsdp.split_factor != mesh.size(-1): + return None + elif type(fsdp) is not Shard: + return None + return { + "ep_shard_dim": None, + "fsdp_shard_dim": fsdp.dim % tensor.ndim, + "ep_size": 1, + "fsdp_size": mesh.size(-2), + "expert_order": "canonical", + "source_kind": "veomni_tp_fsdp_no_extra", + } + + class _ExtraParallelSavePlanner(DefaultSavePlanner): """Record per-tensor source layouts in DCP metadata, including optimizer state.""" - def __init__(self, *args, parallel_state=None, source_model=None, adapter_name=None, **kwargs): + def __init__( + self, *args, parallel_state=None, source_model=None, adapter_name=None, allow_staged_copies=False, **kwargs + ): super().__init__(*args, **kwargs) self.parallel_state = parallel_state + self.allow_staged_copies = allow_staged_copies + self.tp_only_parameters, self.tp_only_mesh = _tp_only_source_parameters(source_model, parallel_state) self.ordinary_fqns = ( _ordinary_expert_fqns(source_model, parallel_state) if isinstance(parallel_state, ParallelState) @@ -355,8 +467,16 @@ def create_local_plan(self): layouts = {} for key, value in self.state_dict.items(): extra_names = getattr(self.parallel_state, "extra_parallel_names", ()) - layout = _expert_layout(value, extra_names) or _no_ep_expert_layout( - key, value, self.parallel_state, self.ordinary_fqns, self.shared_lora_fqns + layout = ( + _expert_layout(value, extra_names) + or _no_ep_expert_layout(key, value, self.parallel_state, self.ordinary_fqns, self.shared_lora_fqns) + or _tp_only_source_layout( + key, + value, + self.tp_only_parameters, + self.tp_only_mesh, + allow_staged_copies=self.allow_staged_copies, + ) ) if layout is not None: layouts[key] = layout @@ -2197,6 +2317,7 @@ def execute_save( dedup_save_to_lowest_rank=save_to_lowest_rank, parallel_state=parallel_state if parallel_state is not None else get_parallel_state(), source_model=source_model, + allow_staged_copies=save_async, ) except BaseException as exc: error = exc From e3e335059ee1cd04f99628b51c55285e5dcf5102 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 16:17:50 +0800 Subject: [PATCH 069/120] [test] remove preflight guard-only cases --- .../test_optimizer_checkpoint_preflight.py | 36 ------------------- 1 file changed, 36 deletions(-) diff --git a/tests/models/test_optimizer_checkpoint_preflight.py b/tests/models/test_optimizer_checkpoint_preflight.py index 0d917d32c5..541243a13f 100644 --- a/tests/models/test_optimizer_checkpoint_preflight.py +++ b/tests/models/test_optimizer_checkpoint_preflight.py @@ -173,42 +173,6 @@ def _require_two_native_devices(): pytest.fail("Assigned accelerator backend does not match the runtime") -@pytest.mark.parametrize("count", [0, 1]) -def test_native_preflight_fails_when_assigned_devices_are_unusable(monkeypatch, count): - monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "declared") - monkeypatch.setitem(globals(), "runner_capabilities", lambda: SimpleNamespace(backend="npu")) - monkeypatch.setattr(torch.accelerator, "is_available", lambda: count > 0) - monkeypatch.setattr(torch.accelerator, "device_count", lambda: count) - with pytest.raises(pytest.fail.Exception, match="two usable devices"): - _require_two_native_devices() - - -def test_native_preflight_skips_without_optional_local_devices(monkeypatch): - monkeypatch.delenv("VEOMNI_CI_CAPABILITIES", raising=False) - monkeypatch.setattr(torch.accelerator, "is_available", lambda: False) - with pytest.raises(pytest.skip.Exception, match="two accelerators"): - _require_two_native_devices() - - -def test_native_preflight_accepts_two_assigned_devices(monkeypatch): - monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "declared") - monkeypatch.setitem(globals(), "runner_capabilities", lambda: SimpleNamespace(backend="npu")) - monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) - monkeypatch.setattr(torch.accelerator, "device_count", lambda: 2) - monkeypatch.setitem(globals(), "get_device_type", lambda: "npu") - _require_two_native_devices() - - -def test_native_preflight_rejects_mismatched_backend(monkeypatch): - monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "declared") - monkeypatch.setitem(globals(), "runner_capabilities", lambda: SimpleNamespace(backend="npu")) - monkeypatch.setattr(torch.accelerator, "is_available", lambda: True) - monkeypatch.setattr(torch.accelerator, "device_count", lambda: 2) - monkeypatch.setitem(globals(), "get_device_type", lambda: "cuda") - with pytest.raises(pytest.fail.Exception, match="backend does not match"): - _require_two_native_devices() - - @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="accelerator")) def test_native_accelerator_optimizer_checkpoint_preflight(): from tests.tools.launch_utils import spawn_with_rendezvous From 3e67bc48b961df128f8afa953070ad76d55a2973 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 17:09:21 +0800 Subject: [PATCH 070/120] [test] schedule TP provenance on its actual backend --- tests/checkpoints/test_tp_fsdp_provenance.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/checkpoints/test_tp_fsdp_provenance.py b/tests/checkpoints/test_tp_fsdp_provenance.py index 4f41d25323..1dee4ddb9b 100644 --- a/tests/checkpoints/test_tp_fsdp_provenance.py +++ b/tests/checkpoints/test_tp_fsdp_provenance.py @@ -221,18 +221,19 @@ def test_tp_fsdp_checkpoint_requires_provenance(tmp_path): def _npu_launch_plan(context): - if os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" and context.capabilities.backend != "npu": - # Let the test fail explicitly even under --ci-batch on a CPU runner. + if context.capabilities.backend != "npu": + # Non-NPU runners need no accelerator lease to skip (or to fail + # explicitly when VEOMNI_REQUIRE_NPU_TESTS=1). return LaunchSpec(device="cpu") return LaunchSpec(nproc=4) @pytest.mark.launch_plan(provider=_npu_launch_plan) def test_tp_fsdp_checkpoint_native_npu(tmp_path): - if not IS_NPU_AVAILABLE: - require_npu = os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" or ( - bool(os.environ.get("VEOMNI_CI_CAPABILITIES")) and runner_capabilities().backend == "npu" - ) + ci_backend = runner_capabilities().backend if os.environ.get("VEOMNI_CI_CAPABILITIES") else None + require_npu = os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" or ci_backend == "npu" + available = ci_backend in (None, "npu") and IS_NPU_AVAILABLE and torch.npu.is_available() + if not available: if require_npu: pytest.fail("NPU CI or an explicit NPU test requires an available NPU") pytest.skip("native TP/FSDP requires NPU") From 2ba4cf6a4376d206689deff497297d04628fdfb4 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 17:11:49 +0800 Subject: [PATCH 071/120] [test] derive checkpoint launch plans from geometry --- .../test_expert_noep_checkpoint.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/checkpoints/test_expert_noep_checkpoint.py b/tests/checkpoints/test_expert_noep_checkpoint.py index 1971206236..03367b79be 100644 --- a/tests/checkpoints/test_expert_noep_checkpoint.py +++ b/tests/checkpoints/test_expert_noep_checkpoint.py @@ -70,8 +70,12 @@ def _step(model, optimizer, value): optimizer.zero_grad(set_to_none=True) +def _noep_world(geometry): + return 4 if geometry == "hsdp" else 2 + + def _noep_worker(rank, directory, geometry, optimizer_kind, phase, endpoint): - world = 4 if geometry == "hsdp" else 2 + world = _noep_world(geometry) dist.init_process_group( "gloo", store=endpoint.connect(), rank=rank, world_size=world, timeout=timedelta(seconds=60) ) @@ -113,9 +117,9 @@ def _noep_worker(rank, directory, geometry, optimizer_kind, phase, endpoint): @pytest.mark.parametrize("geometry", ["fsdp", "fsdp2", "sp", "hsdp"]) @pytest.mark.parametrize("optimizer_kind", ["adamw", "multi"]) -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +@pytest.mark.launch_plan(provider=lambda ctx: LaunchSpec(device="cpu", nproc=_noep_world(ctx.params["geometry"]))) def test_new_noep_expert_checkpoint_fresh_process_roundtrip(tmp_path, geometry, optimizer_kind): - world = 4 if geometry == "hsdp" else 2 + world = _noep_world(geometry) spawn_with_rendezvous(_noep_worker, args=(str(tmp_path), geometry, optimizer_kind, "save"), nprocs=world) model_dir = tmp_path / "checkpoint" / "model" / "ckpt" optimizer_dir = tmp_path / "checkpoint" / "model" / "optimizer" @@ -131,8 +135,12 @@ def test_new_noep_expert_checkpoint_fresh_process_roundtrip(tmp_path, geometry, spawn_with_rendezvous(_noep_worker, args=(str(tmp_path), geometry, optimizer_kind, "resume"), nprocs=world) +def _unproven_world(mode): + return 4 if mode == "tp" else 2 + + def _unproven_source_worker(rank, directory, mode, endpoint): - world = 4 if mode == "tp" else 2 + world = _unproven_world(mode) dist.init_process_group( "gloo", store=endpoint.connect(), rank=rank, world_size=world, timeout=timedelta(seconds=60) ) @@ -169,9 +177,9 @@ def _unproven_source_worker(rank, directory, mode, endpoint): @pytest.mark.parametrize("mode", ["ep", "tp", "custom"]) -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4)) +@pytest.mark.launch_plan(provider=lambda ctx: LaunchSpec(device="cpu", nproc=_unproven_world(ctx.params["mode"]))) def test_ep_tp_and_custom_mesh_are_not_inferred_as_noep(tmp_path, mode): - spawn_with_rendezvous(_unproven_source_worker, args=(str(tmp_path), mode), nprocs=4 if mode == "tp" else 2) + spawn_with_rendezvous(_unproven_source_worker, args=(str(tmp_path), mode), nprocs=_unproven_world(mode)) def _uninitialized_parallel_state_worker(rank, directory, endpoint): From 6dd9407678744953248fb8aad82034344ccc74a2 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:00:39 +0800 Subject: [PATCH 072/120] [test] fail assigned accelerator group shortages --- tests/utils/test_dist_utils_group.py | 53 +++++++++++++++++++++++++--- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/tests/utils/test_dist_utils_group.py b/tests/utils/test_dist_utils_group.py index df06c5baa0..8d7470344c 100644 --- a/tests/utils/test_dist_utils_group.py +++ b/tests/utils/test_dist_utils_group.py @@ -1,11 +1,13 @@ """Ensure error propagation follows the backend of the actual process group.""" +import os +from types import SimpleNamespace from unittest.mock import patch import pytest import torch.distributed as dist -from scripts.ci.accelerator.launch_plan import LaunchSpec +from scripts.ci.accelerator.launch_plan import LaunchSpec, runner_capabilities from tests.tools.launch_utils import spawn_with_rendezvous from veomni.utils.device import get_device_type, get_dist_comm_backend, get_torch_device from veomni.utils.dist_utils import any_rank_failed, raise_if_any_rank_failed @@ -53,10 +55,51 @@ def _accelerator_group_worker(rank, world_size, endpoint): dist.destroy_process_group() +def _require_two_accelerators(): + """Allow optional local skips, but fail when accelerator CI is misconfigured.""" + capabilities = runner_capabilities() if os.environ.get("VEOMNI_CI_CAPABILITIES") else None + assigned = capabilities is not None and capabilities.backend != "cpu" + device_type = get_device_type() + if device_type == "cpu" or not get_torch_device().is_available() or get_torch_device().device_count() < 2: + if assigned: + pytest.fail("Assigned accelerator runtime must provide two usable devices") + pytest.skip("Requires two accelerator devices") + if assigned and device_type != capabilities.backend: + pytest.fail("Assigned accelerator backend does not match the runtime") + + @pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="accelerator")) -@pytest.mark.skipif( - get_device_type() == "cpu" or not get_torch_device().is_available() or get_torch_device().device_count() < 2, - reason="Requires two accelerator devices", -) def test_default_accelerator_group_failure_propagation(): + _require_two_accelerators() spawn_with_rendezvous(_accelerator_group_worker, args=(2,), nprocs=2) + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +@pytest.mark.parametrize( + ("assigned", "device_type", "count", "outcome"), + ( + (None, "cpu", 0, "skip"), + (None, "cuda", 1, "skip"), + ("cuda", "cuda", 0, "fail"), + ("cuda", "cuda", 1, "fail"), + ("npu", "cuda", 2, "fail"), + ("npu", "npu", 2, "pass"), + ), +) +def test_two_accelerator_requirement(monkeypatch, assigned, device_type, count, outcome): + monkeypatch.setattr(__name__ + ".get_device_type", lambda: device_type) + monkeypatch.setattr( + __name__ + ".get_torch_device", + lambda: SimpleNamespace(is_available=lambda: count > 0, device_count=lambda: count), + ) + if assigned: + monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "assigned") + monkeypatch.setattr(__name__ + ".runner_capabilities", lambda: SimpleNamespace(backend=assigned)) + else: + monkeypatch.delenv("VEOMNI_CI_CAPABILITIES", raising=False) + if outcome == "pass": + _require_two_accelerators() + else: + error = pytest.skip.Exception if outcome == "skip" else pytest.fail.Exception + with pytest.raises(error): + _require_two_accelerators() From 9be702695cffa665eb733e5cde430e5c665039e0 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 22:00:01 +0800 Subject: [PATCH 073/120] [test] remove accelerator skip meta tests --- tests/utils/test_dist_utils_group.py | 32 ---------------------------- 1 file changed, 32 deletions(-) diff --git a/tests/utils/test_dist_utils_group.py b/tests/utils/test_dist_utils_group.py index 8d7470344c..5459d832bb 100644 --- a/tests/utils/test_dist_utils_group.py +++ b/tests/utils/test_dist_utils_group.py @@ -1,7 +1,6 @@ """Ensure error propagation follows the backend of the actual process group.""" import os -from types import SimpleNamespace from unittest.mock import patch import pytest @@ -72,34 +71,3 @@ def _require_two_accelerators(): def test_default_accelerator_group_failure_propagation(): _require_two_accelerators() spawn_with_rendezvous(_accelerator_group_worker, args=(2,), nprocs=2) - - -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) -@pytest.mark.parametrize( - ("assigned", "device_type", "count", "outcome"), - ( - (None, "cpu", 0, "skip"), - (None, "cuda", 1, "skip"), - ("cuda", "cuda", 0, "fail"), - ("cuda", "cuda", 1, "fail"), - ("npu", "cuda", 2, "fail"), - ("npu", "npu", 2, "pass"), - ), -) -def test_two_accelerator_requirement(monkeypatch, assigned, device_type, count, outcome): - monkeypatch.setattr(__name__ + ".get_device_type", lambda: device_type) - monkeypatch.setattr( - __name__ + ".get_torch_device", - lambda: SimpleNamespace(is_available=lambda: count > 0, device_count=lambda: count), - ) - if assigned: - monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "assigned") - monkeypatch.setattr(__name__ + ".runner_capabilities", lambda: SimpleNamespace(backend=assigned)) - else: - monkeypatch.delenv("VEOMNI_CI_CAPABILITIES", raising=False) - if outcome == "pass": - _require_two_accelerators() - else: - error = pytest.skip.Exception if outcome == "skip" else pytest.fail.Exception - with pytest.raises(error): - _require_two_accelerators() From 4800b5e57d3161601a04e666fa87da3d6a5a66cb Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 22:58:36 +0800 Subject: [PATCH 074/120] [test] use CPU plan for expert layout tests --- tests/checkpoints/test_expert_shard_layout.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 6893f24b9f..7e667a5748 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -16,7 +16,7 @@ from torch.distributed.tensor import DTensor, Partial, Replicate, Shard from torch.distributed.tensor.placement_types import _StridedShard -from scripts.ci.accelerator.launch_plan import from_parameter +from scripts.ci.accelerator.launch_plan import LaunchSpec, from_parameter from tests.tools.launch_utils import spawn_with_rendezvous from veomni.checkpoint.dcp_checkpointer import ( _EXPERT_LAYOUTS_KEY, @@ -32,6 +32,9 @@ ) +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + def _layout_worker(rank, world, directory, rendezvous): dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=world, timeout=timedelta(seconds=90)) try: From f6a449c31ea087b4985b4c2fc4108c63fd5c5038 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 4 Nov 2026 23:00:10 +0800 Subject: [PATCH 075/120] [test] use CPU plan for expert layout tests --- tests/checkpoints/test_expert_shard_layout.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 79bc0dd0f6..9e8d549c2b 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -35,6 +35,9 @@ ) +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + def _layout_worker(rank, world, directory, rendezvous): dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=world, timeout=timedelta(seconds=90)) try: From 845ccfde55038485ad3d941f15b3e651d94669eb Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 01:37:57 +0800 Subject: [PATCH 076/120] [checkpoint, model] fix: preserve safe RoPE reset and dense layouts --- tests/checkpoints/test_expert_shard_layout.py | 4 +- .../test_multi_optimizer_rope_cache.py | 94 +++++++++++++++++-- .../test_optimizer_checkpoint_preflight.py | 41 ++++++++ veomni/checkpoint/dynamic_rope.py | 25 ++++- veomni/checkpoint/optimizer_resume_layout.py | 18 +++- veomni/models/model_runtime.py | 11 ++- 6 files changed, 178 insertions(+), 15 deletions(-) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 9e8d549c2b..e13bb5c8d2 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -39,7 +39,9 @@ def _layout_worker(rank, world, directory, rendezvous): - dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=world, timeout=timedelta(seconds=90)) + dist.init_process_group( + "gloo", init_method=rendezvous, rank=rank, world_size=world, timeout=timedelta(seconds=180) + ) try: shape, names = ((2, 2), ("ep_fsdp", "ep")) if world == 4 else ((2, 2, 2), ("ep_replicate", "ep_fsdp", "ep")) mesh = init_device_mesh("cpu", shape, mesh_dim_names=names) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index 7d998297c8..ad8c431d3f 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -69,10 +69,10 @@ def _update(model, optimizer, value): optimizer.zero_grad() -def _rope_update(model, optimizer): - inputs = torch.arange(64, dtype=torch.float32).reshape(1, 8, 8) / 64 +def _rope_update(model, optimizer, length=8): + inputs = torch.arange(length * 8, dtype=torch.float32).reshape(1, length, 8) / (length * 8) projected = model.projection(inputs) - cos, sin = model.rope(projected, torch.arange(8).reshape(1, -1)) + cos, sin = model.rope(projected, torch.arange(length).reshape(1, -1)) loss = (projected * cos + sin).square().mean() loss.backward() trace = {"cos": cos.detach().clone(), "sin": sin.detach().clone(), "loss": loss.detach().clone()} @@ -82,6 +82,17 @@ def _rope_update(model, optimizer): return trace +def _rope_transitions(model, optimizer): + result = {} + for length in (8, 32, 24, 4): + trace = _rope_update(model, optimizer, length) + result.update({f"{length}.{key}": value for key, value in trace.items()}) + result[f"{length}.inv_freq"] = model.rope.inv_freq.detach().clone() + result[f"{length}.cached_length"] = torch.tensor(int(model.rope.max_seq_len_cached)) + result.update({f"{length}.param.{name}": value.detach().clone() for name, value in model.named_parameters()}) + return result + + def _fresh_resume(_rank, checkpoint, expected_path): model, optimizer = _build() DistributedCheckpointer.load( @@ -89,7 +100,7 @@ def _fresh_resume(_rank, checkpoint, expected_path): {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=SimpleNamespace(dp_mode="ddp"), ) - trace = _rope_update(model, optimizer) + trace = _rope_transitions(model, optimizer) expected = torch.load(expected_path, map_location="cpu", weights_only=True) actual = {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace} for group in ("model", "optimizer", "trace"): @@ -209,15 +220,64 @@ def test_rank_local_expanded_rope_error_is_coordinated(tmp_path): assert not (tmp_path / "checkpoint").exists() -def test_real_qwen3_dynamic_rope_reset_is_rejected_before_save(tmp_path): +def test_real_qwen3_dynamic_rope_reset_fresh_process_matches_next_update(tmp_path): model, optimizer = _build() _update(model, optimizer, 1.0) model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) assert model.rope.inv_freq is model.rope.original_inv_freq assert not any("inv_freq" in key for key in model.state_dict()) + checkpoint = str(tmp_path / "checkpoint") + DistributedCheckpointer.save( + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + trace = _rope_transitions(model, optimizer) + expected = tmp_path / "expected.pt" + torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace}, expected) + mp.spawn(_fresh_resume, args=(checkpoint, str(expected)), nprocs=1, join=True) + + +def _rank_local_rope_reset(rank, checkpoint, directory, phase, endpoint): + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) + try: + torch.manual_seed(13) + expected = str(directory / f"expected-{rank}.pt") + if phase == "load": + _fresh_resume(rank, checkpoint, expected) + return + model, optimizer = _build() + _update(model, optimizer, 1.0) + if rank == 0: + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) + assert (model.rope.inv_freq is model.rope.original_inv_freq) == (rank == 0) + DistributedCheckpointer.save( + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + trace = _rope_transitions(model, optimizer) + torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace}, expected) + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +def test_rank_local_qwen3_rope_reset_fresh_resume(tmp_path): + checkpoint = str(tmp_path / "checkpoint") + spawn_with_rendezvous(_rank_local_rope_reset, args=(checkpoint, tmp_path, "save"), nprocs=2) + spawn_with_rendezvous(_rank_local_rope_reset, args=(checkpoint, tmp_path, "load"), nprocs=2) + + +def test_qwen3_reset_alias_still_rejects_wrong_frequency_values(tmp_path): + model, optimizer = _build() + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) + model.rope.inv_freq.add_(1) checkpoint = tmp_path / "checkpoint" - with pytest.raises(RuntimeError, match="Cannot prove dynamic RoPE checkpoint cache"): + with pytest.raises(RuntimeError, match="Dynamic RoPE cache.*depends on prior sequence lengths"): DistributedCheckpointer.save( str(checkpoint), {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, @@ -226,6 +286,26 @@ def test_real_qwen3_dynamic_rope_reset_is_rejected_before_save(tmp_path): assert not checkpoint.exists() +def test_qwen3_reset_alias_with_third_buffer_keeps_identity_check(tmp_path): + model, optimizer = _build() + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) + model.rope.register_buffer("other_cache", model.rope.inv_freq, persistent=False) + checkpoint = str(tmp_path / "checkpoint") + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + target, restored = _build() + target.rope.register_buffer("other_cache", target.rope.inv_freq.clone(), persistent=False) + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + + def test_unexpanded_dynamic_rope_fresh_process_matches_training_update(tmp_path): model, optimizer = _build() _update(model, optimizer, 1.0) @@ -235,7 +315,7 @@ def test_unexpanded_dynamic_rope_fresh_process_matches_training_update(tmp_path) {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=SimpleNamespace(dp_mode="ddp"), ) - trace = _rope_update(model, optimizer) + trace = _rope_transitions(model, optimizer) expected = tmp_path / "expected.pt" torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace}, expected) mp.spawn(_fresh_resume, args=(checkpoint, str(expected)), nprocs=1, join=True) diff --git a/tests/models/test_optimizer_checkpoint_preflight.py b/tests/models/test_optimizer_checkpoint_preflight.py index 541243a13f..d1bbed35cb 100644 --- a/tests/models/test_optimizer_checkpoint_preflight.py +++ b/tests/models/test_optimizer_checkpoint_preflight.py @@ -68,6 +68,47 @@ def test_other_checkpoint_managers_keep_their_own_async_policy(monkeypatch): assert _runtime(True, True, manager="custom").optimizer._is_multi_optimizer is True +@pytest.mark.parametrize("active,has_plan", [(False, False), (True, False), (False, True), (True, True)]) +def test_canonical_layout_requires_only_an_active_extra_parallel_plan(monkeypatch, active, has_plan): + from torch.distributed.tensor import Shard + + from veomni.distributed.parallel_plan import ParallelPlan + + model = torch.nn.Linear(2, 2) + if has_plan: + model.get_parallel_plan = lambda: ParallelPlan(extra_parallel_plan={"ep": {"weight": Shard(0)}}) + state = SimpleNamespace( + any_extra_parallel_enabled=active, + extra_parallel_sizes=("ep",), + extra_parallel_enabled=lambda name: active and name == "ep", + ) + monkeypatch.setattr(model_runtime, "use_parallel_state", lambda *_: nullcontext()) + monkeypatch.setattr(model_runtime, "get_parallel_state_by_name", lambda _name: state) + monkeypatch.setattr("veomni.distributed.torch_parallelize.get_parallel_state", lambda: state) + for name in ( + "setup", + "_freeze_model_module", + "_build_parallelized_model", + "_build_model_assets", + "build_checkpoint", + ): + monkeypatch.setattr(model_runtime.VeOmniModelRuntime, name, lambda self: None) + monkeypatch.setattr(model_runtime.VeOmniModelRuntime, "_build_model", lambda self: setattr(self, "model", model)) + reached = [] + monkeypatch.setattr( + model_runtime.VeOmniModelRuntime, "_build_optimizer_checked", lambda self, *_: reached.append(True) + ) + args = SimpleNamespace(optimizer=SimpleNamespace(type="adamw")) + train = SimpleNamespace(checkpoint=SimpleNamespace(load_path=None, extra_parallel_canonical_layout=True)) + if active and has_plan: + with pytest.raises(RuntimeError, match="requires an ExtraParallel model plan"): + model_runtime.VeOmniModelRuntime(args, train=train) + assert not reached + else: + model_runtime.VeOmniModelRuntime(args, train=train) + assert reached == [True] + + def test_omni_composed_wrap_rejects_async_multi_optimizer_before_checkpoint(monkeypatch): from veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime import ModuleRuntime diff --git a/veomni/checkpoint/dynamic_rope.py b/veomni/checkpoint/dynamic_rope.py index 9e30a81dff..68898e4b09 100644 --- a/veomni/checkpoint/dynamic_rope.py +++ b/veomni/checkpoint/dynamic_rope.py @@ -9,6 +9,23 @@ import torch +def _is_native_qwen3_dynamic_reset_alias(module): + """Recognize the native Qwen3 reset alias, not custom mutable caches.""" + cls = type(module) + if cls.__module__ != "transformers.models.qwen3.modeling_qwen3" or cls.__name__ != "Qwen3RotaryEmbedding": + return False + from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding + + buffers = module._buffers + return ( + cls is Qwen3RotaryEmbedding + and module.rope_type == "dynamic" + and {"inv_freq", "original_inv_freq"} <= module._non_persistent_buffers_set + and buffers.get("inv_freq") is not None + and buffers.get("inv_freq") is buffers.get("original_inv_freq") + ) + + def dynamic_rope_checkpoint_contract(model): """Reject dynamic caches that a fresh runtime cannot reconstruct exactly. @@ -31,10 +48,10 @@ def dynamic_rope_checkpoint_contract(model): cached = int(getattr(module, f"{prefix}max_seq_len_cached", module.max_seq_len_cached)) inv = getattr(module, f"{prefix}inv_freq") initial_inv = getattr(module, f"{prefix}original_inv_freq") - # After a dynamic grow/reset, HF can assign the original - # buffer object directly. A fresh instance may not share - # this alias, and DCP does not persist it. - if inv is initial_inv: + # Native Qwen3 resets by reusing the original buffer. Its + # forward reads the buffer and growth replaces it, so object + # identity is not part of its recoverable numerical state. + if inv is initial_inv and not (layer_type is None and _is_native_qwen3_dynamic_reset_alias(module)): raise RuntimeError("dynamic RoPE reset shares nonpersistent cache buffers") kwargs = {} if layer_type is None else {"layer_type": layer_type} expected_inv, expected_scale = ROPE_INIT_FUNCTIONS[kind]( diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index 87317b3f3f..52c4182681 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -159,7 +159,7 @@ def check_aliases(value, path): optimizer_names.setdefault(id(param), set()).update(_get_fqns(optimizer.model, runtime_name)) if names != optimizer_names: raise RuntimeError("MultiOptimizer and checkpoint model have different canonical parameter identities") - from .dynamic_rope import dynamic_rope_checkpoint_contract + from .dynamic_rope import _is_native_qwen3_dynamic_reset_alias, dynamic_rope_checkpoint_contract model_aliases = {} for kind, items in ( @@ -170,7 +170,21 @@ def check_aliases(value, path): model_aliases.setdefault(id(value), set()).update( f"{kind}:{name}" for name in _get_fqns(model, runtime_name) ) - alias_graph = sorted(sorted(aliases) for aliases in model_aliases.values() if len(aliases) > 1) + safe_rope_pairs = set() + for module_name, module in model.named_modules(remove_duplicate=False): + if _is_native_qwen3_dynamic_reset_alias(module): + rope_names = { + f"buffer:{fqn}" + for buffer_name in ("inv_freq", "original_inv_freq") + for fqn in _get_fqns(model, f"{module_name}.{buffer_name}" if module_name else buffer_name) + } + if len(rope_names) == 2: + safe_rope_pairs.add(frozenset(rope_names)) + alias_graph = sorted( + sorted(aliases) + for aliases in model_aliases.values() + if len(aliases) > 1 and frozenset(aliases) not in safe_rope_pairs + ) leaves = [] seen_param_ids, seen_fqns, lr_aliases, lr_owners = set(), set(), {}, {} diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index 234150cd2c..a6b49bc673 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -166,7 +166,16 @@ def __init__( layout_error = None try: if canonical: - _configure_extra_parallel_checkpoint_layout(self.model, self.train_args.checkpoint) + from ..distributed.torch_parallelize import _has_extra_parallel_plan + + _configure_extra_parallel_checkpoint_layout( + self.model, + self.train_args.checkpoint, + require_plan=( + bool(getattr(self.parallel_state, "any_extra_parallel_enabled", False)) + and _has_extra_parallel_plan(self.model) + ), + ) except BaseException as exc: layout_error = exc raise_if_any_rank_failed(layout_error, "configuring ExtraParallel checkpoint layout") From 0c9002dadbe19bbdd9a106c7e9e3ab567da70f58 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 02:51:01 +0800 Subject: [PATCH 077/120] [checkpoint] fix: bind dynamic RoPE table bytes on resume --- docs/usage/checkpoint.md | 21 +- .../test_multi_optimizer_proven_resume.py | 11 + .../test_multi_optimizer_rope_cache.py | 231 +++++++++++++++++- veomni/checkpoint/dynamic_rope.py | 73 +++++- veomni/checkpoint/optimizer_resume_layout.py | 2 +- 5 files changed, 308 insertions(+), 30 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index f06ff841ff..946029e020 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -471,17 +471,24 @@ identity covers parameters and buffers, including nonpersistent buffer aliases. Optimizer adapters and model serialization hooks are trusted to preserve their state and semantics. Identity tracks repeated tensor objects, but does not detect distinct views sharing storage; models that rely on mutable shared views are -outside this checkpoint contract. A reset dynamic RoPE cache can have a -different alias graph from a newly constructed model and is conservatively -rejected before writing until that cache state can be restored safely. +outside this checkpoint contract. The upstream and VeOmni-generated Qwen3 +dynamic RoPE reset alias is permitted when its current and original tables match +a valid reconstruction; other reset aliases remain outside this contract. Dynamic RoPE frequencies and cached maximum lengths are not checkpointed. If the cache grew beyond its original context and has not reset, a fresh process can compute different later outputs. MultiOptimizer checkpoints therefore reject a non-reconstructible dynamic or longrope cache at save. With a dynamic -or longrope target, resume requires a matching cache contract; older checkpoints -without that proof need a verified migration. Older checkpoints do not identify -the source's RoPE configuration, so the operator must also independently verify -that the source and target configurations match. Saving an expanded cache is not +or longrope target, resume requires the exact saved current/reset frequency +bytes, shape, implementation and scaling to match the fresh target. The supported +Qwen3 classes accept an exact CPU-built table moved to the accelerator as well as +an exact table built on that accelerator; close-but-different tables are not +interchangeable. The byte-level contract replaces the earlier configuration-only +contract, so older dynamic or longrope checkpoints need a verified migration. +Unversioned sources lack the source's RoPE configuration, including when the +target is static; the operator must independently verify source and target. +Matching the RoPE table alone does not establish whole-training determinism: the +runtime, hardware, topology, input and other checkpoint state must also match. +Saving an expanded cache is not yet supported. Dynamic and longrope training may run without saving, but a MultiOptimizer DCP save fails if its cache cannot be reconstructed; operators should account for this when planning checkpoint cadence. Ordinary diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index dcaa9b54ca..14e6708abc 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -665,8 +665,11 @@ def _runtime_plan_mismatch_worker(rank, endpoint): from datetime import timedelta import torch.distributed as dist + from torch.distributed.tensor import Shard + import veomni.distributed.torch_parallelize as parallelize_module import veomni.models.model_runtime as runtime_module + from veomni.distributed.parallel_plan import ParallelPlan dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) try: @@ -677,6 +680,7 @@ def setup(self): def _build_model(self): self.model = torch.nn.Linear(2, 2) + self.model.get_parallel_plan = lambda: ParallelPlan(extra_parallel_plan={"ep": {"weight": Shard(0)}}) def _freeze_model_module(self): pass @@ -685,6 +689,13 @@ def _build_parallelized_model(self): pass runtime_module.use_parallel_state = lambda *_args, **_kwargs: contextlib.nullcontext() + state = SimpleNamespace( + any_extra_parallel_enabled=True, + extra_parallel_sizes=("ep",), + extra_parallel_enabled=lambda name: name == "ep", + ) + runtime_module.get_parallel_state_by_name = lambda _name: state + parallelize_module.get_parallel_state = lambda: state checkpoint = SimpleNamespace(extra_parallel_canonical_layout=True) args = SimpleNamespace(optimizer=SimpleNamespace(type="adamw")) with pytest.raises(RuntimeError, match="requires an ExtraParallel model plan"): diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index ad8c431d3f..ff54844895 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -42,7 +42,7 @@ def test_runtime_preflight_allows_dynamic_rope_training_without_checkpoint(): class _TinyQwenRope(torch.nn.Module): - def __init__(self, theta=10000): + def __init__(self, theta=10000, rope_class=Qwen3RotaryEmbedding): super().__init__() self.projection = torch.nn.Linear(8, 8) config = Qwen3Config( @@ -53,15 +53,28 @@ def __init__(self, theta=10000): max_position_embeddings=16, rope_parameters={"rope_type": "dynamic", "factor": 2.0, "rope_theta": theta}, ) - self.rope = Qwen3RotaryEmbedding(config) + self.rope = rope_class(config) -def _build(theta=10000): - model = _TinyQwenRope(theta) +def _build(theta=10000, rope_class=Qwen3RotaryEmbedding): + model = _TinyQwenRope(theta, rope_class) leaf = torch.optim.AdamW(model.parameters(), lr=0.01) return model, MultiOptimizer(model, {"model": leaf}, ["model"]) +def _registry_rope_class(backend): + import importlib + from unittest.mock import patch + + from veomni.models.registry import MODELING_REGISTRY + from veomni.models.transformers import qwen3 + + with patch.object(qwen3, "IS_NPU_AVAILABLE", backend == "npu"): + model_class = MODELING_REGISTRY["qwen3"]("Qwen3Model") + assert model_class.__module__.endswith(f"patched_modeling_qwen3_{backend}") + return importlib.import_module(model_class.__module__).Qwen3RotaryEmbedding + + def _update(model, optimizer, value): for parameter in model.parameters(): parameter.grad = torch.full_like(parameter, value) @@ -84,17 +97,19 @@ def _rope_update(model, optimizer, length=8): def _rope_transitions(model, optimizer): result = {} - for length in (8, 32, 24, 4): + for index, length in enumerate((8, 16, 17, 16, 15, 32, 24, 4)): trace = _rope_update(model, optimizer, length) - result.update({f"{length}.{key}": value for key, value in trace.items()}) - result[f"{length}.inv_freq"] = model.rope.inv_freq.detach().clone() - result[f"{length}.cached_length"] = torch.tensor(int(model.rope.max_seq_len_cached)) - result.update({f"{length}.param.{name}": value.detach().clone() for name, value in model.named_parameters()}) + label = f"{index}:{length}" + result.update({f"{label}.{key}": value for key, value in trace.items()}) + result[f"{label}.inv_freq"] = model.rope.inv_freq.detach().clone() + result[f"{label}.cached_length"] = torch.tensor(int(model.rope.max_seq_len_cached)) + result.update({f"{label}.param.{name}": value.detach().clone() for name, value in model.named_parameters()}) return result -def _fresh_resume(_rank, checkpoint, expected_path): - model, optimizer = _build() +def _fresh_resume(_rank, checkpoint, expected_path, backend=None): + rope_class = Qwen3RotaryEmbedding if backend is None else _registry_rope_class(backend) + model, optimizer = _build(rope_class=rope_class) DistributedCheckpointer.load( checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, @@ -138,6 +153,51 @@ def test_standard_qwen4_rope_128k_does_not_require_dynamic_cache_contract(): assert dynamic_rope_checkpoint_contract(rope) == [] +def test_qwen4_default_rope_keeps_legacy_checkpoint_layout(tmp_path): + import json + + from transformers.models.qwen4_exp.configuration_qwen4_exp import Qwen4ExpTextConfig + from transformers.models.qwen4_exp.modeling_qwen4_exp import Qwen4ExpTextRotaryEmbedding + + from veomni.checkpoint.optimizer_resume_layout import capture_multi_layout + + def build(): + model, optimizer = _build() + config = Qwen4ExpTextConfig( + hidden_size=32, + head_dim=8, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=262144, + rope_parameters={"rope_type": "default", "rope_theta": 10000000, "mrope_section": [1, 1, 2]}, + ) + model.rope = Qwen4ExpTextRotaryEmbedding(config) + return model, optimizer + + model, optimizer = build() + _update(model, optimizer, 1.0) + layout = json.loads(capture_multi_layout(model, optimizer, None, None)) + assert "dynamic_rope_cache_v1" not in layout and "dynamic_rope_cache_v2" not in layout + checkpoint = str(tmp_path / "checkpoint") + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + target, restored = build() + DistributedCheckpointer.load( + checkpoint, {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + for actual, expected in zip(target.parameters(), model.parameters(), strict=True): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + for key in ("step", "exp_avg", "exp_avg_sq"): + torch.testing.assert_close( + restored.optimizers_dict["model"].state[actual][key], + optimizer.optimizers_dict["model"].state[expected][key], + rtol=0, + atol=0, + ) + + def test_expanded_dynamic_rope_cache_is_rejected_before_checkpoint_write(tmp_path): model, optimizer = _build() model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) @@ -197,6 +257,136 @@ def test_changed_dynamic_rope_config_is_rejected_before_model_load(tmp_path): torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) +def _skew_device_recomputation(monkeypatch): + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + original = ROPE_INIT_FUNCTIONS["dynamic"] + + def skewed(config, device, seq_len=None, **kwargs): + values, scale = original(config, device, seq_len=seq_len, **kwargs) + # Model the observed CPU/NPU 1-ULP difference without requiring an NPU. + if isinstance(device, torch.device): + values = torch.nextafter(values, torch.full_like(values, float("inf"))) + return values, scale + + monkeypatch.setitem(ROPE_INIT_FUNCTIONS, "dynamic", skewed) + + +@pytest.mark.parametrize("backend", ["hf", "gpu", "npu"]) +def test_cpu_origin_rope_with_device_rounding_difference_resumes_exactly(tmp_path, monkeypatch, backend): + rope_class = Qwen3RotaryEmbedding if backend == "hf" else _registry_rope_class(backend) + model, optimizer = _build(rope_class=rope_class) + _update(model, optimizer, 1.0) + _skew_device_recomputation(monkeypatch) + checkpoint = str(tmp_path / "checkpoint") + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + target, restored = _build(rope_class=rope_class) + DistributedCheckpointer.load( + checkpoint, {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + expected = _rope_transitions(model, optimizer) + actual = _rope_transitions(target, restored) + assert actual.keys() == expected.keys() + for name in actual: + torch.testing.assert_close(actual[name], expected[name], rtol=0, atol=0) + + +@pytest.mark.parametrize("backend", ["hf", "gpu", "npu"]) +def test_distinct_cpu_and_device_origin_rope_is_rejected_before_restore(tmp_path, monkeypatch, backend): + rope_class = Qwen3RotaryEmbedding if backend == "hf" else _registry_rope_class(backend) + model, optimizer = _build(rope_class=rope_class) + _skew_device_recomputation(monkeypatch) + checkpoint = str(tmp_path / "checkpoint") + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + target, restored = _build(rope_class=rope_class) + shifted = torch.nextafter(target.rope.inv_freq, torch.full_like(target.rope.inv_freq, float("inf"))) + target.rope.inv_freq.copy_(shifted) + target.rope.original_inv_freq.copy_(shifted) + before = target.projection.weight.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) + assert not restored.optimizers_dict["model"].state + + +def test_signed_zero_in_rope_table_is_not_a_canonical_match(tmp_path): + model, optimizer = _build(theta=1e12) + model.rope.inv_freq = model.rope.inv_freq.to(torch.float16) + model.rope.original_inv_freq = model.rope.original_inv_freq.to(torch.float16) + assert model.rope.inv_freq[-1].item() == 0 + model.rope.inv_freq[-1] = -0.0 + model.rope.original_inv_freq[-1] = -0.0 + checkpoint = tmp_path / "checkpoint" + with pytest.raises(RuntimeError, match="Dynamic RoPE cache.*depends on prior sequence lengths"): + DistributedCheckpointer.save( + str(checkpoint), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not checkpoint.exists() + + +def test_changed_qwen3_rope_implementation_is_rejected_before_restore(tmp_path): + model, optimizer = _build() + checkpoint = str(tmp_path / "checkpoint") + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + target, restored = _build(rope_class=_registry_rope_class("gpu")) + torch.testing.assert_close(model.rope.inv_freq, target.rope.inv_freq, rtol=0, atol=0) + before = target.projection.weight.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) + + +def test_old_dynamic_rope_contract_without_table_identity_is_rejected(tmp_path, monkeypatch): + import veomni.checkpoint.optimizer_resume_layout as layout + + model, optimizer = _build() + original = layout._local_layout + + def capture_legacy(*args, **kwargs): + payload = original(*args, **kwargs) + if kwargs.get("capture"): + payload["dynamic_rope_cache_v1"] = [item[:6] for item in payload.pop("dynamic_rope_cache_v2")] + return payload + + checkpoint = str(tmp_path / "checkpoint") + parallel = SimpleNamespace(dp_mode="ddp") + with monkeypatch.context() as legacy: + legacy.setattr(layout, "_local_layout", capture_legacy) + DistributedCheckpointer.save( + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + target, restored = _build() + before = target.projection.weight.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load( + checkpoint, + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + torch.testing.assert_close(target.projection.weight, before, rtol=0, atol=0) + + def _rank_local_rope_error(rank, checkpoint, endpoint): dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) try: @@ -239,6 +429,25 @@ def test_real_qwen3_dynamic_rope_reset_fresh_process_matches_next_update(tmp_pat mp.spawn(_fresh_resume, args=(checkpoint, str(expected)), nprocs=1, join=True) +@pytest.mark.parametrize("backend", ["gpu", "npu"]) +def test_registry_qwen3_dynamic_rope_reset_fresh_process_matches_next_update(tmp_path, backend): + model, optimizer = _build(rope_class=_registry_rope_class(backend)) + _update(model, optimizer, 1.0) + model.rope(torch.zeros((1, 1, 8)), torch.arange(32).view(1, -1)) + model.rope(torch.zeros((1, 1, 8)), torch.arange(4).view(1, -1)) + assert model.rope.inv_freq is model.rope.original_inv_freq + checkpoint = str(tmp_path / "checkpoint") + DistributedCheckpointer.save( + checkpoint, + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + trace = _rope_transitions(model, optimizer) + expected = tmp_path / "expected.pt" + torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(), "trace": trace}, expected) + mp.spawn(_fresh_resume, args=(checkpoint, str(expected), backend), nprocs=1, join=True) + + def _rank_local_rope_reset(rank, checkpoint, directory, phase, endpoint): dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)) try: diff --git a/veomni/checkpoint/dynamic_rope.py b/veomni/checkpoint/dynamic_rope.py index 68898e4b09..024a03cf4e 100644 --- a/veomni/checkpoint/dynamic_rope.py +++ b/veomni/checkpoint/dynamic_rope.py @@ -3,29 +3,44 @@ """Checkpoint preflight for dynamic RoPE's nonpersistent runtime state.""" import hashlib +import importlib import json import math -import torch +from .tensor_bytes import tensor_byte_chunks -def _is_native_qwen3_dynamic_reset_alias(module): - """Recognize the native Qwen3 reset alias, not custom mutable caches.""" +def _is_native_qwen3_dynamic(module): + """Recognize upstream and VeOmni-generated Qwen3 rotary implementations.""" cls = type(module) - if cls.__module__ != "transformers.models.qwen3.modeling_qwen3" or cls.__name__ != "Qwen3RotaryEmbedding": + approved_modules = { + "transformers.models.qwen3.modeling_qwen3", + "veomni.models.transformers.qwen3.generated.patched_modeling_qwen3_gpu", + "veomni.models.transformers.qwen3.generated.patched_modeling_qwen3_npu", + } + if cls.__module__ not in approved_modules or cls.__name__ != "Qwen3RotaryEmbedding": return False - from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding + return importlib.import_module(cls.__module__).Qwen3RotaryEmbedding is cls and module.rope_type == "dynamic" + +def _is_native_qwen3_dynamic_reset_alias(module): + """Recognize the native Qwen3 reset alias, not custom mutable caches.""" buffers = module._buffers return ( - cls is Qwen3RotaryEmbedding - and module.rope_type == "dynamic" + _is_native_qwen3_dynamic(module) and {"inv_freq", "original_inv_freq"} <= module._non_persistent_buffers_set and buffers.get("inv_freq") is not None and buffers.get("inv_freq") is buffers.get("original_inv_freq") ) +def _tensor_digest(value): + digest = hashlib.sha256() + for chunk in tensor_byte_chunks(value, chunk_bytes=4 << 20): + digest.update(chunk) + return digest.hexdigest() + + def dynamic_rope_checkpoint_contract(model): """Reject dynamic caches that a fresh runtime cannot reconstruct exactly. @@ -67,16 +82,35 @@ def dynamic_rope_checkpoint_contract(model): config_hash = hashlib.sha256( json.dumps(config, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() ).hexdigest() + inv_digest = _tensor_digest(inv) + initial_digest = _tensor_digest(initial_inv) + expected_inv = expected_inv.to(dtype=inv.dtype) + matches_canonical = ( + tuple(inv.shape) == tuple(expected_inv.shape) + and inv_digest == _tensor_digest(expected_inv) + and scale == float(expected_scale) + ) + if not matches_canonical and layer_type is None and _is_native_qwen3_dynamic(module): + # HF often builds the original table on CPU before moving it + # to an accelerator. Recompute that origin exactly rather + # than tolerating device-dependent rounding differences. + cpu_inv, cpu_scale = ROPE_INIT_FUNCTIONS[kind](module.config, device="cpu", seq_len=original) + cpu_inv = cpu_inv.to(dtype=inv.dtype) + matches_canonical = ( + tuple(inv.shape) == tuple(cpu_inv.shape) + and inv_digest == _tensor_digest(cpu_inv) + and scale == float(cpu_scale) + ) valid = ( original > 0 and cached == original and not inv.requires_grad and not initial_inv.requires_grad and inv.dtype == initial_inv.dtype - and torch.equal(inv, initial_inv.to(device=inv.device)) - and torch.equal(inv, expected_inv.to(device=inv.device, dtype=inv.dtype)) + and tuple(inv.shape) == tuple(initial_inv.shape) + and inv_digest == initial_digest + and matches_canonical and math.isfinite(scale) - and scale == float(expected_scale) and (kind != "longrope" or not hasattr(module, f"{prefix}long_inv_freq")) ) except Exception as exc: @@ -86,5 +120,22 @@ def dynamic_rope_checkpoint_contract(model): f"Dynamic RoPE cache for {name or ''} depends on prior sequence lengths; " "its nonpersistent state cannot yet be restored exactly" ) - descriptors.append((name, layer_type, kind, original, str(inv.dtype), config_hash)) + # DCP omits these nonpersistent tables: bind the exact current and + # reset bytes to the layout so restore rejects another origin even + # when both individually match a valid canonical reconstruction. + descriptors.append( + ( + name, + layer_type, + kind, + original, + str(inv.dtype), + config_hash, + f"{type(module).__module__}.{type(module).__qualname__}", + tuple(inv.shape), + scale.hex(), + inv_digest, + initial_digest, + ) + ) return descriptors diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index 52c4182681..fdf1655790 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -333,7 +333,7 @@ def mesh_identity(mesh): if rope_contract: # Earlier checkpoints did not prove that transient cache history was # reconstructible by a fresh process. - layout["dynamic_rope_cache_v1"] = rope_contract + layout["dynamic_rope_cache_v2"] = rope_contract return layout if capture else _layout_digest(layout) From 1a96e9b1dfc59408ba6b9876251c9e60423bb596 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 03:18:52 +0800 Subject: [PATCH 078/120] [checkpoint] reject nonfinite RoPE tables without changing FP8 identity --- .../test_multi_optimizer_rope_cache.py | 51 +++++++++++++++++++ veomni/checkpoint/dynamic_rope.py | 15 ++++++ 2 files changed, 66 insertions(+) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index ff54844895..6bd4dd3450 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -336,6 +336,57 @@ def test_signed_zero_in_rope_table_is_not_a_canonical_match(tmp_path): assert not checkpoint.exists() +@pytest.mark.parametrize("dtype", [torch.float32, torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz]) +def test_nonfinite_canonical_rope_is_rejected(tmp_path, monkeypatch, dtype): + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + model, optimizer = _build() + nonfinite_table = model.rope.inv_freq.clone() + nonfinite_table[0] = float("nan") + model.rope.inv_freq = nonfinite_table.to(dtype) + model.rope.original_inv_freq = nonfinite_table.to(dtype) + original = ROPE_INIT_FUNCTIONS["dynamic"] + + def nonfinite(config, device, seq_len=None, **kwargs): + value, scale = original(config, device, seq_len=seq_len, **kwargs) + value[0] = float("nan") + return value, scale + + monkeypatch.setitem(ROPE_INIT_FUNCTIONS, "dynamic", nonfinite) + checkpoint = tmp_path / "checkpoint" + with pytest.raises(RuntimeError, match="Cannot prove|Dynamic RoPE cache"): + DistributedCheckpointer.save( + str(checkpoint), + {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + assert not checkpoint.exists() + + +@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz]) +def test_finite_float8_rope_tables_keep_exact_checkpoint_identity(tmp_path, dtype): + def build(): + model, optimizer = _build() + model.rope.inv_freq = model.rope.inv_freq.to(dtype) + model.rope.original_inv_freq = model.rope.original_inv_freq.to(dtype) + return model, optimizer + + model, optimizer = build() + checkpoint = str(tmp_path / "checkpoint") + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + target, restored = build() + DistributedCheckpointer.load( + checkpoint, {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, parallel_state=parallel + ) + expected, actual = _rope_transitions(model, optimizer), _rope_transitions(target, restored) + assert expected.keys() == actual.keys() + for name in expected: + torch.testing.assert_close(expected[name], actual[name], rtol=0, atol=0) + + def test_changed_qwen3_rope_implementation_is_rejected_before_restore(tmp_path): model, optimizer = _build() checkpoint = str(tmp_path / "checkpoint") diff --git a/veomni/checkpoint/dynamic_rope.py b/veomni/checkpoint/dynamic_rope.py index 024a03cf4e..cfc64be82a 100644 --- a/veomni/checkpoint/dynamic_rope.py +++ b/veomni/checkpoint/dynamic_rope.py @@ -7,6 +7,8 @@ import json import math +import torch + from .tensor_bytes import tensor_byte_chunks @@ -35,6 +37,19 @@ def _is_native_qwen3_dynamic_reset_alias(module): def _tensor_digest(value): + value = value.detach().cpu() + if not value.is_floating_point(): + raise ValueError("Dynamic RoPE frequencies must be finite floating-point values") + try: + finite = torch.isfinite(value).all() + except NotImplementedError: + if not str(value.dtype).startswith("torch.float8_"): + raise + # Every finite FP8 value is exactly representable in FP32. Check there + # only when the CPU backend lacks isfinite for the FP8 storage dtype. + finite = torch.isfinite(value.float()).all() + if not bool(finite): + raise ValueError("Dynamic RoPE frequencies must be finite floating-point values") digest = hashlib.sha256() for chunk in tensor_byte_chunks(value, chunk_bytes=4 << 20): digest.update(chunk) From 80f61cfc379686df5952c990c2498b985918909a Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 04:18:31 +0800 Subject: [PATCH 079/120] [checkpoint] limit RoPE finiteness checks to supported types --- .../test_multi_optimizer_rope_cache.py | 33 ++----------------- veomni/checkpoint/dynamic_rope.py | 12 +------ 2 files changed, 4 insertions(+), 41 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index 6bd4dd3450..1ee475b465 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -336,15 +336,12 @@ def test_signed_zero_in_rope_table_is_not_a_canonical_match(tmp_path): assert not checkpoint.exists() -@pytest.mark.parametrize("dtype", [torch.float32, torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz]) -def test_nonfinite_canonical_rope_is_rejected(tmp_path, monkeypatch, dtype): +def test_nonfinite_canonical_rope_is_rejected(tmp_path, monkeypatch): from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS model, optimizer = _build() - nonfinite_table = model.rope.inv_freq.clone() - nonfinite_table[0] = float("nan") - model.rope.inv_freq = nonfinite_table.to(dtype) - model.rope.original_inv_freq = nonfinite_table.to(dtype) + model.rope.inv_freq[0] = float("nan") + model.rope.original_inv_freq[0] = float("nan") original = ROPE_INIT_FUNCTIONS["dynamic"] def nonfinite(config, device, seq_len=None, **kwargs): @@ -363,30 +360,6 @@ def nonfinite(config, device, seq_len=None, **kwargs): assert not checkpoint.exists() -@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2fnuz]) -def test_finite_float8_rope_tables_keep_exact_checkpoint_identity(tmp_path, dtype): - def build(): - model, optimizer = _build() - model.rope.inv_freq = model.rope.inv_freq.to(dtype) - model.rope.original_inv_freq = model.rope.original_inv_freq.to(dtype) - return model, optimizer - - model, optimizer = build() - checkpoint = str(tmp_path / "checkpoint") - parallel = SimpleNamespace(dp_mode="ddp") - DistributedCheckpointer.save( - checkpoint, {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, parallel_state=parallel - ) - target, restored = build() - DistributedCheckpointer.load( - checkpoint, {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, parallel_state=parallel - ) - expected, actual = _rope_transitions(model, optimizer), _rope_transitions(target, restored) - assert expected.keys() == actual.keys() - for name in expected: - torch.testing.assert_close(expected[name], actual[name], rtol=0, atol=0) - - def test_changed_qwen3_rope_implementation_is_rejected_before_restore(tmp_path): model, optimizer = _build() checkpoint = str(tmp_path / "checkpoint") diff --git a/veomni/checkpoint/dynamic_rope.py b/veomni/checkpoint/dynamic_rope.py index cfc64be82a..bcaa7457cb 100644 --- a/veomni/checkpoint/dynamic_rope.py +++ b/veomni/checkpoint/dynamic_rope.py @@ -38,17 +38,7 @@ def _is_native_qwen3_dynamic_reset_alias(module): def _tensor_digest(value): value = value.detach().cpu() - if not value.is_floating_point(): - raise ValueError("Dynamic RoPE frequencies must be finite floating-point values") - try: - finite = torch.isfinite(value).all() - except NotImplementedError: - if not str(value.dtype).startswith("torch.float8_"): - raise - # Every finite FP8 value is exactly representable in FP32. Check there - # only when the CPU backend lacks isfinite for the FP8 storage dtype. - finite = torch.isfinite(value.float()).all() - if not bool(finite): + if not value.is_floating_point() or not bool(torch.isfinite(value).all()): raise ValueError("Dynamic RoPE frequencies must be finite floating-point values") digest = hashlib.sha256() for chunk in tensor_byte_chunks(value, chunk_bytes=4 << 20): From f1b3e033912f4d59e3d02dcb0ba347597e7a9842 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 04:18:32 +0800 Subject: [PATCH 080/120] [checkpoint] document reserved runtime state keys --- tests/checkpoints/test_multi_optimizer_proven_resume.py | 3 ++- veomni/models/checkpoint/manager.py | 2 -- veomni/models/model_runtime.py | 4 +++- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 14e6708abc..a6e2670032 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -1744,13 +1744,14 @@ def build(shared): def test_model_checkpoint_extra_state_records_optimizer_global_step(): - from veomni.models.checkpoint.manager import ModelCheckpointManager + from veomni.models.checkpoint.manager import ModelCheckpointManager, _scheduler_recipe _, optimizer = _model_and_optimizer() manager = object.__new__(ModelCheckpointManager) manager.runtime = SimpleNamespace(optimizer=optimizer, lr_scheduler=None, extra_state=lambda: {}) saved = manager._extra_state(SimpleNamespace(global_step=3)) assert saved["_veomni_optimizer_global_step"] == 3 + assert saved["_veomni_scheduler_recipe"] == {"status": "supported", "value": _scheduler_recipe(None, optimizer)} def test_proven_optimizer_load_sets_identity_receipt(tmp_path): diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index a50a0012d0..14abcf775c 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -305,8 +305,6 @@ def _extra_state(self, state: "TrainerState") -> Dict[str, Any]: } extra_state.update(self.runtime.extra_state()) if getattr(self.runtime.optimizer, "_is_multi_optimizer", False) is True: - if any(key in extra_state for key in ("_veomni_optimizer_global_step", "_veomni_scheduler_recipe")): - raise RuntimeError("Reserved MultiOptimizer checkpoint step key is already in use") extra_state["_veomni_optimizer_global_step"] = int(state.global_step) try: extra_state["_veomni_scheduler_recipe"] = { diff --git a/veomni/models/model_runtime.py b/veomni/models/model_runtime.py index a6b49bc673..9074761f26 100644 --- a/veomni/models/model_runtime.py +++ b/veomni/models/model_runtime.py @@ -739,7 +739,9 @@ def extra_state(self) -> Dict[str, Any]: Models with extra state (e.g. the DiT condition model's noise/timestep generator) contribute it here. The checkpoint manager merges the result - into the model's extra_state blob. Default: nothing. + into the model's extra_state blob. Keys prefixed with ``_veomni_`` are + reserved for framework checkpoint metadata and must not be returned by + custom runtimes. Default: nothing. """ return {} From d075345a1b6bc4ae068e4d70605a9632eadfff33 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 10:54:11 +0800 Subject: [PATCH 081/120] [docs] describe file-backed checkpoint abort hook --- .../rank_local_adamw_checkpoint_adapter.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/docs/design/rank_local_adamw_checkpoint_adapter.md b/docs/design/rank_local_adamw_checkpoint_adapter.md index 53fb98f9f5..eb654867e5 100644 --- a/docs/design/rank_local_adamw_checkpoint_adapter.md +++ b/docs/design/rank_local_adamw_checkpoint_adapter.md @@ -117,3 +117,20 @@ for the complete recovery constraints. Test an extension with initialized and unused parameters, save, fresh-process restore, and the next update at its actual dtype and placement; an adapter registration or a single forward pass is not an acceptance test. + +## Optional file-backed state abort hook + +An optimizer that owns file-backed state may provide the zero-argument +`checkpoint_abort_file_backed_state()` hook. If the final distributed barrier +in `ModelCheckpointManager.load()` fails after loading state, the manager +calls this hook on the optimizer and on each `MultiOptimizer` leaf that defines +it. The hook owns its optimizer's cleanup or invalidation: it should make any +partially restored file-backed state unusable until the optimizer is rebuilt +or explicitly recovered. It must not perform collectives and should be safe +to call more than once. The manager ignores ordinary cleanup exceptions so the +original barrier failure is propagated; it does not report the checkpoint as +successfully loaded. Optimizers without file-backed state need no hook. + +The manager calls this hook only on its final-barrier failure path; it is not a +general notification for every earlier load error. File-backed optimizers and +checkpointers remain responsible for cleanup at their own failure sites. From 8c97fbc182fb3b4def5ecf3eb5fd2615af78c0e1 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:23:27 +0800 Subject: [PATCH 082/120] [checkpoint] unify scheduler identity validation --- .../test_multi_optimizer_proven_resume.py | 14 ++++++ veomni/checkpoint/optimizer_resume_layout.py | 30 ++++++++----- veomni/models/checkpoint/manager.py | 44 ++----------------- 3 files changed, 37 insertions(+), 51 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index a6e2670032..3346834735 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -1991,6 +1991,20 @@ def test_multi_checkpoint_requires_explicit_scheduler_declaration(tmp_path): assert not (tmp_path / "model").exists() +def test_scheduler_recipe_rejects_schedules_outside_the_checkpoint_contract(): + from veomni.models.checkpoint.manager import _scheduler_recipe + + _, optimizer = _model_and_optimizer() + leaf = optimizer.optimizers_dict["weight"] + step = torch.optim.lr_scheduler.StepLR(leaf, step_size=1) + with pytest.raises(RuntimeError, match="scheduler implementation or binding is unsupported"): + _scheduler_recipe(step, leaf) + + scheduler = torch.optim.lr_scheduler.LambdaLR(leaf, lambda value, *, factor=0.5: value * factor) + with pytest.raises(RuntimeError, match="scheduler captures unsupported mutable state"): + _scheduler_recipe(scheduler, leaf) + + def test_multi_checkpoint_accepts_explicit_no_scheduler(tmp_path): model, optimizer = _model_and_optimizer() model.weight.grad = torch.ones_like(model.weight) diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index fdf1655790..be910af399 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -65,18 +65,10 @@ def _lr_layout(group, key, aliases): return {"type": type(value).__name__} -def _scheduler_layout(leaf): +def _scheduler_functions(scheduler): + """Describe LambdaLR callables using the existing checkpoint identity format.""" import math - from ..optim.lr_scheduler import checkpoint_scheduler_binding - - scheduler = checkpoint_scheduler_binding(leaf) - if scheduler is None: - return None - if type(scheduler) is not torch.optim.lr_scheduler.LambdaLR: - raise RuntimeError("MultiOptimizer requires a supported scheduler") - if scheduler.optimizer is not leaf: - raise RuntimeError("MultiOptimizer scheduler binding changed") functions = [] for function in scheduler.lr_lambdas: if not isinstance(function, types.FunctionType): @@ -99,7 +91,23 @@ def normalize(code): "closure": [_immutable_value(cell.cell_contents) for cell in function.__closure__ or ()], } ) - return {"type": f"{type(scheduler).__module__}.{type(scheduler).__qualname__}", "functions": functions} + return functions + + +def _scheduler_layout(leaf): + from ..optim.lr_scheduler import checkpoint_scheduler_binding + + scheduler = checkpoint_scheduler_binding(leaf) + if scheduler is None: + return None + if type(scheduler) is not torch.optim.lr_scheduler.LambdaLR: + raise RuntimeError("MultiOptimizer requires a supported scheduler") + if scheduler.optimizer is not leaf: + raise RuntimeError("MultiOptimizer scheduler binding changed") + return { + "type": f"{type(scheduler).__module__}.{type(scheduler).__qualname__}", + "functions": _scheduler_functions(scheduler), + } def _local_layout( diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py index 14abcf775c..8532979076 100644 --- a/veomni/models/checkpoint/manager.py +++ b/veomni/models/checkpoint/manager.py @@ -29,14 +29,10 @@ :class:`~veomni.trainer.callbacks.global_state_callback.GlobalStateCallback`. """ -import hashlib -import marshal -import math -import types from typing import TYPE_CHECKING, Any, Dict, Optional import torch.distributed as dist -from torch.optim.lr_scheduler import LambdaLR, StepLR +from torch.optim.lr_scheduler import LambdaLR from ...checkpoint import CheckpointerBase, build_checkpointer, layout from ...utils import helper @@ -45,6 +41,7 @@ def _scheduler_recipe(scheduler, optimizer): """Stable schedule identity; LambdaLR state_dict omits its callable.""" + from ...checkpoint.optimizer_resume_layout import _scheduler_functions from ...optim.lr_scheduler import MultiLRScheduler if scheduler is None: @@ -57,42 +54,9 @@ def _scheduler_recipe(scheduler, optimizer): } if getattr(optimizer, "_is_multi_optimizer", False) is True: raise RuntimeError("MultiOptimizer requires a scheduler for each optimizer leaf") - if type(scheduler) not in (LambdaLR, StepLR) or scheduler.optimizer is not optimizer: + if type(scheduler) is not LambdaLR or scheduler.optimizer is not optimizer: raise RuntimeError("MultiOptimizer scheduler implementation or binding is unsupported") - if type(scheduler) is StepLR: - return {"type": "StepLR", "step_size": scheduler.step_size, "gamma": scheduler.gamma} - - def immutable(value): - if value is None or type(value) in (bool, int, float, str): - return value - if isinstance(value, tuple): - return [immutable(item) for item in value] - if isinstance(value, dict) and all(isinstance(key, str) for key in value): - return {key: immutable(item) for key, item in sorted(value.items())} - raise RuntimeError("MultiOptimizer scheduler captures unsupported mutable state") - - def normalize(code): - constants = tuple(normalize(item) if isinstance(item, types.CodeType) else item for item in code.co_consts) - return code.replace(co_filename="", co_firstlineno=0, co_consts=constants) - - functions = [] - for function in scheduler.lr_lambdas: - if not isinstance(function, types.FunctionType): - raise RuntimeError("MultiOptimizer scheduler callable is unsupported") - if any( - name in function.__globals__ and (name != "math" or function.__globals__[name] is not math) - for name in function.__code__.co_names - ): - raise RuntimeError("MultiOptimizer scheduler depends on unsupported global state") - functions.append( - { - "code": hashlib.sha256(marshal.dumps(normalize(function.__code__))).hexdigest(), - "defaults": immutable(function.__defaults__), - "kwdefaults": immutable(function.__kwdefaults__), - "closure": [immutable(cell.cell_contents) for cell in function.__closure__ or ()], - } - ) - return {"type": "LambdaLR", "functions": functions} + return {"type": "LambdaLR", "functions": _scheduler_functions(scheduler)} def _validate_scheduler_checkpoint(saved, current, global_step): From 230f18cd07f3c6ae894fb2fe9986fe46ce849239 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:39:36 +0800 Subject: [PATCH 083/120] [checkpoint] Coordinate EP optimizer type preflight --- .../test_extra_parallel_optimizer_keys.py | 22 +++++++++++++------ veomni/checkpoint/dcp_checkpointer.py | 11 +++++++--- 2 files changed, 23 insertions(+), 10 deletions(-) diff --git a/tests/checkpoints/test_extra_parallel_optimizer_keys.py b/tests/checkpoints/test_extra_parallel_optimizer_keys.py index 33ce9aa7ff..93245b245b 100644 --- a/tests/checkpoints/test_extra_parallel_optimizer_keys.py +++ b/tests/checkpoints/test_extra_parallel_optimizer_keys.py @@ -178,9 +178,13 @@ def __getitem__(self, _): return SimpleNamespace(size=lambda: 2) class FakeOptimizer: - _is_multi_optimizer = not (case == "wrong_type" and rank == 0) + _is_multi_optimizer = not (case.startswith("wrong_type") and rank == 0) + called = False def state_dict(self): + self.called = True + if case == "wrong_type_collective": + dist.barrier() return {"state.experts.weight.exp_avg": torch.ones(2, 3)} if rank == 0 else {} dist.init_process_group( @@ -191,15 +195,19 @@ def state_dict(self): spec = SimpleNamespace(placement=Shard(0), para_fsdp_mesh=FakeMesh()) model = SimpleNamespace(_fqn2spec_info={"experts.weight": spec}) parallel = SimpleNamespace(dp_mode="fsdp2") - optimizer = dcp_checkpointer.OptimizerState(model, FakeOptimizer(), parallel_state=parallel) + leaf = FakeOptimizer() + optimizer = dcp_checkpointer.OptimizerState(model, leaf, parallel_state=parallel) checkpoint = Path(root) / "rejected" - if rank == 0 and case == "wrong_type": - exception, message = AssertionError, "not a MultiOptimizer" + if rank == 0 and case.startswith("wrong_type"): + message = "not a MultiOptimizer" + elif case.startswith("wrong_type"): + message = "validating ExtraParallel optimizer type" else: - exception = RuntimeError message = "no proven replica layout" if rank == 0 else "preparing ExtraParallel optimizer checkpoint state" - with pytest.raises(exception, match=message): + with pytest.raises(RuntimeError, match=message): dcp.save({"optimizer": optimizer}, checkpoint_id=checkpoint) + if case.startswith("wrong_type"): + assert not leaf.called assert not (checkpoint / ".metadata").exists() value = torch.ones(()) dist.all_reduce(value) @@ -209,6 +217,6 @@ def state_dict(self): @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -@pytest.mark.parametrize("case", ["plain_state", "wrong_type"]) +@pytest.mark.parametrize("case", ["plain_state", "wrong_type", "wrong_type_collective"]) def test_asymmetric_optimizer_preflight_rejects_every_rank(tmp_path, case): torch.multiprocessing.spawn(_asymmetric_optimizer_preflight, args=(str(tmp_path), case), nprocs=2, join=True) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index e0b39e93a8..bf72419202 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -434,11 +434,16 @@ def state_dict(self): logger.info_rank0( "Getting optimizer state_dict from OptimizerState wrapper, would restore ExtraParallel dim for Experts module" ) + # Validate locally before any rank enters state_dict(), which can + # contain collectives. A single rank's type error must not race + # another rank's optimizer collective against our error reduction. + error = None + if not getattr(self.optimizer, "_is_multi_optimizer", False): + error = RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") + raise_if_any_rank_failed(error, "validating ExtraParallel optimizer type", group=self.process_group) + optim_sd_with_extra_parallel_dim, error = None, None try: - assert self.optimizer._is_multi_optimizer, ( - "ExtraParallel is enabled but optimizer is not a MultiOptimizer instance" - ) vanilla_optim_sd = self.optimizer.state_dict() optim_sd_with_extra_parallel_dim = self.get_state_dict_with_extra_parallel_dim_preprocess( vanilla_optim_sd, "restore" From ed055f74ab0b0a35789c4789c4323734b1c7dda9 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:44:54 +0800 Subject: [PATCH 084/120] [checkpoint] Coordinate optimizer type lookup errors --- .../test_extra_parallel_optimizer_keys.py | 19 +++++++++++++------ veomni/checkpoint/dcp_checkpointer.py | 7 +++++-- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/tests/checkpoints/test_extra_parallel_optimizer_keys.py b/tests/checkpoints/test_extra_parallel_optimizer_keys.py index 93245b245b..ec911ae75f 100644 --- a/tests/checkpoints/test_extra_parallel_optimizer_keys.py +++ b/tests/checkpoints/test_extra_parallel_optimizer_keys.py @@ -178,12 +178,17 @@ def __getitem__(self, _): return SimpleNamespace(size=lambda: 2) class FakeOptimizer: - _is_multi_optimizer = not (case.startswith("wrong_type") and rank == 0) called = False + @property + def _is_multi_optimizer(self): + if case == "type_property_error" and rank == 0: + raise ValueError("injected optimizer type lookup error") + return not (case.startswith("wrong_type") and rank == 0) + def state_dict(self): self.called = True - if case == "wrong_type_collective": + if case in ("wrong_type_collective", "type_property_error"): dist.barrier() return {"state.experts.weight.exp_avg": torch.ones(2, 3)} if rank == 0 else {} @@ -200,13 +205,15 @@ def state_dict(self): checkpoint = Path(root) / "rejected" if rank == 0 and case.startswith("wrong_type"): message = "not a MultiOptimizer" - elif case.startswith("wrong_type"): + elif rank == 0 and case == "type_property_error": + message = "optimizer type lookup error" + elif case.startswith("wrong_type") or case == "type_property_error": message = "validating ExtraParallel optimizer type" else: message = "no proven replica layout" if rank == 0 else "preparing ExtraParallel optimizer checkpoint state" - with pytest.raises(RuntimeError, match=message): + with pytest.raises((RuntimeError, ValueError), match=message): dcp.save({"optimizer": optimizer}, checkpoint_id=checkpoint) - if case.startswith("wrong_type"): + if case.startswith("wrong_type") or case == "type_property_error": assert not leaf.called assert not (checkpoint / ".metadata").exists() value = torch.ones(()) @@ -217,6 +224,6 @@ def state_dict(self): @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -@pytest.mark.parametrize("case", ["plain_state", "wrong_type", "wrong_type_collective"]) +@pytest.mark.parametrize("case", ["plain_state", "wrong_type", "wrong_type_collective", "type_property_error"]) def test_asymmetric_optimizer_preflight_rejects_every_rank(tmp_path, case): torch.multiprocessing.spawn(_asymmetric_optimizer_preflight, args=(str(tmp_path), case), nprocs=2, join=True) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index bf72419202..27d5f7b13b 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -438,8 +438,11 @@ def state_dict(self): # contain collectives. A single rank's type error must not race # another rank's optimizer collective against our error reduction. error = None - if not getattr(self.optimizer, "_is_multi_optimizer", False): - error = RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") + try: + if not getattr(self.optimizer, "_is_multi_optimizer", False): + raise RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") + except BaseException as exc: + error = exc raise_if_any_rank_failed(error, "validating ExtraParallel optimizer type", group=self.process_group) optim_sd_with_extra_parallel_dim, error = None, None From 3f7002e3c86c13eb60b04616e103443030b2ebec Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:45:33 +0800 Subject: [PATCH 085/120] [checkpoint] Coordinate EP preflight and align layouts --- tests/checkpoints/test_expert_shard_layout.py | 28 +- veomni/checkpoint/dcp_checkpointer.py | 392 +++++++++--------- 2 files changed, 217 insertions(+), 203 deletions(-) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 7e667a5748..ed2eec2e74 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -836,7 +836,7 @@ def test_plain_optimizer_shards_are_not_treated_as_replicas(world, hsdp): spawn_with_rendezvous(_plain_optimizer_unproven_worker, args=(hsdp,), nprocs=world) -def _sparse_plain_optimizer_failure_worker(rank, endpoint): +def _sparse_plain_optimizer_failure_worker(rank, case, endpoint): from veomni.checkpoint.dcp_checkpointer import OptimizerState from veomni.distributed.parallel_plan import SpecInfo @@ -850,14 +850,29 @@ def _sparse_plain_optimizer_failure_worker(rank, endpoint): model._fqn2spec_info = {"experts.weight": SpecInfo("ep", Shard(0), "experts.weight", mesh, 1)} class SparseOptimizer: - _is_multi_optimizer = True + called = False + + @property + def _is_multi_optimizer(self): + if case == "type_property_error" and rank == 0: + raise ValueError("injected optimizer type lookup error") + return not (case.startswith("wrong_type") and rank == 0) def state_dict(self): + self.called = True + if case in ("wrong_type_collective", "type_property_error"): + dist.barrier() return {"state.experts.weight.exp_avg": torch.ones(4, 4)} if rank == 0 else {} - state = OptimizerState(model, SparseOptimizer(), parallel_state=parallel) - with pytest.raises(RuntimeError, match="no proven replica layout|failed on another rank"): + leaf = SparseOptimizer() + state = OptimizerState(model, leaf, parallel_state=parallel) + with pytest.raises( + (RuntimeError, ValueError), + match="no proven replica layout|not a MultiOptimizer|failed on another rank|optimizer type lookup error", + ): state.state_dict() + if case.startswith("wrong_type") or case == "type_property_error": + assert not leaf.called completed = torch.tensor(1) dist.all_reduce(completed) assert completed.item() == 4 @@ -867,8 +882,9 @@ def state_dict(self): @pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) @pytest.mark.parametrize("world", [4]) -def test_sparse_plain_optimizer_failure_exits_all_ranks(world): - spawn_with_rendezvous(_sparse_plain_optimizer_failure_worker, args=(), nprocs=world) +@pytest.mark.parametrize("case", ["plain_state", "wrong_type", "wrong_type_collective", "type_property_error"]) +def test_sparse_plain_optimizer_failure_exits_all_ranks(world, case): + spawn_with_rendezvous(_sparse_plain_optimizer_failure_worker, args=(case,), nprocs=world) def _subgroup_sparse_optimizer_failure_worker(rank, endpoint): diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 5e6d6dd86e..ec774658e9 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -946,6 +946,179 @@ def _dropped_extra_parallel_shape(loaded_tensor, ep_shard_dim, fsdp_shard_dim, e return tuple(shape), tuple(stride) +def drop_extra_parallel_dim_v2( + loaded_tensor: torch.Tensor, + device_mesh: DeviceMesh, + fsdp_shard_dim: int = 1, + *, + extra_parallel_name: Optional[str] = None, +): + """ + Drop ExtraParallel dims after loading from DCP so that ExtraParallel-FSDP would not be confused. + + ``device_mesh`` is the FSDP sub-mesh (ExtraParallel dim already excluded): + 1D ``(ep_fsdp,)`` for plain FSDP, or 2D ``(ep_replicate, ep_fsdp)`` for HSDP. + ``fsdp_shard_dim`` is the tensor dim FSDP shards along (1 by default, 0 for + the Muon zero-comm layout). The number of placements on the loaded DTensor + reflects the full saved mesh: + + * 1 placement: pure ExtraParallel, no FSDP -> return the plain local tensor. + * 2 placements: ``(ep_fsdp, ep)`` -> keep FSDP ``Shard(fsdp_shard_dim)`` on + the 1D sub-mesh. + * 3 placements: ``(ep_replicate, ep_fsdp, ep)`` (HSDP) -> keep + ``[Replicate(), Shard(fsdp_shard_dim)]`` on the 2D sub-mesh. + """ + + if not -loaded_tensor.ndim <= fsdp_shard_dim < loaded_tensor.ndim: + raise RuntimeError("FSDP shard dimension is outside the tensor rank") + fsdp_shard_dim %= loaded_tensor.ndim + num_placements = len(loaded_tensor.placements) + if num_placements == 1: + tensor_to_put = loaded_tensor.to_local() + elif num_placements == 2: + ep_placement = loaded_tensor.placements[-1] + if not isinstance(ep_placement, Shard): + raise RuntimeError(f"ExtraParallel placement must be Shard, got {ep_placement}") + ep_shard_dim = ep_placement.dim + shape, stride = _dropped_extra_parallel_shape(loaded_tensor, ep_shard_dim, fsdp_shard_dim, extra_parallel_name) + tensor_to_put = _from_local_preserve_device( + loaded_tensor._local_tensor, + device_mesh=device_mesh, + placements=[Shard(fsdp_shard_dim)], + shape=shape, + stride=stride, + ) + elif num_placements == 3: + ep_placement = loaded_tensor.placements[-1] + if not isinstance(ep_placement, Shard): + raise RuntimeError(f"ExtraParallel placement must be Shard, got {ep_placement}") + ep_shard_dim = ep_placement.dim + shape, stride = _dropped_extra_parallel_shape(loaded_tensor, ep_shard_dim, fsdp_shard_dim, extra_parallel_name) + tensor_to_put = _from_local_preserve_device( + loaded_tensor._local_tensor, + device_mesh=device_mesh, + placements=[Replicate(), Shard(fsdp_shard_dim)], + shape=shape, + stride=stride, + ) + else: + raise RuntimeError( + "Expect ExtraParallel parameters from checkpoints to be DTensor with 1-dim (no FSDP), " + f"2-dim (ExtraParallel+FSDP) or 3-dim (ExtraParallel+HSDP), got {loaded_tensor}" + ) + + return tensor_to_put + + +def _validate_extra_parallel_tensor_v2(tensor, mesh, fsdp_mesh, ep_dim, fsdp_dim, global_shape): + if not isinstance(tensor, DTensor): + raise RuntimeError("The v2 ExtraParallel checkpoint layout requires an FSDP DTensor") + ndim = tensor.ndim + if not (-ndim <= ep_dim < ndim and -ndim <= fsdp_dim < ndim): + raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") + ep_dim, fsdp_dim = ep_dim % ndim, fsdp_dim % ndim + if tensor.device_mesh != fsdp_mesh: + raise RuntimeError("ExtraParallel source tensor does not use the declared FSDP mesh") + if tuple(tensor.placements) != (Replicate(),) * (fsdp_mesh.ndim - 1) + (Shard(fsdp_dim),): + raise RuntimeError("ExtraParallel source tensor does not use the declared FSDP placement") + if global_shape is None or len(global_shape) != ndim: + raise RuntimeError("FSDP-only ExtraParallel tensors require the original global_shape") + + expected = list(global_shape) + ep_size, ep_rank = mesh.size(-1), mesh.get_local_rank(-1) + ep_chunk = (expected[ep_dim] + ep_size - 1) // ep_size + expected[ep_dim] = max(min(expected[ep_dim] - ep_rank * ep_chunk, ep_chunk), 0) + if tuple(tensor.shape) != tuple(expected): + raise RuntimeError("ExtraParallel source tensor shape does not match the declared global shape") + local_size, local_rank = fsdp_mesh.size(-1), fsdp_mesh.get_local_rank(-1) + local_chunk = (expected[fsdp_dim] + local_size - 1) // local_size + expected[fsdp_dim] = max(min(expected[fsdp_dim] - local_rank * local_chunk, local_chunk), 0) + if tuple(tensor.to_local().shape) != tuple(expected): + raise RuntimeError("ExtraParallel local tensor shape does not match its declared FSDP placement") + + +def restore_extra_parallel_dim_v2( + orgin_tensor: torch.Tensor, + fsdp_mesh: DeviceMesh, + extra_parallel_fsdp_mesh: DeviceMesh, + ep_shard_dim: int = 0, + fsdp_shard_dim: int = 1, + global_shape: Optional[tuple[int, ...]] = None, +): + """ + Restore ExtraParallel dim so that DCP can be aware about ExtraParallel ranks + + The ExtraParallel (e.g. ``ep``) dim and the FSDP dim shard the tensor along + DIFFERENT axes: ``ep_shard_dim`` (the EP plan's ``Shard.dim``, dim 0 for + experts) and ``fsdp_shard_dim`` (1 by default, 0 for the Muon zero-comm + layout). They must be mapped to the matching mesh dims explicitly. + + args: + orgin_tensor (torch.Tensor): The orgin tensor (FSDP-local shard). + fsdp_mesh (DeviceMesh): The full ExtraParallel mesh, i.e. 2D + ``(ep_fsdp, ep)`` for plain FSDP or 3D + ``(ep_replicate, ep_fsdp, ep)`` for HSDP. + extra_parallel_fsdp_mesh (DeviceMesh): The FSDP sub-mesh (ExtraParallel + dim excluded), used for the pure-ExtraParallel (no FSDP) path. + ep_shard_dim (int): Tensor dim the ExtraParallel dim shards along. + fsdp_shard_dim (int): Tensor dim FSDP shards along. + global_shape (tuple[int, ...], optional): The model parameter's shape + before ExtraParallel and FSDP sharding. It is required for an + FSDP-only DTensor: a local slice alone cannot prove global shape. + + Note: + When both dimensions shard the same tensor axis, the FSDP placement + uses ``_StridedShard`` on the FSDP-first runtime mesh. If the global + axis is uneven across EP, use an EP-first mesh view instead: current + Torch strided metadata can report incorrect offsets for that case. + """ + assert fsdp_mesh.ndim in (2, 3), f"global_mesh.ndim must be 2 or 3, got {fsdp_mesh.ndim}" + ndim = orgin_tensor.ndim + if not (-ndim <= ep_shard_dim < ndim and -ndim <= fsdp_shard_dim < ndim): + raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") + ep_shard_dim %= ndim + fsdp_shard_dim %= ndim + + if isinstance(orgin_tensor, DTensor): + _validate_extra_parallel_tensor_v2( + orgin_tensor, fsdp_mesh, extra_parallel_fsdp_mesh, ep_shard_dim, fsdp_shard_dim, global_shape + ) + # ExtraParallel+FSDP2. mesh order is (ep_fsdp, ep) or, for HSDP, + # (ep_replicate, ep_fsdp, ep): ep_fsdp -> Shard(fsdp_shard_dim), + # ep -> Shard(ep_shard_dim), ep_replicate -> Replicate(). + same_axis = ep_shard_dim == fsdp_shard_dim and fsdp_mesh.size(-1) > 1 + if same_axis and global_shape[ep_shard_dim] % fsdp_mesh.size(-1): + checkpoint_mesh = _extra_parallel_first_mesh(fsdp_mesh) + placements = [Shard(ep_shard_dim), Shard(fsdp_shard_dim)] + if checkpoint_mesh.ndim == 3: + placements.insert(0, Replicate()) + else: + checkpoint_mesh = fsdp_mesh + fsdp_placement = ( + _StridedShard(fsdp_shard_dim, split_factor=fsdp_mesh.size(-1)) if same_axis else Shard(fsdp_shard_dim) + ) + placements = [fsdp_placement, Shard(ep_shard_dim)] + if checkpoint_mesh.ndim == 3: + placements.insert(0, Replicate()) + shape = tuple(global_shape) + stride = [1] * len(shape) + for index in range(len(shape) - 2, -1, -1): + stride[index] = stride[index + 1] * shape[index + 1] + dtensor = _from_local_preserve_device( + orgin_tensor._local_tensor, + device_mesh=checkpoint_mesh, + placements=placements, + shape=shape, + stride=tuple(stride), + ) + elif torch.is_tensor(orgin_tensor): + raise RuntimeError("The v2 ExtraParallel checkpoint layout requires an FSDP DTensor") + else: + raise RuntimeError(f"origin_tensor - {orgin_tensor} is not a tensor!") + + return dtensor + + def _apply_extra_parallel_dim( state_dict: Dict[str, Any], extra_parallel_fqn2spec_info: Dict[str, Any], @@ -1216,12 +1389,20 @@ def state_dict(self): logger.info_rank0( "Getting optimizer state_dict from OptimizerState wrapper, would restore ExtraParallel dim for Experts module" ) - assert self.optimizer._is_multi_optimizer, ( - "ExtraParallel is enabled but optimizer is not a MultiOptimizer instance" - ) - vanilla_optim_sd = self.optimizer.state_dict() + # Validate locally before any rank enters state_dict(), which can + # contain collectives. A single rank's type error must not race + # another rank's optimizer collective against our error reduction. + error = None + try: + if not getattr(self.optimizer, "_is_multi_optimizer", False): + raise RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "validating ExtraParallel optimizer type", group=self.process_group) + optim_sd_with_extra_parallel_dim, error = None, None try: + vanilla_optim_sd = self.optimizer.state_dict() optim_sd_with_extra_parallel_dim = self.get_state_dict_with_extra_parallel_dim_preprocess( vanilla_optim_sd, "restore" ) @@ -1314,24 +1495,21 @@ def drop_extra_parallel_dim(loaded_tensor: torch.Tensor, device_mesh: DeviceMesh """ num_placements = len(loaded_tensor.placements) - if num_placements in (2, 3) and any(not isinstance(p, Replicate) for p in loaded_tensor.placements[:-1]): + has_fsdp_shard = any(not isinstance(p, Replicate) for p in loaded_tensor.placements[:-1]) + if num_placements in (2, 3) and has_fsdp_shard: ndim = loaded_tensor.ndim if not -ndim <= fsdp_shard_dim < ndim: raise RuntimeError("FSDP shard dimension is outside the tensor rank") fsdp_shard_dim %= ndim - if num_placements == 1 or all(isinstance(p, Replicate) for p in loaded_tensor.placements[:-1]): + if num_placements == 1 or (num_placements in (2, 3) and not has_fsdp_shard): tensor_to_put = loaded_tensor.to_local() elif num_placements == 2: tensor_to_put = _from_local_preserve_device( - loaded_tensor._local_tensor, - device_mesh=device_mesh, - placements=[Shard(fsdp_shard_dim)], + loaded_tensor._local_tensor, device_mesh=device_mesh, placements=[Shard(fsdp_shard_dim)] ) elif num_placements == 3: tensor_to_put = _from_local_preserve_device( - loaded_tensor._local_tensor, - device_mesh=device_mesh, - placements=[Replicate(), Shard(fsdp_shard_dim)], + loaded_tensor._local_tensor, device_mesh=device_mesh, placements=[Replicate(), Shard(fsdp_shard_dim)] ) else: raise RuntimeError( @@ -1378,10 +1556,8 @@ def restore_extra_parallel_dim( if isinstance(orgin_tensor, DTensor): ndim = orgin_tensor.ndim - if not -ndim <= ep_shard_dim < ndim: + if not (-ndim <= ep_shard_dim < ndim and -ndim <= fsdp_shard_dim < ndim): raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") - if not -ndim <= fsdp_shard_dim < ndim: - raise RuntimeError("FSDP shard dimension is outside the tensor rank") ep_shard_dim %= ndim fsdp_shard_dim %= ndim # ExtraParallel+FSDP2. mesh order is (ep_fsdp, ep) or, for HSDP, @@ -1396,16 +1572,11 @@ def restore_extra_parallel_dim( placements = [Replicate(), fsdp_placement, Shard(ep_shard_dim)] else: placements = [fsdp_placement, Shard(ep_shard_dim)] - dtensor = _from_local_preserve_device( - orgin_tensor._local_tensor, - device_mesh=fsdp_mesh, - placements=placements, - ) + dtensor = _from_local_preserve_device(orgin_tensor._local_tensor, device_mesh=fsdp_mesh, placements=placements) elif torch.is_tensor(orgin_tensor): - # A plain model tensor has no FSDP shard. The caller must ensure every - # FSDP/HSDP replica for the same ExtraParallel coordinate has identical - # values; from_local does not check or synchronize them. Optimizer - # states without this guarantee are rejected by the caller. + # Plain model tensors must be replicated across FSDP coordinates for + # the same ExtraParallel rank. Optimizer shards without that proof are + # rejected by the caller. ndim = orgin_tensor.ndim if not -ndim <= ep_shard_dim < ndim: raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") @@ -1419,179 +1590,6 @@ def restore_extra_parallel_dim( return dtensor -def drop_extra_parallel_dim_v2( - loaded_tensor: torch.Tensor, - device_mesh: DeviceMesh, - fsdp_shard_dim: int = 1, - *, - extra_parallel_name: Optional[str] = None, -): - """ - Drop ExtraParallel dims after loading from DCP so that ExtraParallel-FSDP would not be confused. - - ``device_mesh`` is the FSDP sub-mesh (ExtraParallel dim already excluded): - 1D ``(ep_fsdp,)`` for plain FSDP, or 2D ``(ep_replicate, ep_fsdp)`` for HSDP. - ``fsdp_shard_dim`` is the tensor dim FSDP shards along (1 by default, 0 for - the Muon zero-comm layout). The number of placements on the loaded DTensor - reflects the full saved mesh: - - * 1 placement: pure ExtraParallel, no FSDP -> return the plain local tensor. - * 2 placements: ``(ep_fsdp, ep)`` -> keep FSDP ``Shard(fsdp_shard_dim)`` on - the 1D sub-mesh. - * 3 placements: ``(ep_replicate, ep_fsdp, ep)`` (HSDP) -> keep - ``[Replicate(), Shard(fsdp_shard_dim)]`` on the 2D sub-mesh. - """ - - if not -loaded_tensor.ndim <= fsdp_shard_dim < loaded_tensor.ndim: - raise RuntimeError("FSDP shard dimension is outside the tensor rank") - fsdp_shard_dim %= loaded_tensor.ndim - num_placements = len(loaded_tensor.placements) - if num_placements == 1: - tensor_to_put = loaded_tensor.to_local() - elif num_placements == 2: - ep_placement = loaded_tensor.placements[-1] - if not isinstance(ep_placement, Shard): - raise RuntimeError(f"ExtraParallel placement must be Shard, got {ep_placement}") - ep_shard_dim = ep_placement.dim - shape, stride = _dropped_extra_parallel_shape(loaded_tensor, ep_shard_dim, fsdp_shard_dim, extra_parallel_name) - tensor_to_put = _from_local_preserve_device( - loaded_tensor._local_tensor, - device_mesh=device_mesh, - placements=[Shard(fsdp_shard_dim)], - shape=shape, - stride=stride, - ) - elif num_placements == 3: - ep_placement = loaded_tensor.placements[-1] - if not isinstance(ep_placement, Shard): - raise RuntimeError(f"ExtraParallel placement must be Shard, got {ep_placement}") - ep_shard_dim = ep_placement.dim - shape, stride = _dropped_extra_parallel_shape(loaded_tensor, ep_shard_dim, fsdp_shard_dim, extra_parallel_name) - tensor_to_put = _from_local_preserve_device( - loaded_tensor._local_tensor, - device_mesh=device_mesh, - placements=[Replicate(), Shard(fsdp_shard_dim)], - shape=shape, - stride=stride, - ) - else: - raise RuntimeError( - "Expect ExtraParallel parameters from checkpoints to be DTensor with 1-dim (no FSDP), " - f"2-dim (ExtraParallel+FSDP) or 3-dim (ExtraParallel+HSDP), got {loaded_tensor}" - ) - - return tensor_to_put - - -def _validate_extra_parallel_tensor_v2(tensor, mesh, fsdp_mesh, ep_dim, fsdp_dim, global_shape): - if not isinstance(tensor, DTensor): - raise RuntimeError("The v2 ExtraParallel checkpoint layout requires an FSDP DTensor") - ndim = tensor.ndim - if not (-ndim <= ep_dim < ndim and -ndim <= fsdp_dim < ndim): - raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") - ep_dim, fsdp_dim = ep_dim % ndim, fsdp_dim % ndim - if tensor.device_mesh != fsdp_mesh: - raise RuntimeError("ExtraParallel source tensor does not use the declared FSDP mesh") - if tuple(tensor.placements) != (Replicate(),) * (fsdp_mesh.ndim - 1) + (Shard(fsdp_dim),): - raise RuntimeError("ExtraParallel source tensor does not use the declared FSDP placement") - if global_shape is None or len(global_shape) != ndim: - raise RuntimeError("FSDP-only ExtraParallel tensors require the original global_shape") - - expected = list(global_shape) - ep_size, ep_rank = mesh.size(-1), mesh.get_local_rank(-1) - ep_chunk = (expected[ep_dim] + ep_size - 1) // ep_size - expected[ep_dim] = max(min(expected[ep_dim] - ep_rank * ep_chunk, ep_chunk), 0) - if tuple(tensor.shape) != tuple(expected): - raise RuntimeError("ExtraParallel source tensor shape does not match the declared global shape") - local_size, local_rank = fsdp_mesh.size(-1), fsdp_mesh.get_local_rank(-1) - local_chunk = (expected[fsdp_dim] + local_size - 1) // local_size - expected[fsdp_dim] = max(min(expected[fsdp_dim] - local_rank * local_chunk, local_chunk), 0) - if tuple(tensor.to_local().shape) != tuple(expected): - raise RuntimeError("ExtraParallel local tensor shape does not match its declared FSDP placement") - - -def restore_extra_parallel_dim_v2( - orgin_tensor: torch.Tensor, - fsdp_mesh: DeviceMesh, - extra_parallel_fsdp_mesh: DeviceMesh, - ep_shard_dim: int = 0, - fsdp_shard_dim: int = 1, - global_shape: Optional[tuple[int, ...]] = None, -): - """ - Restore ExtraParallel dim so that DCP can be aware about ExtraParallel ranks - - The ExtraParallel (e.g. ``ep``) dim and the FSDP dim shard the tensor along - DIFFERENT axes: ``ep_shard_dim`` (the EP plan's ``Shard.dim``, dim 0 for - experts) and ``fsdp_shard_dim`` (1 by default, 0 for the Muon zero-comm - layout). They must be mapped to the matching mesh dims explicitly. - - args: - orgin_tensor (torch.Tensor): The orgin tensor (FSDP-local shard). - fsdp_mesh (DeviceMesh): The full ExtraParallel mesh, i.e. 2D - ``(ep_fsdp, ep)`` for plain FSDP or 3D - ``(ep_replicate, ep_fsdp, ep)`` for HSDP. - extra_parallel_fsdp_mesh (DeviceMesh): The FSDP sub-mesh (ExtraParallel - dim excluded), used for the pure-ExtraParallel (no FSDP) path. - ep_shard_dim (int): Tensor dim the ExtraParallel dim shards along. - fsdp_shard_dim (int): Tensor dim FSDP shards along. - global_shape (tuple[int, ...], optional): The model parameter's shape - before ExtraParallel and FSDP sharding. It is required for an - FSDP-only DTensor: a local slice alone cannot prove global shape. - - Note: - When both dimensions shard the same tensor axis, the FSDP placement - uses ``_StridedShard`` on the FSDP-first runtime mesh. If the global - axis is uneven across EP, use an EP-first mesh view instead: current - Torch strided metadata can report incorrect offsets for that case. - """ - assert fsdp_mesh.ndim in (2, 3), f"global_mesh.ndim must be 2 or 3, got {fsdp_mesh.ndim}" - ndim = orgin_tensor.ndim - if not (-ndim <= ep_shard_dim < ndim and -ndim <= fsdp_shard_dim < ndim): - raise RuntimeError("ExtraParallel shard dimension is outside the tensor rank") - ep_shard_dim %= ndim - fsdp_shard_dim %= ndim - - if isinstance(orgin_tensor, DTensor): - _validate_extra_parallel_tensor_v2( - orgin_tensor, fsdp_mesh, extra_parallel_fsdp_mesh, ep_shard_dim, fsdp_shard_dim, global_shape - ) - # ExtraParallel+FSDP2. mesh order is (ep_fsdp, ep) or, for HSDP, - # (ep_replicate, ep_fsdp, ep): ep_fsdp -> Shard(fsdp_shard_dim), - # ep -> Shard(ep_shard_dim), ep_replicate -> Replicate(). - same_axis = ep_shard_dim == fsdp_shard_dim and fsdp_mesh.size(-1) > 1 - if same_axis and global_shape[ep_shard_dim] % fsdp_mesh.size(-1): - checkpoint_mesh = _extra_parallel_first_mesh(fsdp_mesh) - placements = [Shard(ep_shard_dim), Shard(fsdp_shard_dim)] - if checkpoint_mesh.ndim == 3: - placements.insert(0, Replicate()) - else: - checkpoint_mesh = fsdp_mesh - fsdp_placement = ( - _StridedShard(fsdp_shard_dim, split_factor=fsdp_mesh.size(-1)) if same_axis else Shard(fsdp_shard_dim) - ) - placements = [fsdp_placement, Shard(ep_shard_dim)] - if checkpoint_mesh.ndim == 3: - placements.insert(0, Replicate()) - shape = tuple(global_shape) - stride = [1] * len(shape) - for index in range(len(shape) - 2, -1, -1): - stride[index] = stride[index + 1] * shape[index + 1] - dtensor = _from_local_preserve_device( - orgin_tensor._local_tensor, - device_mesh=checkpoint_mesh, - placements=placements, - shape=shape, - stride=tuple(stride), - ) - elif torch.is_tensor(orgin_tensor): - raise RuntimeError("The v2 ExtraParallel checkpoint layout requires an FSDP DTensor") - else: - raise RuntimeError(f"origin_tensor - {orgin_tensor} is not a tensor!") - - return dtensor - - def _local_rank() -> int: """This process's rank within its node, as set by the elastic launcher. From cccc1dfc42c0b64fa4100689bff7ed1c4991bf91 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 11:54:17 +0800 Subject: [PATCH 086/120] [checkpoint] Validate nested scheduler identity --- docs/usage/checkpoint.md | 9 + .../test_multi_optimizer_proven_resume.py | 180 ++++++++++++++++++ veomni/checkpoint/optimizer_resume_layout.py | 81 ++++++-- 3 files changed, 258 insertions(+), 12 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 946029e020..f85f31020a 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -162,6 +162,15 @@ step and generation alone do not prove exact dataloader or RNG recovery. If a restored run has not established a same-step save receipt, an immediate export to the already completed step fails rather than overwriting that DCP. +MultiOptimizer checkpoint identity supports VeOmni's built-in constant, linear +and cosine schedules. Custom `LambdaLR` functions must use immutable captures +and the supported pure builtins or standard `math` members; global state, +dynamic lookup, mutable closures and unrecognized attribute access are rejected +at save and resume. Older custom scheduler checkpoints that cannot meet this +contract require a verified migration or weights-only restart. Identity checks +assume a trusted Python runtime whose `math` and builtins were not monkeypatched +before VeOmni was imported; they are not a code sandbox. + With `stage_dir` the model half waits: a staged save has overwritten nothing until the copy to the destination, so it invalidates the destination there instead of at the start. A staged save that fails leaves the previous checkpoint diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 3346834735..e59e9cd2fb 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -2005,6 +2005,186 @@ def test_scheduler_recipe_rejects_schedules_outside_the_checkpoint_contract(): _scheduler_recipe(scheduler, leaf) +@pytest.mark.parametrize("kind", ["function", "comprehension"]) +def test_scheduler_recipe_rejects_global_in_nested_function(monkeypatch, kind): + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + monkeypatch.setitem(globals(), "_TEST_NESTED_LR_FACTOR", 0.5) + + def schedule(step): + def nested(): + return _TEST_NESTED_LR_FACTOR # noqa: F821 + + return step * nested() + + def comprehension(step): + return sum(value * _TEST_NESTED_LR_FACTOR for value in range(step)) # noqa: F821 + + # Changing this global changes the next learning-rate update even though + # the outer code object never names it. Its identity is therefore unsafe. + with pytest.raises(RuntimeError, match="unsupported global state"): + _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule if kind == "function" else comprehension])) + + +def test_scheduler_recipe_accepts_nested_pure_function(): + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + def schedule(step): + def nested(value): + return value * 0.5 + + return nested(step) + + recipe = _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule])) + assert len(recipe) == 1 + assert schedule(2) == 1.0 + + +def test_scheduler_recipe_accepts_local_temporary_cell(): + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + def schedule(step): + value = step * 0.5 + return (lambda: value)() + + assert len(_scheduler_functions(SimpleNamespace(lr_lambdas=[schedule]))) == 1 + assert schedule(2) == 1.0 + + +def test_scheduler_recipe_accepts_math_attribute_with_unrelated_global(monkeypatch): + import math as math_module + + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + monkeypatch.setitem(globals(), "math", math_module) + monkeypatch.setitem(globals(), "cos", "unrelated") + + def schedule(step): + def nested(value): + return math.cos(value) # noqa: F821 + + return nested(step) + + assert len(_scheduler_functions(SimpleNamespace(lr_lambdas=[schedule]))) == 1 + + +def test_scheduler_recipe_rejects_changed_math_member(monkeypatch): + import math as math_module + + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + monkeypatch.setitem(globals(), "math", math_module) + + def schedule(step): + return math.cos(step) # noqa: F821 + + _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule])) + monkeypatch.setattr(math_module, "cos", lambda value: 0.25) + with pytest.raises(RuntimeError, match="unsupported.*state|unsupported.*math"): + _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule])) + + +def test_scheduler_recipe_rejects_tuple_subclass_closure(): + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + class StatefulTuple(tuple): + def __getitem__(self, index): + return 0.5 + + factor = StatefulTuple((1.0,)) + + def schedule(step): + return step * factor[0] + + with pytest.raises(RuntimeError, match="unsupported mutable state"): + _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule])) + + +def test_scheduler_recipe_rejects_nonlocal_mutation(): + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + state = 1.0 + + def schedule(step): + nonlocal state + state += 1.0 + return step * state + + with pytest.raises(RuntimeError, match="unsupported dynamic state"): + _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule])) + + +@pytest.mark.parametrize("builtin_name", ["eval", "globals", "getattr", "__import__"]) +def test_scheduler_recipe_rejects_dynamic_builtin_in_nested_function(builtin_name): + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + def schedule_eval(step): + return lambda: eval("step") # noqa: S307 + + def schedule_globals(step): + return lambda: globals() + + def schedule_getattr(step): + return lambda: getattr(step, str(step)) + + def schedule_import(step): + return lambda: __import__("math") + + schedules = { + "eval": schedule_eval, + "globals": schedule_globals, + "getattr": schedule_getattr, + "__import__": schedule_import, + } + with pytest.raises(RuntimeError, match="unsupported.*state|unsupported.*builtin"): + _scheduler_functions(SimpleNamespace(lr_lambdas=[schedules[builtin_name]])) + + +def test_nested_global_scheduler_is_rejected_before_checkpoint_load(tmp_path, monkeypatch): + model, optimizer = _model_and_optimizer() + scheduler = build_lr_scheduler(optimizer, train_steps=10) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": scheduler.state_dict()}, + "_veomni_scheduler_object": scheduler, + }, + parallel_state=parallel, + ) + + target, restored = _model_and_optimizer() + target_scheduler = build_lr_scheduler(restored, train_steps=10) + monkeypatch.setitem(globals(), "_TEST_NESTED_LR_FACTOR", 0.5) + + def schedule(step): + def nested(): + return _TEST_NESTED_LR_FACTOR # noqa: F821 + + return step * nested() + + target_scheduler["weight"].lr_lambdas[0] = schedule + before = {name: value.detach().clone() for name, value in target.named_parameters()} + scheduler_before = target_scheduler.state_dict() + with pytest.raises(RuntimeError, match="unsupported global state"): + DistributedCheckpointer.load( + str(tmp_path), + { + "model": target, + "optimizer": restored, + "extra_state": {}, + "_veomni_scheduler_object": target_scheduler, + }, + parallel_state=parallel, + ) + for name, value in target.named_parameters(): + torch.testing.assert_close(value, before[name], rtol=0, atol=0) + assert target_scheduler.state_dict() == scheduler_before + assert all(not leaf.state for leaf in restored.optimizers_dict.values()) + + def test_multi_checkpoint_accepts_explicit_no_scheduler(tmp_path): model, optimizer = _model_and_optimizer() model.weight.grad = torch.ones_like(model.weight) diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index be910af399..7ca4a394fb 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -3,9 +3,12 @@ """Identity of a MultiOptimizer's rank-local parameter and scheduler layout.""" +import builtins +import dis import hashlib import json import marshal +import math import types import torch @@ -17,11 +20,21 @@ MULTI_LAYOUT_MARKER = "__veomni_multi_optimizer_layout_v1__" +_SCHEDULER_SAFE_BUILTINS = { + name: getattr(builtins, name) + for name in ("abs", "bool", "float", "int", "len", "max", "min", "pow", "range", "round", "sum") +} +_SCHEDULER_MATH_MEMBERS = { + name: value + for name, value in vars(math).items() + if not name.startswith("_") and type(value) in (float, int, types.BuiltinFunctionType) +} + def _immutable_value(value): if value is None or type(value) in (bool, int, float, str): return value - if isinstance(value, tuple): + if type(value) is tuple: return [_immutable_value(item) for item in value] raise RuntimeError("MultiOptimizer scheduler captures unsupported mutable state") @@ -67,25 +80,69 @@ def _lr_layout(group, key, aliases): def _scheduler_functions(scheduler): """Describe LambdaLR callables using the existing checkpoint identity format.""" - import math + + def normalize(code, globals_dict, builtins_dict): + # Nested code objects share this function's global namespace. A + # referenced global can change the next LR without changing code. + # LOAD_ATTR names are not globals; inspect actual global lookups instead. + instructions = list(dis.get_instructions(code)) + for index, instruction in enumerate(instructions): + opcode = instruction.opname + if opcode in ("STORE_DEREF", "DELETE_DEREF") and instruction.argval in code.co_freevars: + raise RuntimeError("MultiOptimizer scheduler uses unsupported dynamic state") + if opcode in ( + "LOAD_NAME", + "STORE_NAME", + "DELETE_NAME", + "STORE_GLOBAL", + "DELETE_GLOBAL", + "STORE_ATTR", + "DELETE_ATTR", + "STORE_SUBSCR", + "DELETE_SUBSCR", + "IMPORT_NAME", + "IMPORT_FROM", + "IMPORT_STAR", + "LOAD_BUILD_CLASS", + ): + raise RuntimeError("MultiOptimizer scheduler uses unsupported dynamic state") + if opcode in ("LOAD_ATTR", "LOAD_METHOD"): + previous = instructions[index - 1] if index else None + if previous is None or previous.opname != "LOAD_GLOBAL" or previous.argval != "math": + raise RuntimeError("MultiOptimizer scheduler uses unsupported attribute access") + expected = _SCHEDULER_MATH_MEMBERS.get(instruction.argval) + actual = vars(math).get(instruction.argval) + if ( + expected is None + or type(expected) is not type(actual) + or (actual is not expected and (type(expected) not in (int, float) or actual != expected)) + ): + raise RuntimeError("MultiOptimizer scheduler uses unsupported math state") + continue + if opcode != "LOAD_GLOBAL": + continue + name = instruction.argval + if name in globals_dict: + if name != "math" or globals_dict[name] is not math: + raise RuntimeError("MultiOptimizer scheduler refers to unsupported global state") + elif name not in _SCHEDULER_SAFE_BUILTINS or builtins_dict.get(name) is not _SCHEDULER_SAFE_BUILTINS[name]: + raise RuntimeError("MultiOptimizer scheduler uses an unsupported builtin") + constants = tuple( + normalize(item, globals_dict, builtins_dict) if isinstance(item, types.CodeType) else item + for item in code.co_consts + ) + return code.replace(co_filename="", co_firstlineno=0, co_consts=constants) functions = [] for function in scheduler.lr_lambdas: if not isinstance(function, types.FunctionType): raise RuntimeError("MultiOptimizer scheduler uses an unsupported callable") - if any( - name in function.__globals__ and (name != "math" or function.__globals__[name] is not math) - for name in function.__code__.co_names - ): - raise RuntimeError("MultiOptimizer scheduler refers to unsupported global state") - - def normalize(code): - constants = tuple(normalize(item) if isinstance(item, types.CodeType) else item for item in code.co_consts) - return code.replace(co_filename="", co_firstlineno=0, co_consts=constants) functions.append( { - "code": hashlib.sha256(marshal.dumps(normalize(function.__code__))).hexdigest(), + "code": hashlib.sha256( + marshal.dumps(normalize(function.__code__, function.__globals__, function.__builtins__)) + ).hexdigest(), "defaults": _immutable_value(function.__defaults__), "kwdefaults": _immutable_value(function.__kwdefaults__), "closure": [_immutable_value(cell.cell_contents) for cell in function.__closure__ or ()], From 6384c62b1f4b52f243f8247e5baaaf7a2fb6e077 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 5 Nov 2026 23:28:40 +0800 Subject: [PATCH 087/120] [checkpoint] fix standalone SwapAdamW recovery guards --- docs/usage/checkpoint.md | 6 +- .../test_swap_proven_multi_resume.py | 156 +++++++++++ veomni/checkpoint/dcp_checkpointer.py | 248 ++++++++++++------ veomni/optim/swap_adamw.py | 16 +- 4 files changed, 337 insertions(+), 89 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 513bfa4995..6e0b7f61bf 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -482,8 +482,10 @@ dtype and does not provide FP32 master weights; it is not equivalent to an FP32-master optimizer. Its float32 checkpoint step is accepted only while the integer step is exactly representable (at most 2**24). For a standalone SwapAdamW, use `DistributedCheckpointer` when parameters may -be frozen during recovery: it restores their saved optimizer history before -they are unfrozen. If a frozen parameter has saved history, custom optimizer +be frozen during recovery: it restores saved optimizer history for frozen +parameters still owned by the optimizer before they are unfrozen. Source and +target parameter groups must match even for an owner with no saved moments; a mismatch +requires the original groups or a verified migration. If a frozen parameter has saved history, custom optimizer load-state hooks and trainable-only model checkpoints are rejected; model hooks see the original frozen flags. Raw DCP `OptimizerState` does not inspect checkpoint owner metadata and can silently omit that history; it retains its original behavior. diff --git a/tests/checkpoints/test_swap_proven_multi_resume.py b/tests/checkpoints/test_swap_proven_multi_resume.py index 4036601273..a6aac299a6 100644 --- a/tests/checkpoints/test_swap_proven_multi_resume.py +++ b/tests/checkpoints/test_swap_proven_multi_resume.py @@ -406,6 +406,52 @@ def test_swap_rejects_inexact_checkpoint_step_before_saving(tmp_path): assert not list(tmp_path.iterdir()) +def test_single_swap_state_dict_rejects_inexact_checkpoint_step(tmp_path): + model = torch.nn.Linear(3, 2) + optimizer = SwapAdamW(model.parameters(), pin_memory=False) + param = next(model.parameters()) + optimizer.state[param]["step"] = 2**24 + assert optimizer.state_dict()["state"][0]["step"].item() == 2**24 + optimizer.state[param]["step"] = 2**24 + 1 + with pytest.raises(RuntimeError, match="cannot be represented exactly"): + optimizer.state_dict() + with pytest.raises(Exception, match="cannot be represented exactly"): + DistributedCheckpointer.save( + str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + + +def test_single_swap_failed_partial_restore_blocks_update(tmp_path, monkeypatch): + import veomni.checkpoint.dcp_checkpointer as implementation + + model = torch.nn.Linear(3, 2) + optimizer = SwapAdamW(model.parameters(), pin_memory=False) + _step(model, optimizer) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target = torch.nn.Linear(3, 2) + resumed = SwapAdamW(target.parameters(), pin_memory=False) + original_load = implementation.dcp.load + + def fail_after_partial_write(*args, **kwargs): + if "optimizer" in kwargs["state_dict"] and not isinstance(kwargs["state_dict"]["optimizer"], dict): + resumed._host_states[target.weight]["exp_avg"].fill_(0.25) + raise RuntimeError("injected partial optimizer read failure") + return original_load(*args, **kwargs) + + monkeypatch.setattr(implementation.dcp, "load", fail_after_partial_write) + with pytest.raises(RuntimeError, match="partial optimizer read failure"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": resumed}, parallel_state=parallel) + with pytest.raises(RuntimeError, match="failed checkpoint restore"): + resumed.step() + with pytest.raises(RuntimeError, match="failed checkpoint restore"): + resumed.state_dict() + with pytest.raises(BaseException, match="failed checkpoint restore"): + DistributedCheckpointer.save( + str(tmp_path / "invalid"), {"model": target, "optimizer": resumed}, parallel_state=parallel + ) + + def test_single_swap_preserves_frozen_history_on_resume(tmp_path): model = torch.nn.Linear(3, 2) source = SwapAdamW(model.parameters(), lr=0.01, pin_memory=False) @@ -446,6 +492,71 @@ def test_single_swap_absent_frozen_history_stays_absent(tmp_path): assert target.bias not in resumed._host_states +def test_single_swap_rejects_saved_history_missing_from_runtime_groups(tmp_path): + model = torch.nn.Linear(3, 2) + optimizer = SwapAdamW(model.parameters(), pin_memory=False) + _step(model, optimizer) + model.bias.requires_grad_(False) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target = torch.nn.Linear(3, 2) + target.bias.requires_grad_(False) + resumed = SwapAdamW((p for p in target.parameters() if p.requires_grad), pin_memory=False) + before = {name: value.clone() for name, value in target.state_dict().items()} + with pytest.raises(RuntimeError, match="checkpoint parameter groups differ"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": resumed}, parallel_state=parallel) + for name, value in target.state_dict().items(): + torch.testing.assert_close(value, before[name], rtol=0, atol=0) + assert target.bias not in resumed.state + + +def test_single_swap_rejects_omitted_owner_without_history(tmp_path): + model = torch.nn.Linear(3, 2) + model.bias.requires_grad_(False) + optimizer = SwapAdamW(model.parameters(), pin_memory=False) + _step(model, optimizer) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target = torch.nn.Linear(3, 2) + target.bias.requires_grad_(False) + resumed = SwapAdamW((p for p in target.parameters() if p.requires_grad), pin_memory=False) + before = {name: value.clone() for name, value in target.state_dict().items()} + with pytest.raises(RuntimeError, match="checkpoint parameter groups differ"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": resumed}, parallel_state=parallel) + for name, value in target.state_dict().items(): + torch.testing.assert_close(value, before[name], rtol=0, atol=0) + + +def test_single_swap_rejects_reordered_trainable_groups(tmp_path): + model = torch.nn.Linear(3, 2) + optimizer = SwapAdamW([{"params": [model.weight]}, {"params": [model.bias]}], pin_memory=False) + _step(model, optimizer) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target = torch.nn.Linear(3, 2) + resumed = SwapAdamW([{"params": [target.bias]}, {"params": [target.weight]}], pin_memory=False) + before = {name: value.clone() for name, value in target.state_dict().items()} + with pytest.raises(RuntimeError, match="checkpoint parameter groups differ"): + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": resumed}, parallel_state=parallel) + for name, value in target.state_dict().items(): + torch.testing.assert_close(value, before[name], rtol=0, atol=0) + + +def test_single_swap_accepts_unowned_frozen_parameter_without_history(tmp_path): + model = torch.nn.Linear(3, 2) + model.bias.requires_grad_(False) + optimizer = SwapAdamW((p for p in model.parameters() if p.requires_grad), pin_memory=False) + _step(model, optimizer) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": optimizer}, parallel_state=parallel) + target = torch.nn.Linear(3, 2) + target.bias.requires_grad_(False) + resumed = SwapAdamW((p for p in target.parameters() if p.requires_grad), pin_memory=False) + DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": resumed}, parallel_state=parallel) + assert target.bias not in resumed.state + torch.testing.assert_close(target.weight, model.weight, rtol=0, atol=0) + + def test_single_swap_frozen_history_preserves_model_load_hook_view(tmp_path): model = torch.nn.Linear(3, 2) optimizer = SwapAdamW(model.parameters(), pin_memory=False) @@ -522,3 +633,48 @@ def test_mismatched_swap_leaves_fail_before_checkpoint(tmp_path): from tests.tools.launch_utils import spawn_with_rendezvous spawn_with_rendezvous(_mismatched_leaf_worker, nprocs=2, args=(str(tmp_path / "ckpt"),)) + + +def _single_swap_save_failure_worker(rank, checkpoint, failure, save_async, endpoint): + from datetime import timedelta + + import torch.distributed as dist + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) + try: + model = torch.nn.Linear(3, 2) + optimizer = SwapAdamW(model.parameters(), pin_memory=False) + if rank == 0: + if failure == "step": + optimizer.state[model.weight]["step"] = 2**24 + 1 + elif failure == "poison": + optimizer._veomni_dcp_restore_failed = True + else: + + def fail(_optimizer): + raise RuntimeError("injected optimizer state hook failure") + + optimizer.register_state_dict_pre_hook(fail) + with pytest.raises(RuntimeError): + DistributedCheckpointer.save( + checkpoint, + {"model": model, "optimizer": optimizer}, + save_async=save_async, + parallel_state=SimpleNamespace(dp_mode="ddp"), + ) + value = torch.ones(1) + dist.all_reduce(value) + assert value.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) +@pytest.mark.parametrize("failure", ["step", "poison", "hook"]) +@pytest.mark.parametrize("save_async", [False, True]) +def test_single_swap_save_failure_is_coordinated(tmp_path, failure, save_async): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous( + _single_swap_save_failure_worker, nprocs=2, args=(str(tmp_path / "ckpt"), failure, save_async) + ) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 0fcc889b1d..ca98a52bff 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -38,7 +38,12 @@ load, ) from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner -from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE, Metadata, TensorStorageMetadata +from torch.distributed.checkpoint.metadata import ( + STATE_DICT_TYPE, + BytesStorageMetadata, + Metadata, + TensorStorageMetadata, +) from torch.distributed.checkpoint.state_dict import ( StateDictOptions, get_model_state_dict, @@ -94,6 +99,26 @@ _EXTRA_STATE_FORMAT = "extra_state_rank_{}.pt" +class _PreparedLoadState: + """Expose a rank-coordinated Stateful snapshot to DCP exactly once.""" + + def __init__(self, owner, snapshot, group): + self.owner = owner + self.snapshot = snapshot + self.group = group + + def state_dict(self): + return self.snapshot + + def load_state_dict(self, state_dict): + error = None + try: + self.owner.load_state_dict(state_dict) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "applying checkpoint state", group=self.group) + + _EXPERT_LAYOUTS_KEY = "__veomni_expert_layouts__" _UNKNOWN_EXPERT_LAYOUTS_KEY = "__veomni_unknown_expert_layouts__" _LEGACY_VEOMNI_TP1_EP0_FSDP1 = "veomni_tp1_ep0_fsdp1" @@ -2086,17 +2111,65 @@ def _single_swap_frozen_resume(model, optimizer, reader, *, group=None, trainabl if type(optimizer) is not SwapAdamW: yield () return - has_frozen = any(not p.requires_grad for g in optimizer.param_groups for p in g["params"]) - if not any_rank_failed(has_frozen, group=group): - yield () + names, owned, error = None, None, None + try: + names = _dcp_parameter_name_map(model) + owned = {id(p): p for g in optimizer.param_groups for p in g["params"]} + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "inspecting SwapAdamW parameter ownership", group=group) + needs_inspection = any(not p.requires_grad or id(p) not in owned for p in names.values()) + needs_inspection = any_rank_failed(needs_inspection, group=group) + metadata, group_indices, error = None, None, None + try: + if needs_inspection: + _validate_fresh_swap_optimizer(optimizer) + metadata = reader.read_metadata().state_dict_metadata + group_indices = sorted( + int(match.group(1)) + for key in metadata + if (match := re.fullmatch(r"optimizer\.param_groups\.(\d+)\.params", key)) + ) + if not group_indices or group_indices != list(range(len(group_indices))): + raise RuntimeError("SwapAdamW checkpoint parameter groups cannot be proven") + if any( + not isinstance(metadata[f"optimizer.param_groups.{index}.params"], BytesStorageMetadata) + for index in group_indices + ): + raise RuntimeError("SwapAdamW checkpoint parameter groups have an unsupported format") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "inspecting SwapAdamW checkpoint groups", group=group) + if dist.is_initialized(): + layouts = [None] * dist.get_world_size(group) + dist.all_gather_object(layouts, group_indices, group=group) + if any(layout != layouts[0] for layout in layouts[1:]): + raise RuntimeError("SwapAdamW checkpoint group metadata differs across ranks") + saved_groups = {"optimizer": {"param_groups": [{"params": None} for _ in group_indices]}} + dcp.load(state_dict=saved_groups, storage_reader=reader, process_group=group) + error = None + try: + names_by_id = {id(param): name for name, param in names.items()} + current_groups = [[names_by_id[id(param)] for param in entry["params"]] for entry in optimizer.param_groups] + source_groups = [entry["params"] for entry in saved_groups["optimizer"]["param_groups"]] + if source_groups != current_groups: + raise RuntimeError( + "SwapAdamW checkpoint parameter groups differ; restore the original groups or use a verified migration" + ) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "validating SwapAdamW checkpoint groups", group=group) + if not needs_inspection: + try: + yield () + except BaseException: + optimizer._veomni_dcp_restore_failed = True + raise return changed, allocated, error = [], [], None try: - _validate_fresh_swap_optimizer(optimizer) - metadata = reader.read_metadata().state_dict_metadata - owned = {id(p): p for g in optimizer.param_groups for p in g["params"]} - for name, param in _dcp_parameter_name_map(model).items(): - if id(param) not in owned or param.requires_grad: + for name, param in names.items(): + if id(param) in owned and param.requires_grad: continue fields = { field: metadata.get(f"optimizer.state.{name}.{field}") for field in ("step", *optimizer._state_keys) @@ -2106,6 +2179,11 @@ def _single_swap_frozen_resume(model, optimizer, reader, *, group=None, trainabl continue if present != set(fields): raise RuntimeError(f"SwapAdamW frozen state for {name} is incomplete") + if id(param) not in owned: + raise RuntimeError( + f"SwapAdamW checkpoint has state for unowned parameter {name}; " + "restore the original optimizer groups or use a verified migration" + ) if trainable_only: raise RuntimeError("Frozen SwapAdamW history requires a full-model checkpoint") if optimizer._optimizer_load_state_dict_pre_hooks or optimizer._optimizer_load_state_dict_post_hooks: @@ -2266,6 +2344,20 @@ def save( _reject_unproven_extra_parallel_optimizer( state, parallel_state, trainable_only=trainable_only, allow_proven_optimizer=multi, saving=True ) + from ..optim.swap_adamw import SwapAdamW + + single_swap_snapshot = None + if type(optimizer) is SwapAdamW: + if save_async: + cls._drain_slot(OPTIMIZER_DIRNAME) + error = None + try: + single_swap_snapshot = OptimizerState( + model=state["model"], optimizer=optimizer, parallel_state=parallel_state + ).state_dict() + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "preparing SwapAdamW checkpoint state") scheduler_object = state.get("_veomni_scheduler_object") if multi else None source_model = state.get("model") extra = state.get("extra_state") if multi else None @@ -2355,10 +2447,7 @@ def save( cls._save_extra_state(checkpoint_dir=write_root, state=state) except BaseException as exc: error = exc - if multi: - raise_if_any_rank_failed(error, "preparing MultiOptimizer checkpoint") - elif error is not None: - raise error + raise_if_any_rank_failed(error, "preparing checkpoint files") extra_digest = None if multi: error = None @@ -2431,6 +2520,8 @@ def prepare_writer_state(factory, what): group=None, ) if multi + else single_swap_snapshot + if single_swap_snapshot is not None else OptimizerState(model=state["model"], optimizer=optimizer, parallel_state=parallel_state), cls._create_storage_writer(os.path.join(write_root, OPTIMIZER_DIRNAME)), ), @@ -2704,111 +2795,96 @@ def _optimizer_state() -> "OptimizerState": state["_veomni_optimizer_identity_verified"] = True state["_veomni_optimizer_checkpoint_generation"] = generation return state - optimizer_path = fused_dir if fused_dir is not None else optimizer_dir(path, module) - reader, error = None, None + model_state, optimizer_state, model_reader, reader, error = None, None, None, None, None try: - reader = cls._create_storage_reader(optimizer_path) if wants_optimizer else None + model_state = _model_state() + optimizer_state = _optimizer_state() if wants_optimizer else None + model_reader = cls._create_storage_reader(fused_dir or weights_dir(path, module)) + if wants_optimizer: + reader = ( + model_reader if fused_dir is not None else cls._create_storage_reader(optimizer_dir(path, module)) + ) except BaseException as exc: error = exc - raise_if_any_rank_failed(error, "constructing optimizer checkpoint reader", group=process_group) + raise_if_any_rank_failed(error, "preparing distributed checkpoint load", group=process_group) + if dist.is_initialized(): + decisions = [None] * dist.get_world_size(group=process_group) + local = (fused_dir, wants_optimizer, multi) + dist.all_gather_object(decisions, local, group=process_group) + if any(decision != local for decision in decisions): + raise RuntimeError("Checkpoint load format differs across ranks") + + def prepared(owner, description): + snapshot, error = None, None + try: + snapshot = owner.state_dict() + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, description, group=process_group) + return _PreparedLoadState(owner, snapshot, process_group) + + def load_extra_coordinated(checkpoint_dir): + error = None + try: + cls._load_extra_state(checkpoint_dir=checkpoint_dir, state=state) + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "loading checkpoint extra state", group=process_group) + with _single_swap_frozen_resume( state["model"], optimizer, reader, group=process_group, trainable_only=trainable_only ) as frozen_swap_owners: - has_frozen_swap_owners = any_rank_failed(bool(frozen_swap_owners), group=process_group) + has_frozen = any_rank_failed(bool(frozen_swap_owners), group=process_group) error = None - if has_frozen_swap_owners and ( + if has_frozen and ( optimizer._optimizer_load_state_dict_pre_hooks or optimizer._optimizer_load_state_dict_post_hooks ): error = RuntimeError("Frozen SwapAdamW history does not support optimizer load-state hooks") raise_if_any_rank_failed(error, "checking frozen optimizer recovery hooks", group=process_group) - split_fused = fused_dir is not None and has_frozen_swap_owners + split_fused = fused_dir is not None and has_frozen if fused_dir is not None and not split_fused: - load_state: Dict[str, Any] = {"model": _model_state()} + model_state = prepared(model_state, "preparing model checkpoint state") + load_state = {"model": model_state} if wants_optimizer: - load_state["optimizer"] = _optimizer_state() + load_state["optimizer"] = prepared(optimizer_state, "preparing optimizer checkpoint state") dcp.load( state_dict=load_state, - storage_reader=cls._create_storage_reader(fused_dir), + storage_reader=model_reader, process_group=process_group, planner=_ModelStrictLoadPlanner( strict_model=not trainable_only, legacy_expert_source_layout=legacy_expert_source_layout ), ) - cls._load_extra_state(checkpoint_dir=fused_dir, state=state) + load_extra_coordinated(fused_dir) logger.info_rank0(f"Loaded pre-split checkpoint from {fused_dir}") return state model_root = model_dir(path, module) - model_state, model_reader, planner, error = None, None, None, None - try: - model_state = _model_state() - model_reader = cls._create_storage_reader(fused_dir if split_fused else weights_dir(path, module)) - planner = _ModelStrictLoadPlanner( - strict_model=not trainable_only, legacy_expert_source_layout=legacy_expert_source_layout - ) - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "preparing model recovery", group=process_group) error = None try: with _single_swap_frozen_model_view(frozen_swap_owners): - if has_frozen_swap_owners: - model_values, local_error = None, None - try: - model_values = model_state.state_dict() - except BaseException as exc: - local_error = exc - raise_if_any_rank_failed(local_error, "constructing model recovery state", group=process_group) - dcp.load( - state_dict={"model": model_values}, - storage_reader=model_reader, - process_group=process_group, - planner=planner, - ) - model_state.load_state_dict(model_values) - else: - dcp.load( - state_dict={"model": model_state}, - storage_reader=model_reader, - process_group=process_group, - planner=planner, - ) + model_state = prepared(model_state, "preparing model checkpoint state") + dcp.load( + state_dict={"model": model_state}, + storage_reader=model_reader, + process_group=process_group, + planner=_ModelStrictLoadPlanner( + strict_model=not trainable_only, legacy_expert_source_layout=legacy_expert_source_layout + ), + ) except BaseException as exc: error = exc raise_if_any_rank_failed(error, "loading model before optimizer recovery", group=process_group) if wants_optimizer: - optimizer_state = _optimizer_state() - if has_frozen_swap_owners: - optimizer_values, error = None, None - try: - optimizer_values = optimizer_state.state_dict() - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "constructing optimizer recovery state", group=process_group) - dcp.load( - state_dict={"optimizer": optimizer_values}, - storage_reader=reader, - process_group=process_group, - planner=_ModelStrictLoadPlanner(strict_model=False), - ) - error = None - try: - optimizer_state.load_state_dict(optimizer_values) - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "applying optimizer recovery state", group=process_group) - else: - dcp.load( - state_dict={"optimizer": optimizer_state}, - storage_reader=reader, - process_group=process_group, - planner=_ModelStrictLoadPlanner(strict_model=False), - ) - - cls._load_extra_state(checkpoint_dir=fused_dir if split_fused else model_root, state=state) - + optimizer_state = prepared(optimizer_state, "preparing optimizer checkpoint state") + dcp.load( + state_dict={"optimizer": optimizer_state}, + storage_reader=reader, + process_group=process_group, + planner=_ModelStrictLoadPlanner(strict_model=False), + ) + load_extra_coordinated(fused_dir if split_fused else model_root) logger.info_rank0(f"Loaded checkpoint from {model_root}") - return state @classmethod diff --git a/veomni/optim/swap_adamw.py b/veomni/optim/swap_adamw.py index 6ddb996ca6..aae576b895 100644 --- a/veomni/optim/swap_adamw.py +++ b/veomni/optim/swap_adamw.py @@ -320,6 +320,8 @@ def state_dict(self): stager deep-copies them before writing, so a save transiently needs about twice the host optimizer-state memory (on top of the host buffers). """ + if getattr(self, "_veomni_dcp_restore_failed", False): + raise RuntimeError("SwapAdamW cannot save after a failed checkpoint restore; rebuild the optimizer") staged = getattr(self, "_veomni_dcp_staged_state", None) if staged is None: self._bind_host_states() @@ -364,10 +366,22 @@ def state_dict(self): state_dict = result finally: self.state = current + # ``Optimizer.state_dict`` hands back the live inner state dicts; copy # them before swapping the int step for the tensor stock AdamW expects. + def serialize_step(step): + if torch.is_tensor(step): + if step.ndim != 0: + raise RuntimeError("SwapAdamW checkpoint step must be a scalar") + step = step.item() + if isinstance(step, bool) or not isinstance(step, (int, float)) or not math.isfinite(step): + raise RuntimeError("SwapAdamW checkpoint step must be a finite integer") + if not 0 <= step <= 2**24 or int(step) != step: + raise RuntimeError("SwapAdamW checkpoint step cannot be represented exactly as float32") + return torch.tensor(float(step), dtype=torch.float32, device="cpu") + state_dict["state"] = { - key: ({**value, "step": torch.tensor(float(value["step"]))} if "step" in value else value) + key: ({**value, "step": serialize_step(value["step"])} if "step" in value else value) for key, value in state_dict["state"].items() } return state_dict From 82fbfc9650b1cf56f50801c601c7c6d6dcc1d0f8 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:22:09 +0800 Subject: [PATCH 088/120] [checkpoint] constrain scheduler identity to trusted runtime --- docs/usage/checkpoint.md | 13 +++-- .../test_multi_optimizer_proven_resume.py | 20 ++++---- veomni/checkpoint/optimizer_resume_layout.py | 51 +++++++++++-------- 3 files changed, 49 insertions(+), 35 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 6e0b7f61bf..62ec6b0b24 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -166,10 +166,15 @@ MultiOptimizer checkpoint identity supports VeOmni's built-in constant, linear and cosine schedules. Custom `LambdaLR` functions must use immutable captures and the supported pure builtins or standard `math` members; global state, dynamic lookup, mutable closures and unrecognized attribute access are rejected -at save and resume. Older custom scheduler checkpoints that cannot meet this -contract require a verified migration or weights-only restart. Identity checks -assume a trusted Python runtime whose `math` and builtins were not monkeypatched -before VeOmni was imported; they are not a code sandbox. +at save and resume. Set construction, frozenset constants and f-string object +formatting are also rejected because their results can differ between processes. +Custom scheduler authors must ensure that the schedule is deterministic across +processes and recovery: the identity check does not prove arbitrary Python +functions deterministic. Older custom scheduler checkpoints that cannot meet +this contract require a verified migration or weights-only restart. Identity checks +assume a trusted Python runtime whose standard `math` members and scheduler +builtin bindings remain unchanged throughout training, saving and recovery. +They do not detect monkeypatches to these objects and are not a code sandbox. With `stage_dir` the model half waits: a staged save has overwritten nothing until the copy to the destination, so it invalidates the destination there diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index e59e9cd2fb..03ec1a6004 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -2068,19 +2068,21 @@ def nested(value): assert len(_scheduler_functions(SimpleNamespace(lr_lambdas=[schedule]))) == 1 -def test_scheduler_recipe_rejects_changed_math_member(monkeypatch): - import math as math_module - +@pytest.mark.parametrize("kind", ["set", "frozenset_constant", "object_format"]) +def test_scheduler_recipe_rejects_process_dependent_computation(kind): from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions - monkeypatch.setitem(globals(), "math", math_module) + def unordered(step): + return sum(float(value) for value in {"0.125", "1e16", "-1e16"}) - def schedule(step): - return math.cos(step) # noqa: F821 + def unordered_constant(step): + return step in {1, 2, 3} + + def object_format(step): + return int(f"{(lambda: 0)}"[-5], 16) / 20 - _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule])) - monkeypatch.setattr(math_module, "cos", lambda value: 0.25) - with pytest.raises(RuntimeError, match="unsupported.*state|unsupported.*math"): + schedule = {"set": unordered, "frozenset_constant": unordered_constant, "object_format": object_format}[kind] + with pytest.raises(RuntimeError, match="unsupported.*state"): _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule])) diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index 062e94e756..7870a0536a 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -3,7 +3,6 @@ """Identity of a MultiOptimizer's rank-local parameter and scheduler layout.""" -import builtins import dis import hashlib import json @@ -20,15 +19,14 @@ MULTI_LAYOUT_MARKER = "__veomni_multi_optimizer_layout_v1__" -_SCHEDULER_SAFE_BUILTINS = { - name: getattr(builtins, name) - for name in ("abs", "bool", "float", "int", "len", "max", "min", "pow", "range", "round", "sum") -} -_SCHEDULER_MATH_MEMBERS = { - name: value +_SCHEDULER_SAFE_BUILTINS = frozenset( + ("abs", "bool", "float", "int", "len", "max", "min", "pow", "range", "round", "sum") +) +_SCHEDULER_MATH_MEMBERS = frozenset( + name for name, value in vars(math).items() if not name.startswith("_") and type(value) in (float, int, types.BuiltinFunctionType) -} +) def _immutable_value(value): @@ -81,7 +79,12 @@ def _lr_layout(group, key, aliases): def _scheduler_functions(scheduler): """Describe LambdaLR callables using the existing checkpoint identity format.""" - def normalize(code, globals_dict, builtins_dict): + def has_unordered_constant(value): + if isinstance(value, frozenset): + return True + return isinstance(value, tuple) and any(has_unordered_constant(item) for item in value) + + def normalize(code, globals_dict): # Nested code objects share this function's global namespace. A # referenced global can change the next LR without changing code. # LOAD_ATTR names are not globals; inspect actual global lookups instead. @@ -104,19 +107,24 @@ def normalize(code, globals_dict, builtins_dict): "IMPORT_FROM", "IMPORT_STAR", "LOAD_BUILD_CLASS", + # Hash-randomized set iteration can change floating-point LR + # results across processes without changing the recipe hash. + "BUILD_SET", + "SET_ADD", + "SET_UPDATE", + # Formatting arbitrary objects can include process-local ids. + "FORMAT_VALUE", + "FORMAT_SIMPLE", + "FORMAT_WITH_SPEC", + "CONVERT_VALUE", + "BUILD_STRING", ): raise RuntimeError("MultiOptimizer scheduler uses unsupported dynamic state") if opcode in ("LOAD_ATTR", "LOAD_METHOD"): previous = instructions[index - 1] if index else None if previous is None or previous.opname != "LOAD_GLOBAL" or previous.argval != "math": raise RuntimeError("MultiOptimizer scheduler uses unsupported attribute access") - expected = _SCHEDULER_MATH_MEMBERS.get(instruction.argval) - actual = vars(math).get(instruction.argval) - if ( - expected is None - or type(expected) is not type(actual) - or (actual is not expected and (type(expected) not in (int, float) or actual != expected)) - ): + if instruction.argval not in _SCHEDULER_MATH_MEMBERS: raise RuntimeError("MultiOptimizer scheduler uses unsupported math state") continue if opcode != "LOAD_GLOBAL": @@ -125,11 +133,12 @@ def normalize(code, globals_dict, builtins_dict): if name in globals_dict: if name != "math" or globals_dict[name] is not math: raise RuntimeError("MultiOptimizer scheduler refers to unsupported global state") - elif name not in _SCHEDULER_SAFE_BUILTINS or builtins_dict.get(name) is not _SCHEDULER_SAFE_BUILTINS[name]: + elif name not in _SCHEDULER_SAFE_BUILTINS: raise RuntimeError("MultiOptimizer scheduler uses an unsupported builtin") + if any(has_unordered_constant(item) for item in code.co_consts): + raise RuntimeError("MultiOptimizer scheduler uses unsupported unordered state") constants = tuple( - normalize(item, globals_dict, builtins_dict) if isinstance(item, types.CodeType) else item - for item in code.co_consts + normalize(item, globals_dict) if isinstance(item, types.CodeType) else item for item in code.co_consts ) return code.replace(co_filename="", co_firstlineno=0, co_consts=constants) @@ -140,9 +149,7 @@ def normalize(code, globals_dict, builtins_dict): functions.append( { - "code": hashlib.sha256( - marshal.dumps(normalize(function.__code__, function.__globals__, function.__builtins__)) - ).hexdigest(), + "code": hashlib.sha256(marshal.dumps(normalize(function.__code__, function.__globals__))).hexdigest(), "defaults": _immutable_value(function.__defaults__), "kwdefaults": _immutable_value(function.__kwdefaults__), "closure": [_immutable_value(cell.cell_contents) for cell in function.__closure__ or ()], From be814f4d9651c6e6b94aa0890c82f9c5603aa423 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 00:24:35 +0800 Subject: [PATCH 089/120] [checkpoint] fix: validate versioned expert ownership --- docs/usage/checkpoint.md | 12 +- tests/checkpoints/test_expert_shard_layout.py | 113 ++++++++++++++++-- veomni/checkpoint/dcp_checkpointer.py | 75 ++++++++---- 3 files changed, 170 insertions(+), 30 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 2fa4cf5151..5253149097 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -241,9 +241,10 @@ The separately namespaced EP-outer layout is opt-in through `SpecInfo.checkpoint_layout_v2`. A plan enabling it must retain the original `global_shape` and must use FSDP DTensor parameters. Every non-persistent ExtraParallel shard in a checkpoint must use the canonical -layout: mixing legacy and canonical shards is rejected. The v2 tensor namespace -identifies canonical global values independently of source geometry, so DCP can -reshard a model checkpoint across supported FSDP shard axes. Use +layout: mixing legacy and canonical shards is rejected. For tensors whose +source layout has been proved, the v2 namespace represents canonical global +values independently of source geometry, so DCP can reshard a model checkpoint +across supported FSDP shard axes. Use `DistributedCheckpointer` for save and resume: it rejects old/unversioned layouts before copying model tensors. Migrating a legacy checkpoint requires verified source layout information. Canonical v2 optimizer save and resume are @@ -251,6 +252,11 @@ rejected until an owner-aware codec can preserve rank-specific optimizer states; this opt-in therefore does not yet support resuming training. The standard VeOmni HF export strips the namespace. +The namespace marker alone does not prove that an ordinary or shared expert's +replicas were synchronized. Older v2 checkpoints whose planner metadata omits +an expert's source layout are rejected on resume and export; verify and migrate +them from their original training layout before use. + The standard, non-v2 writer uses canonical EP-first coordinates for same-axis EP/FSDP shards and records source expert layouts in DCP planner metadata. Older readers may ignore that metadata and interpret these shards using the old diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index ed2eec2e74..5b04010e6a 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -161,6 +161,60 @@ def _lora_model(config=None): return VeOmniLoraModel(model, config) +@pytest.mark.parametrize("root", ["model", "optimizer"]) +def test_versioned_owner_lookup_requires_exact_logical_fqn(root): + from veomni.checkpoint.dcp_checkpointer import _key_belongs_to_fqns, _logical_extra_parallel_key + + marker = _extra_parallel_layout_key({}) + state_dict = {f"{root}.{marker}": "ep-outer-v1"} + fqn = "layers.0.experts.gate_proj.lora_A.weight" + tail = fqn if root == "model" else f"state.{fqn}.exp_avg" + key = _logical_extra_parallel_key(f"{root}.{marker}__{tail}", state_dict) + assert _key_belongs_to_fqns(key, {fqn}) + assert not _key_belongs_to_fqns(key, {f"other.{fqn}"}) + assert _logical_extra_parallel_key(f"{root}.{marker}", state_dict) is None + nested = _logical_extra_parallel_key(f"{root}.{marker}__{marker}__{tail}", state_dict) + assert not _key_belongs_to_fqns(nested, {fqn}) + unversioned = f"{root}.{marker}__{tail}" + assert _logical_extra_parallel_key(unversioned, {}) == unversioned + with pytest.raises(RuntimeError, match="mixes versioned and legacy"): + _logical_extra_parallel_key(f"{root}.{tail}", state_dict) + + +@pytest.mark.parametrize("fqn", ["experts.weight", "layers.0.experts.weight"]) +def test_versioned_experts_without_proven_owner_are_recorded_as_unproven(fqn): + marker = _extra_parallel_layout_key({}) + key = f"model.{marker}__{fqn}" + planner = _ExtraParallelSavePlanner() + planner.set_up_planner( + {"model": {marker: "ep-outer-v1", f"{marker}__{fqn}": torch.ones(2, 2)}}, is_coordinator=True + ) + plan = planner.create_local_plan() + assert plan.planner_data[_EXPERT_LAYOUTS_KEY][key] == {"expert_order": "unproven"} + + +@pytest.mark.parametrize("fqn", ["experts.weight", "layers.0.experts.weight"]) +@pytest.mark.parametrize("layout_metadata", [False, True]) +def test_v2_checkpoint_missing_expert_proof_rejects_load_and_export(tmp_path, fqn, layout_metadata): + class MissingLayoutPlanner(_ExtraParallelSavePlanner): + def create_local_plan(self): + plan = super().create_local_plan() + plan.planner_data[_EXPERT_LAYOUTS_KEY].clear() + return plan + + marker = _extra_parallel_layout_key({}) + name = f"{marker}__{fqn}" + state = {"model": {marker: "ep-outer-v1", name: torch.ones(2, 2)}} + planner = MissingLayoutPlanner() if layout_metadata else DefaultSavePlanner() + dcp.save(state, checkpoint_id=tmp_path, planner=planner) + target = {"model": {marker: "", name: torch.zeros(2, 2)}} + with pytest.raises(dcp.CheckpointException, match="metadata omits its layout"): + dcp.load(target, checkpoint_id=tmp_path, planner=_ModelStrictLoadPlanner(strict_model=True)) + assert target["model"][name].count_nonzero() == 0 + with pytest.raises(RuntimeError, match="metadata omits its layout"): + dcp_to_torch_state_dict(tmp_path) + + def _lora_export_worker(rank, world, shard_dim, negative_ep, use_v2, directory, endpoint): from veomni.distributed.parallel_plan import SpecInfo from veomni.lora import VeOmniLoraConfig @@ -718,10 +772,53 @@ def _ep_shared_lora_worker(rank, root, mode, endpoint): ) selected = get_lora_state_dict(model, config=model.get_lora_config()) expected = {name: value.full_tensor().detach().to(torch.bfloat16).cpu() for name, value in selected.items()} - if mode == "plan_conflict": + if mode in ("plan_conflict", "v2_plan_conflict"): base.get_parallel_plan = lambda: ParallelPlan( {"ep": {"experts.gate_proj.lora_A.default.weight": Shard(0)}} ) + if mode.startswith("v2_"): + from veomni.checkpoint.dcp_checkpointer import _namespace_extra_parallel_layout_state + + marker = _extra_parallel_layout_key({}) + state = {"model": _namespace_extra_parallel_layout_state(selected, marker)} + planner = _ExtraParallelSavePlanner(parallel_state=parallel, source_model=model, adapter_name="default") + dcp.save(state, checkpoint_id=root, planner=planner) + metadata = dcp.FileSystemReader(root).read_metadata() + records = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] + orders = [] + for name in selected: + if "experts" not in name.split("."): + continue + key = f"model.{marker}__{name}" + assert key in records + orders.append(records[key]["expert_order"]) + assert orders + assert ("unproven" in orders) if mode == "v2_plan_conflict" else all(o == "canonical" for o in orders) + target = { + "model": _namespace_extra_parallel_layout_state( + {k: torch.zeros_like(v) for k, v in selected.items()}, marker + ) + } + if mode == "v2_plan_conflict": + with pytest.raises(dcp.CheckpointException, match="noncanonical expert order"): + dcp.load(target, checkpoint_id=root, planner=_ModelStrictLoadPlanner(strict_model=True)) + if rank == 0: + with pytest.raises(RuntimeError, match="noncanonical expert order"): + dcp_to_torch_state_dict(root) + dist.barrier() + return + dcp.load(target, checkpoint_id=root, planner=_ModelStrictLoadPlanner(strict_model=True)) + for name, value in selected.items(): + torch.testing.assert_close( + target["model"][f"{marker}__{name}"].to_local(), value.to_local(), rtol=0, atol=0 + ) + expected_fp32 = {name: value.full_tensor().detach().cpu() for name, value in selected.items()} + if rank == 0: + exported = dcp_to_torch_state_dict(root) + for name, value in expected_fp32.items(): + torch.testing.assert_close(exported[name], value, rtol=0, atol=0) + dist.barrier() + return if mode != "valid": with pytest.raises(RuntimeError, match="noncanonical expert order|failed on another rank"): save_lora_adapter_with_dcp(model, root, parallel_state=parallel) @@ -741,7 +838,7 @@ def _ep_shared_lora_worker(rank, root, mode, endpoint): @pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) @pytest.mark.parametrize("world", [2]) -@pytest.mark.parametrize("mode", ["valid", "plan_conflict", "unproven_bias"]) +@pytest.mark.parametrize("mode", ["valid", "plan_conflict", "unproven_bias", "v2_valid", "v2_plan_conflict"]) def test_ep_replicated_shared_expert_lora_export(world, mode, tmp_path): spawn_with_rendezvous(_ep_shared_lora_worker, args=(str(tmp_path / "adapter"), mode), nprocs=world) @@ -997,8 +1094,9 @@ def _named_extra_expert_export_worker(rank, root, endpoint): dcp.load(scalar_target, checkpoint_id=scalar_path, planner=_ModelStrictLoadPlanner(strict_model=False)) torch.testing.assert_close(scalar_target["optimizer"]["experts"]["step"], torch.tensor(1.0), rtol=0, atol=0) - # The old per-tensor planner did not recognize ``custom``, but the - # canonical v2 namespace itself proves this checkpoint's source order. + # The old planner did not recognize ``custom``. A v2 marker describes + # the EP conversion but does not prove all shared expert replicas, so a + # missing per-tensor layout must not be silently accepted. marker = _extra_parallel_layout_key({}) canonical_mesh = _extra_parallel_first_mesh(mesh) namespaced = DTensor.from_local( @@ -1016,12 +1114,13 @@ def _named_extra_expert_export_worker(rank, root, endpoint): planner=_LegacyNamedExtraSavePlanner(), ) target = {"model": {marker: "", f"{marker}__{key}": torch.zeros_like(namespaced)}} - dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=True)) - torch.testing.assert_close(target["model"][f"{marker}__{key}"].full_tensor(), expected, rtol=0, atol=0) + with pytest.raises(dcp.CheckpointException, match="metadata omits its layout"): + dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=True)) result = [None] if rank == 0: try: - torch.testing.assert_close(dcp_to_torch_state_dict(path)[key], expected, rtol=0, atol=0) + with pytest.raises(RuntimeError, match="metadata omits its layout"): + dcp_to_torch_state_dict(path) except Exception as exc: result[0] = repr(exc) dist.broadcast_object_list(result, src=0) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index ec774658e9..5e8876efff 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -185,6 +185,25 @@ def _ordinary_expert_fqns(model, parallel_state): return result - ambiguous +def _logical_extra_parallel_key(key, state_dict): + # The save planner's metadata keeps the serialized key, but ownership is + # expressed in logical FQNs. Canonical v2 checkpoints namespace both model + # and optimizer keys, so strip only the verified layout marker for lookup. + root, sep, rest = key.partition(".") + if not sep or root not in ("model", "optimizer"): + return key + marker = _extra_parallel_layout_key({}) + marker_value = state_dict.get(f"{root}.{marker}") + if isinstance(marker_value, str) and marker_value == "ep-outer-v1": + if rest == marker: + return None + prefix = marker + "__" + if not rest.startswith(prefix): + raise RuntimeError("ExtraParallel checkpoint mixes versioned and legacy tensor keys") + return f"{root}.{rest[len(prefix) :]}" + return key + + def _key_belongs_to_fqns(key, fqns): if key.startswith("model.") and (key in fqns or key.removeprefix("model.") in fqns): return True @@ -467,12 +486,17 @@ def create_local_plan(self): plan = super().create_local_plan() layouts = {} for key, value in self.state_dict.items(): + logical_key = _logical_extra_parallel_key(key, self.state_dict) + if logical_key is None: + continue extra_names = getattr(self.parallel_state, "extra_parallel_names", ()) layout = ( _expert_layout(value, extra_names) - or _no_ep_expert_layout(key, value, self.parallel_state, self.ordinary_fqns, self.shared_lora_fqns) + or _no_ep_expert_layout( + logical_key, value, self.parallel_state, self.ordinary_fqns, self.shared_lora_fqns + ) or _tp_only_source_layout( - key, + logical_key, value, self.tp_only_parameters, self.tp_only_mesh, @@ -482,7 +506,7 @@ def create_local_plan(self): if layout is not None: layouts[key] = layout elif _unidentified_same_axis_shards(value) or ( - "experts" in key.split(".") and torch.is_tensor(value) and value.ndim > 0 + "experts" in logical_key.split(".") and torch.is_tensor(value) and value.ndim > 0 ): layouts[key] = {"expert_order": "unproven"} return replace(plan, planner_data={**(plan.planner_data or {}), _EXPERT_LAYOUTS_KEY: layouts}) @@ -578,16 +602,14 @@ def _validate_expert_layouts(self): "Migrate model and optimizer state using the original training layout before loading." ) - # Checkpoints produced by the earlier version of this feature already - # used canonical coordinates, but did not record individual layouts. - if planner_data.get("__veomni_ep_first_layout__") == 1: - return - - # Earlier namespaced v2 checkpoints used an EP-first mesh but predate - # the per-tensor planner metadata. The marker and namespace were - # already checked against the target above. + # The v2 namespace proves how sharded EP tensors were serialized, but + # not whether an ordinary/shared expert was a synchronized replica. + # Each saved expert needs its own layout proof even in a v2 checkpoint. marker = _extra_parallel_layout_key({}) - if any(key.split(".")[-1] == marker for key in self.metadata.state_dict_metadata): + has_v2_marker = any(key.split(".")[-1] == marker for key in self.metadata.state_dict_metadata) + # Preserve the earlier unversioned EP-first contract. A v2 marker must + # not bypass the per-tensor proof even if both markers are present. + if not has_v2_marker and planner_data.get("__veomni_ep_first_layout__") == 1: return for key, value in self.state_dict.items(): @@ -597,11 +619,18 @@ def _validate_expert_layouts(self): saved = self.metadata.state_dict_metadata.get(key) if saved is None: continue - if source is not None and "experts" in key.split(".") and bool(getattr(saved, "size", ())): + root = key.partition(".")[0] + root_marker = f"{root}.{marker}" + source_marker = {root_marker: "ep-outer-v1"} if root_marker in self.metadata.state_dict_metadata else {} + logical_key = _logical_extra_parallel_key(key, source_marker) + is_expert = logical_key is not None and "experts" in logical_key.split(".") + if (source is not None or has_v2_marker) and is_expert and bool(getattr(saved, "size", ())): raise RuntimeError( - f"Cannot determine source expert order for {key!r}: version-1 metadata omits its layout. " + f"Cannot determine source expert order for {key!r}: checkpoint metadata omits its layout. " "Verify and migrate the original training layout before loading." ) + if has_v2_marker: + continue # For legacy VeOmni EP weights, the faulty same-axis layout only # split the expert dimension. A source chunk that splits matrix # rows (old dim-1 FSDP), or spans all experts, cannot have that @@ -611,7 +640,7 @@ def _validate_expert_layouts(self): # letting it re-save would incorrectly mark that data as # canonical. The stacked VeOmni expert axis is dimension 0. shape = getattr(saved, "size", ()) - if "experts" not in key.split(".") or not shape: + if not is_expert or not shape: continue ep_dim = 0 else: @@ -2471,8 +2500,10 @@ def _validate_expert_export_provenance(metadata: Metadata, model_keys) -> None: f"Checkpoint tensor {key!r} has noncanonical expert order; migrate it before export" ) - marker_key = "model." + _extra_parallel_layout_key({}) - if marker_key in metadata.state_dict_metadata or planner_data.get("__veomni_ep_first_layout__") == 1: + marker = _extra_parallel_layout_key({}) + marker_key = "model." + marker + has_v2_marker = marker_key in metadata.state_dict_metadata + if not has_v2_marker and planner_data.get("__veomni_ep_first_layout__") == 1: return # An old expert-only chunk may use either EP-first or FSDP-first order. @@ -2480,11 +2511,15 @@ def _validate_expert_export_provenance(metadata: Metadata, model_keys) -> None: for key in model_keys: tensor_metadata = metadata.state_dict_metadata[key] shape = getattr(tensor_metadata, "size", ()) - if "experts" not in key.split(".") or not shape or key in layouts: + logical_key = key + if has_v2_marker: + fqn = _export_extra_parallel_layout_key(key.removeprefix("model.")) + logical_key = f"model.{fqn}" if fqn is not None else "" + if "experts" not in logical_key.split(".") or not shape or key in layouts: continue - if recorded is not None: + if recorded is not None or has_v2_marker: raise RuntimeError( - f"Cannot determine source expert order for {key!r}: version-1 metadata omits its layout. " + f"Cannot determine source expert order for {key!r}: checkpoint metadata omits its layout. " "Verify and migrate the original training layout before exporting." ) if any( From 37b473eb05dcf04d3ea528f7ad1a2e01e1fd7b73 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 10:28:44 +0800 Subject: [PATCH 090/120] [checkpoint] simplify extra-parallel optimizer preflight --- .../test_extra_parallel_optimizer_keys.py | 38 +++++++------------ veomni/checkpoint/dcp_checkpointer.py | 12 +----- 2 files changed, 16 insertions(+), 34 deletions(-) diff --git a/tests/checkpoints/test_extra_parallel_optimizer_keys.py b/tests/checkpoints/test_extra_parallel_optimizer_keys.py index ec911ae75f..2304fe2402 100644 --- a/tests/checkpoints/test_extra_parallel_optimizer_keys.py +++ b/tests/checkpoints/test_extra_parallel_optimizer_keys.py @@ -162,7 +162,7 @@ def test_plain_extra_parallel_values_survive_dcp_roundtrip(tmp_path): torch.multiprocessing.spawn(_plain_extra_parallel_roundtrip, args=(str(tmp_path),), nprocs=2, join=True) -def _asymmetric_optimizer_preflight(rank, root, case): +def _asymmetric_optimizer_preflight(rank, root): from datetime import timedelta from pathlib import Path from types import SimpleNamespace @@ -178,18 +178,9 @@ def __getitem__(self, _): return SimpleNamespace(size=lambda: 2) class FakeOptimizer: - called = False - - @property - def _is_multi_optimizer(self): - if case == "type_property_error" and rank == 0: - raise ValueError("injected optimizer type lookup error") - return not (case.startswith("wrong_type") and rank == 0) + _is_multi_optimizer = True def state_dict(self): - self.called = True - if case in ("wrong_type_collective", "type_property_error"): - dist.barrier() return {"state.experts.weight.exp_avg": torch.ones(2, 3)} if rank == 0 else {} dist.init_process_group( @@ -203,18 +194,9 @@ def state_dict(self): leaf = FakeOptimizer() optimizer = dcp_checkpointer.OptimizerState(model, leaf, parallel_state=parallel) checkpoint = Path(root) / "rejected" - if rank == 0 and case.startswith("wrong_type"): - message = "not a MultiOptimizer" - elif rank == 0 and case == "type_property_error": - message = "optimizer type lookup error" - elif case.startswith("wrong_type") or case == "type_property_error": - message = "validating ExtraParallel optimizer type" - else: - message = "no proven replica layout" if rank == 0 else "preparing ExtraParallel optimizer checkpoint state" + message = "no proven replica layout" if rank == 0 else "preparing ExtraParallel optimizer checkpoint state" with pytest.raises((RuntimeError, ValueError), match=message): dcp.save({"optimizer": optimizer}, checkpoint_id=checkpoint) - if case.startswith("wrong_type") or case == "type_property_error": - assert not leaf.called assert not (checkpoint / ".metadata").exists() value = torch.ones(()) dist.all_reduce(value) @@ -224,6 +206,14 @@ def state_dict(self): @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -@pytest.mark.parametrize("case", ["plain_state", "wrong_type", "wrong_type_collective", "type_property_error"]) -def test_asymmetric_optimizer_preflight_rejects_every_rank(tmp_path, case): - torch.multiprocessing.spawn(_asymmetric_optimizer_preflight, args=(str(tmp_path), case), nprocs=2, join=True) +def test_asymmetric_optimizer_preflight_rejects_every_rank(tmp_path): + torch.multiprocessing.spawn(_asymmetric_optimizer_preflight, args=(str(tmp_path),), nprocs=2, join=True) + + +def test_extra_parallel_rejects_non_multi_optimizer_locally(): + from types import SimpleNamespace + + model = SimpleNamespace(_fqn2spec_info={}) + state = dcp_checkpointer.OptimizerState(model, object(), parallel_state=SimpleNamespace(dp_mode="fsdp2")) + with pytest.raises(RuntimeError, match="optimizer is not a MultiOptimizer"): + state.state_dict() diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 27d5f7b13b..22f1080e0a 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -434,16 +434,8 @@ def state_dict(self): logger.info_rank0( "Getting optimizer state_dict from OptimizerState wrapper, would restore ExtraParallel dim for Experts module" ) - # Validate locally before any rank enters state_dict(), which can - # contain collectives. A single rank's type error must not race - # another rank's optimizer collective against our error reduction. - error = None - try: - if not getattr(self.optimizer, "_is_multi_optimizer", False): - raise RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "validating ExtraParallel optimizer type", group=self.process_group) + if not getattr(self.optimizer, "_is_multi_optimizer", False): + raise RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") optim_sd_with_extra_parallel_dim, error = None, None try: From 95d085540166d8096466d01171573e9bb5cb346e Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 10:29:18 +0800 Subject: [PATCH 091/120] [checkpoint] simplify extra-parallel optimizer preflight --- tests/checkpoints/test_expert_shard_layout.py | 34 ++++++++----------- veomni/checkpoint/dcp_checkpointer.py | 12 ++----- 2 files changed, 16 insertions(+), 30 deletions(-) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index 5b04010e6a..ff70fa6458 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -933,7 +933,7 @@ def test_plain_optimizer_shards_are_not_treated_as_replicas(world, hsdp): spawn_with_rendezvous(_plain_optimizer_unproven_worker, args=(hsdp,), nprocs=world) -def _sparse_plain_optimizer_failure_worker(rank, case, endpoint): +def _sparse_plain_optimizer_failure_worker(rank, endpoint): from veomni.checkpoint.dcp_checkpointer import OptimizerState from veomni.distributed.parallel_plan import SpecInfo @@ -947,29 +947,15 @@ def _sparse_plain_optimizer_failure_worker(rank, case, endpoint): model._fqn2spec_info = {"experts.weight": SpecInfo("ep", Shard(0), "experts.weight", mesh, 1)} class SparseOptimizer: - called = False - - @property - def _is_multi_optimizer(self): - if case == "type_property_error" and rank == 0: - raise ValueError("injected optimizer type lookup error") - return not (case.startswith("wrong_type") and rank == 0) + _is_multi_optimizer = True def state_dict(self): - self.called = True - if case in ("wrong_type_collective", "type_property_error"): - dist.barrier() return {"state.experts.weight.exp_avg": torch.ones(4, 4)} if rank == 0 else {} leaf = SparseOptimizer() state = OptimizerState(model, leaf, parallel_state=parallel) - with pytest.raises( - (RuntimeError, ValueError), - match="no proven replica layout|not a MultiOptimizer|failed on another rank|optimizer type lookup error", - ): + with pytest.raises(RuntimeError, match="no proven replica layout|failed on another rank"): state.state_dict() - if case.startswith("wrong_type") or case == "type_property_error": - assert not leaf.called completed = torch.tensor(1) dist.all_reduce(completed) assert completed.item() == 4 @@ -979,9 +965,17 @@ def state_dict(self): @pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) @pytest.mark.parametrize("world", [4]) -@pytest.mark.parametrize("case", ["plain_state", "wrong_type", "wrong_type_collective", "type_property_error"]) -def test_sparse_plain_optimizer_failure_exits_all_ranks(world, case): - spawn_with_rendezvous(_sparse_plain_optimizer_failure_worker, args=(case,), nprocs=world) +def test_sparse_plain_optimizer_failure_exits_all_ranks(world): + spawn_with_rendezvous(_sparse_plain_optimizer_failure_worker, args=(), nprocs=world) + + +def test_extra_parallel_rejects_non_multi_optimizer_locally(): + from veomni.checkpoint.dcp_checkpointer import OptimizerState + + model = SimpleNamespace(_fqn2spec_info={}) + state = OptimizerState(model, object(), parallel_state=SimpleNamespace(dp_mode="fsdp2")) + with pytest.raises(RuntimeError, match="optimizer is not a MultiOptimizer"): + state.state_dict() def _subgroup_sparse_optimizer_failure_worker(rank, endpoint): diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 5e8876efff..6587328fe7 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -1418,16 +1418,8 @@ def state_dict(self): logger.info_rank0( "Getting optimizer state_dict from OptimizerState wrapper, would restore ExtraParallel dim for Experts module" ) - # Validate locally before any rank enters state_dict(), which can - # contain collectives. A single rank's type error must not race - # another rank's optimizer collective against our error reduction. - error = None - try: - if not getattr(self.optimizer, "_is_multi_optimizer", False): - raise RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") - except BaseException as exc: - error = exc - raise_if_any_rank_failed(error, "validating ExtraParallel optimizer type", group=self.process_group) + if not getattr(self.optimizer, "_is_multi_optimizer", False): + raise RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") optim_sd_with_extra_parallel_dim, error = None, None try: From a64fd963547d5c91b63cdc0b57ecc0d99cd6dd0d Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 10:43:05 +0800 Subject: [PATCH 092/120] [checkpoint] simplify multi-optimizer leaf coordination --- .../rank_local_adamw_checkpoint_adapter.md | 10 +++--- docs/usage/checkpoint.md | 5 +-- .../test_swap_proven_multi_resume.py | 36 ------------------- veomni/checkpoint/optimizer_resume_layout.py | 12 ------- veomni/optim/optimizer.py | 7 ---- 5 files changed, 8 insertions(+), 62 deletions(-) diff --git a/docs/design/rank_local_adamw_checkpoint_adapter.md b/docs/design/rank_local_adamw_checkpoint_adapter.md index eb654867e5..e3e2417495 100644 --- a/docs/design/rank_local_adamw_checkpoint_adapter.md +++ b/docs/design/rank_local_adamw_checkpoint_adapter.md @@ -62,14 +62,14 @@ the saved checkpoint. Change these only with an explicitly verified migration. ### Save and load call sites -[`_local_layout`](../../veomni/checkpoint/optimizer_resume_layout.py#L105) +[`_local_layout`](../../veomni/checkpoint/optimizer_resume_layout.py#L177) records `algorithm_identity`, each parameter's `state_dtype`, and `state_is_offloaded` in the checkpoint identity; a mismatch rejects resume before applying state. During save, -[`_validate_source_optimizer_states`](../../veomni/checkpoint/multi_optimizer_checkpoint.py#L89) +[`_validate_source_optimizer_states`](../../veomni/checkpoint/multi_optimizer_checkpoint.py#L121) uses `state_dtype` to check initialized fields against their declared precision. During load, -[`_prepare_multi_optimizer_saved_states`](../../veomni/checkpoint/multi_optimizer_checkpoint.py#L408) +[`_prepare_multi_optimizer_saved_states`](../../veomni/checkpoint/multi_optimizer_checkpoint.py#L466) checks stored field dtypes and uses `state_is_offloaded` to allocate missing destination state on CPU when the adapter specifies offload. [`rank_local_adamw_context`](../../veomni/optim/checkpoint_adapter.py#L55) @@ -87,7 +87,7 @@ must also be validated for an offloading implementation. ## Framework call sites and restore behavior -1. [`MultiOptimizer._sparse_state_dict()`](../../veomni/optim/optimizer.py#L563) +1. [`MultiOptimizer._sparse_state_dict()`](../../veomni/optim/optimizer.py#L568) enters the context around PyTorch's optimizer-state serialization. Registered adapter implementations must preserve live state and serialize all initialized fields and group options @@ -98,7 +98,7 @@ must also be validated for an offloading implementation. load checks the checkpoint's identity, field names, dtypes and shapes before preparing initialized state for the saved owners. Rank-local ownership and the supported parallel geometry are not automatically resharded. -3. [`MultiOptimizer._dcp_load_state_dict()`](../../veomni/optim/optimizer.py#L601) +3. [`MultiOptimizer._dcp_load_state_dict()`](../../veomni/optim/optimizer.py#L618) enters the context again while PyTorch applies the optimizer state. The trusted adapter must restore every validated field without changing its value or dtype. An optimizer may diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 62ec6b0b24..5fa7a8cb5f 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -452,8 +452,9 @@ and uninitialized states, but increases checkpoint storage for replicated optimizer shards and does not support elastic resharding. All ranks must use the same trusted model and configuration, including optimizer leaf ordering. Owner-specific freezing that removes a whole leaf from only some ranks is not -supported; differing leaf names or order are rejected before checkpoint collectives. Replica -values, trainability, dtype, shape and geometry +supported. Leaf names and order are not checked across ranks before checkpoint +collectives; divergent configurations can hang or produce an invalid checkpoint. +Replica values, trainability, dtype, shape and geometry are trusted to be consistent across ranks. Source and target checkpoint layout metadata are checked during restore; tensor bytes are not hashed for replica consistency. Trainable-only MultiOptimizer checkpoints do hash each rank's local diff --git a/tests/checkpoints/test_swap_proven_multi_resume.py b/tests/checkpoints/test_swap_proven_multi_resume.py index a6aac299a6..4d274a94db 100644 --- a/tests/checkpoints/test_swap_proven_multi_resume.py +++ b/tests/checkpoints/test_swap_proven_multi_resume.py @@ -599,42 +599,6 @@ def test_single_swap_frozen_history_rejects_optimizer_load_hook_before_loading(t torch.testing.assert_close(target.weight, before, rtol=0, atol=0) -def _mismatched_leaf_worker(rank, checkpoint, endpoint): - from datetime import timedelta - - import torch.distributed as dist - - dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30)) - try: - model = torch.nn.Linear(3, 2) - if rank == 0: - leaves = { - "weight": SwapAdamW([model.weight], pin_memory=False), - "bias": SwapAdamW([model.bias], pin_memory=False), - } - else: - leaves = {"both": SwapAdamW(model.parameters(), pin_memory=False)} - optimizer = MultiOptimizer(model, leaves, list(leaves)) - with pytest.raises(RuntimeError, match="leaf order differs"): - DistributedCheckpointer.save( - checkpoint, - {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None}, - parallel_state=SimpleNamespace(dp_mode="ddp"), - ) - value = torch.ones(1) - dist.all_reduce(value) - assert value.item() == 2 - finally: - dist.destroy_process_group() - - -@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -def test_mismatched_swap_leaves_fail_before_checkpoint(tmp_path): - from tests.tools.launch_utils import spawn_with_rendezvous - - spawn_with_rendezvous(_mismatched_leaf_worker, nprocs=2, args=(str(tmp_path / "ckpt"),)) - - def _single_swap_save_failure_worker(rank, checkpoint, failure, save_async, endpoint): from datetime import timedelta diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index 7870a0536a..cffa48d92c 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -439,16 +439,6 @@ def frozen_parameter_digest(model, *, enabled, group=None, chunk_bytes=4 << 20): return digest.hexdigest() -def _validate_multi_leaf_order(optimizer, *, group=None): - """All ranks must enter per-leaf checkpoint collectives in the same order.""" - if not dist.is_initialized(): - return - orders = [None] * dist.get_world_size(group) - dist.all_gather_object(orders, tuple(optimizer.key_names), group=group) - if any(order != orders[0] for order in orders[1:]): - raise RuntimeError(f"MultiOptimizer leaf order differs across checkpoint ranks: {orders}") - - def capture_multi_layout(model, optimizer, scheduler_snapshot, scheduler_object, *, group=None, frozen_digest=None): """Validate the live save layout once, before creating checkpoint files.""" encoded, error = None, None @@ -467,7 +457,6 @@ def capture_multi_layout(model, optimizer, scheduler_snapshot, scheduler_object, except BaseException as exc: error = exc raise_if_any_rank_failed(error, "capturing MultiOptimizer checkpoint layout", group=group) - _validate_multi_leaf_order(optimizer, group=group) return encoded @@ -517,7 +506,6 @@ def coordinated_multi_layout( except BaseException as exc: error = exc raise_if_any_rank_failed(error, "inspecting MultiOptimizer layout", group=group) - _validate_multi_leaf_order(optimizer, group=group) if dist.is_initialized(): world = dist.get_world_size(group) digests = [None] * world diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index d36c76ebf7..579e8c41a6 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -579,10 +579,6 @@ def _sparse_state_dict( by a scheduler even when no parameter has received a gradient. """ merged: Dict[str, Any] = {} - if coordinate: - from ..checkpoint.optimizer_resume_layout import _validate_multi_leaf_order - - _validate_multi_leaf_order(self, group=group) for name in self.key_names: error, sd = None, None try: @@ -620,9 +616,6 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None: ) def _dcp_load_state_dict(self, state_dict: Dict[str, Any], *, group=None) -> None: - from ..checkpoint.optimizer_resume_layout import _validate_multi_leaf_order - - _validate_multi_leaf_order(self, group=group) for name in self.key_names: error = None try: From 8424b853ad78701c1cd89b02b4f29f252cfdb651 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 12:46:05 +0800 Subject: [PATCH 093/120] [checkpoint] support deterministic scheduler formatting --- docs/usage/checkpoint.md | 6 +- .../test_multi_optimizer_proven_resume.py | 113 +++++++++++++++++- veomni/checkpoint/optimizer_resume_layout.py | 6 - 3 files changed, 111 insertions(+), 14 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index 5fa7a8cb5f..361409cf13 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -166,9 +166,9 @@ MultiOptimizer checkpoint identity supports VeOmni's built-in constant, linear and cosine schedules. Custom `LambdaLR` functions must use immutable captures and the supported pure builtins or standard `math` members; global state, dynamic lookup, mutable closures and unrecognized attribute access are rejected -at save and resume. Set construction, frozenset constants and f-string object -formatting are also rejected because their results can differ between processes. -Custom scheduler authors must ensure that the schedule is deterministic across +at save and resume. Set construction and frozenset constants are also rejected +because their iteration order can differ between processes. Custom scheduler +authors must ensure that formatting and the overall schedule are deterministic across processes and recovery: the identity check does not prove arbitrary Python functions deterministic. Older custom scheduler checkpoints that cannot meet this contract require a verified migration or weights-only restart. Identity checks diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 03ec1a6004..2b217b39ee 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -2068,7 +2068,7 @@ def nested(value): assert len(_scheduler_functions(SimpleNamespace(lr_lambdas=[schedule]))) == 1 -@pytest.mark.parametrize("kind", ["set", "frozenset_constant", "object_format"]) +@pytest.mark.parametrize("kind", ["set", "frozenset_constant"]) def test_scheduler_recipe_rejects_process_dependent_computation(kind): from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions @@ -2078,14 +2078,35 @@ def unordered(step): def unordered_constant(step): return step in {1, 2, 3} - def object_format(step): - return int(f"{(lambda: 0)}"[-5], 16) / 20 - - schedule = {"set": unordered, "frozenset_constant": unordered_constant, "object_format": object_format}[kind] + schedule = {"set": unordered, "frozenset_constant": unordered_constant}[kind] with pytest.raises(RuntimeError, match="unsupported.*state"): _scheduler_functions(SimpleNamespace(lr_lambdas=[schedule])) +def test_scheduler_recipe_accepts_deterministic_format_and_records_its_identity(): + from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions + + def build_schedule(two_decimals): + if two_decimals: + + def schedule(step): + return float(f"{step / 300:.2f}") + else: + + def schedule(step): + return float(f"{step / 300:.3f}") + + return schedule + + two_decimals = build_schedule(True) + three_decimals = build_schedule(False) + assert two_decimals.__qualname__ == three_decimals.__qualname__ + two = _scheduler_functions(SimpleNamespace(lr_lambdas=[two_decimals])) + three = _scheduler_functions(SimpleNamespace(lr_lambdas=[three_decimals])) + assert two_decimals(13) != three_decimals(13) + assert two != three + + def test_scheduler_recipe_rejects_tuple_subclass_closure(): from veomni.checkpoint.optimizer_resume_layout import _scheduler_functions @@ -2212,6 +2233,88 @@ def test_multi_checkpoint_accepts_explicit_no_scheduler(tmp_path): torch.testing.assert_close(target_optimizer.state_dict(), optimizer.state_dict(), rtol=0, atol=0) +def test_deterministic_formatted_scheduler_recovers_next_update(tmp_path): + import weakref + + from veomni.optim.lr_scheduler import _CHECKPOINT_SCHEDULER_BINDINGS, MultiLRScheduler + + def build_formatted_scheduler(optimizer, *, precision=4): + schedules = {} + for name, leaf in optimizer.optimizers_dict.items(): + if precision == 4: + + def lr_lambda(step): + return float(f"{1 / (step + 1):.4f}") + else: + + def lr_lambda(step): + return float(f"{1 / (step + 1):.3f}") + + schedule = torch.optim.lr_scheduler.LambdaLR(leaf, lr_lambda) + _CHECKPOINT_SCHEDULER_BINDINGS[leaf] = weakref.ref(schedule) + schedules[name] = schedule + return MultiLRScheduler(schedules) + + torch.manual_seed(12) + model, optimizer = _model_and_optimizer() + scheduler = build_formatted_scheduler(optimizer) + for parameter in model.parameters(): + parameter.grad = torch.ones_like(parameter) + optimizer.step() + optimizer.zero_grad() + scheduler.step() + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + { + "model": model, + "optimizer": optimizer, + "extra_state": {"lr_scheduler": scheduler.state_dict()}, + "_veomni_scheduler_object": scheduler, + }, + parallel_state=parallel, + ) + changed, changed_optimizer = _model_and_optimizer() + changed_scheduler = build_formatted_scheduler(changed_optimizer, precision=3) + before = [parameter.detach().clone() for parameter in changed.parameters()] + with pytest.raises(RuntimeError, match="(?i)(scheduler|identity|layout)"): + DistributedCheckpointer.load( + str(tmp_path), + { + "model": changed, + "optimizer": changed_optimizer, + "extra_state": {}, + "_veomni_scheduler_object": changed_scheduler, + }, + parallel_state=parallel, + ) + for parameter, original in zip(changed.parameters(), before, strict=True): + torch.testing.assert_close(parameter, original, rtol=0, atol=0) + assert all(not leaf.state for leaf in changed_optimizer.optimizers_dict.values()) + + target, target_optimizer = _model_and_optimizer() + target_scheduler = build_formatted_scheduler(target_optimizer) + restored = { + "model": target, + "optimizer": target_optimizer, + "extra_state": {}, + "_veomni_scheduler_object": target_scheduler, + } + DistributedCheckpointer.load(str(tmp_path), restored, parallel_state=parallel) + target_scheduler.load_state_dict(restored["extra_state"]["lr_scheduler"]) + for source_parameter, target_parameter in zip(model.parameters(), target.parameters(), strict=True): + source_parameter.grad = torch.full_like(source_parameter, 0.25) + target_parameter.grad = torch.full_like(target_parameter, 0.25) + optimizer.step() + target_optimizer.step() + scheduler.step() + target_scheduler.step() + assert scheduler.state_dict() == target_scheduler.state_dict() + for source_parameter, target_parameter in zip(model.parameters(), target.parameters(), strict=True): + torch.testing.assert_close(source_parameter, target_parameter, rtol=0, atol=0) + torch.testing.assert_close(optimizer.state_dict(), target_optimizer.state_dict(), rtol=0, atol=0) + + def test_direct_multi_checkpoint_applies_returned_scheduler_state(tmp_path): model, optimizer = _model_and_optimizer() scheduler = build_lr_scheduler(optimizer, train_steps=10, lr_decay_style="linear") diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index cffa48d92c..ecb3b060bb 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -112,12 +112,6 @@ def normalize(code, globals_dict): "BUILD_SET", "SET_ADD", "SET_UPDATE", - # Formatting arbitrary objects can include process-local ids. - "FORMAT_VALUE", - "FORMAT_SIMPLE", - "FORMAT_WITH_SPEC", - "CONVERT_VALUE", - "BUILD_STRING", ): raise RuntimeError("MultiOptimizer scheduler uses unsupported dynamic state") if opcode in ("LOAD_ATTR", "LOAD_METHOD"): From ec34eefe316c44a8ab2776b449d97df4a1bb472f Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 12:52:56 +0800 Subject: [PATCH 094/120] [docs] refresh checkpoint configuration references --- docs/seed_omni/mixins/metric_meter.md | 10 +++++----- docs/seed_omni/mixins/offline_encoding.md | 12 ++++++------ docs/seed_omni/usage/hf_layout.md | 8 ++++---- docs/usage/arguments.md | 1 + 4 files changed, 16 insertions(+), 15 deletions(-) diff --git a/docs/seed_omni/mixins/metric_meter.md b/docs/seed_omni/mixins/metric_meter.md index 8c5d1edf4e..3cc5bda968 100644 --- a/docs/seed_omni/mixins/metric_meter.md +++ b/docs/seed_omni/mixins/metric_meter.md @@ -89,8 +89,8 @@ class VeOmniMixin(BaseMixin, TrainingMixin, InferenceModuleMixin, MetricMeter): | Function | Called by | Effect | |----------|-----------|--------| | [`MetricMeterMixin.metric_meter_add(method)`](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L120) | [`execute_train_node`](../../../veomni/models/seed_omni/accelerated/utils/executor.py#L66), once per training node, right after `pre_forward` and before the endpoint runs, inside the module's `ParallelState` scope | [Appends](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L136) the lengths stashed for the running `method` to the module's per-step buffer, then [clears](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L137) the whole stash, so copies a shared hook left under other call-sites cannot leak into a later node. Over gradient accumulation it runs once per micro-batch, so the buffer holds the whole global step. | -| [`MetricMeterMixin.metric_meter_collect() -> (flops, seqlens)`](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L139) | [`OmniModelRuntime.metric_meter_collect`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L590) | Returns `(estimate_flops(buffer), buffer)` and empties the buffer for the next step. | -| [`OmniModelRuntime.metric_meter_collect() -> {name: (flops, seqlens)}`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L590) | [`OmniStepMetricsCallback.on_step_end`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L76) | Drains every module with `isinstance(omni_module, MetricMeterMixin)`. Keys are module names, so they are identical on every rank even when a rank's batch did not use a module (it reports `(0.0, [])`). | +| [`MetricMeterMixin.metric_meter_collect() -> (flops, seqlens)`](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L139) | [`OmniModelRuntime.metric_meter_collect`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L608) | Returns `(estimate_flops(buffer), buffer)` and empties the buffer for the next step. | +| [`OmniModelRuntime.metric_meter_collect() -> {name: (flops, seqlens)}`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L608) | [`OmniStepMetricsCallback.on_step_end`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L76) | Drains every module with `isinstance(omni_module, MetricMeterMixin)`. Keys are module names, so they are identical on every rank even when a rank's batch did not use a module (it reports `(0.0, [])`). | | [`OmniEnvironMeter.add(micro_batch)`](../../../veomni/utils/omni_helper.py#L123) | [`OmniStepMetricsCallback.on_step_begin`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L66), once per micro-batch | Counts samples (`len(conversation_list)`) and, for multi-source, gathers `ds_idx`. Never looks at tokens. | | [`OmniEnvironMeter.step(delta_time, global_step, module_metrics)`](../../../veomni/utils/omni_helper.py#L137) | [`OmniStepMetricsCallback.on_step_end`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L73) | [One DP all-reduce](../../../veomni/utils/omni_helper.py#L169) of FLOPs, sample count and per-module token sums, then the metrics listed below. | @@ -117,7 +117,7 @@ and stored as `trainer.environ_meter`. - [`on_step_end`](../../../veomni/trainer/omni/omni_trainer.py#L570) → [`OmniStepMetricsCallback.on_step_end`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L69) - [`delta_time = time.time() - start_time`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L72) - - [`OmniModelRuntime.metric_meter_collect()`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L590) → + - [`OmniModelRuntime.metric_meter_collect()`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L608) → per metered module [`(estimate_flops(buffer), buffer)`](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L150); buffer reset. - [`OmniEnvironMeter.step(delta_time, global_step, module_metrics)`](../../../veomni/utils/omni_helper.py#L137) - [`trainer.step_env_metrics = ...`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L107), @@ -162,8 +162,8 @@ reported per module under `trace//` and there is no top-level `consume_tokens` is part of [`OmniEnvironMeter.state_dict()`](../../../veomni/utils/omni_helper.py#L111), which `GlobalStateCallback` -[saves](../../../veomni/trainer/callbacks/global_state_callback.py#L142) and -[restores](../../../veomni/trainer/callbacks/global_state_callback.py#L300), so +[saves](../../../veomni/trainer/callbacks/global_state_callback.py#L253) and +[restores](../../../veomni/trainer/callbacks/global_state_callback.py#L496), so the cumulative counters survive a resume. ## Rules for a correct implementation diff --git a/docs/seed_omni/mixins/offline_encoding.md b/docs/seed_omni/mixins/offline_encoding.md index d0ae083ae0..479c4c461d 100644 --- a/docs/seed_omni/mixins/offline_encoding.md +++ b/docs/seed_omni/mixins/offline_encoding.md @@ -7,7 +7,7 @@ it declares no method. - Mixin: [`OfflineEncodingMixin`](../../../veomni/models/seed_omni/mixins/offline_encoding_mixin.py#L20) - Run selector: [`train.training_task`](../../../veomni/arguments/omni_arguments_types.py#L702) -- Meta build: [`ModuleRuntime.reads_offline_cache`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L216) +- Meta build: [`ModuleRuntime.reads_offline_cache`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L234) ## Why cache a module's encoder @@ -87,7 +87,7 @@ so the cache looks complete and the error only shows in `offline_training`. `offline_embedding` freezes every module, so it never writes or restores module weights; encoder weights always come from each module's `model_path`. A saved step is therefore only the cursor, and its manifest -[says so](../../../veomni/trainer/callbacks/global_state_callback.py#L214), so +[says so](../../../veomni/trainer/callbacks/global_state_callback.py#L324), so `train.checkpoint.load_path: auto` finds it like any other step. Give the encoding run and the training run separate `train.checkpoint.output_dir`s: `auto` resumes the newest complete step in the directory, and @@ -120,11 +120,11 @@ This follows the DiT trainer's condition model: either the module is loaded in full, or it is built on meta. - **`offline_training`.** - [`ModuleRuntime.reads_offline_cache`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L216) + [`ModuleRuntime.reads_offline_cache`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L234) is true when the module's class mixes in `OfflineEncodingMixin`. The class is resolved from the module config with `get_model_class`, because the init device has to be chosen before the module exists. Then - [`_build_model`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L318) + [`_build_model`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L336) passes `init_device="meta"`, and the [training build stops](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L160) right after `requires_grad_(False)`: no parallel wrap, no weight load, no @@ -136,7 +136,7 @@ full, or it is built on meta. detaches cache readers before the composed wrap so they stay on meta and are not materialised with the rest. - **`offline_embedding`.** Every module, cached or not, is built and - [frozen](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L171), + [frozen](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L173), then loaded and wrapped like any other: the run trains nothing, so no module gets an optimizer or a checkpoint manager, and [`OmniTrainer`](../../../veomni/trainer/omni/omni_trainer.py#L276) allows the @@ -323,7 +323,7 @@ A missing module or method still fails before the first request runs. ## Current scope - **Cached modules are not checkpointed.** They are frozen in both cache tasks, - and [fully frozen modules have no checkpoint manager](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L631), + and [fully frozen modules have no checkpoint manager](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L642), so a checkpoint never holds a meta module. - **Every cacheable module is read from the cache.** `offline_training` builds every module that mixes in `OfflineEncodingMixin` on meta; one run cannot diff --git a/docs/seed_omni/usage/hf_layout.md b/docs/seed_omni/usage/hf_layout.md index 8d4189b2ed..8bb2d4d868 100644 --- a/docs/seed_omni/usage/hf_layout.md +++ b/docs/seed_omni/usage/hf_layout.md @@ -149,15 +149,15 @@ Training from an upstream root: (skipped for `train.training_task: offline_embedding`, which exports nothing). - [`OmniConfig.load_checkpoint_sidecars`](../../../veomni/models/seed_omni/configuration_omni.py#L426) records the `HFSource`, which each `ModuleRuntime` receives. - `ModuleRuntime._build_model` builds the module on meta and - [attaches the key converter](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L347); - [`weights_path`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L233) is the upstream root. + [attaches the key converter](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L365); + [`weights_path`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L251) is the upstream root. - The FSDP2 / DDP wrap (per module, or [composed](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L306)) loads weights: - [`load_model_weights`](../../../veomni/models/checkpoint/weights.py#L410) and the [rank-0 broadcast loader](../../../veomni/models/checkpoint/weights.py#L1208) skip keys the converter declares, without reading them; - [`HFSourceKeyConverter`](../../../veomni/models/seed_omni/utils/hf_layout.py#L331) renames the module's own keys and hands them to the module's converter. -HF export (`CheckpointCallback` → [`OmniModelRuntime.save_hf_or_lora`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L619)): +HF export (`CheckpointCallback` → [`OmniModelRuntime.save_hf_or_lora`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L631)): -- [`_save_hf_source_layout`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L632) exports LoRA modules' adapters per module, skips frozen modules, and makes sure each trained module's DCP exists; +- [`_save_hf_source_layout`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L644) exports LoRA modules' adapters per module, skips frozen modules, and makes sure each trained module's DCP exists; - [`save_hf_source_checkpoint`](../../../veomni/models/seed_omni/utils/hf_layout.py#L683) renames keys back (raising on a key with no source key or a different shape), casts to the source dtypes, fills tied duplicates, copies every other source tensor on rank 0, and writes `global_step_N/hf_ckpt/` with the source's shard files and non-weight files. Bare `OmniModel.from_pretrained` and eager inference: diff --git a/docs/usage/arguments.md b/docs/usage/arguments.md index 1f315d92b5..54e32a5ef2 100644 --- a/docs/usage/arguments.md +++ b/docs/usage/arguments.md @@ -928,6 +928,7 @@ and is not intended to be captured by `torch.compile`. | hf_save_steps | `int` | `0` | Steps between HuggingFace weight saves. `0` to disable. | | hf_save_epochs | `int` | `0` | Epochs between HuggingFace weight saves. `0` to disable. | | save_hf_weights | `bool` | `True` | Save HuggingFace-format weights to the last checkpoint directory. | +| extra_parallel_canonical_layout | `bool` | `False` | Opt in to versioned canonical EP/PLE model checkpoint layout. See [layout compatibility and migration restrictions](checkpoint.md#versioned-extraparallel-layout). | ### InferArguments From 5a5520d9a89742d3ecf869e619caef3541ed6cf5 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:32:27 +0800 Subject: [PATCH 095/120] [test] fix SeedOmni offload config fixture --- tests/seed_omni/runtime/test_module_weight_load.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/seed_omni/runtime/test_module_weight_load.py b/tests/seed_omni/runtime/test_module_weight_load.py index bbc2138f4b..a3dfb384c4 100644 --- a/tests/seed_omni/runtime/test_module_weight_load.py +++ b/tests/seed_omni/runtime/test_module_weight_load.py @@ -101,6 +101,8 @@ def test_parallelize_forwards_module_skip_decision( forward_prefetch=False, offload=False, offload_pin_memory=False, + cpu_offload_extra_parallel_names=[], + cpu_offload_exclude_module_patterns=[], low_precision_reduce_scatter_comm=False, expert_weight_fp8_all_gather=expert_weight_fp8_all_gather, expert_fsdp_shard_dim=0, From a656e438a63012e78daf9b63edfc85c5a2912bbf Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:35:35 +0800 Subject: [PATCH 096/120] [docs] fix checkpoint and offline cache anchors --- docs/design/rank_local_adamw_checkpoint_adapter.md | 2 +- docs/seed_omni/mixins/offline_encoding.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/design/rank_local_adamw_checkpoint_adapter.md b/docs/design/rank_local_adamw_checkpoint_adapter.md index e3e2417495..be85913be1 100644 --- a/docs/design/rank_local_adamw_checkpoint_adapter.md +++ b/docs/design/rank_local_adamw_checkpoint_adapter.md @@ -62,7 +62,7 @@ the saved checkpoint. Change these only with an explicitly verified migration. ### Save and load call sites -[`_local_layout`](../../veomni/checkpoint/optimizer_resume_layout.py#L177) +[`_local_layout`](../../veomni/checkpoint/optimizer_resume_layout.py#L171) records `algorithm_identity`, each parameter's `state_dtype`, and `state_is_offloaded` in the checkpoint identity; a mismatch rejects resume before applying state. During save, diff --git a/docs/seed_omni/mixins/offline_encoding.md b/docs/seed_omni/mixins/offline_encoding.md index 479c4c461d..e3b5541be2 100644 --- a/docs/seed_omni/mixins/offline_encoding.md +++ b/docs/seed_omni/mixins/offline_encoding.md @@ -353,7 +353,7 @@ A missing module or method still fails before the first request runs. [it is frozen and never wrapped, trained or saved](../../../tests/seed_omni/runtime/test_module_runtime.py#L209), and [an `offline_embedding` run freezes every module](../../../tests/seed_omni/runtime/test_module_runtime.py#L241). - [`tests/seed_omni/runtime/test_omni_model_runtime.py`](../../../tests/seed_omni/runtime/test_omni_model_runtime.py): - [under `fsdp_scope: model`, cache readers stay on meta outside the composed wrap](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L417), and + [under `fsdp_scope: model`, cache readers stay on meta outside the composed wrap](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L440), and [under a module-level input cast, a node's item edits reach the batch only from its hooks or its return](../../../tests/seed_omni/runtime/test_omni_model_runtime.py#L304). - [`tests/seed_omni/test_offline_cache_writer.py`](../../../tests/seed_omni/test_offline_cache_writer.py): [the `seedomni_cached` transform](../../../tests/seed_omni/test_offline_cache_writer.py#L33), the From a2708c3df48a22a900dc2ee89f6ff4b8175d5bf7 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 16:31:00 +0800 Subject: [PATCH 097/120] [test] retain historical checkpoint rejection coverage --- tests/checkpoints/test_expert_shard_layout.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index dd0deeb428..a1316742f5 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -466,6 +466,33 @@ def create_local_plan(self): dcp_to_torch_state_dict(tmp_path) +def test_existing_v2_unproven_expert_checkpoint_rejects_load_and_export(tmp_path): + """Historical invalid checkpoints remain unreadable after save-side checks.""" + marker = _extra_parallel_layout_key({}) + name = f"{marker}__layers.0.experts.weight" + + class HistoricalUnprovenPlanner(DefaultSavePlanner): + def create_global_plan(self, plans): + global_plans, metadata = super().create_global_plan(plans) + metadata.planner_data[_EXPERT_LAYOUTS_KEY] = { + "version": 1, + "tensors": {f"model.{name}": {"expert_order": "unproven"}}, + } + return global_plans, metadata + + dcp.save( + {"model": {marker: "ep-outer-v1", name: torch.ones(2, 2)}}, + checkpoint_id=tmp_path, + planner=HistoricalUnprovenPlanner(), + ) + target = {"model": {marker: "", name: torch.zeros(2, 2)}} + with pytest.raises(dcp.CheckpointException, match="noncanonical expert order"): + dcp.load(target, checkpoint_id=tmp_path, planner=_ModelStrictLoadPlanner(strict_model=True)) + assert target["model"][name].count_nonzero() == 0 + with pytest.raises(RuntimeError, match="noncanonical expert order"): + dcp_to_torch_state_dict(tmp_path) + + def _lora_export_worker(rank, world, shard_dim, directory, ep_dim, canonical, endpoint): from veomni.distributed.parallel_plan import SpecInfo from veomni.lora import VeOmniLoraConfig @@ -1272,6 +1299,60 @@ def test_raw_multi_optimizer_is_rejected_on_all_ranks(world): spawn_with_rendezvous(_sparse_plain_optimizer_failure_worker, args=(), nprocs=world) +def _inconsistent_optimizer_preflight_worker(rank, case, endpoint): + from veomni.checkpoint.dcp_checkpointer import OptimizerState + from veomni.distributed.parallel_plan import SpecInfo + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=4, timeout=timedelta(seconds=30)) + try: + mesh = init_device_mesh("cpu", (2, 2), mesh_dim_names=("ep_fsdp", "ep")) + parallel = SimpleNamespace( + dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh} + ) + model = torch.nn.Module() + model._fqn2spec_info = {"experts.weight": SpecInfo("ep", Shard(0), "experts.weight", mesh, 1)} + + class SparseOptimizer: + called = False + + @property + def _is_multi_optimizer(self): + if case == "type_property_error" and rank == 0: + raise ValueError("injected optimizer type lookup error") + return not (case == "wrong_type_collective" and rank == 0) + + def state_dict(self): + self.called = True + dist.barrier() + return {"state.experts.weight.exp_avg": torch.ones(4, 4)} if rank == 0 else {} + + leaf = SparseOptimizer() + state = OptimizerState(model, leaf, parallel_state=parallel) + expected_type = ValueError if case == "type_property_error" and rank == 0 else RuntimeError + expected_message = ( + "injected optimizer type lookup error" + if expected_type is ValueError + else "optimizer is not a MultiOptimizer" + if rank == 0 + else "Raw DCP MultiOptimizer" + ) + with pytest.raises(expected_type, match=expected_message): + state.state_dict() + assert not leaf.called + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 4 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu")) +@pytest.mark.parametrize("world", [4]) +@pytest.mark.parametrize("case", ["wrong_type_collective", "type_property_error"]) +def test_inconsistent_optimizer_preflight_exits_all_ranks(world, case): + spawn_with_rendezvous(_inconsistent_optimizer_preflight_worker, args=(case,), nprocs=world) + + def test_extra_parallel_rejects_non_multi_optimizer_locally(): from veomni.checkpoint.dcp_checkpointer import OptimizerState From be8195a6b847201c5383d9592090aba5fe009592 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Fri, 6 Nov 2026 16:58:23 +0800 Subject: [PATCH 098/120] [checkpoint] coordinate optimizer policy preflight --- tests/checkpoints/test_expert_shard_layout.py | 10 ++ .../test_multi_optimizer_proven_resume.py | 73 ++++++++++++++ veomni/checkpoint/dcp_checkpointer.py | 99 ++++++++++--------- 3 files changed, 135 insertions(+), 47 deletions(-) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index a1316742f5..b668db6b09 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -1300,6 +1300,7 @@ def test_raw_multi_optimizer_is_rejected_on_all_ranks(world): def _inconsistent_optimizer_preflight_worker(rank, case, endpoint): + import veomni.checkpoint.dcp_checkpointer as checkpoint_module from veomni.checkpoint.dcp_checkpointer import OptimizerState from veomni.distributed.parallel_plan import SpecInfo @@ -1328,6 +1329,14 @@ def state_dict(self): leaf = SparseOptimizer() state = OptimizerState(model, leaf, parallel_state=parallel) + coordinator = checkpoint_module.raise_if_any_rank_failed + coordinated = [] + + def record_coordination(error, what, group=None): + coordinated.append(what) + return coordinator(error, what, group=group) + + checkpoint_module.raise_if_any_rank_failed = record_coordination expected_type = ValueError if case == "type_property_error" and rank == 0 else RuntimeError expected_message = ( "injected optimizer type lookup error" @@ -1338,6 +1347,7 @@ def state_dict(self): ) with pytest.raises(expected_type, match=expected_message): state.state_dict() + assert coordinated == ["validating ExtraParallel optimizer checkpoint"] assert not leaf.called completed = torch.tensor(1) dist.all_reduce(completed) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index 2b217b39ee..b0f29fb2f8 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -98,6 +98,79 @@ def test_corrupt_optimizer_inventory_fails_on_all_ranks(): spawn_with_rendezvous(_corrupt_inventory_worker, nprocs=2) +def _asymmetric_optimizer_policy_worker(rank, checkpoint, operation, failure, endpoint): + from datetime import timedelta + from unittest.mock import patch + + import torch.distributed as dist + + import veomni.checkpoint.dcp_checkpointer as checkpoint_module + from veomni.optim.swap_adamw import SwapAdamW + + dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=6)) + try: + model = torch.nn.Linear(2, 2) + + class FaultAdamW(torch.optim.AdamW): + @property + def _is_multi_optimizer(self): + if failure == "type_property" and rank == 0: + raise OSError("injected optimizer type failure") + return False + + leaf = FaultAdamW(model.parameters()) + optimizer = leaf + if failure == "multi_mismatch" and rank == 0: + optimizer = MultiOptimizer(model, {"all": leaf}, ["all"]) + if failure == "presence_mismatch" and rank == 0: + optimizer = None + if failure == "swap_mismatch" and rank == 0: + optimizer = SwapAdamW(model.parameters(), pin_memory=False) + state = { + "model": model, + "optimizer": optimizer, + "_veomni_scheduler_object": None, + "_veomni_optimizer_identity_verified": True, + "_veomni_optimizer_checkpoint_generation": "stale", + } + original = checkpoint_module._requires_rank_local_optimizer_checkpoint + + def inspect_optimizer(value): + if failure == "rank_local_check" and rank == 0: + raise OSError("injected optimizer precision policy failure") + return original(value) + + with patch.object(checkpoint_module, "_requires_rank_local_optimizer_checkpoint", inspect_optimizer): + with pytest.raises((OSError, RuntimeError), match="injected|failed on another rank|optimizer.*differ"): + getattr(DistributedCheckpointer, operation)( + checkpoint, state, parallel_state=SimpleNamespace(dp_mode="ddp") + ) + if operation == "load": + assert "_veomni_optimizer_identity_verified" not in state + assert "_veomni_optimizer_checkpoint_generation" not in state + assert not Path(checkpoint).exists() + completed = torch.tensor(1) + dist.all_reduce(completed) + assert completed.item() == 2 + finally: + dist.destroy_process_group() + + +@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu")) +@pytest.mark.parametrize("operation", ["save", "load"]) +@pytest.mark.parametrize( + "failure", ["rank_local_check", "type_property", "multi_mismatch", "presence_mismatch", "swap_mismatch"] +) +def test_asymmetric_optimizer_policy_fails_before_checkpoint_io(tmp_path, operation, failure): + from tests.tools.launch_utils import spawn_with_rendezvous + + spawn_with_rendezvous( + _asymmetric_optimizer_policy_worker, + args=(str(tmp_path / "checkpoint"), operation, failure), + nprocs=2, + ) + + def _load_preflight_failure_worker(rank, checkpoint, failure, endpoint): from datetime import timedelta from unittest.mock import patch diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 045cb77c36..38b12adb7a 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -1468,6 +1468,34 @@ def _requires_rank_local_optimizer_checkpoint(optimizer): ) +def _checkpoint_optimizer_policy(state, *, loading, group=None): + """Agree on the optimizer path before entering any checkpoint collective.""" + optimizer, multi, swap, error = None, False, False, None + try: + from ..optim.swap_adamw import SwapAdamW + + if not isinstance(state, dict) or "model" not in state: + if loading and state is None: + raise ValueError("State dict must be provided to load a distributed checkpoint.") + raise ValueError(f"Model must be provided to {'load' if loading else 'save'} a distributed checkpoint.") + optimizer = state.get("optimizer") + multi = getattr(optimizer, "_is_multi_optimizer", False) is True + swap = type(optimizer) is SwapAdamW + if not multi and _requires_rank_local_optimizer_checkpoint(optimizer): + raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") + except BaseException as exc: + error = exc + raise_if_any_rank_failed(error, "checking checkpoint optimizer policy", group=group) + + if any_rank_failed(multi, group=group) and any_rank_failed(not multi, group=group): + raise RuntimeError("Checkpoint optimizer types differ across ranks") + if any_rank_failed(optimizer is None, group=group) and any_rank_failed(optimizer is not None, group=group): + raise RuntimeError("Checkpoint optimizer presence differs across ranks") + if any_rank_failed(swap, group=group) and any_rank_failed(not swap, group=group): + raise RuntimeError("Checkpoint optimizer offload types differ across ranks") + return optimizer, multi + + class OptimizerState(Stateful): """A wrapper around an optimizer to make it stateful. @@ -1526,44 +1554,35 @@ def __init__(self, model, optimizer, parallel_state=None, *, load: bool = False, self._load = load def state_dict(self): - if _requires_rank_local_optimizer_checkpoint(self.optimizer): - raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") - if self.should_extra_parallel_aware and self._layout_marker_key is not None: - raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") - if getattr(self.optimizer, "_is_multi_optimizer", False): - raise RuntimeError( - "Raw DCP MultiOptimizer does not prove parameter groups or scheduler identity; " - "use DistributedCheckpointer with the live scheduler." - ) if self.should_extra_parallel_aware: - logger.info_rank0( - "Getting optimizer state_dict from OptimizerState wrapper, would restore ExtraParallel dim for Experts module" - ) - if not getattr(self.optimizer, "_is_multi_optimizer", False): - raise RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") - - optim_sd_with_extra_parallel_dim, error = None, None + # Reject unsupported raw DCP on every rank before any optimizer + # can enter its own collectives, even if a local check fails. + error = None try: - vanilla_optim_sd = self.optimizer.state_dict() - optim_sd_with_extra_parallel_dim = self.get_state_dict_with_extra_parallel_dim_preprocess( - vanilla_optim_sd, "restore" + if _requires_rank_local_optimizer_checkpoint(self.optimizer): + raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") + if self._layout_marker_key is not None: + raise RuntimeError( + "Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec" + ) + if not getattr(self.optimizer, "_is_multi_optimizer", False): + raise RuntimeError("ExtraParallel is enabled but optimizer is not a MultiOptimizer instance") + raise RuntimeError( + "Raw DCP MultiOptimizer does not prove parameter groups or scheduler identity; " + "use DistributedCheckpointer with the live scheduler." ) except BaseException as exc: error = exc - raise_if_any_rank_failed( - error, "preparing ExtraParallel optimizer checkpoint state", group=self.process_group - ) - if _needs_extra_parallel_layout_marker(self.extra_parallel_fqn2spec_info): - optim_sd_with_extra_parallel_dim = _namespace_extra_parallel_layout_state( - optim_sd_with_extra_parallel_dim, self._layout_marker_key - ) - return optim_sd_with_extra_parallel_dim + raise_if_any_rank_failed(error, "validating ExtraParallel optimizer checkpoint", group=self.process_group) + raise RuntimeError("ExtraParallel optimizer checkpoint preflight unexpectedly succeeded") + if _requires_rank_local_optimizer_checkpoint(self.optimizer): + raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") if getattr(self.optimizer, "_is_multi_optimizer", False): - if self._load: - return self.optimizer.state_dict() - return self.optimizer._sparse_state_dict() - + raise RuntimeError( + "Raw DCP MultiOptimizer does not prove parameter groups or scheduler identity; " + "use DistributedCheckpointer with the live scheduler." + ) return get_optimizer_state_dict(model=self.model, optimizers=self.optimizer) def load_state_dict(self, state_dict): @@ -2375,16 +2394,10 @@ def save( return: None """ - if not isinstance(state, dict) or "model" not in state: - raise ValueError("Model must be provided to save a distributed checkpoint.") - optimizer = state.get("optimizer") - multi = getattr(optimizer, "_is_multi_optimizer", False) is True + optimizer, multi = _checkpoint_optimizer_policy(state, loading=False) # An async write may outlive a staged copy, so the modes cannot combine. if stage_dir and save_async: raise ValueError("stage_dir cannot be combined with save_async") - single_exact = not multi and _requires_rank_local_optimizer_checkpoint(optimizer) - if single_exact: - raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") _reject_unproven_extra_parallel_optimizer( state, parallel_state, trainable_only=trainable_only, allow_proven_optimizer=multi, saving=True ) @@ -2670,15 +2683,7 @@ def load( if isinstance(state, dict): state.pop("_veomni_optimizer_identity_verified", None) state.pop("_veomni_optimizer_checkpoint_generation", None) - if not isinstance(state, dict) or "model" not in state: - if state is None: - raise ValueError("State dict must be provided to load a distributed checkpoint.") - raise ValueError("Model must be provided to load a distributed checkpoint.") - optimizer = state.get("optimizer") - multi = getattr(optimizer, "_is_multi_optimizer", False) is True - single_exact = not multi and _requires_rank_local_optimizer_checkpoint(optimizer) - if single_exact: - raise RuntimeError("Exact-precision optimizers require MultiOptimizer for sparse DCP recovery") + optimizer, multi = _checkpoint_optimizer_policy(state, loading=True, group=process_group) _reject_unproven_extra_parallel_optimizer( state, parallel_state, group=process_group, trainable_only=trainable_only, allow_proven_optimizer=multi ) From 1fefa138947b1e856f1b577316bae430e54e5f29 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 05:17:13 +0800 Subject: [PATCH 099/120] fix: compose expert precision with parallel patch paths --- tests/ci/test_npu_moe_precision.py | 32 +++++++ veomni/distributed/torch_parallelize.py | 31 +++---- .../patched_modeling_qwen4_exp_gpu.diff | 88 +++++++++---------- .../patched_modeling_qwen4_exp_gpu.py | 6 +- .../patched_modeling_qwen4_exp_npu.diff | 88 +++++++++---------- .../patched_modeling_qwen4_exp_npu.py | 6 +- .../qwen4_exp_gpu_patch_gen_config.py | 6 +- 7 files changed, 140 insertions(+), 117 deletions(-) diff --git a/tests/ci/test_npu_moe_precision.py b/tests/ci/test_npu_moe_precision.py index 9abe58d103..375446088a 100644 --- a/tests/ci/test_npu_moe_precision.py +++ b/tests/ci/test_npu_moe_precision.py @@ -42,6 +42,38 @@ def test_valid_or_disabled_configuration_needs_no_process_group(): _validate_config(MixedPrecisionConfig(npu_expert_fp32_wgrad=False), mode="ddp", tp=2) +@pytest.mark.parametrize("backend", ["gpu", "npu"]) +def test_qwen4_expert_forwards_fp32_wgrad_option_only_when_enabled(backend): + import importlib + + module = importlib.import_module( + f"veomni.models.transformers.qwen4_exp.generated.patched_modeling_qwen4_exp_{backend}" + ) + config = SimpleNamespace( + num_experts=2, + hidden_size=4, + moe_intermediate_size=3, + hidden_act="silu", + ops_config=SimpleNamespace(moe_implementation="eager"), + ) + experts = module.Qwen4ExpTextExperts(config) + calls = [] + + def capture(*args, **kwargs): + calls.append(kwargs) + return args[0] + + experts.veomni_moe = capture + inputs = torch.zeros(2, 4) + indices = torch.zeros(2, 1, dtype=torch.int64) + weights = torch.ones(2, 1) + assert experts(inputs, indices, weights) is inputs + assert calls.pop() == {"num_experts": 2} + experts._veomni_fp32_expert_wgrad = True + assert experts(inputs, indices, weights) is inputs + assert calls.pop() == {"num_experts": 2, "fp32_wgrad": True} + + @pytest.mark.parametrize( "mixed_changes,parallel_changes", [ diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py index 2582f2cf22..9f995f2b01 100644 --- a/veomni/distributed/torch_parallelize.py +++ b/veomni/distributed/torch_parallelize.py @@ -485,6 +485,18 @@ def parallelize_model_fsdp2( if not parallel_state.fsdp_enabled or parallel_state.dp_mode != "fsdp2": raise ValueError("FP32 expert weight gradients require FSDP2") + expert_fp32_wgrad = mixed_precision.npu_expert_fp32_wgrad + if expert_fp32_wgrad: + if not IS_NPU_AVAILABLE or get_device_type() != "npu": + raise ValueError("FP32 expert weight gradients require an NPU") + from .npu_moe_precision import validate_qwen4_experts + + expert_ids = validate_qwen4_experts(model) + basic_modules = list(_fsdp_target_classes(model, basic_modules, True)) + else: + expert_ids = set() + wrapped_expert_ids = set() + if expert_weight_fp8_all_gather is not False: from .fsdp2.fp8_all_gather import resolve_fp8_all_gather @@ -504,18 +516,8 @@ def parallelize_model_fsdp2( elif low_precision_reduce_scatter_comm: logger.info_rank0("Parameter dtype matches reduce dtype; using the native PyTorch collective.") - expert_fp32_wgrad = mixed_precision.npu_expert_fp32_wgrad - if expert_fp32_wgrad: - if not IS_NPU_AVAILABLE or get_device_type() != "npu": - raise ValueError("FP32 expert weight gradients require an NPU") - from .npu_moe_precision import validate_qwen4_experts - - expert_ids = validate_qwen4_experts(model) - else: - expert_ids = set() - wrapped_expert_ids = set() - - target_classes = _fsdp_target_classes(model, basic_modules, expert_fp32_wgrad) + model_no_split_modules = getattr(model, "_no_split_modules", None) or [] + target_classes = set(model_no_split_modules) | set(basic_modules or []) # Make a list of tuples that contains target classes' name and module # Note that all target classes should include all ExtraParallel modules. @@ -906,9 +908,6 @@ def configure_fp32_expert(expert): layer_mod._fsdp_modules.append(layer_mod) logger.info_rank0(f"{layer_fqn=}, {layer_mod._fsdp_modules=}") - if expert_fp32_wgrad and wrapped_expert_ids != expert_ids: - raise RuntimeError("Not every Qwen4-Exp expert was wrapped for FP32 weight gradients") - # Shard root WITHOUT explicit `reshard_after_forward` so FSDP2's # auto-no-reshard for the root kicks in (`_fsdp_state.py:_lazy_init`). # This keeps the root's params (lm_head + embeddings + final norm) @@ -918,6 +917,8 @@ def configure_fp32_expert(expert): # `setStorage … storage of size 0` when the saved reference points # to a freed buffer. Decoder layers reshard normally (their calls # above pass `reshard_after_forward` explicitly). + if expert_fp32_wgrad and wrapped_expert_ids != expert_ids: + raise RuntimeError("Not every Qwen4-Exp expert was wrapped for FP32 weight gradients") root_fsdp_kwargs = {k: v for k, v in fsdp_kwargs.items() if k != "reshard_after_forward"} fully_shard(model, **root_fsdp_kwargs) # Persistent ExtraParallel parameters are ignored by every FSDP group and diff --git a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.diff b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.diff index 749c4abf2a..5df1185094 100644 --- a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.diff +++ b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.diff @@ -1370,12 +1370,12 @@ self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) self.act_fn = ACT2FN[config.hidden_act] -+ self.use_swiglu_mlp = config.hidden_act in {"silu", "swish"} ++ self._veomni_fp32_expert_wgrad, self.use_swiglu_mlp = False, config.hidden_act in {"silu", "swish"} + self.veomni_moe = VeomniOp("moe_experts", "standard", config.ops_config.moe_implementation) def forward( self, -@@ -874,25 +1434,31 @@ +@@ -874,25 +1434,29 @@ top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: @@ -1410,7 +1410,6 @@ + self.num_experts, + ) + unused = self.gate_up_proj.new_empty(0) -+ expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} + return self.veomni_moe( + hidden_states, + top_k_weights, @@ -1419,14 +1418,13 @@ + unused, + self.down_proj, + self.gate_up_proj, -+ num_experts=self.num_experts, -+ **expert_options, ++ **dict(num_experts=self.num_experts, **({"fp32_wgrad": True} if self._veomni_fp32_expert_wgrad else {})), + ) + # --- Patch.1 --- class Qwen4ExpTextTopKRouter(nn.Module): -@@ -916,6 +1482,12 @@ +@@ -916,6 +1480,12 @@ return router_logits, router_scores, router_indices @@ -1439,7 +1437,7 @@ class Qwen4ExpTextSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() -@@ -924,11 +1496,33 @@ +@@ -924,11 +1494,33 @@ self.shared_expert = Qwen4ExpTextMLP(config, intermediate_size=config.shared_expert_intermediate_size) self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) @@ -1475,7 +1473,7 @@ expert_output = self.experts(hidden_states_reshaped, selected_experts, routing_weights) shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states_reshaped)) * shared_expert_output -@@ -938,13 +1532,21 @@ +@@ -938,13 +1530,21 @@ return expert_output @@ -1500,7 +1498,7 @@ self.input_mix_weight_down = nn.Linear(hc_hidden_size, config.hc_lowrank, bias=False) self.input_mix_weight_up = nn.Linear(config.hc_lowrank, hc_hidden_size, bias=False) self.block_inject_weight = nn.Linear(hc_hidden_size, self.hc_count, bias=False) if use_combine else None -@@ -970,9 +1572,13 @@ +@@ -970,9 +1570,13 @@ _MASK64 = (1 << 64) - 1 @@ -1514,7 +1512,7 @@ _PRIME_1 = 10007 -@@ -1015,8 +1621,14 @@ +@@ -1015,8 +1619,14 @@ return prime @@ -1530,7 +1528,7 @@ super().__init__() self.layer_idx = layer_idx self.ngram_size = config.ngram_size -@@ -1028,7 +1640,6 @@ +@@ -1028,7 +1638,6 @@ self.ngram_vocab_size_base = config.ngram_vocab_size_base head_dim_per_ngram = embedding_dim // self.ngram_heads self.seed = config.seed @@ -1538,7 +1536,7 @@ self.eos_token_id = config.eos_token_id[0] if isinstance(config.eos_token_id, list) else config.eos_token_id self.head_vocab_sizes = [] -@@ -1046,9 +1657,29 @@ +@@ -1046,9 +1655,29 @@ ) self.ngram_heads_vocab_sizes = nn.Buffer(torch.tensor(self.head_vocab_sizes, dtype=torch.long)) self.ngram_heads_offsets = nn.Buffer(torch.tensor(self.head_offsets, dtype=torch.long)) @@ -1571,7 +1569,7 @@ def _shift_right_ignore_eos(self, token_ids: torch.Tensor, shift: int) -> torch.Tensor: if shift == 0: -@@ -1066,19 +1697,183 @@ +@@ -1066,19 +1695,183 @@ valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0) return torch.where(valid, shifted, token_ids.new_full((), self.eos_token_id)) @@ -1762,7 +1760,7 @@ if ( not past_key_values.has_previous_state(self.layer_idx, state_idx=2) and input_ids.shape[1] < self.context_len -@@ -1090,28 +1885,48 @@ +@@ -1090,28 +1883,48 @@ input_ids_to_cache, self.layer_idx, state_idx=2, conv_kernel_size=self.context_len ) @@ -1819,7 +1817,7 @@ class Qwen4ExpTextPLELayer(nn.Module): -@@ -1122,8 +1937,8 @@ +@@ -1122,8 +1935,8 @@ The returned tensor has shape `(batch_size, sequence_length, hc_count * hidden_size)`. """ @@ -1830,7 +1828,7 @@ self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.hc_count = config.hc_count -@@ -1135,9 +1950,15 @@ +@@ -1135,9 +1948,15 @@ self.short_conv_state_len = (conv_kernel_size - 1) * conv_dilation self.key_proj = nn.Linear(ple_embed_dim, hc_hidden_size, bias=False) self.value_proj = nn.Linear(ple_embed_dim, self.hidden_size, bias=False) @@ -1849,7 +1847,7 @@ self.conv1d = nn.Conv1d( hc_hidden_size, hc_hidden_size, -@@ -1147,24 +1968,91 @@ +@@ -1147,24 +1966,91 @@ bias=False, ) @@ -1957,7 +1955,7 @@ def forward( self, -@@ -1172,8 +2060,25 @@ +@@ -1172,8 +2058,25 @@ input_ids: torch.Tensor, past_key_values: Cache | None, conv_mask: torch.Tensor | None = None, @@ -1984,7 +1982,7 @@ key_normed = self.norm_key(self.key_proj(embeddings)).unflatten(-1, (self.hc_count, self.hidden_size)) value = self.value_proj(embeddings) query_normed = self.norm_query(hidden_states).unflatten(-1, (self.hc_count, self.hidden_size)) -@@ -1185,8 +2090,21 @@ +@@ -1185,8 +2088,21 @@ if conv_mask is not None: gated_value = apply_mask_to_padding_states(gated_value, conv_mask) gated_value_normed = apply_mask_to_padding_states(gated_value_normed, conv_mask) @@ -2007,7 +2005,7 @@ class Qwen4ExpTextDecoderLayer(GradientCheckpointingLayer): -@@ -1204,6 +2122,15 @@ +@@ -1204,6 +2120,15 @@ self.attn_hyper_connection = Qwen4ExpTextGatedResidual(config) self.mlp_hyper_connection = Qwen4ExpTextGatedResidual(config) @@ -2023,7 +2021,7 @@ def forward( self, hidden_states: torch.Tensor, -@@ -1212,17 +2139,42 @@ +@@ -1212,17 +2137,42 @@ conv_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, ple_input_ids: torch.LongTensor | None = None, @@ -2069,7 +2067,7 @@ ) else: hidden_states, _ = self.self_attn( -@@ -1230,6 +2182,9 @@ +@@ -1230,6 +2180,9 @@ position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, @@ -2079,7 +2077,7 @@ **kwargs, ) -@@ -1308,6 +2263,12 @@ +@@ -1308,6 +2261,12 @@ rope_deltas: torch.LongTensor | None = None router_logits: tuple[torch.FloatTensor] | None = None @@ -2092,7 +2090,7 @@ @auto_docstring -@@ -1352,20 +2313,73 @@ +@@ -1352,20 +2311,73 @@ """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -2168,7 +2166,7 @@ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1) -@@ -1381,31 +2395,45 @@ +@@ -1381,31 +2393,45 @@ else: text_position_ids = None @@ -2220,7 +2218,7 @@ # CODEPATH: @ArthurZucker fix flagging for no reason here if self.config.ple_layer_ids and conv_mask is not None: eos_token_id = self.config.eos_token_id -@@ -1416,16 +2444,40 @@ +@@ -1416,16 +2442,40 @@ position_embeddings = self.rotary_emb(hidden_states, position_ids) hidden_states = hidden_states.repeat(1, 1, self.config.hc_count) @@ -2263,7 +2261,7 @@ hidden_states = self.hyper_connection_mixer(hidden_states) -@@ -1435,12 +2487,6 @@ +@@ -1435,12 +2485,6 @@ ) def reverse_embedding(self, inputs_embeds: torch.Tensor) -> torch.Tensor: @@ -2276,7 +2274,7 @@ with torch.no_grad(): input_ids = ( (inputs_embeds[:, :, None, :] == self.embed_tokens.weight[None, None, :, :]).all(dim=3).nonzero()[:, 2] -@@ -1448,11 +2494,15 @@ +@@ -1448,11 +2492,15 @@ try: input_ids = input_ids.view(inputs_embeds.shape[:2]) except RuntimeError: @@ -2296,7 +2294,7 @@ return input_ids -@@ -1718,21 +2768,14 @@ +@@ -1718,21 +2766,14 @@ return x @@ -2323,7 +2321,7 @@ def __init__(self, config: Qwen4ExpVisionConfig) -> None: super().__init__() self.dim = config.hidden_size -@@ -1745,13 +2788,16 @@ +@@ -1745,13 +2786,16 @@ self.config = config self.attention_dropout = 0.0 self.is_causal = False @@ -2341,7 +2339,7 @@ **kwargs, ) -> torch.Tensor: seq_length = hidden_states.shape[0] -@@ -1759,20 +2805,29 @@ +@@ -1759,20 +2803,29 @@ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) ) cos, sin = position_embeddings @@ -2380,7 +2378,7 @@ self, query_states, key_states, -@@ -1787,28 +2842,6 @@ +@@ -1787,28 +2840,6 @@ is_causal=False, **kwargs, ) @@ -2409,7 +2407,7 @@ attn_output = attn_output.reshape(seq_length, -1).contiguous() attn_output = self.proj(attn_output) -@@ -1845,6 +2878,12 @@ +@@ -1845,6 +2876,12 @@ return hidden_states @@ -2422,7 +2420,7 @@ @auto_docstring class Qwen4ExpVisionModel(Qwen4ExpPreTrainedModel): config: Qwen4ExpVisionConfig -@@ -1884,30 +2923,127 @@ +@@ -1884,30 +2921,127 @@ self.post_init() @@ -2572,7 +2570,7 @@ @merge_with_config_defaults @capture_outputs -@@ -1922,33 +3058,145 @@ +@@ -1922,33 +3056,145 @@ Returns: `torch.Tensor`: hidden_states. """ @@ -2733,7 +2731,7 @@ position_embeddings=position_embeddings, **kwargs, ) -@@ -1959,6 +3207,39 @@ +@@ -1959,6 +3205,39 @@ last_hidden_state=hidden_states, pooler_output=merged_hidden_states, ) @@ -2773,7 +2771,7 @@ @auto_docstring -@@ -1970,11 +3251,12 @@ +@@ -1970,11 +3249,12 @@ def __init__(self, config): super().__init__(config) @@ -2791,7 +2789,7 @@ self.post_init() def get_vision_position_ids( -@@ -2260,53 +3542,226 @@ +@@ -2260,53 +3540,226 @@ mm_token_type_ids: torch.IntTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen4ExpModelOutputWithPast: @@ -3049,7 +3047,7 @@ outputs = self.language_model( input_ids=None, position_ids=position_ids, -@@ -2316,11 +3771,7 @@ +@@ -2316,11 +3769,7 @@ ple_input_ids=ple_input_ids, **kwargs, ) @@ -3062,7 +3060,7 @@ @auto_docstring -@@ -2337,6 +3788,27 @@ +@@ -2337,6 +3786,27 @@ aux_loss: torch.FloatTensor | None = None @@ -3090,7 +3088,7 @@ @auto_docstring class Qwen4ExpForConditionalGeneration(Qwen4ExpPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} -@@ -2346,11 +3818,23 @@ +@@ -2346,11 +3816,23 @@ _fsdp_plan = {"lm_head": "keep_full_weight"} @@ -3115,7 +3113,7 @@ self.post_init() @auto_docstring -@@ -2396,47 +3880,10 @@ +@@ -2396,47 +3878,10 @@ mm_token_type_ids: torch.IntTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -3167,7 +3165,7 @@ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, -@@ -2448,33 +3895,64 @@ +@@ -2448,33 +3893,64 @@ attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, @@ -3247,7 +3245,7 @@ loss=loss, aux_loss=aux_loss, logits=logits, -@@ -2483,6 +3961,7 @@ +@@ -2483,6 +3959,7 @@ attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, router_logits=outputs.router_logits, @@ -3255,7 +3253,7 @@ ) def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): -@@ -2696,6 +4175,32 @@ +@@ -2696,6 +4173,32 @@ } return causal_mask_mapping diff --git a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py index ed297ddd3e..2968109c9b 100644 --- a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py +++ b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py @@ -1425,7 +1425,7 @@ def __init__(self, config): self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) self.act_fn = ACT2FN[config.hidden_act] - self.use_swiglu_mlp = config.hidden_act in {"silu", "swish"} + self._veomni_fp32_expert_wgrad, self.use_swiglu_mlp = False, config.hidden_act in {"silu", "swish"} self.veomni_moe = VeomniOp("moe_experts", "standard", config.ops_config.moe_implementation) def forward( @@ -1446,7 +1446,6 @@ def forward( self.num_experts, ) unused = self.gate_up_proj.new_empty(0) - expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} return self.veomni_moe( hidden_states, top_k_weights, @@ -1455,8 +1454,7 @@ def forward( unused, self.down_proj, self.gate_up_proj, - num_experts=self.num_experts, - **expert_options, + **dict(num_experts=self.num_experts, **({"fp32_wgrad": True} if self._veomni_fp32_expert_wgrad else {})), ) # --- Patch.1 --- diff --git a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.diff b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.diff index 2532c559bf..6cf6b4329d 100644 --- a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.diff +++ b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.diff @@ -1370,12 +1370,12 @@ self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) self.act_fn = ACT2FN[config.hidden_act] -+ self.use_swiglu_mlp = config.hidden_act in {"silu", "swish"} ++ self._veomni_fp32_expert_wgrad, self.use_swiglu_mlp = False, config.hidden_act in {"silu", "swish"} + self.veomni_moe = VeomniOp("moe_experts", "standard", config.ops_config.moe_implementation) def forward( self, -@@ -874,25 +1434,31 @@ +@@ -874,25 +1434,29 @@ top_k_index: torch.Tensor, top_k_weights: torch.Tensor, ) -> torch.Tensor: @@ -1410,7 +1410,6 @@ + self.num_experts, + ) + unused = self.gate_up_proj.new_empty(0) -+ expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} + return self.veomni_moe( + hidden_states, + top_k_weights, @@ -1419,14 +1418,13 @@ + unused, + self.down_proj, + self.gate_up_proj, -+ num_experts=self.num_experts, -+ **expert_options, ++ **dict(num_experts=self.num_experts, **({"fp32_wgrad": True} if self._veomni_fp32_expert_wgrad else {})), + ) + # --- Patch.1 --- class Qwen4ExpTextTopKRouter(nn.Module): -@@ -916,6 +1482,12 @@ +@@ -916,6 +1480,12 @@ return router_logits, router_scores, router_indices @@ -1439,7 +1437,7 @@ class Qwen4ExpTextSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() -@@ -924,11 +1496,33 @@ +@@ -924,11 +1494,33 @@ self.shared_expert = Qwen4ExpTextMLP(config, intermediate_size=config.shared_expert_intermediate_size) self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) @@ -1475,7 +1473,7 @@ expert_output = self.experts(hidden_states_reshaped, selected_experts, routing_weights) shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states_reshaped)) * shared_expert_output -@@ -938,13 +1532,21 @@ +@@ -938,13 +1530,21 @@ return expert_output @@ -1500,7 +1498,7 @@ self.input_mix_weight_down = nn.Linear(hc_hidden_size, config.hc_lowrank, bias=False) self.input_mix_weight_up = nn.Linear(config.hc_lowrank, hc_hidden_size, bias=False) self.block_inject_weight = nn.Linear(hc_hidden_size, self.hc_count, bias=False) if use_combine else None -@@ -970,9 +1572,13 @@ +@@ -970,9 +1570,13 @@ _MASK64 = (1 << 64) - 1 @@ -1514,7 +1512,7 @@ _PRIME_1 = 10007 -@@ -1015,8 +1621,14 @@ +@@ -1015,8 +1619,14 @@ return prime @@ -1530,7 +1528,7 @@ super().__init__() self.layer_idx = layer_idx self.ngram_size = config.ngram_size -@@ -1028,7 +1640,6 @@ +@@ -1028,7 +1638,6 @@ self.ngram_vocab_size_base = config.ngram_vocab_size_base head_dim_per_ngram = embedding_dim // self.ngram_heads self.seed = config.seed @@ -1538,7 +1536,7 @@ self.eos_token_id = config.eos_token_id[0] if isinstance(config.eos_token_id, list) else config.eos_token_id self.head_vocab_sizes = [] -@@ -1046,9 +1657,29 @@ +@@ -1046,9 +1655,29 @@ ) self.ngram_heads_vocab_sizes = nn.Buffer(torch.tensor(self.head_vocab_sizes, dtype=torch.long)) self.ngram_heads_offsets = nn.Buffer(torch.tensor(self.head_offsets, dtype=torch.long)) @@ -1571,7 +1569,7 @@ def _shift_right_ignore_eos(self, token_ids: torch.Tensor, shift: int) -> torch.Tensor: if shift == 0: -@@ -1066,19 +1697,183 @@ +@@ -1066,19 +1695,183 @@ valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0) return torch.where(valid, shifted, token_ids.new_full((), self.eos_token_id)) @@ -1762,7 +1760,7 @@ if ( not past_key_values.has_previous_state(self.layer_idx, state_idx=2) and input_ids.shape[1] < self.context_len -@@ -1090,28 +1885,48 @@ +@@ -1090,28 +1883,48 @@ input_ids_to_cache, self.layer_idx, state_idx=2, conv_kernel_size=self.context_len ) @@ -1819,7 +1817,7 @@ class Qwen4ExpTextPLELayer(nn.Module): -@@ -1122,8 +1937,8 @@ +@@ -1122,8 +1935,8 @@ The returned tensor has shape `(batch_size, sequence_length, hc_count * hidden_size)`. """ @@ -1830,7 +1828,7 @@ self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.hc_count = config.hc_count -@@ -1135,9 +1950,15 @@ +@@ -1135,9 +1948,15 @@ self.short_conv_state_len = (conv_kernel_size - 1) * conv_dilation self.key_proj = nn.Linear(ple_embed_dim, hc_hidden_size, bias=False) self.value_proj = nn.Linear(ple_embed_dim, self.hidden_size, bias=False) @@ -1849,7 +1847,7 @@ self.conv1d = nn.Conv1d( hc_hidden_size, hc_hidden_size, -@@ -1147,24 +1968,91 @@ +@@ -1147,24 +1966,91 @@ bias=False, ) @@ -1957,7 +1955,7 @@ def forward( self, -@@ -1172,8 +2060,25 @@ +@@ -1172,8 +2058,25 @@ input_ids: torch.Tensor, past_key_values: Cache | None, conv_mask: torch.Tensor | None = None, @@ -1984,7 +1982,7 @@ key_normed = self.norm_key(self.key_proj(embeddings)).unflatten(-1, (self.hc_count, self.hidden_size)) value = self.value_proj(embeddings) query_normed = self.norm_query(hidden_states).unflatten(-1, (self.hc_count, self.hidden_size)) -@@ -1185,8 +2090,21 @@ +@@ -1185,8 +2088,21 @@ if conv_mask is not None: gated_value = apply_mask_to_padding_states(gated_value, conv_mask) gated_value_normed = apply_mask_to_padding_states(gated_value_normed, conv_mask) @@ -2007,7 +2005,7 @@ class Qwen4ExpTextDecoderLayer(GradientCheckpointingLayer): -@@ -1204,6 +2122,15 @@ +@@ -1204,6 +2120,15 @@ self.attn_hyper_connection = Qwen4ExpTextGatedResidual(config) self.mlp_hyper_connection = Qwen4ExpTextGatedResidual(config) @@ -2023,7 +2021,7 @@ def forward( self, hidden_states: torch.Tensor, -@@ -1212,17 +2139,42 @@ +@@ -1212,17 +2137,42 @@ conv_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, ple_input_ids: torch.LongTensor | None = None, @@ -2069,7 +2067,7 @@ ) else: hidden_states, _ = self.self_attn( -@@ -1230,6 +2182,9 @@ +@@ -1230,6 +2180,9 @@ position_embeddings, attention_mask=attention_mask, past_key_values=past_key_values, @@ -2079,7 +2077,7 @@ **kwargs, ) -@@ -1308,6 +2263,12 @@ +@@ -1308,6 +2261,12 @@ rope_deltas: torch.LongTensor | None = None router_logits: tuple[torch.FloatTensor] | None = None @@ -2092,7 +2090,7 @@ @auto_docstring -@@ -1352,20 +2313,73 @@ +@@ -1352,20 +2311,73 @@ """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -2168,7 +2166,7 @@ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1) -@@ -1381,31 +2395,45 @@ +@@ -1381,31 +2393,45 @@ else: text_position_ids = None @@ -2220,7 +2218,7 @@ # CODEPATH: @ArthurZucker fix flagging for no reason here if self.config.ple_layer_ids and conv_mask is not None: eos_token_id = self.config.eos_token_id -@@ -1416,16 +2444,40 @@ +@@ -1416,16 +2442,40 @@ position_embeddings = self.rotary_emb(hidden_states, position_ids) hidden_states = hidden_states.repeat(1, 1, self.config.hc_count) @@ -2263,7 +2261,7 @@ hidden_states = self.hyper_connection_mixer(hidden_states) -@@ -1435,12 +2487,6 @@ +@@ -1435,12 +2485,6 @@ ) def reverse_embedding(self, inputs_embeds: torch.Tensor) -> torch.Tensor: @@ -2276,7 +2274,7 @@ with torch.no_grad(): input_ids = ( (inputs_embeds[:, :, None, :] == self.embed_tokens.weight[None, None, :, :]).all(dim=3).nonzero()[:, 2] -@@ -1448,11 +2494,15 @@ +@@ -1448,11 +2492,15 @@ try: input_ids = input_ids.view(inputs_embeds.shape[:2]) except RuntimeError: @@ -2296,7 +2294,7 @@ return input_ids -@@ -1718,21 +2768,14 @@ +@@ -1718,21 +2766,14 @@ return x @@ -2323,7 +2321,7 @@ def __init__(self, config: Qwen4ExpVisionConfig) -> None: super().__init__() self.dim = config.hidden_size -@@ -1745,13 +2788,16 @@ +@@ -1745,13 +2786,16 @@ self.config = config self.attention_dropout = 0.0 self.is_causal = False @@ -2341,7 +2339,7 @@ **kwargs, ) -> torch.Tensor: seq_length = hidden_states.shape[0] -@@ -1759,20 +2805,29 @@ +@@ -1759,20 +2803,29 @@ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) ) cos, sin = position_embeddings @@ -2380,7 +2378,7 @@ self, query_states, key_states, -@@ -1787,28 +2842,6 @@ +@@ -1787,28 +2840,6 @@ is_causal=False, **kwargs, ) @@ -2409,7 +2407,7 @@ attn_output = attn_output.reshape(seq_length, -1).contiguous() attn_output = self.proj(attn_output) -@@ -1845,6 +2878,12 @@ +@@ -1845,6 +2876,12 @@ return hidden_states @@ -2422,7 +2420,7 @@ @auto_docstring class Qwen4ExpVisionModel(Qwen4ExpPreTrainedModel): config: Qwen4ExpVisionConfig -@@ -1884,30 +2923,127 @@ +@@ -1884,30 +2921,127 @@ self.post_init() @@ -2572,7 +2570,7 @@ @merge_with_config_defaults @capture_outputs -@@ -1922,33 +3058,145 @@ +@@ -1922,33 +3056,145 @@ Returns: `torch.Tensor`: hidden_states. """ @@ -2733,7 +2731,7 @@ position_embeddings=position_embeddings, **kwargs, ) -@@ -1959,6 +3207,39 @@ +@@ -1959,6 +3205,39 @@ last_hidden_state=hidden_states, pooler_output=merged_hidden_states, ) @@ -2773,7 +2771,7 @@ @auto_docstring -@@ -1970,11 +3251,12 @@ +@@ -1970,11 +3249,12 @@ def __init__(self, config): super().__init__(config) @@ -2791,7 +2789,7 @@ self.post_init() def get_vision_position_ids( -@@ -2260,53 +3542,226 @@ +@@ -2260,53 +3540,226 @@ mm_token_type_ids: torch.IntTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | Qwen4ExpModelOutputWithPast: @@ -3049,7 +3047,7 @@ outputs = self.language_model( input_ids=None, position_ids=position_ids, -@@ -2316,11 +3771,7 @@ +@@ -2316,11 +3769,7 @@ ple_input_ids=ple_input_ids, **kwargs, ) @@ -3062,7 +3060,7 @@ @auto_docstring -@@ -2337,6 +3788,27 @@ +@@ -2337,6 +3786,27 @@ aux_loss: torch.FloatTensor | None = None @@ -3090,7 +3088,7 @@ @auto_docstring class Qwen4ExpForConditionalGeneration(Qwen4ExpPreTrainedModel, GenerationMixin): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} -@@ -2346,11 +3818,23 @@ +@@ -2346,11 +3816,23 @@ _fsdp_plan = {"lm_head": "keep_full_weight"} @@ -3115,7 +3113,7 @@ self.post_init() @auto_docstring -@@ -2396,47 +3880,10 @@ +@@ -2396,47 +3878,10 @@ mm_token_type_ids: torch.IntTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -3167,7 +3165,7 @@ outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, -@@ -2448,33 +3895,64 @@ +@@ -2448,33 +3893,64 @@ attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, @@ -3247,7 +3245,7 @@ loss=loss, aux_loss=aux_loss, logits=logits, -@@ -2483,6 +3961,7 @@ +@@ -2483,6 +3959,7 @@ attentions=outputs.attentions, rope_deltas=outputs.rope_deltas, router_logits=outputs.router_logits, @@ -3255,7 +3253,7 @@ ) def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs): -@@ -2696,6 +4175,32 @@ +@@ -2696,6 +4173,32 @@ } return causal_mask_mapping diff --git a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py index 5f41c7842c..efa283ba04 100644 --- a/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py +++ b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py @@ -1425,7 +1425,7 @@ def __init__(self, config): self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) self.act_fn = ACT2FN[config.hidden_act] - self.use_swiglu_mlp = config.hidden_act in {"silu", "swish"} + self._veomni_fp32_expert_wgrad, self.use_swiglu_mlp = False, config.hidden_act in {"silu", "swish"} self.veomni_moe = VeomniOp("moe_experts", "standard", config.ops_config.moe_implementation) def forward( @@ -1446,7 +1446,6 @@ def forward( self.num_experts, ) unused = self.gate_up_proj.new_empty(0) - expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} return self.veomni_moe( hidden_states, top_k_weights, @@ -1455,8 +1454,7 @@ def forward( unused, self.down_proj, self.gate_up_proj, - num_experts=self.num_experts, - **expert_options, + **dict(num_experts=self.num_experts, **({"fp32_wgrad": True} if self._veomni_fp32_expert_wgrad else {})), ) # --- Patch.1 --- diff --git a/veomni/models/transformers/qwen4_exp/qwen4_exp_gpu_patch_gen_config.py b/veomni/models/transformers/qwen4_exp/qwen4_exp_gpu_patch_gen_config.py index 0fa8140376..519eed3170 100644 --- a/veomni/models/transformers/qwen4_exp/qwen4_exp_gpu_patch_gen_config.py +++ b/veomni/models/transformers/qwen4_exp/qwen4_exp_gpu_patch_gen_config.py @@ -1054,7 +1054,7 @@ def __init__(self, config): self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) self.act_fn = ACT2FN[config.hidden_act] - self.use_swiglu_mlp = config.hidden_act in {"silu", "swish"} + self._veomni_fp32_expert_wgrad, self.use_swiglu_mlp = False, config.hidden_act in {"silu", "swish"} self.veomni_moe = VeomniOp("moe_experts", "standard", config.ops_config.moe_implementation) def forward( @@ -1075,7 +1075,6 @@ def forward( self.num_experts, ) unused = self.gate_up_proj.new_empty(0) - expert_options = {"fp32_wgrad": True} if getattr(self, "_veomni_fp32_expert_wgrad", False) else {} return self.veomni_moe( hidden_states, top_k_weights, @@ -1084,8 +1083,7 @@ def forward( unused, self.down_proj, self.gate_up_proj, - num_experts=self.num_experts, - **expert_options, + **dict(num_experts=self.num_experts, **({"fp32_wgrad": True} if self._veomni_fp32_expert_wgrad else {})), ) # --- Patch.1 --- From 759e6fbc1a0132719d6b99a0b894c40cf2f50aef Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 05:22:42 +0800 Subject: [PATCH 100/120] fix: coordinate cursor and optimizer checkpoint recovery --- tests/checkpoints/test_dcp_checkpointer.py | 5 ++- .../test_multi_optimizer_trainer_step.py | 43 ++++++++++++++++++- .../callbacks/global_state_callback.py | 41 ++++++++++++------ 3 files changed, 73 insertions(+), 16 deletions(-) diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py index 2896c2eec8..b9b95bd85f 100644 --- a/tests/checkpoints/test_dcp_checkpointer.py +++ b/tests/checkpoints/test_dcp_checkpointer.py @@ -532,9 +532,9 @@ def test_load_passes_partial_planner_to_dcp(self, mock_dcp, mock_gps): mock_gps.return_value = SimpleNamespace(dp_mode="fsdp2") from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer - model = MagicMock() + model = torch.nn.Linear(2, 2) model._fqn2spec_info = None - optimizer = MagicMock() + optimizer = torch.optim.AdamW(model.parameters()) state = {"model": model, "optimizer": optimizer, "extra_state": {}} @@ -2488,6 +2488,7 @@ def test_a_rank_that_cannot_remove_the_markers_fails_the_group(self, tmp_path): with ( patch("veomni.checkpoint.dcp_checkpointer.dist.is_initialized", return_value=True), patch("veomni.checkpoint.dcp_checkpointer.dist.get_rank", return_value=1), + patch("veomni.checkpoint.dcp_checkpointer._checkpoint_optimizer_policy", return_value=(None, False)), patch("veomni.checkpoint.dcp_checkpointer._reject_unproven_extra_parallel_optimizer"), patch("veomni.utils.dist_utils.any_rank_failed", return_value=True), patch.object(DistributedCheckpointer, "execute_save") as execute_save, diff --git a/tests/checkpoints/test_multi_optimizer_trainer_step.py b/tests/checkpoints/test_multi_optimizer_trainer_step.py index dab475dd57..0ced70675f 100644 --- a/tests/checkpoints/test_multi_optimizer_trainer_step.py +++ b/tests/checkpoints/test_multi_optimizer_trainer_step.py @@ -8,6 +8,7 @@ from unittest.mock import patch import pytest +import torch import torch.distributed as dist from scripts.ci.accelerator.launch_plan import LaunchSpec @@ -215,6 +216,23 @@ def _resume_worker(rank, tmp_path, mode, endpoint): trainer = _trainer(tmp_path, optimizer_step=step) trainer.args.train.global_rank = rank trainer.args.train.world_size = 2 + if mode in ("asymmetric_loader", "asymmetric_loader_error"): + trainer.args.train_steps = 3 + + class Loader: + def __init__(self): + self.iterations = 0 + + def load_state_dict(self, state): + assert state == {"cursor": 2} + + def __iter__(self): + self.iterations += 1 + if mode == "asymmetric_loader_error": + raise ValueError("injected iterator failure") + return iter(()) + + trainer.train_dataloader = Loader() if rank == 0 else None callback = GlobalStateCallback(trainer) state = ( None @@ -224,6 +242,7 @@ def _resume_worker(rank, tmp_path, mode, endpoint): "_veomni_optimizer_generations": { "model": "older" if mode == "trainer_generation" and rank == 1 else "generation" }, + "train_dataloader": {"cursor": 2}, } ) with ( @@ -231,9 +250,12 @@ def _resume_worker(rank, tmp_path, mode, endpoint): patch("veomni.checkpoint.legacy_v0_1_12.apply_legacy_global_state", return_value=None), patch("veomni.trainer.callbacks.global_state_callback.set_device_rng_state"), ): - if mode == "matching": + if mode in ("matching", "asymmetric_loader"): callback.load_global_state() assert trainer.state.global_step == 3 + elif mode == "asymmetric_loader_error": + with pytest.raises((ValueError, RuntimeError), match="iterator failure|restored data iterator"): + callback.load_global_state() else: with pytest.raises( RuntimeError, @@ -242,13 +264,30 @@ def _resume_worker(rank, tmp_path, mode, endpoint): ): callback.load_global_state() assert trainer.state.global_step == 0 + if mode in ("asymmetric_loader", "asymmetric_loader_error"): + if rank == 0: + assert trainer.train_dataloader.iterations == 1 + value = torch.ones(1) + dist.all_reduce(value) + assert value.item() == 2 finally: clear_parallel_state() dist.destroy_process_group() @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -@pytest.mark.parametrize("mode", ["missing", "optimizer_step", "trainer_step", "trainer_generation", "matching"]) +@pytest.mark.parametrize( + "mode", + [ + "missing", + "optimizer_step", + "trainer_step", + "trainer_generation", + "matching", + "asymmetric_loader", + "asymmetric_loader_error", + ], +) def test_multi_optimizer_trainer_resume_is_rank_coordinated(tmp_path, mode): spawn_with_rendezvous(_resume_worker, args=(str(tmp_path), mode), nprocs=2) diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index 4e4bfa171e..569af4df47 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -271,18 +271,13 @@ def save_global_state(self, state: TrainerState) -> None: it clears. """ args: "VeOmniArguments" = self.trainer.args - step_root, payload, generations, error = None, None, {}, None + step_root, path_error = None, None try: step_root = layout.step_dir(args.train.checkpoint.save_path, state.global_step) - payload = self.state_dict(state) - generations = _saved_multi_optimizer_generations(self.trainer, state.global_step) - except Exception as exc: # noqa: BLE001 - all ranks must fail before writing - error = exc - raise_if_any_rank_failed(error, "preparing trainer checkpoint state") - if generations: - payload["_veomni_optimizer_generations"] = generations - loader_payload = {key: payload[key] for key in _LOADER_KEYS if key in payload} - extra_payload = {key: value for key, value in payload.items() if key not in _LOADER_KEYS} + except Exception as exc: # noqa: BLE001 - all ranks must agree before invalidating a manifest + path_error = exc + raise_if_any_rank_failed(path_error, "preparing the trainer checkpoint path") + assert step_root is not None # A restarted run reaching this step again overwrites the cursor files, so # the manifest an earlier attempt left has to go before the first of them @@ -300,11 +295,29 @@ def save_global_state(self, state: TrainerState) -> None: remove_error = e raise_if_any_rank_failed(remove_error, "removing the stale checkpoint manifest") + payload, state_error = None, None + try: + payload = self.state_dict(state) + except Exception as exc: # noqa: BLE001 - synchronize a rank-local cursor failure + state_error = exc + raise_if_any_rank_failed(state_error, "collecting trainer checkpoint state") + assert payload is not None + loader_payload = {key: payload[key] for key in _LOADER_KEYS if key in payload} + extra_payload = {key: value for key, value in payload.items() if key not in _LOADER_KEYS} + # Each rank writes its own files, so a full disk or a bad pickle starts out # visible to that rank alone. Reduce before the manifest: the manifest # claims every rank's cursor is down, and rank 0 only knows about its own. # A rank that raised here would otherwise leave its peers recording a step # whose state is incomplete, or waiting in a collective it never reaches. + generations, generation_error = {}, None + try: + generations = _saved_multi_optimizer_generations(self.trainer, state.global_step) + except Exception as exc: # noqa: BLE001 - all ranks must fail before writing + generation_error = exc + raise_if_any_rank_failed(generation_error, "preparing trainer checkpoint state") + if generations: + extra_payload["_veomni_optimizer_generations"] = generations write_error: Optional[Exception] = None try: for path, blob in ( @@ -506,8 +519,12 @@ def restore_rng(): set_device_rng_state(global_state.get("device_rng_state")) restore_local(restore_rng, "restoring training RNG") - if self.trainer.start_step == 0 and self.trainer.train_dataloader is not None: - restore_local(lambda: iter(self.trainer.train_dataloader), "initializing the restored data iterator") + restore_local( + lambda: iter(self.trainer.train_dataloader) + if self.trainer.start_step == 0 and self.trainer.train_dataloader is not None + else None, + "initializing the restored data iterator", + ) if reuse_owners: for checkpoint in reuse_owners: From ad278e3a39f95d302df5bd53a0e70a77661c6f2e Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 05:29:54 +0800 Subject: [PATCH 101/120] ci: isolate selective expert test entries --- scripts/ci/accelerator/suites/gpu.json | 20 ++++++++++---------- scripts/ci/accelerator/suites/npu.json | 20 ++++++++++---------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/scripts/ci/accelerator/suites/gpu.json b/scripts/ci/accelerator/suites/gpu.json index cfa780909c..023bbebf75 100644 --- a/scripts/ci/accelerator/suites/gpu.json +++ b/scripts/ci/accelerator/suites/gpu.json @@ -38,6 +38,16 @@ "tests/parallel/ulysses/test_slice_input_tensor.py" ] }, + { + "name": "gpu-selective-expert-cpu-clip-unit", + "suite": "unit", + "args": [ + "-v", + "-s", + "-x", + "tests/utils/test_fsdp2_clip_grad_norm.py" + ] + }, { "name": "04-unit", "suite": "unit", @@ -644,16 +654,6 @@ "tests/models/test_omni_fsdp.py" ] }, - { - "name": "gpu-selective-expert-cpu-clip-unit", - "suite": "unit", - "args": [ - "-v", - "-s", - "-x", - "tests/utils/test_fsdp2_clip_grad_norm.py" - ] - }, { "name": "swap-adamw-unit", "suite": "unit", diff --git a/scripts/ci/accelerator/suites/npu.json b/scripts/ci/accelerator/suites/npu.json index 45e8885b50..ba7dfc3e33 100644 --- a/scripts/ci/accelerator/suites/npu.json +++ b/scripts/ci/accelerator/suites/npu.json @@ -48,6 +48,16 @@ "tests/parallel/ulysses/test_wan_t2v_ulysses.py" ] }, + { + "name": "npu-selective-expert-cpu-clip-unit", + "suite": "unit", + "args": [ + "-v", + "-s", + "-x", + "tests/utils/test_fsdp2_clip_grad_norm.py" + ] + }, { "name": "npu-unit-04", "suite": "unit", @@ -438,16 +448,6 @@ "tests/distributed/test_fsdp_equivalence.py" ] }, - { - "name": "npu-selective-expert-cpu-clip-unit", - "suite": "unit", - "args": [ - "-v", - "-s", - "-x", - "tests/utils/test_fsdp2_clip_grad_norm.py" - ] - }, { "name": "swap-adamw-unit", "suite": "unit", From ed45ea520cb42bd562523ec80bfcc0c66817ede6 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 05:43:17 +0800 Subject: [PATCH 102/120] fix: coordinate invalid checkpoint snapshots --- veomni/trainer/callbacks/global_state_callback.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index 569af4df47..303b401d8f 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -298,7 +298,9 @@ def save_global_state(self, state: TrainerState) -> None: payload, state_error = None, None try: payload = self.state_dict(state) - except Exception as exc: # noqa: BLE001 - synchronize a rank-local cursor failure + if not isinstance(payload, dict): + raise TypeError("Trainer checkpoint state must be a dictionary") + except BaseException as exc: # noqa: BLE001 - synchronize a rank-local cursor failure state_error = exc raise_if_any_rank_failed(state_error, "collecting trainer checkpoint state") assert payload is not None From 9a8f1fe3fc491d484e72897a79c3392da775ec23 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 05:45:42 +0800 Subject: [PATCH 103/120] fix: coordinate fatal optimizer generation failures --- .../test_multi_optimizer_trainer_step.py | 20 ++++++++++++++++--- .../callbacks/global_state_callback.py | 2 +- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/tests/checkpoints/test_multi_optimizer_trainer_step.py b/tests/checkpoints/test_multi_optimizer_trainer_step.py index 0ced70675f..c2a3d5faa1 100644 --- a/tests/checkpoints/test_multi_optimizer_trainer_step.py +++ b/tests/checkpoints/test_multi_optimizer_trainer_step.py @@ -293,6 +293,8 @@ def test_multi_optimizer_trainer_resume_is_rank_coordinated(tmp_path, mode): def _save_failure_worker(rank, tmp_path, mode, endpoint): + from contextlib import nullcontext + from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state from veomni.trainer.callbacks.base import TrainerState @@ -309,8 +311,20 @@ def _save_failure_worker(rank, tmp_path, mode, endpoint): callback = GlobalStateCallback(trainer) if mode == "cursor" and rank == 1: callback.state_dict = lambda _: (_ for _ in ()).throw(ValueError("injected cursor failure")) - with pytest.raises( - (RuntimeError, ValueError), match="trainer checkpoint state|injected cursor failure|generation" + inject = ( + patch( + "veomni.trainer.callbacks.global_state_callback._saved_multi_optimizer_generations", + side_effect=SystemExit("injected generation failure"), + ) + if mode == "generation_fatal" and rank == 1 + else nullcontext() + ) + with ( + inject, + pytest.raises( + (RuntimeError, ValueError, SystemExit), + match="trainer checkpoint state|injected cursor failure|generation", + ), ): callback.save_global_state(TrainerState(global_step=3)) dist.barrier() @@ -321,6 +335,6 @@ def _save_failure_worker(rank, tmp_path, mode, endpoint): @pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2)) -@pytest.mark.parametrize("mode", ["cursor", "generation"]) +@pytest.mark.parametrize("mode", ["cursor", "generation", "generation_fatal"]) def test_trainer_save_failure_is_coordinated_before_write(tmp_path, mode): spawn_with_rendezvous(_save_failure_worker, args=(tmp_path, mode), nprocs=2) diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py index 303b401d8f..1f4f1fcec3 100644 --- a/veomni/trainer/callbacks/global_state_callback.py +++ b/veomni/trainer/callbacks/global_state_callback.py @@ -315,7 +315,7 @@ def save_global_state(self, state: TrainerState) -> None: generations, generation_error = {}, None try: generations = _saved_multi_optimizer_generations(self.trainer, state.global_step) - except Exception as exc: # noqa: BLE001 - all ranks must fail before writing + except BaseException as exc: # noqa: BLE001 - all ranks must fail before writing generation_error = exc raise_if_any_rank_failed(generation_error, "preparing trainer checkpoint state") if generations: From 6f867feea453a304e280309e64a3e76cbbee1354 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 05:51:00 +0800 Subject: [PATCH 104/120] docs: clarify optimizer placement resume identity --- docs/design/selective_expert_cpu_offload.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/design/selective_expert_cpu_offload.md b/docs/design/selective_expert_cpu_offload.md index 2c1a8a7438..ae4de582ff 100644 --- a/docs/design/selective_expert_cpu_offload.md +++ b/docs/design/selective_expert_cpu_offload.md @@ -21,8 +21,12 @@ model: The exclusion list keeps complete matching expert modules resident on the accelerator. Patterns match complete module FQNs using the parallel plan's `*` wildcard rules; every exclusion must match a selected expert module. An empty -list offloads all selected experts. Keeping 8 or 16 layers resident changes this -list, not the implementation or checkpoint layout. +list offloads all selected experts. The list changes placement, while model +tensor shapes and parallel shard geometry stay unchanged. With an owner-aware +MultiOptimizer checkpoint, each optimizer-owned parameter's device type is part +of the optimizer identity: preserve its CPU/accelerator placement when resuming. +The remaining owner, mesh, dtype and scheduler checks still apply. This feature +does not provide a validated optimizer-state migration across placements. ## Placement and update lifecycle From 43926979f09f825e2c7eabeabb881f5278b03619 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 06:22:07 +0800 Subject: [PATCH 105/120] fix: compose selective offload guard with optimizer setup --- tests/optim/test_selective_expert_muon_guard.py | 16 ++++++++++++++++ veomni/optim/optimizer.py | 10 +++++----- 2 files changed, 21 insertions(+), 5 deletions(-) create mode 100644 tests/optim/test_selective_expert_muon_guard.py diff --git a/tests/optim/test_selective_expert_muon_guard.py b/tests/optim/test_selective_expert_muon_guard.py new file mode 100644 index 0000000000..7a8a4e3f37 --- /dev/null +++ b/tests/optim/test_selective_expert_muon_guard.py @@ -0,0 +1,16 @@ +import pytest +import torch + +from veomni.optim import optimizer as optimizer_module + + +def test_muon_rejects_only_trainable_selectively_offloaded_parameters(monkeypatch): + model = torch.nn.Linear(4, 4, bias=False) + parameter = model.weight + parameter._veomni_optimizer_compute_device = "npu" + with pytest.raises(ValueError, match="DistributedMuon does not support selectively CPU-offloaded"): + optimizer_module.build_optimizer(model, optimizer_type="muon") + sentinel = object() + monkeypatch.setattr(optimizer_module, "_build_muon_with_adamw", lambda *args, **kwargs: sentinel) + parameter.requires_grad_(False) + assert optimizer_module.build_optimizer(model, optimizer_type="muon") is sentinel diff --git a/veomni/optim/optimizer.py b/veomni/optim/optimizer.py index 3a85371183..7a0512511a 100644 --- a/veomni/optim/optimizer.py +++ b/veomni/optim/optimizer.py @@ -492,12 +492,12 @@ def build_optimizer( selects the fused kernel from the backend (``FUSED_ADAMW_DEVICE_TYPES``) and takes the elementwise path elsewhere. """ - if optimizer_type == "muon" and any( - parameter.requires_grad and getattr(parameter, "_veomni_optimizer_compute_device", None) is not None - for parameter in model.parameters() - ): - raise ValueError("DistributedMuon does not support selectively CPU-offloaded parameters") if optimizer_type == "muon": + if any( + parameter.requires_grad and getattr(parameter, "_veomni_optimizer_compute_device", None) is not None + for parameter in model.parameters() + ): + raise ValueError("DistributedMuon does not support selectively CPU-offloaded parameters") if optimizer_config is not None: muon_kwargs = {**_collect_muon_kwargs(optimizer_config), **(muon_kwargs or {})} if param_groups is not None: From 8f1b7d390b8f77204c17562900cec2ce0ca1d918 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 06:24:40 +0800 Subject: [PATCH 106/120] test: run selective optimizer guard in existing CI suite --- tests/optim/test_selective_expert_muon_guard.py | 16 ---------------- tests/utils/test_fsdp2_clip_grad_norm.py | 14 ++++++++++++++ 2 files changed, 14 insertions(+), 16 deletions(-) delete mode 100644 tests/optim/test_selective_expert_muon_guard.py diff --git a/tests/optim/test_selective_expert_muon_guard.py b/tests/optim/test_selective_expert_muon_guard.py deleted file mode 100644 index 7a8a4e3f37..0000000000 --- a/tests/optim/test_selective_expert_muon_guard.py +++ /dev/null @@ -1,16 +0,0 @@ -import pytest -import torch - -from veomni.optim import optimizer as optimizer_module - - -def test_muon_rejects_only_trainable_selectively_offloaded_parameters(monkeypatch): - model = torch.nn.Linear(4, 4, bias=False) - parameter = model.weight - parameter._veomni_optimizer_compute_device = "npu" - with pytest.raises(ValueError, match="DistributedMuon does not support selectively CPU-offloaded"): - optimizer_module.build_optimizer(model, optimizer_type="muon") - sentinel = object() - monkeypatch.setattr(optimizer_module, "_build_muon_with_adamw", lambda *args, **kwargs: sentinel) - parameter.requires_grad_(False) - assert optimizer_module.build_optimizer(model, optimizer_type="muon") is sentinel diff --git a/tests/utils/test_fsdp2_clip_grad_norm.py b/tests/utils/test_fsdp2_clip_grad_norm.py index c3e366340c..033d800f8d 100644 --- a/tests/utils/test_fsdp2_clip_grad_norm.py +++ b/tests/utils/test_fsdp2_clip_grad_norm.py @@ -149,3 +149,17 @@ def test_selective_expert_offload_rejects_incompatible_configuration(settings): with pytest.raises(ValueError): FSDPConfig(**settings) + + +def test_muon_rejects_only_trainable_selectively_offloaded_parameters(monkeypatch): + from veomni.optim import optimizer as optimizer_module + + model = torch.nn.Linear(4, 4, bias=False) + parameter = model.weight + parameter._veomni_optimizer_compute_device = "npu" + with pytest.raises(ValueError, match="DistributedMuon does not support selectively CPU-offloaded"): + optimizer_module.build_optimizer(model, optimizer_type="muon") + sentinel = object() + monkeypatch.setattr(optimizer_module, "_build_muon_with_adamw", lambda *args, **kwargs: sentinel) + parameter.requires_grad_(False) + assert optimizer_module.build_optimizer(model, optimizer_type="muon") is sentinel From 1b0f552b6f11f544817e1ced1058b52c204e982e Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 15:15:01 +0800 Subject: [PATCH 107/120] Clarify selective expert offload resume requirements --- docs/design/selective_expert_cpu_offload.md | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/docs/design/selective_expert_cpu_offload.md b/docs/design/selective_expert_cpu_offload.md index ae4de582ff..388ca9dfe9 100644 --- a/docs/design/selective_expert_cpu_offload.md +++ b/docs/design/selective_expert_cpu_offload.md @@ -22,17 +22,15 @@ The exclusion list keeps complete matching expert modules resident on the accelerator. Patterns match complete module FQNs using the parallel plan's `*` wildcard rules; every exclusion must match a selected expert module. An empty list offloads all selected experts. The list changes placement, while model -tensor shapes and parallel shard geometry stay unchanged. With an owner-aware -MultiOptimizer checkpoint, each optimizer-owned parameter's device type is part -of the optimizer identity: preserve its CPU/accelerator placement when resuming. -The remaining owner, mesh, dtype and scheduler checks still apply. This feature -does not provide a validated optimizer-state migration across placements. +tensor shapes and parallel shard geometry stay unchanged. Preserve each +optimizer-owned parameter's CPU/accelerator placement when resuming. This +feature does not provide a validated optimizer-state migration across placements. ## Placement and update lifecycle [`VeOmniModelRuntime._build_parallelized_model()`](../../veomni/models/model_runtime.py#L401) forwards these settings to -[`parallelize_model_fsdp2()`](../../veomni/distributed/torch_parallelize.py#L432). +[`parallelize_model_fsdp2()`](../../veomni/distributed/torch_parallelize.py#L427). The latter applies CPU offload only to non-excluded EP FSDP groups, loads the model, and moves those groups to CPU before optimizer construction. FSDP stages their weights for forward/backward and accumulates From 6d520d99d67e10c8ac14f0636b0bd4c2441dc6eb Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Sat, 7 Nov 2026 15:33:42 +0800 Subject: [PATCH 108/120] Initialize checkpoint manager config in failure fixture --- tests/checkpoints/test_multi_optimizer_proven_resume.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index b0f29fb2f8..e42f2cedab 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -276,6 +276,7 @@ def test_failed_final_load_barrier_does_not_certify_multi_optimizer(): leaf = next(iter(optimizer.optimizers_dict.values())) leaf.checkpoint_abort_file_backed_state = MagicMock() manager = object.__new__(ModelCheckpointManager) + manager.config = SimpleNamespace(dcp_evict_source_cache=False) manager.runtime = SimpleNamespace( model=model, optimizer=optimizer, lr_scheduler=None, load_extra_state=lambda _: None ) From 135f718cea427df023cc394bdafe5901ebf67220 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 9 Nov 2026 15:35:33 +0800 Subject: [PATCH 109/120] [docs] clarify legacy HF export collective --- veomni/utils/save_safetensor_utils.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/veomni/utils/save_safetensor_utils.py b/veomni/utils/save_safetensor_utils.py index 4c95178c35..687bdda8b5 100644 --- a/veomni/utils/save_safetensor_utils.py +++ b/veomni/utils/save_safetensor_utils.py @@ -292,26 +292,28 @@ def save_hf_safetensor( """Save model weights in HuggingFace safetensors format. This function is self-contained w.r.t. synchronization: it calls ``synchronize()`` at - entry to flush pending GPU operations before reading tensor data. All ranks - coordinate completion and errors before returning; only rank 0 writes in legacy - mode. Callers do not need to add external synchronization around this function. + entry to flush pending device operations before reading tensor data. When + distributed is initialized, all ranks must call this function. The legacy + path coordinates the rank-0 export result on the default world group before + returning, even when a per-module parallel state is provided. Supports two modes: - - Distributed mode (PyTorch >= 2.9, ckpt_manager="dcp", non-LoRA): Uses HuggingFaceStorageWriter - for efficient distributed save directly from the live FSDP model. Must be called on all ranks. - - Legacy mode: Loads from checkpoint and converts to safetensors on rank 0. + - Distributed mode (PyTorch >= 2.9, ckpt_manager="dcp", non-NPU): Uses HuggingFaceStorageWriter + for efficient distributed save directly from the live FSDP model. + - Legacy mode, including NPU: Loads from checkpoint and converts to safetensors on rank 0. Args: save_hf_safetensor_path: Output path for saved HuggingFace safetensors. - ckpt_manager: Checkpoint manager type. Used for routing (distributed when "dcp") - and passed to legacy ``ckpt_to_state_dict``. + ckpt_manager: Checkpoint manager type. Used for routing and passed to legacy + ``ckpt_to_state_dict``. On NPU, "dcp" also uses the legacy path. model_assets: Model assets (e.g., config, tokenizer) to save alongside weights. save_checkpoint_path: [Legacy only] Path to the distributed checkpoint for conversion. output_dir: [Legacy only] Output directory passed to ``ckpt_to_state_dict``. is_rank_0: [Legacy only] Whether the current process is global rank 0. - Legacy save is rank-0 only; non-rank-0 processes return immediately. - Required by non-dcp checkpoint managers (e.g., omnistore). + Only rank 0 writes in legacy mode; when distributed is initialized, + all ranks must call this function and participate in the collective. + Set to true on the global rank 0 whenever the legacy path is selected. model: [Distributed only] Live FSDP model for distributed save. fqn_to_index_mapping: [Distributed only] Maps FQNs to safetensors file indices for multi-file output. From a47732c27c9dac2909775ed8b25712a60f580524 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 9 Nov 2026 15:35:33 +0800 Subject: [PATCH 110/120] [docs] clarify SeedOmni FP32 expert gradient scope --- docs/design/op_selection.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/design/op_selection.md b/docs/design/op_selection.md index 1a8894ddf2..9dbc2c983e 100644 --- a/docs/design/op_selection.md +++ b/docs/design/op_selection.md @@ -626,8 +626,10 @@ The current integration requires Qwen4-Exp SiLU `fused_npu`, FSDP2, TP=PP=1, FP32 reduction, native communication, and `torch.compile` disabled. Each rank validates its local precision settings and constructed expert modules before FSDP wrapping. As with the rest of the trainer configuration, all ranks must -use the same settings. Composed Omni wraps and custom model parallelization -reject the opt-in. MoE LoRA expert wrappers also reject the opt-in because +use the same settings. SeedOmni `ModuleRuntime` rejects the opt-in when enabled +in its resolved per-module accelerator configuration, regardless of `fsdp_scope`. +Composed Omni wraps and custom model parallelization also reject the opt-in. MoE +LoRA expert wrappers reject the opt-in because their separate kernel is not covered. DPO reference-model isolation and resume configuration checks are not wired by this option yet. Keep the setting fixed across checkpoint save and resume when comparing training trajectories. From 5c2837caac0017e0c50e50a2e313bfcf1846c9e5 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 9 Nov 2026 16:07:11 +0800 Subject: [PATCH 111/120] [docs] clarify checkpoint resume and source links --- docs/seed_omni/mixins/metric_meter.md | 8 ++++---- docs/seed_omni/mixins/offline_encoding.md | 2 +- docs/seed_omni/usage/hf_layout.md | 4 ++-- docs/usage/checkpoint.md | 15 ++++++++++++++- 4 files changed, 21 insertions(+), 8 deletions(-) diff --git a/docs/seed_omni/mixins/metric_meter.md b/docs/seed_omni/mixins/metric_meter.md index 1f9bb502a3..0d2b955ee8 100644 --- a/docs/seed_omni/mixins/metric_meter.md +++ b/docs/seed_omni/mixins/metric_meter.md @@ -89,8 +89,8 @@ class VeOmniMixin(BaseMixin, TrainingMixin, InferenceModuleMixin, MetricMeter): | Function | Called by | Effect | |----------|-----------|--------| | [`MetricMeterMixin.metric_meter_add(method)`](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L120) | [`execute_train_node`](../../../veomni/models/seed_omni/accelerated/utils/executor.py#L66), once per training node, right after `pre_forward` and before the endpoint runs, inside the module's `ParallelState` scope | [Appends](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L136) the lengths stashed for the running `method` to the module's per-step buffer, then [clears](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L137) the whole stash, so copies a shared hook left under other call-sites cannot leak into a later node. Over gradient accumulation it runs once per micro-batch, so the buffer holds the whole global step. | -| [`MetricMeterMixin.metric_meter_collect() -> (flops, seqlens)`](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L139) | [`OmniModelRuntime.metric_meter_collect`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L608) | Returns `(estimate_flops(buffer), buffer)` and empties the buffer for the next step. | -| [`OmniModelRuntime.metric_meter_collect() -> {name: (flops, seqlens)}`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L608) | [`OmniStepMetricsCallback.on_step_end`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L76) | Drains every module with `isinstance(omni_module, MetricMeterMixin)`. Keys are module names, so they are identical on every rank even when a rank's batch did not use a module (it reports `(0.0, [])`). | +| [`MetricMeterMixin.metric_meter_collect() -> (flops, seqlens)`](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L139) | [`OmniModelRuntime.metric_meter_collect`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L610) | Returns `(estimate_flops(buffer), buffer)` and empties the buffer for the next step. | +| [`OmniModelRuntime.metric_meter_collect() -> {name: (flops, seqlens)}`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L610) | [`OmniStepMetricsCallback.on_step_end`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L76) | Drains every module with `isinstance(omni_module, MetricMeterMixin)`. Keys are module names, so they are identical on every rank even when a rank's batch did not use a module (it reports `(0.0, [])`). | | [`OmniEnvironMeter.add(micro_batch)`](../../../veomni/utils/omni_helper.py#L123) | [`OmniStepMetricsCallback.on_step_begin`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L66), once per micro-batch | Counts samples (`len(conversation_list)`) and, for multi-source, gathers `ds_idx`. Never looks at tokens. | | [`OmniEnvironMeter.step(delta_time, global_step, module_metrics)`](../../../veomni/utils/omni_helper.py#L137) | [`OmniStepMetricsCallback.on_step_end`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L73) | [One DP all-reduce](../../../veomni/utils/omni_helper.py#L169) of FLOPs, sample count and per-module token sums, then the metrics listed below. | @@ -117,7 +117,7 @@ and stored as `trainer.environ_meter`. - [`on_step_end`](../../../veomni/trainer/omni/omni_trainer.py#L570) → [`OmniStepMetricsCallback.on_step_end`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L69) - [`delta_time = time.time() - start_time`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L72) - - [`OmniModelRuntime.metric_meter_collect()`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L608) → + - [`OmniModelRuntime.metric_meter_collect()`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L610) → per metered module [`(estimate_flops(buffer), buffer)`](../../../veomni/models/seed_omni/mixins/metric_meter_mixin.py#L150); buffer reset. - [`OmniEnvironMeter.step(delta_time, global_step, module_metrics)`](../../../veomni/utils/omni_helper.py#L137) - [`trainer.step_env_metrics = ...`](../../../veomni/trainer/callbacks/omni_callbacks/step_metrics_callback.py#L107), @@ -163,7 +163,7 @@ reported per module under `trace//` and there is no top-level [`OmniEnvironMeter.state_dict()`](../../../veomni/utils/omni_helper.py#L111), which `GlobalStateCallback` [saves](../../../veomni/trainer/callbacks/global_state_callback.py#L253) and -[restores](../../../veomni/trainer/callbacks/global_state_callback.py#L496), so +[restores](../../../veomni/trainer/callbacks/global_state_callback.py#L511), so the cumulative counters survive a resume. ## Rules for a correct implementation diff --git a/docs/seed_omni/mixins/offline_encoding.md b/docs/seed_omni/mixins/offline_encoding.md index cc2c1d957f..a5d2d1f27b 100644 --- a/docs/seed_omni/mixins/offline_encoding.md +++ b/docs/seed_omni/mixins/offline_encoding.md @@ -87,7 +87,7 @@ so the cache looks complete and the error only shows in `offline_training`. `offline_embedding` freezes every module, so it never writes or restores module weights; encoder weights always come from each module's `model_path`. A saved step is therefore only the cursor, and its manifest -[says so](../../../veomni/trainer/callbacks/global_state_callback.py#L324), so +[says so](../../../veomni/trainer/callbacks/global_state_callback.py#L348), so `train.checkpoint.load_path: auto` finds it like any other step. Give the encoding run and the training run separate `train.checkpoint.output_dir`s: `auto` resumes the newest complete step in the directory, and diff --git a/docs/seed_omni/usage/hf_layout.md b/docs/seed_omni/usage/hf_layout.md index 8458375ff2..d03becf69b 100644 --- a/docs/seed_omni/usage/hf_layout.md +++ b/docs/seed_omni/usage/hf_layout.md @@ -155,9 +155,9 @@ Training from an upstream root: - [`load_model_weights`](../../../veomni/models/checkpoint/weights.py#L410) and the [rank-0 broadcast loader](../../../veomni/models/checkpoint/weights.py#L1208) skip keys the converter declares, without reading them; - [`HFSourceKeyConverter`](../../../veomni/models/seed_omni/utils/hf_layout.py#L331) renames the module's own keys and hands them to the module's converter. -HF export (`CheckpointCallback` → [`OmniModelRuntime.save_hf_or_lora`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L631)): +HF export (`CheckpointCallback` → [`OmniModelRuntime.save_hf_or_lora`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L633)): -- [`_save_hf_source_layout`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L644) exports LoRA modules' adapters per module, skips frozen modules, and makes sure each trained module's DCP exists; +- [`_save_hf_source_layout`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L646) exports LoRA modules' adapters per module, skips frozen modules, and makes sure each trained module's DCP exists; - [`save_hf_source_checkpoint`](../../../veomni/models/seed_omni/utils/hf_layout.py#L683) renames keys back (raising on a key with no source key or a different shape), casts to the source dtypes, fills tied duplicates, copies every other source tensor on rank 0, and writes `global_step_N/hf_ckpt/` with the source's shard files and non-weight files. Bare `OmniModel.from_pretrained` and eager inference: diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index ce6400a378..d8a53621a1 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -23,6 +23,19 @@ A step directory is split by *who owns the state*. `model/`, `loader/` and nothing else. The two exports are inference artifacts that resume never reads. `/` appears only in a [multi-module job](#multi-module-jobs-seedomni). +With `train.checkpoint.manager: dcp`, `load_path: auto` discovers the newest +structurally complete step, including a model-only step without optimizer +state. Discovery does not prove that the step is valid for training resume. +A normal trainer resume requests the optimizer; +MultiOptimizer recovery rejects missing or unproven optimizer identity rather +than silently replacing its state. To start a separate weights-only run from a +full-model checkpoint, export verified model weights to HF format, check the +exported dtypes and layout, supply them as `model.model_path`, and leave +`load_path` unset with a new checkpoint output directory. LoRA requires its +base model and adapter separately. This starts new optimizer, scheduler and data +state rather than continuing the old run. If only BF16 model weights were +saved, exporting them as FP32 cannot recover the master's discarded low bits. + ## What each path holds | Path | Holds | Written by | @@ -554,7 +567,7 @@ resume: they cannot establish scheduler and group provenance. The standalone `MultiOptimizer.state_dict()` returns a snapshot and copies tensor states on their current devices, which can temporarily require as much extra device memory as the optimizer tensors themselves. That raw snapshot supports the -built-in AdamW, AnyPrecisionAdamW and DistributedMuon leaves, as well as +built-in AdamW, AnyPrecisionAdamW, SwapAdamW and DistributedMuon leaves, as well as registered adapters. Custom serialization hooks and adapters must preserve state faithfully. It does not prove the identity needed for raw restoration. Use the high-level `DistributedCheckpointer` for save and restore; its save path avoids this copy. From 158d83eac45662674802529a2db89ee10ed7d253 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 9 Nov 2026 18:22:07 +0800 Subject: [PATCH 112/120] [tests] fix: resolve checkpoint export CLI path --- tests/checkpoints/test_expert_shard_layout.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index b668db6b09..7b8bf76aae 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -3,6 +3,7 @@ """CPU/Gloo regressions for canonical EP-first DCP expert coordinates.""" +import os from datetime import timedelta from pathlib import Path from types import SimpleNamespace @@ -170,10 +171,13 @@ def wrap(dim, *, empty=False): assert exported["layer.experts.weight"].dtype == torch.float32 assert exported["layer.bias"].dtype == torch.bfloat16 cli_dir = directory + "known_legacy_hf" + repo_root = Path(__file__).resolve().parents[2] + cli_env = os.environ.copy() + cli_env["PYTHONPATH"] = os.pathsep.join((str(repo_root), cli_env.get("PYTHONPATH", ""))) subprocess.run( [ sys.executable, - "scripts/merge_dcp_to_hf.py", + str(repo_root / "scripts" / "merge_dcp_to_hf.py"), "--load-dir", path, "--save-dir", @@ -187,6 +191,8 @@ def wrap(dim, *, empty=False): ], check=True, timeout=60, + cwd=repo_root, + env=cli_env, ) exported = load_file(cli_dir + "/model.safetensors") torch.testing.assert_close(exported["layer.experts.weight"], global_weight, rtol=0, atol=0) From c051d9c45aea141d7c7c5913bd0a6a30303e7d91 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Mon, 9 Nov 2026 19:50:36 +0800 Subject: [PATCH 113/120] [checkpoint] clarify existing save destination --- tests/checkpoints/test_multi_export_receipt.py | 3 +++ veomni/trainer/callbacks/checkpoint_callback.py | 5 ++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/checkpoints/test_multi_export_receipt.py b/tests/checkpoints/test_multi_export_receipt.py index fc99a1e0c2..442360a0eb 100644 --- a/tests/checkpoints/test_multi_export_receipt.py +++ b/tests/checkpoints/test_multi_export_receipt.py @@ -265,6 +265,9 @@ def stat_with_error(path, *args, **kwargs): if target_exists == "stat_error": expected_error = "manifest stat error" if rank == 1 else "failed on another rank" assert expected_error in str(caught.value) + else: + assert repr(layout.step_dir(str(target_path), 1)) in str(caught.value) + assert "choose a new checkpoint save_path" in str(caught.value) for runtime in restored: runtime.checkpoint.save_dcp.assert_not_called() if rank == 0: diff --git a/veomni/trainer/callbacks/checkpoint_callback.py b/veomni/trainer/callbacks/checkpoint_callback.py index bd33c4bd82..4faecf140a 100644 --- a/veomni/trainer/callbacks/checkpoint_callback.py +++ b/veomni/trainer/callbacks/checkpoint_callback.py @@ -84,7 +84,10 @@ def _check_multi_destination(trainer, step): except FileNotFoundError: pass else: - raise RuntimeError("Cannot overwrite a completed checkpoint transaction with MultiOptimizer state") + raise RuntimeError( + f"Cannot overwrite a completed checkpoint transaction with MultiOptimizer state at {step_root!r}; " + "choose a new checkpoint save_path before retrying" + ) except BaseException as exc: error = exc raise_if_any_rank_failed(error, "checking MultiOptimizer checkpoint destination") From cfafa433cae2c7fdb5d6cca805b04579e65003dd Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 10 Nov 2026 12:34:14 +0800 Subject: [PATCH 114/120] [test] guard invalid prefixed dynamic rope cache --- .../test_multi_optimizer_rope_cache.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/checkpoints/test_multi_optimizer_rope_cache.py b/tests/checkpoints/test_multi_optimizer_rope_cache.py index 1ee475b465..7fbea95aeb 100644 --- a/tests/checkpoints/test_multi_optimizer_rope_cache.py +++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py @@ -22,6 +22,35 @@ pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +def test_dynamic_rope_prefixed_cache_without_shared_fallback_fails_closed(): + from transformers import Gemma3TextConfig + from transformers.models.gemma3.modeling_gemma3 import Gemma3RotaryEmbedding + + from veomni.checkpoint.dynamic_rope import dynamic_rope_checkpoint_contract + + config = Gemma3TextConfig( + hidden_size=32, + head_dim=8, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=1, + max_position_embeddings=16, + layer_types=["full_attention"], + rope_parameters={"full_attention": {"rope_type": "dynamic", "rope_theta": 10000.0, "factor": 2.0}}, + ) + rope = Gemma3RotaryEmbedding(config) + rope.full_attention_max_seq_len_cached = 16 + assert len(dynamic_rope_checkpoint_contract(rope)) == 1 + del rope.max_seq_len_cached + # Transformers 5.16's dynamic-rope decorator eagerly reads the shared + # fallback even with a prefixed cache. Do not certify an unrestorable state. + with pytest.raises(RuntimeError, match="Cannot prove dynamic RoPE checkpoint cache") as exc: + dynamic_rope_checkpoint_contract(rope) + assert isinstance(exc.value.__cause__, AttributeError) + with pytest.raises(AttributeError, match="max_seq_len_cached"): + rope(torch.zeros((1, 16, 8)), torch.arange(16).reshape(1, -1), layer_type="full_attention") + + def test_runtime_preflight_allows_dynamic_rope_training_without_checkpoint(): from veomni.models.model_runtime import VeOmniModelRuntime From f111b9115091ee2e3a762a7737b655ea11db5f35 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:41:01 +0800 Subject: [PATCH 115/120] [ci, tests] require assigned NPU and fix precision fixture --- tests/ci/test_npu_moe_fp32_cold.py | 87 ++++++++++++++++--- .../test_optimizer_checkpoint_preflight.py | 7 +- 2 files changed, 82 insertions(+), 12 deletions(-) diff --git a/tests/ci/test_npu_moe_fp32_cold.py b/tests/ci/test_npu_moe_fp32_cold.py index 8bfe2ccc83..b738c80c60 100644 --- a/tests/ci/test_npu_moe_fp32_cold.py +++ b/tests/ci/test_npu_moe_fp32_cold.py @@ -15,12 +15,13 @@ """FSDP2 + expert-parallel optimizer/DCP cold-recovery regression on Ascend.""" import os +from types import SimpleNamespace import pytest import torch from torch.distributed.checkpoint import FileSystemReader -from scripts.ci.accelerator.launch_plan import LaunchSpec +from scripts.ci.accelerator.launch_plan import LaunchSpec, runner_capabilities from tests.ci._npu_moe_fp32_cold_worker import run, run_selective from veomni.checkpoint.layout import weights_dir from veomni.utils.device import IS_NPU_AVAILABLE @@ -86,13 +87,80 @@ def _assert_ep2_checkpoint_layout(checkpoint_dir): assert {(tuple(chunk.offsets), tuple(chunk.sizes)) for chunk in metadata.chunks} == chunks -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2)) +def _require_two_native_npus(): + """Fail an assigned or explicitly required NPU test when its lease is unusable.""" + capabilities = runner_capabilities() if os.environ.get("VEOMNI_CI_CAPABILITIES") else None + required = os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" + if capabilities is not None and capabilities.backend != "npu": + if required: + pytest.fail("VEOMNI_REQUIRE_NPU_TESTS=1 requires an NPU CI assignment") + pytest.skip("requires an NPU CI assignment") + npu = getattr(torch, "npu", None) + if not IS_NPU_AVAILABLE or npu is None or not npu.is_available() or npu.device_count() < 2: + if required or (capabilities is not None and capabilities.backend == "npu"): + pytest.fail("Assigned or required NPU runtime must provide two usable devices") + pytest.skip("requires two native Ascend devices") + + +def _native_npu_launch(context): + """Keep an unavailable required test runnable so its guard reports failure.""" + caps = context.capabilities + if caps.backend == "npu" and len(caps.devices) >= 2: + return LaunchSpec(nproc=2, device="accelerator") + return LaunchSpec(device="cpu") + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +@pytest.mark.parametrize( + "backend,required,available,count,outcome", + [ + (None, False, False, 0, "skip"), + (None, True, False, 0, "fail"), + ("npu", False, False, 0, "fail"), + ("npu", False, True, 1, "fail"), + ("npu", False, True, 2, "pass"), + ("cuda", False, True, 2, "skip"), + ("cuda", True, True, 2, "fail"), + ], +) +def test_native_npu_guard_respects_ci_assignment(monkeypatch, backend, required, available, count, outcome): + monkeypatch.setitem(globals(), "IS_NPU_AVAILABLE", available) + monkeypatch.setattr( + torch, "npu", SimpleNamespace(is_available=lambda: available, device_count=lambda: count), raising=False + ) + if backend is None: + monkeypatch.delenv("VEOMNI_CI_CAPABILITIES", raising=False) + else: + monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "test") + monkeypatch.setitem(globals(), "runner_capabilities", lambda: SimpleNamespace(backend=backend)) + monkeypatch.setenv("VEOMNI_REQUIRE_NPU_TESTS", "1" if required else "0") + if outcome == "fail": + with pytest.raises(pytest.fail.Exception): + _require_two_native_npus() + elif outcome == "skip": + with pytest.raises(pytest.skip.Exception): + _require_two_native_npus() + else: + _require_two_native_npus() + + +@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +def test_native_npu_launch_keeps_missing_devices_runnable(): + assert _native_npu_launch(SimpleNamespace(capabilities=SimpleNamespace(backend="cpu", devices=()))) == LaunchSpec( + device="cpu" + ) + assert _native_npu_launch( + SimpleNamespace(capabilities=SimpleNamespace(backend="npu", devices=("0",))) + ) == LaunchSpec(device="cpu") + assert _native_npu_launch( + SimpleNamespace(capabilities=SimpleNamespace(backend="npu", devices=("0", "1"))) + ) == LaunchSpec(nproc=2, device="accelerator") + + +@pytest.mark.launch_plan(provider=_native_npu_launch) @pytest.mark.parametrize("ep", [1, 2]) def test_native_fsdp_expert_fp32_dcp_cold_resume(tmp_path, ep): - if not IS_NPU_AVAILABLE or torch.npu.device_count() < 2: - if os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1": - pytest.fail("VEOMNI_REQUIRE_NPU_TESTS=1 requires two native NPU devices") - pytest.skip("requires two native Ascend devices") + _require_two_native_npus() from tests.tools.launch_utils import spawn_with_rendezvous for phase in ("fresh", "resume"): @@ -116,12 +184,9 @@ def test_native_fsdp_expert_fp32_dcp_cold_resume(tmp_path, ep): _assert_ep2_checkpoint_layout(tmp_path / "ep2-dcp") -@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2)) +@pytest.mark.launch_plan(provider=_native_npu_launch) def test_native_fp32_expert_selective_offload_cold_resume(tmp_path): - if not IS_NPU_AVAILABLE or torch.npu.device_count() < 2: - if os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1": - pytest.fail("VEOMNI_REQUIRE_NPU_TESTS=1 requires two native NPU devices") - pytest.skip("requires two native Ascend devices") + _require_two_native_npus() from tests.tools.launch_utils import spawn_with_rendezvous results = {} diff --git a/tests/models/test_optimizer_checkpoint_preflight.py b/tests/models/test_optimizer_checkpoint_preflight.py index d1bbed35cb..5f44d3faf6 100644 --- a/tests/models/test_optimizer_checkpoint_preflight.py +++ b/tests/models/test_optimizer_checkpoint_preflight.py @@ -130,6 +130,7 @@ def test_omni_composed_wrap_rejects_async_multi_optimizer_before_checkpoint(monk @pytest.mark.parametrize("multi", [None, False, True]) def test_omni_module_scope_checks_async_before_checkpoint(monkeypatch, multi): + from veomni.arguments import MixedPrecisionConfig from veomni.models.seed_omni.accelerated.omni_module.omni_module_runtime import ModuleRuntime monkeypatch.setattr(ModuleRuntime, "setup", lambda self: None) @@ -145,7 +146,11 @@ def test_omni_module_scope_checks_async_before_checkpoint(monkeypatch, multi): ) reached_checkpoint = [] monkeypatch.setattr(ModuleRuntime, "build_checkpoint", lambda self: reached_checkpoint.append(True)) - accelerator = SimpleNamespace(fsdp_config=SimpleNamespace(fsdp_scope="module", fsdp_mode="fsdp2")) + accelerator = SimpleNamespace( + fsdp_config=SimpleNamespace( + fsdp_scope="module", fsdp_mode="fsdp2", mixed_precision=MixedPrecisionConfig(enable=False) + ) + ) train = SimpleNamespace(training_task="sft", checkpoint=SimpleNamespace(save_async=True)) kwargs = dict(module_config=SimpleNamespace(), global_accelerator=accelerator, train_args=train) if multi: From f222247cf4a5793c583d169f9f51d25694fd3687 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Tue, 10 Nov 2026 14:44:18 +0800 Subject: [PATCH 116/120] [checkpoint] bind model storage views for resume --- docs/usage/checkpoint.md | 12 ++ .../test_multi_optimizer_proven_resume.py | 178 ++++++++++++++++++ veomni/checkpoint/optimizer_resume_layout.py | 94 ++++++++- 3 files changed, 283 insertions(+), 1 deletion(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index d8a53621a1..4d9f9d60d3 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -36,6 +36,18 @@ base model and adapter separately. This starts new optimizer, scheduler and data state rather than continuing the old run. If only BF16 model weights were saved, exporting them as FP32 cannot recover the master's discarded low bits. +MultiOptimizer recovery also binds overlapping parameter and buffer storage +views, including their local shape, strides and relative offsets. A checkpoint +created before this storage-view proof existed is rejected for optimizer resume, +even when its target has no shared views; the old metadata cannot prove that the +source was also unshared. Use a separately verified migration or a new +weights-only run rather than silently assuming identical update behavior. +This proof currently requires Linux procfs for nonempty CPU model tensors and +rejects shared CPU mappings, quantized, non-strided, and meta model tensors. +Accelerator tensors must use framework-owned storage; externally imported or +shared accelerator allocations are not proved by this check. File-backed +optimizer moments are separate from model tensors and are not restricted here. + ## What each path holds | Path | Holds | Written by | diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index e42f2cedab..ab1fd12353 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -1950,6 +1950,184 @@ def build(shared): torch.testing.assert_close(target.p, before, rtol=0, atol=0) +@pytest.mark.parametrize("source_shared", [False, True]) +def test_distinct_tensor_storage_alias_must_match_before_resume(tmp_path, source_shared): + def build(shared): + model = torch.nn.Module() + model.p = torch.nn.Parameter(torch.ones(3)) + view = model.p.detach() + model.register_buffer("view", view if shared else view.clone()) + optimizer = MultiOptimizer(model, {"leaf": torch.optim.AdamW([model.p])}, ["leaf"]) + return model, optimizer + + source, optimizer = build(source_shared) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + target, restored = build(not source_shared) + before = (target.p.detach().clone(), target.view.detach().clone()) + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + torch.testing.assert_close(target.p, before[0], rtol=0, atol=0) + torch.testing.assert_close(target.view, before[1], rtol=0, atol=0) + + +def test_matching_storage_views_resume_and_update(tmp_path): + def build(): + model = torch.nn.Module() + model.p = torch.nn.Parameter(torch.arange(4, dtype=torch.float32)) + model.register_buffer("view", model.p.detach()[::2]) + optimizer = MultiOptimizer(model, {"leaf": torch.optim.AdamW([model.p], lr=0.01)}, ["leaf"]) + return model, optimizer + + def step(model, optimizer): + model.p.grad = torch.ones_like(model.p) + optimizer.step() + optimizer.zero_grad() + + source, optimizer = build() + step(source, optimizer) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + target, restored = build() + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + step(source, optimizer) + step(target, restored) + torch.testing.assert_close(target.p, source.p, rtol=0, atol=0) + torch.testing.assert_close(target.view, source.view, rtol=0, atol=0) + + +def test_storage_view_geometry_and_disjoint_views_are_distinguished(): + from veomni.checkpoint.optimizer_resume_layout import _model_storage_alias_graph + + base = torch.arange(8, dtype=torch.float32) + parameter = torch.nn.Parameter(base[:4]) + overlap = base[1:4:2] + shifted = base[:3:2] + disjoint = base[4:] + + def graph(view): + return _model_storage_alias_graph( + {id(parameter): {"parameter:p"}, id(view): {"buffer:view"}}, + {id(parameter): parameter, id(view): view}, + ) + + assert graph(overlap) and graph(shifted) + assert graph(overlap) != graph(shifted) + assert graph(disjoint)["storage_groups"] + independent = base[4:].clone() + assert not graph(independent)["storage_groups"] + assert graph(disjoint) != graph(independent) + + +def test_storage_view_proof_binds_internal_overlap_and_quantization(): + from veomni.checkpoint.optimizer_resume_layout import _model_storage_alias_graph + + def graph(value): + return _model_storage_alias_graph({id(value): {"buffer:view"}}, {id(value): value}) + + assert graph(torch.ones(1).expand(3)) != graph(torch.ones(3)) + with pytest.raises(RuntimeError, match="quantized"): + graph(torch.quantize_per_tensor(torch.ones(3), 0.1, 0, torch.quint8)) + + +def test_storage_view_proof_binds_empty_shared_storage(): + from veomni.checkpoint.optimizer_resume_layout import _model_storage_alias_graph + + first = torch.empty(0) + shared = first.detach() + independent = torch.empty(0) + + def graph(second): + return _model_storage_alias_graph( + {id(first): {"buffer:first"}, id(second): {"buffer:second"}}, + {id(first): first, id(second): second}, + ) + + assert graph(shared)["storage_groups"] + assert not graph(independent)["storage_groups"] + assert graph(shared) != graph(independent) + + +def test_storage_view_proof_ignores_unrelated_non_ascii_mapping(tmp_path): + import mmap + + from veomni.checkpoint.optimizer_resume_layout import _model_storage_alias_graph + + path = tmp_path / "模型.bin" + path.write_bytes(bytes(64)) + with path.open("r+b") as file, mmap.mmap(file.fileno(), 64, access=mmap.ACCESS_COPY): + value = torch.ones(3) + _model_storage_alias_graph({id(value): {"buffer:view"}}, {id(value): value}) + + +@pytest.mark.parametrize("method", ["from_file", "frombuffer"]) +def test_storage_view_proof_rejects_shared_cpu_mappings(tmp_path, method): + import mmap + + from veomni.checkpoint.optimizer_resume_layout import _model_storage_alias_graph + + path = tmp_path / "mapping.bin" + path.write_bytes(bytes(64)) + with path.open("r+b") as file: + if method == "from_file": + value = torch.from_file(str(path), shared=True, size=16, dtype=torch.float32) + else: + mapping = mmap.mmap(file.fileno(), 64, access=mmap.ACCESS_WRITE) + value = torch.frombuffer(mapping, dtype=torch.float32) + with pytest.raises(RuntimeError, match="shared CPU model tensor mappings"): + _model_storage_alias_graph({id(value): {"buffer:view"}}, {id(value): value}) + if method == "frombuffer": + del value + mapping.close() + + +def test_old_layout_without_storage_view_proof_is_rejected(tmp_path, monkeypatch): + import veomni.checkpoint.optimizer_resume_layout as layout + + source, optimizer = _model_and_optimizer() + original = layout._local_layout + + def legacy_layout(*args, **kwargs): + value = original(*args, **kwargs) + if isinstance(value, dict): + value.pop("model_storage_aliases_v1") + return value + + parallel = SimpleNamespace(dp_mode="ddp") + with monkeypatch.context() as patch: + patch.setattr(layout, "_local_layout", legacy_layout) + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + target, restored = _model_and_optimizer() + before = target.weight.detach().clone() + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load( + str(tmp_path), + {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + torch.testing.assert_close(target.weight, before, rtol=0, atol=0) + + def _save_destination_stat_error_worker(rank, checkpoint, location, endpoint): import errno import os diff --git a/veomni/checkpoint/optimizer_resume_layout.py b/veomni/checkpoint/optimizer_resume_layout.py index ecb3b060bb..69ce22ac59 100644 --- a/veomni/checkpoint/optimizer_resume_layout.py +++ b/veomni/checkpoint/optimizer_resume_layout.py @@ -3,6 +3,7 @@ """Identity of a MultiOptimizer's rank-local parameter and scheduler layout.""" +import bisect import dis import hashlib import json @@ -61,6 +62,93 @@ def _group_option_value(value): return _immutable_value(value) +def _model_storage_alias_graph(names, tensors, safe_rope_pairs=()): + """Describe local view geometry and shared storage without recording addresses. + + CPU shared mappings are rejected: different virtual addresses can reference + the same underlying pages. Accelerator storage must be framework-owned. + """ + from torch.distributed.tensor import DTensor + + intervals, empty_storages, geometry = {}, {}, [] + for object_id, value in tensors.items(): + local = value.to_local() if isinstance(value, DTensor) else value + if local.layout != torch.strided or local.device.type == "meta" or local.is_quantized: + raise RuntimeError("Cannot prove model tensor storage aliases for non-strided, meta or quantized tensors") + storage = local.untyped_storage() + base, size = storage.data_ptr(), storage.nbytes() + if size and not base: + raise RuntimeError("Cannot prove model tensor storage aliases without a storage address") + for name in sorted(names[object_id]): + geometry.append( + ( + name, + local.device.type, + str(local.dtype), + tuple(local.shape), + tuple(local.stride()), + local.storage_offset(), + size, + local.is_conj(), + local.is_neg(), + ) + ) + if size: + intervals.setdefault(str(local.device), []).append((base, base + size, name)) + else: + storage_identity = getattr(storage, "_cdata", None) + if not storage_identity: + raise RuntimeError("Cannot prove empty model tensor storage identity") + empty_storages.setdefault((str(local.device), storage_identity), []).append((0, 0, name)) + + cpu_intervals = intervals.get("cpu", ()) + if cpu_intervals: + try: + with open("/proc/self/maps", "rb") as maps: + mappings = [] + for line in maps: + address, permissions, *_ = line.split(maxsplit=2) + start, end = (int(value, 16) for value in address.split(b"-")) + mappings.append((start, end, permissions)) + except (OSError, ValueError) as exc: + raise RuntimeError("Cannot prove CPU model tensor storage mapping") from exc + starts = [item[0] for item in mappings] + for start, end, _ in cpu_intervals: + position = start + index = bisect.bisect_right(starts, start) - 1 + while position < end: + if index < 0 or index >= len(mappings): + raise RuntimeError("Cannot prove CPU model tensor storage mapping") + mapped_start, mapped_end, permissions = mappings[index] + if not mapped_start <= position < mapped_end: + raise RuntimeError("Cannot prove CPU model tensor storage mapping") + if len(permissions) < 4 or permissions[3:4] != b"p": + raise RuntimeError("Cannot prove aliases for shared CPU model tensor mappings") + position = min(end, mapped_end) + index += 1 + + groups = [] + + def append_group(group): + group_names = frozenset(item[2] for item in group) + if len(group) > 1 and group_names not in safe_rope_pairs: + origin = min(item[0] for item in group) + groups.append(sorted((name, start - origin, end - start) for start, end, name in group)) + + for device_intervals in intervals.values(): + group, group_end = [], None + for item in sorted(device_intervals): + if group_end is not None and item[0] >= group_end: + append_group(group) + group, group_end = [], None + group.append(item) + group_end = item[1] if group_end is None else max(group_end, item[1]) + append_group(group) + for group in empty_storages.values(): + append_group(group) + return {"geometry": sorted(geometry), "storage_groups": sorted(groups)} + + def _lr_layout(group, key, aliases): if key not in group: return None @@ -228,7 +316,7 @@ def check_aliases(value, path): raise RuntimeError("MultiOptimizer and checkpoint model have different canonical parameter identities") from .dynamic_rope import _is_native_qwen3_dynamic_reset_alias, dynamic_rope_checkpoint_contract - model_aliases = {} + model_aliases, model_tensors = {}, {} for kind, items in ( ("parameter", model.named_parameters(remove_duplicate=False)), ("buffer", model.named_buffers(remove_duplicate=False)), @@ -237,6 +325,7 @@ def check_aliases(value, path): model_aliases.setdefault(id(value), set()).update( f"{kind}:{name}" for name in _get_fqns(model, runtime_name) ) + model_tensors[id(value)] = value safe_rope_pairs = set() for module_name, module in model.named_modules(remove_duplicate=False): if _is_native_qwen3_dynamic_reset_alias(module): @@ -391,6 +480,9 @@ def mesh_identity(mesh): layout = { "leaves": leaves, "model_aliases": alias_graph, + # Always bind this proof. Older layouts did not capture storage-sharing + # views and therefore cannot establish an exact resume identity. + "model_storage_aliases_v1": _model_storage_alias_graph(model_aliases, model_tensors, safe_rope_pairs), "model_trainability": sorted(model_trainability.items()), "extra_parallel_specs": extra_parallel_specs, "scheduler_snapshot": _scheduler_snapshot_value(scheduler_snapshot), From 5926c9911a9cecaa83f2ca9fb4849a297a308bc5 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 11 Nov 2026 17:21:16 +0800 Subject: [PATCH 117/120] [docs] fix HF export and cursor test anchors --- docs/seed_omni/mixins/offline_encoding.md | 2 +- docs/seed_omni/usage/hf_layout.md | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/seed_omni/mixins/offline_encoding.md b/docs/seed_omni/mixins/offline_encoding.md index f860a5d4bb..27473a82cb 100644 --- a/docs/seed_omni/mixins/offline_encoding.md +++ b/docs/seed_omni/mixins/offline_encoding.md @@ -371,7 +371,7 @@ A missing module or method still fails before the first request runs. - [`tests/checkpoints/test_checkpoint_callback.py`](../../../tests/checkpoints/test_checkpoint_callback.py): [the writer's position is saved and restored with the cursor, and the cursor-only step is complete](../../../tests/checkpoints/test_checkpoint_callback.py#L632). - [`tests/checkpoints/test_dcp_checkpointer.py`](../../../tests/checkpoints/test_dcp_checkpointer.py): - [a step without `model/` is found only when its manifest says it carries no model state](../../../tests/checkpoints/test_dcp_checkpointer.py#L2749). + [a step without `model/` is found only when its manifest says it carries no model state](../../../tests/checkpoints/test_dcp_checkpointer.py#L2752). - [`tests/seed_omni/trainer/test_step_metrics_callback.py`](../../../tests/seed_omni/trainer/test_step_metrics_callback.py): [a step without an lr scheduler logs no lr](../../../tests/seed_omni/trainer/test_step_metrics_callback.py#L82). - [`tests/seed_omni/model/test_graph.py`](../../../tests/seed_omni/model/test_graph.py): diff --git a/docs/seed_omni/usage/hf_layout.md b/docs/seed_omni/usage/hf_layout.md index 25d5f66603..4245e8136f 100644 --- a/docs/seed_omni/usage/hf_layout.md +++ b/docs/seed_omni/usage/hf_layout.md @@ -155,9 +155,9 @@ Training from an upstream root: - [`load_model_weights`](../../../veomni/models/checkpoint/weights.py#L410) and the [rank-0 broadcast loader](../../../veomni/models/checkpoint/weights.py#L1208) skip keys the converter declares, without reading them; - [`HFSourceKeyConverter`](../../../veomni/models/seed_omni/utils/hf_layout.py#L331) renames the module's own keys and hands them to the module's converter. -HF export (`CheckpointCallback` → [`OmniModelRuntime.save_hf_or_lora`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L619)): +HF export (`CheckpointCallback` → [`OmniModelRuntime.save_hf_or_lora`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L621)): -- [`_save_hf_source_layout`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L632) exports LoRA modules' adapters per module, skips frozen modules, and makes sure each trained module's DCP exists; +- [`_save_hf_source_layout`](../../../veomni/models/seed_omni/accelerated/omni_model/omni_model_runtime.py#L634) exports LoRA modules' adapters per module, skips frozen modules, and makes sure each trained module's DCP exists; - [`save_hf_source_checkpoint`](../../../veomni/models/seed_omni/utils/hf_layout.py#L683) renames keys back (raising on a key with no source key or a different shape), casts to the source dtypes, fills tied duplicates, copies every other source tensor on rank 0, and writes `global_step_N/hf_ckpt/` with the source's shard files and non-weight files. Bare `OmniModel.from_pretrained` and eager inference: From 1687b828722b41f65375182ed538520b419816b0 Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Wed, 11 Nov 2026 17:29:00 +0800 Subject: [PATCH 118/120] [test] exercise checkpoint storage aliases on accelerators --- .../test_multi_optimizer_proven_resume.py | 128 +++++++++++++++++- 1 file changed, 127 insertions(+), 1 deletion(-) diff --git a/tests/checkpoints/test_multi_optimizer_proven_resume.py b/tests/checkpoints/test_multi_optimizer_proven_resume.py index ab1fd12353..3b9e103df7 100644 --- a/tests/checkpoints/test_multi_optimizer_proven_resume.py +++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py @@ -2,13 +2,14 @@ # Licensed under the Apache License, Version 2.0. """CPU regressions for proven MultiOptimizer checkpoint recovery.""" +import os from pathlib import Path from types import SimpleNamespace import pytest import torch -from scripts.ci.accelerator.launch_plan import LaunchSpec +from scripts.ci.accelerator.launch_plan import LaunchSpec, runner_capabilities from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer from veomni.optim.lr_scheduler import build_lr_scheduler from veomni.optim.muon import DistributedMuon @@ -2012,6 +2013,131 @@ def step(model, optimizer): torch.testing.assert_close(target.view, source.view, rtol=0, atol=0) +def _storage_view_roundtrip(tmp_path, device, matching): + def build(shared): + model = torch.nn.Module() + model.p = torch.nn.Parameter(torch.arange(8, dtype=torch.float32, device=device)) + backing = model.p.detach() if shared else model.p.detach().clone() + model.register_buffer("view", backing[1::2]) + optimizer = MultiOptimizer(model, {"leaf": torch.optim.AdamW([model.p], lr=0.01)}, ["leaf"]) + return model, optimizer + + def step(model, optimizer): + model.p.grad = torch.ones_like(model.p) + optimizer.step() + optimizer.zero_grad() + + source, optimizer = build(True) + step(source, optimizer) + parallel = SimpleNamespace(dp_mode="ddp") + DistributedCheckpointer.save( + str(tmp_path), + {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}, + parallel_state=parallel, + ) + target, restored = build(matching) + assert target.view.shape == source.view.shape + assert target.view.stride() == source.view.stride() + assert target.view.storage_offset() == source.view.storage_offset() + assert target.view.untyped_storage().nbytes() == source.view.untyped_storage().nbytes() + assert (target.view.untyped_storage().data_ptr() == target.p.untyped_storage().data_ptr()) is matching + before = (target.p.detach().clone(), target.view.detach().clone()) + state = {"model": target, "optimizer": restored, "_veomni_scheduler_object": None} + if not matching: + from unittest.mock import patch + + import torch.distributed.checkpoint as dcp + + from veomni.checkpoint.dcp_checkpointer import _ProvenModelState + + original_load = dcp.load + + def guard_model_load(*args, **kwargs): + requested = kwargs.get("state_dict", args[0] if args else {}) + assert not isinstance(requested.get("model"), _ProvenModelState), "model tensor load must not start" + return original_load(*args, **kwargs) + + with patch.object(dcp, "load", side_effect=guard_model_load): + with pytest.raises(RuntimeError, match="layout|identity"): + DistributedCheckpointer.load(str(tmp_path), state, parallel_state=parallel) + torch.testing.assert_close(target.p, before[0], rtol=0, atol=0) + torch.testing.assert_close(target.view, before[1], rtol=0, atol=0) + assert not restored.optimizers_dict["leaf"].state + return + DistributedCheckpointer.load(str(tmp_path), state, parallel_state=parallel) + assert target.view.untyped_storage().data_ptr() == target.p.untyped_storage().data_ptr() + torch.testing.assert_close(target.p, source.p, rtol=0, atol=0) + torch.testing.assert_close(target.view, source.view, rtol=0, atol=0) + torch.testing.assert_close(restored.state_dict(), optimizer.state_dict(), rtol=0, atol=0) + step(source, optimizer) + step(target, restored) + torch.testing.assert_close(target.p, source.p, rtol=0, atol=0) + torch.testing.assert_close(target.view, source.view, rtol=0, atol=0) + torch.testing.assert_close(restored.state_dict(), optimizer.state_dict(), rtol=0, atol=0) + + +@pytest.mark.parametrize("matching", [False, True]) +def test_storage_view_checkpoint_reference_cpu(tmp_path, matching): + _storage_view_roundtrip(tmp_path, torch.device("cpu"), matching) + + +def _storage_view_accelerator_plan(context): + if context.capabilities.backend in ("cuda", "npu") and context.capabilities.devices: + return LaunchSpec(device="accelerator") + return LaunchSpec(device="cpu") + + +@pytest.mark.parametrize( + ("backend", "devices", "expected"), + [("cpu", (), "cpu"), ("npu", (), "cpu"), ("npu", ("0",), "accelerator"), ("cuda", ("0",), "accelerator")], +) +def test_storage_view_accelerator_plan_respects_owned_devices(backend, devices, expected): + context = SimpleNamespace(capabilities=SimpleNamespace(backend=backend, devices=devices)) + assert _storage_view_accelerator_plan(context).device == expected + + +@pytest.mark.launch_plan(provider=_storage_view_accelerator_plan) +@pytest.mark.parametrize("matching", [False, True]) +def test_storage_view_checkpoint_native_accelerator(tmp_path, matching): + from veomni.utils.device import IS_NPU_AVAILABLE, get_device_type + + ci_backend = runner_capabilities().backend if os.environ.get("VEOMNI_CI_CAPABILITIES") else None + require_npu = os.environ.get("VEOMNI_REQUIRE_NPU_TESTS") == "1" or ci_backend == "npu" + if require_npu and ci_backend not in (None, "npu"): + pytest.fail("A working npu device must be assigned to this test") + backend = "npu" if require_npu else (ci_backend or get_device_type()) + available = ( + IS_NPU_AVAILABLE and torch.npu.is_available() + if backend == "npu" + else backend == "cuda" and torch.cuda.is_available() + ) + if not available: + if require_npu or ci_backend == "cuda": + pytest.fail(f"A working {backend} device is required for storage-view recovery") + pytest.skip("requires an available CUDA or NPU device") + _storage_view_roundtrip(tmp_path, torch.device(backend, 0), matching) + + +def test_required_storage_view_checkpoint_rejects_missing_npu(tmp_path, monkeypatch): + import veomni.utils.device as device + + monkeypatch.setenv("VEOMNI_REQUIRE_NPU_TESTS", "1") + monkeypatch.setattr(device, "IS_NPU_AVAILABLE", False) + with pytest.raises(pytest.fail.Exception, match="working npu device"): + test_storage_view_checkpoint_native_accelerator(tmp_path, True) + + +def test_required_storage_view_checkpoint_rejects_cuda_assignment(tmp_path, monkeypatch): + import veomni.utils.device as device + + monkeypatch.setenv("VEOMNI_REQUIRE_NPU_TESTS", "1") + monkeypatch.setenv("VEOMNI_CI_CAPABILITIES", "unused") + monkeypatch.setattr(device, "IS_NPU_AVAILABLE", True) + monkeypatch.setitem(globals(), "runner_capabilities", lambda: SimpleNamespace(backend="cuda")) + with pytest.raises(pytest.fail.Exception, match="npu device must be assigned"): + test_storage_view_checkpoint_native_accelerator(tmp_path, True) + + def test_storage_view_geometry_and_disjoint_views_are_distinguished(): from veomni.checkpoint.optimizer_resume_layout import _model_storage_alias_graph From f05a537dadce5c0ee0ec864b9ab45ce0c3129dcf Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 12 Nov 2026 17:43:07 +0800 Subject: [PATCH 119/120] Normalize persistent shard axes in checkpoint preflight --- docs/usage/checkpoint.md | 10 ++++++---- .../test_extra_parallel_layout_v2.py | 10 ++++++++++ veomni/checkpoint/dcp_checkpointer.py | 18 +++++++++++++++--- 3 files changed, 31 insertions(+), 7 deletions(-) diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md index d52cc32606..ccb2700a8f 100644 --- a/docs/usage/checkpoint.md +++ b/docs/usage/checkpoint.md @@ -296,10 +296,12 @@ partial placements, or any active ExtraParallel dimension remain unproven. All ranks in a checkpoint operation must use the same logical model and checkpoint configuration, including the layout version and corresponding sharding specifications. Normal rank-local shard contents and mesh coordinates -may differ. When any rank enables the canonical v2 layout, a preflight checks -that every rank agrees on the logical layout and optimizer policy before any -checkpoint is modified. Legacy-only configurations still rely on this agreement -as an input contract. Rank-local layout validation errors are propagated to peers. +may differ. When any rank enables the canonical v2 layout, a preflight compares +the v2 flag, shard specifications, optimizer presence and trainable-only +selection before any checkpoint is modified. This is not a complete proof of +all runtime provenance or replica equality. Legacy-only configurations still +rely on agreement as an input contract. Rank-local layout validation errors are +propagated to peers. Raw `torch.distributed.checkpoint.load(..., planner=DefaultLoadPlanner(allow_partial_load=True))` is not a supported migration API: it can return success without reading any tensor when diff --git a/tests/checkpoints/test_extra_parallel_layout_v2.py b/tests/checkpoints/test_extra_parallel_layout_v2.py index 0e43833177..c528f7d1f7 100644 --- a/tests/checkpoints/test_extra_parallel_layout_v2.py +++ b/tests/checkpoints/test_extra_parallel_layout_v2.py @@ -656,6 +656,16 @@ def signature(spec): persistent = SpecInfo("ep", Shard(0), "weight", mesh, 0, 1, (8, 8), False) persistent_unused_fields = SpecInfo("ep", Shard(0), "weight", mesh, 1, 1, (8, 8), True) assert signature(persistent) == signature(persistent_unused_fields) + persistent_unknown_shape = SpecInfo("ep", Shard(-1), "weight", mesh, 0, -1, None, False) + persistent_positive = SpecInfo("ep", Shard(1), "weight", mesh, 0, 1, None, False) + assert _checkpoint_extra_parallel_signature( + {"weight": persistent_unknown_shape}, + optimizer_present=False, + trainable_only=False, + runtime_ndims={"weight": 2}, + ) == _checkpoint_extra_parallel_signature( + {"weight": persistent_positive}, optimizer_present=False, trainable_only=False, runtime_ndims={"weight": 2} + ) assert signature(first) != _checkpoint_extra_parallel_signature( {"weight": first}, optimizer_present=False, trainable_only=True ) diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index f820cccf48..07fe16c05e 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -774,6 +774,13 @@ def _reject_unproven_extra_parallel_optimizer(state, parallel_state, *, group=No if trainable_only else () ), + runtime_ndims={ + name: tensor.ndim + for name, tensor in ( + list(state["model"].named_parameters(remove_duplicate=False)) + + list(state["model"].named_buffers(remove_duplicate=False)) + ) + }, ) except BaseException as exc: error = exc @@ -806,7 +813,9 @@ def _reject_unproven_extra_parallel_optimizer(state, parallel_state, *, group=No raise_if_any_rank_failed(error, "checking canonical ExtraParallel runtime tensors", group=group) -def _checkpoint_extra_parallel_signature(specs, *, optimizer_present, trainable_only, trainability=()): +def _checkpoint_extra_parallel_signature( + specs, *, optimizer_present, trainable_only, trainability=(), runtime_ndims=None +): """Describe logical layout, excluding rank-local coordinates and tensor contents.""" entries = [] for name, spec in sorted((specs or {}).items()): @@ -822,8 +831,11 @@ def _checkpoint_extra_parallel_signature(specs, *, optimizer_present, trainable_ if persistent_axis is not None: # Persistent shards do not use the ordinary FSDP axis or v2 marker. fsdp_axis, layout_v2 = None, None - if shape: - ndim = len(shape) + ndim = len(shape) if shape is not None else (runtime_ndims or {}).get(name) + if ndim: + for shard_axis in (axis, fsdp_axis, persistent_axis): + if shard_axis is not None and not -ndim <= shard_axis < ndim: + raise RuntimeError(f"ExtraParallel checkpoint axis is outside the tensor rank: {name}") axis = axis % ndim if axis is not None else None fsdp_axis = fsdp_axis % ndim if fsdp_axis is not None else None persistent_axis = persistent_axis % ndim if persistent_axis is not None else None From e14b678782125768e3977216a6d7c6db420bd45f Mon Sep 17 00:00:00 2001 From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com> Date: Thu, 12 Nov 2026 17:51:23 +0800 Subject: [PATCH 120/120] Check expert provenance before checkpoint overwrite --- tests/checkpoints/test_expert_shard_layout.py | 35 ++++++++++ .../test_extra_parallel_layout_v2.py | 7 ++ veomni/checkpoint/dcp_checkpointer.py | 65 +++++++++++++++---- 3 files changed, 95 insertions(+), 12 deletions(-) diff --git a/tests/checkpoints/test_expert_shard_layout.py b/tests/checkpoints/test_expert_shard_layout.py index ff70fa6458..421700f47f 100644 --- a/tests/checkpoints/test_expert_shard_layout.py +++ b/tests/checkpoints/test_expert_shard_layout.py @@ -536,6 +536,41 @@ def get_parallel_plan(self): result[0] = repr(exc) dist.broadcast_object_list(result, src=0) assert result[0] is None, result[0] + + # Versioned PLE plus an ordinary expert: inconsistent expert provenance + # must fail before replacing an existing checkpoint's completion marker. + import hashlib + from pathlib import Path + + from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer + + model.shard_expert_alias = False + ple_mesh = parallel.extra_parallel_fsdp_device_mesh["ple"] + fsdp_mesh = ple_mesh[ple_mesh.mesh_dim_names[:-1]] + model.ple.weight = torch.nn.Parameter(distribute_tensor(model.ple.weight.detach(), fsdp_mesh, [Shard(1)])) + model._fqn2spec_info["ple.weight"].checkpoint_layout_v2 = True + path = f"{root}-guard" + DistributedCheckpointer.save(path, {"model": model}, parallel_state=parallel) + + def fingerprint(): + return { + str(p.relative_to(path)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in Path(path).rglob("*") + if p.is_file() + } + + before = fingerprint() + if rank == 1: + model._fqn2spec_info["experts.alias"].para_name = "ep" + with pytest.raises(RuntimeError, match="layout or optimizer policy differs across ranks"): + DistributedCheckpointer.save(path, {"model": model}, parallel_state=parallel) + assert fingerprint() == before + if rank == 1: + model._fqn2spec_info["experts.alias"].para_name = "ple" + model.shard_expert_alias = True + with pytest.raises(RuntimeError, match="layout or optimizer policy differs across ranks"): + DistributedCheckpointer.save(path, {"model": model}, parallel_state=parallel) + assert fingerprint() == before finally: clear_parallel_state() dist.destroy_process_group() diff --git a/tests/checkpoints/test_extra_parallel_layout_v2.py b/tests/checkpoints/test_extra_parallel_layout_v2.py index c528f7d1f7..da326f686c 100644 --- a/tests/checkpoints/test_extra_parallel_layout_v2.py +++ b/tests/checkpoints/test_extra_parallel_layout_v2.py @@ -656,6 +656,13 @@ def signature(spec): persistent = SpecInfo("ep", Shard(0), "weight", mesh, 0, 1, (8, 8), False) persistent_unused_fields = SpecInfo("ep", Shard(0), "weight", mesh, 1, 1, (8, 8), True) assert signature(persistent) == signature(persistent_unused_fields) + replicated_expert = SpecInfo("ep", Replicate(), "experts.weight", mesh) + replicated_other = SpecInfo("ple", Replicate(), "experts.weight", mesh) + assert _checkpoint_extra_parallel_signature( + {"experts.weight": replicated_expert}, optimizer_present=False, trainable_only=False + ) != _checkpoint_extra_parallel_signature( + {"experts.weight": replicated_other}, optimizer_present=False, trainable_only=False + ) persistent_unknown_shape = SpecInfo("ep", Shard(-1), "weight", mesh, 0, -1, None, False) persistent_positive = SpecInfo("ep", Shard(1), "weight", mesh, 0, 1, None, False) assert _checkpoint_extra_parallel_signature( diff --git a/veomni/checkpoint/dcp_checkpointer.py b/veomni/checkpoint/dcp_checkpointer.py index 07fe16c05e..d709088b2f 100644 --- a/veomni/checkpoint/dcp_checkpointer.py +++ b/veomni/checkpoint/dcp_checkpointer.py @@ -760,27 +760,41 @@ def _reject_unproven_extra_parallel_optimizer(state, parallel_state, *, group=No if has_v2: raise RuntimeError("Canonical ExtraParallel optimizer state requires an owner-aware checkpoint codec") if has_v2: + named_params = list(state["model"].named_parameters(remove_duplicate=False)) + if ps is None: + ps = parallel_state if parallel_state is not None else get_parallel_state() + aliases = {} + for name, param in named_params: + aliases.setdefault(id(param), []).append(name) + expert_aliases = { + name + for names in aliases.values() + if any("experts" in alias.split(".") for alias in names) + for name in names + } + ordinary_experts = ( + tuple(sorted(_ordinary_expert_fqns(state["model"], ps))) + if isinstance(ps, ParallelState) and any(size != 1 for size in ps.extra_parallel_sizes.values()) + else () + ) logical_signature = _checkpoint_extra_parallel_signature( specs, optimizer_present=state.get("optimizer") is not None, trainable_only=trainable_only, trainability=( - tuple( - sorted( - (name, param.requires_grad) - for name, param in state["model"].named_parameters(remove_duplicate=False) - ) - ) + tuple(sorted((name, param.requires_grad) for name, param in named_params)) if trainable_only else () ), runtime_ndims={ name: tensor.ndim - for name, tensor in ( - list(state["model"].named_parameters(remove_duplicate=False)) - + list(state["model"].named_buffers(remove_duplicate=False)) - ) + for name, tensor in (named_params + list(state["model"].named_buffers(remove_duplicate=False))) }, + expert_aliases=expert_aliases, + ordinary_experts=ordinary_experts, + shared_lora=tuple(sorted(_shared_lora_fqns(state["model"], ordinary_experts))) + if ordinary_experts + else (), ) except BaseException as exc: error = exc @@ -814,7 +828,15 @@ def _reject_unproven_extra_parallel_optimizer(state, parallel_state, *, group=No def _checkpoint_extra_parallel_signature( - specs, *, optimizer_present, trainable_only, trainability=(), runtime_ndims=None + specs, + *, + optimizer_present, + trainable_only, + trainability=(), + runtime_ndims=None, + expert_aliases=(), + ordinary_experts=(), + shared_lora=(), ): """Describe logical layout, excluding rank-local coordinates and tensor contents.""" entries = [] @@ -822,6 +844,18 @@ def _checkpoint_extra_parallel_signature( mesh = getattr(spec, "para_fsdp_mesh", None) shape = getattr(spec, "global_shape", None) placement = getattr(spec, "placement", None) + if isinstance(placement, Replicate) and ("experts" in name.split(".") or name in expert_aliases): + mesh_names = getattr(mesh, "mesh_dim_names", None) + entries.append( + ( + name, + spec.para_name, + "Replicate", + tuple(mesh_names) if mesh_names is not None else None, + tuple(mesh.shape) if mesh is not None else None, + ) + ) + continue if not isinstance(placement, Shard): continue axis = placement.dim @@ -854,7 +888,14 @@ def _checkpoint_extra_parallel_signature( tuple(mesh.shape) if mesh is not None else None, ) ) - return bool(optimizer_present), bool(trainable_only), tuple(trainability), tuple(entries) + return ( + bool(optimizer_present), + bool(trainable_only), + tuple(trainability), + tuple(ordinary_experts), + tuple(shared_lora), + tuple(entries), + ) def _check_extra_parallel_signature_across_ranks(has_v2, signature, *, group=None):