[checkpoint, optim] fix: validate multi-optimizer resume identity

docs/usage/checkpoint.md

@@ -235,6 +235,82 @@ 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 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
 
 - [Checkpoint conversion](checkpoint_conversion.md) — DCP shards → HuggingFace safetensors (`scripts/merge_dcp_to_hf.py`).

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())
@@ -2367,6 +2358,11 @@ 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.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,
         ):