diff --git a/.agents/knowledge/constraints.md b/.agents/knowledge/constraints.md
index f6f0b73e86..e43d7dd4bd 100644
--- a/.agents/knowledge/constraints.md
+++ b/.agents/knowledge/constraints.md
@@ -103,7 +103,7 @@ Core entry points:
    - Token routing: `veomni/distributed/moe/moe_layer.py` — `preprocess()` computes dispatch counts, `token_pre_all2all()` / `tokens_post_all2all()` exchange tokens between EP ranks via `all_to_all` / `all_to_all_async` in `moe/comm.py`.
    - Expert computation: `EPGroupGemm` runs fused expert MLP on grouped tokens per rank.
    - Device mesh: `init_parallel_state_from_config()` builds `[ep × ep_fsdp]` submesh; accessed via `ParallelState.extra_parallel_mesh("ep")`, `ep_group`, `ep_rank`.
-   - In FSDP2: expert modules get `fully_shard()` on the `ep_fsdp` submesh with `Shard(1)` by default. `expert_fsdp_shard_dim=0` explicitly selects whole-expert FSDP shards after validating the EP plan and even local expert divisibility; `None` preserves the existing placement/Muon policy. Record the selected FSDP dimension in `SpecInfo`. If EP/FSDP split the same tensor dimension, DCP must use `_StridedShard(dim, split_factor=ep_size)` on the FSDP mesh axis: physical sharding is EP-first even though the mesh lists FSDP first. Two ordinary `Shard(0)` placements silently permute global expert IDs. VeOmni DCP save/load planners record per-tensor source EP/FSDP dimensions, sizes and expert order in `planner_data`; load compatibility depends on saved coordinates, not the target shard dimension. Unmarked legacy chunks that split only the expert axis are ambiguous and must be migrated using the original training layout; do not bypass the planners for production resumes.
+   - In FSDP2: expert modules get `fully_shard()` on the `ep_fsdp` submesh with `Shard(1)` by default. `expert_fsdp_shard_dim=0` explicitly selects whole-expert FSDP shards after validating the EP plan and even local expert divisibility; `None` preserves the existing placement/Muon policy. Record the selected FSDP dimension in `SpecInfo`. If EP/FSDP split the same tensor dimension on an FSDP-first mesh, DCP must use `_StridedShard(dim, split_factor=ep_size)` on the FSDP axis: physical sharding is EP-first. Two ordinary `Shard(0)` placements on that mesh silently permute global expert IDs. For the namespaced v2 layout, an EP-uneven global axis instead uses an EP-first mesh view with ordinary shards, since Torch strided metadata can give incorrect offsets for that case. VeOmni DCP save/load planners record per-tensor source EP/FSDP dimensions, sizes and expert order in `planner_data`; load compatibility depends on saved coordinates, not the target shard dimension. Unmarked legacy chunks that split only the expert axis are ambiguous and must be migrated using the original training layout; do not bypass the planners for production resumes.
 
    - LoRA adapter export must pass both `SpecInfo.placement.dim` and `SpecInfo.fsdp_shard_dim` to `restore_extra_parallel_dim`, just like regular model checkpoints. Assuming dim-1 FSDP for whole-expert shards corrupts global adapter shapes and expert order during safetensors consolidation.
 
diff --git a/docs/design/op_selection.md b/docs/design/op_selection.md
index 6d94e1ca27..6be875cedb 100644
--- a/docs/design/op_selection.md
+++ b/docs/design/op_selection.md
@@ -752,6 +752,45 @@ 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. 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. 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.
+
+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/docs/design/rank_local_adamw_checkpoint_adapter.md b/docs/design/rank_local_adamw_checkpoint_adapter.md
new file mode 100644
index 0000000000..be85913be1
--- /dev/null
+++ b/docs/design/rank_local_adamw_checkpoint_adapter.md
@@ -0,0 +1,136 @@
+# 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.
+
+### Save and load call sites
+
+[`_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,
+[`_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#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)
+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)
+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
+
+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
+   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()`](../../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
+   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 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.
+
+## 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.
diff --git a/docs/design/selective_expert_cpu_offload.md b/docs/design/selective_expert_cpu_offload.md
new file mode 100644
index 0000000000..4b2e2a4112
--- /dev/null
+++ b/docs/design/selective_expert_cpu_offload.md
@@ -0,0 +1,66 @@
+# 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. The list changes placement, while model
+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#L476)
+forwards these settings to
+[`parallelize_model_fsdp2()`](../../veomni/distributed/torch_parallelize.py#L435).
+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.
+- `model.optimizer.state_offload: cpu` cannot be combined with this selective
+  parameter offload; the configuration is rejected before training.
+- 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 9eea12101b..6a10adaf40 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -131,9 +131,11 @@ seed_omni/models/bagel.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
+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/docs/seed_omni/mixins/metric_meter.md b/docs/seed_omni/mixins/metric_meter.md
index 7fc1f6d892..e47e03a6e4 100644
--- a/docs/seed_omni/mixins/metric_meter.md
+++ b/docs/seed_omni/mixins/metric_meter.md
@@ -162,8 +162,8 @@ reported per module under `trace/<module>/` 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#L295), so
+[saves](../../../veomni/trainer/callbacks/global_state_callback.py#L253) and
+[restores](../../../veomni/trainer/callbacks/global_state_callback.py#L505), so
 the cumulative counters survive a resume.
 
 ## Rules for a correct implementation
diff --git a/docs/seed_omni/mixins/module_preprocessor.md b/docs/seed_omni/mixins/module_preprocessor.md
index 48e5a405a1..ae55dfc564 100644
--- a/docs/seed_omni/mixins/module_preprocessor.md
+++ b/docs/seed_omni/mixins/module_preprocessor.md
@@ -114,19 +114,20 @@ class VisionPreprocessor(ModulePreprocessorBase):
 | [`from_pretrained(module_path, *, config_overrides=None, **kwargs)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L99) | **May override.** Build and return the CPU preprocessor from the module's checkpoint subfolder without instantiating the `nn.Module`. Default returns `None`. | [`OmniProcessor.from_config`](../../../veomni/models/seed_omni/processing_omni.py#L109) and [`bind_module_assets`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L128) | Returning `None` means the module contributes no CPU preprocessor (for example a pure AR backbone). |
 | [`bind_dummy_inputs(config, dtype=None)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L119) | **May override.** Pre-compute dummy tensor(s) and token lengths from the module's resolved `config` and `dtype` for the `inference=False` branch. Default is a no-op. | [`OmniProcessor.bind_dummy_inputs`](../../../veomni/models/seed_omni/processing_omni.py#L88) during [`OmniTrainer._build_data_transform`](../../../veomni/trainer/omni/omni_trainer.py#L288) | Attaches FSDP-anchor dummy shapes before DataLoader workers fork/spawn. |
 
-### What the framework provides (do not override)
+### Framework default behavior
 
 | 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#L239) | 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. |
+| [`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#L239) | By default, delegates `batch["conversation_list"]` to `preprocess_conversations` in `config.module_names` order. Packed preprocessors may override this path. |
+| [`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#L379) | 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
 
 - **Training sample & micro-batch**:
   - [`OmniTrainer._build_data_transform`](../../../veomni/trainer/omni/omni_trainer.py#L288) builds [`OmniProcessor.from_config`](../../../veomni/models/seed_omni/processing_omni.py#L109), calls [`OmniProcessor.bind_dummy_inputs`](../../../veomni/models/seed_omni/processing_omni.py#L88), and passes `processor` into [`process_seedomni_example`](../../../veomni/data/seed_omni/seedomni_transform.py#L183) / [`process_seedomni_cached_example`](../../../veomni/data/seed_omni/seedomni_transform.py#L254).
-  - For each dataset sample, `seedomni_transform` calls [`OmniProcessor.preprocess(conversation_list)`](../../../veomni/models/seed_omni/processing_omni.py#L216) → [`OmniProcessor.preprocess_batch(..., inference=False)`](../../../veomni/models/seed_omni/processing_omni.py#L239), running each module's [`ModulePreprocessorBase.__call__`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L85) → [`preprocess_conversations`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L89) to mutate items in place, populate `_omni_token_len` / `labels`, and append per-sample `is_dummy=True` placeholders.
+  - On the default conversation-list path, for each dataset sample, `seedomni_transform` calls [`OmniProcessor.preprocess(conversation_list)`](../../../veomni/models/seed_omni/processing_omni.py#L216) → [`OmniProcessor.preprocess_batch(..., inference=False)`](../../../veomni/models/seed_omni/processing_omni.py#L239), running each module's [`ModulePreprocessorBase.__call__`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L85) → [`preprocess_conversations`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L89) to mutate items in place, populate `_omni_token_len` / `labels`, and append per-sample `is_dummy=True` placeholders.
   - When `train.dyn_bsz: true`, [`get_length_fn`](../../../veomni/data/dataset.py#L444) (`count_mode="total"` and `physical_token_cap`, including `is_dummy=True` items) and [`get_length_by_labels_fn`](../../../veomni/data/dataset.py#L480) (`count_mode="effective"`, excluding `is_dummy=True` items) read those metadata fields to batch samples by token budget, and [`SeedOmniCollator.__call__`](../../../veomni/data/seed_omni/collator.py#L102) groups the selected sample dicts into a micro-batch.
+  - **Packed training exception**: Qwen3-VL can enable `packed_preprocess`; its text preprocessor consumes `conversation_list` and emits packed tensors. `OmniProcessor.preprocess` retains the resulting batch dict, and `SeedOmniCollator` concatenates its packed tensors and reconstructs sample boundaries.
 - **Inference request**:
   - [`OmniInferencer._build_model`](../../../veomni/trainer/omni/omni_inferencer.py#L120) builds [`OmniProcessor.from_config`](../../../veomni/models/seed_omni/processing_omni.py#L109).
   - [`OmniInferencer._run`](../../../veomni/trainer/omni/omni_inferencer.py#L237) calls [`OmniProcessor.__call__(..., inference=True)`](../../../veomni/models/seed_omni/processing_omni.py#L167) → [`OmniProcessor.preprocess`](../../../veomni/models/seed_omni/processing_omni.py#L216) → [`OmniProcessor.preprocess_batch(..., inference=True)`](../../../veomni/models/seed_omni/processing_omni.py#L239).
diff --git a/docs/seed_omni/mixins/offline_encoding.md b/docs/seed_omni/mixins/offline_encoding.md
index 21d54fb293..bd84ccfc45 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#L638).
 - [`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#L2706).
+  [a step without `model/` is found only when its manifest says it carries no model state](../../../tests/checkpoints/test_dcp_checkpointer.py#L2798).
 - [`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 a0def1c4e6..0ca728d548 100644
--- a/docs/seed_omni/usage/hf_layout.md
+++ b/docs/seed_omni/usage/hf_layout.md
@@ -151,16 +151,16 @@ 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) resolves through `HFSource.weights_path`.
+  [attaches the key converter](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L372);
+  [`weights_path`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L255) resolves through `HFSource.weights_path`.
 - 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#L391) 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#L633)):
 
-- [`_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_checkpoint`](../../../veomni/models/seed_omni/utils/hf_layout.py#L730) extracts each module's save state (via `layout.extract_module_save_state`), 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.
+- [`_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#L737) extracts each module's save state (via `layout.extract_module_save_state`), 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 953a58f93f..5901372fb3 100644
--- a/docs/usage/arguments.md
+++ b/docs/usage/arguments.md
@@ -524,7 +524,7 @@ The default `mode=None` follows TorchTitan's main path by using the `inductor` b
 | muon_expert_zero_comm | `bool` | `False` | Use whole-expert `Shard(0)` when the FSDP+ExtraParallel topology permits zero-communication expert Muon updates. |
 | muon_ns_implementation | `Literal["std", "gram", "gram_quack"]` | `"gram_quack"` | Newton–Schulz backend: standard, pure-PyTorch Gram-NS, or Gram-NS with quack kernels (default; falls back to `gram` if unavailable). |
 | muon_gram_ns_reset_iterations | `List[int]` | `[2]` | Restart indices for Gram Newton–Schulz (ignored by `std`). |
-| state_offload | `Literal["none", "cpu"]` | `"none"` | Optimizer-managed state offload. `none` disables it (training state may still be offloaded by `model.accelerator.fsdp_config.offload`); `cpu` keeps the moment states in host buffers and streams them to the device in batches, lowering peak device memory (only with `type="adamw"`). |
+| state_offload | `Literal["none", "cpu"]` | `"none"` | Optimizer-managed state offload. `none` disables it (training state may still be offloaded by `model.accelerator.fsdp_config.offload`); `cpu` keeps the moment states in host buffers and streams them to the device in batches, lowering peak device memory (only with `type="adamw"`). It cannot be combined with selective ExtraParallel parameter CPU offload. |
 | state_offload_device_memory_fraction | `float` | `0.8` | With `state_offload="cpu"`, fraction of currently free device memory used to size one streamed batch of optimizer states. It is recomputed every step; a single parameter whose states exceed it cannot be swapped in and `step()` raises. |
 | state_offload_pin_memory | `bool` | `True` | With `state_offload="cpu"`, pin the host-side optimizer state buffers. Disable if page-locked memory is scarce. A failed pinned allocation degrades to pageable buffers (rank-0 warning) instead of failing the run. |
 
@@ -931,6 +931,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
 
diff --git a/docs/usage/basic_modules.md b/docs/usage/basic_modules.md
index bf5f22b789..68d099a72e 100644
--- a/docs/usage/basic_modules.md
+++ b/docs/usage/basic_modules.md
@@ -390,6 +390,13 @@ states in host buffers — pinned by default, pageable when `state_offload_pin_m
 is `false` or a pinned allocation fails — and streams them to the device in
 batches, lowering peak device memory.
 
+This optimizer's moments use the parameter dtype and it does not create FP32
+master weights. BF16 parameters therefore have BF16 moments and BF16 parameter
+writes; select an optimizer with the desired precision policy when FP32 masters
+are required.
+`state_offload: cpu` cannot be combined with selective ExtraParallel parameter
+CPU offload (`cpu_offload_extra_parallel_names`).
+
 Muon-specific hyperparameters live under `model.optimizer.muon_*` (e.g. `muon_lr`, `muon_momentum`, `muon_adjust_lr_fn`); `lr` / `weight_decay` / `betas` continue to drive the AdamW sibling group.
 
 Muon-specific knobs (only consulted when `optimizer.type == "muon"`):
diff --git a/docs/usage/checkpoint.md b/docs/usage/checkpoint.md
index 1b18feafd2..cc771a60c5 100644
--- a/docs/usage/checkpoint.md
+++ b/docs/usage/checkpoint.md
@@ -24,6 +24,31 @@ 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.
 `<module>/` 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.
+
+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 |
@@ -134,7 +159,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#L254), and the manifest
+[`model_state: false`](../../veomni/checkpoint/layout.py#L185), 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.
@@ -147,6 +172,35 @@ new, partial data. Each half clears its own, right before it writes it back:
 the step is correctly rejected throughout, including in the window where one
 half has been invalidated and the other has not.
 
+For a `MultiOptimizer`, a completed checkpoint is not overwritten in place:
+the save fails before changing its manifest or shards. When rewinding to an
+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.
+
+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 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
+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
 instead of at the start. A staged save that fails leaves the previous checkpoint
@@ -216,8 +270,317 @@ such a run, restart it, or convert the weights out first with
 [Checkpoint conversion](checkpoint_conversion.md) and resume from the converted
 export as a fresh run.
 
+## 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 versioned EP-outer tensor namespace is opt-in through
+`SpecInfo.checkpoint_layout_v2`. Ordinary ExtraParallel saves also use canonical
+EP-first shard coordinates and record per-tensor source layout. A versioned
+plan must retain the original `global_shape` and 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 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 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.
+
+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
+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.
+
+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.
+New v2 saves also reject expert layouts marked as unproven before publishing
+DCP metadata.
+
+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
+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 for EP/FSDP as well as ordinary
+FSDP or unsharded targets. An old expert-only FSDP checkpoint without layout
+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
+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
+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.
+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
+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.
+
+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
+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
+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.
+  Preserve the original axis notation too: the current optimizer identity
+  treats `Shard(0)` and `Shard(-2)`, or axis `1` and `-1`, as different even
+  when they describe equivalent two-dimensional shards.
+- 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.
+
+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
+checkpoint uses `MultiOptimizer` is rejected before training begins. 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, SwapAdamW 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. 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. 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
+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, and hooks must not mutate checkpoint
+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. 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
+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.
+When a proven `MultiOptimizer` includes `SwapAdamW`, restore first reads its
+moments into separate CPU buffers. Budget approximately one additional copy of
+the saved SwapAdamW moment shards in host memory while loading, plus DCP's own
+temporary storage. Those buffers are discarded after restore; they do not add
+accelerator-resident optimizer state. A failed restore marks the optimizer
+unusable until it is rebuilt. SwapAdamW stores its moments in the parameter
+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 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.
+`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 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. 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 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
+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
+sharded. Such known expert tensors require a verified migration with the
+original training layout. Offline export can recognize VeOmni's standard
+`experts` names; custom names require the operator to verify source layout.
+
+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, 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.
+
+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`).
+- [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/scripts/ci/accelerator/suites/gpu.json b/scripts/ci/accelerator/suites/gpu.json
index f9b198dfe0..beb9246d53 100644
--- a/scripts/ci/accelerator/suites/gpu.json
+++ b/scripts/ci/accelerator/suites/gpu.json
@@ -43,6 +43,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",
@@ -459,6 +469,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 63adc56c05..838012df7e 100644
--- a/scripts/ci/accelerator/suites/npu.json
+++ b/scripts/ci/accelerator/suites/npu.json
@@ -53,6 +53,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",
@@ -215,6 +225,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/scripts/merge_dcp_to_hf.py b/scripts/merge_dcp_to_hf.py
index e1595c2389..b5dc43adb5 100644
--- a/scripts/merge_dcp_to_hf.py
+++ b/scripts/merge_dcp_to_hf.py
@@ -8,10 +8,11 @@
 
 import torch
 from safetensors.torch import save_file
+from torch.distributed.checkpoint import FileSystemReader
 from transformers import AutoConfig, AutoProcessor
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME
 
-from veomni.checkpoint.dcp_checkpointer import _get_sharding_plan, _process_shard
+from veomni.checkpoint.dcp_checkpointer import _get_sharding_plan, _normalize_key, _process_shard
 from veomni.checkpoint.layout import weights_dir
 from veomni.utils import helper
 
@@ -46,6 +47,8 @@ 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``.
 
@@ -59,15 +62,22 @@ def save_lora_adapter_weights(
     logger.info(f"Saving LoRA adapter to {output_dir}")
 
     # ``shard_size=None`` forces a single shard so we get ``adapter_model.safetensors`` directly.
-    all_keys, _total_size, _all_dcp_keys = _get_sharding_plan(checkpoint_path, shard_size=None, save_dtype=save_dtype)
-    lora_keys = {hf_k: dcp_k for hf_k, dcp_k in all_keys.items() if _is_lora_key(hf_k)}
+    lora_keys, _total_size, _all_dcp_keys = _get_sharding_plan(
+        checkpoint_path,
+        shard_size=None,
+        save_dtype=save_dtype,
+        legacy_expert_source_layout=legacy_expert_source_layout,
+        include_hf_key=_is_lora_key,
+    )
     if not lora_keys:
         raise RuntimeError(
             f"LoRA conversion requested but no LoRA keys (.lora_A./.lora_B./...) found under {checkpoint_path}"
         )
 
     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 +111,8 @@ 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)
@@ -111,7 +123,9 @@ def save_model_weights(
 
     # Plan shards from metadata
     logger.info("Analyzing DCP metadata and planning shards...")
-    shards, total_size, all_dcp_keys = _get_sharding_plan(checkpoint_path, shard_size, save_dtype)
+    shards, total_size, all_dcp_keys = _get_sharding_plan(
+        checkpoint_path, shard_size, save_dtype, legacy_expert_source_layout=legacy_expert_source_layout
+    )
 
     logger.info(f"Found {len(all_dcp_keys)} model tensors, total size: ~{total_size / 1e9:.2f}GB")
     logger.info(f"Split into {len(shards)} shards")
@@ -135,7 +149,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 +192,13 @@ 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 +220,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 +273,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 +301,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
@@ -270,9 +316,9 @@ def main():
 
     mode = args.mode
     if mode == "auto":
-        _shards_for_detection, _size, _dcp_keys = _get_sharding_plan(load_dir, shard_size=None, save_dtype="bfloat16")
-        # _shards_for_detection is a single {hf_key: dcp_key} dict when shard_size is None
-        detected_lora = _detect_lora(_shards_for_detection.keys())
+        metadata = FileSystemReader(load_dir).read_metadata()
+        hf_keys = (normalized for key in metadata.state_dict_metadata if (normalized := _normalize_key(key)))
+        detected_lora = _detect_lora(hf_keys)
         mode = "lora" if detected_lora else "full"
         logger.info(
             f"Auto-detected mode: {mode} "
@@ -283,12 +329,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/arguments/test_checkpoint_config.py b/tests/arguments/test_checkpoint_config.py
new file mode 100644
index 0000000000..4c7570f253
--- /dev/null
+++ b/tests/arguments/test_checkpoint_config.py
@@ -0,0 +1,52 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# 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
+
+
+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
+
+
+@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/tests/arguments/test_model_arguments.py b/tests/arguments/test_model_arguments.py
index 3a163fa90b..9b6da6ec36 100644
--- a/tests/arguments/test_model_arguments.py
+++ b/tests/arguments/test_model_arguments.py
@@ -74,6 +74,15 @@ def test_state_offload_rejects_fsdp_cpu_offload():
         )
 
 
+def test_state_offload_rejects_selective_extra_parallel_cpu_offload():
+    with pytest.raises(ValueError, match="cannot be combined with selective ExtraParallel CPU offload"):
+        ModelArguments(
+            model_path="weights",
+            optimizer=OptimizerConfig(state_offload="cpu"),
+            accelerator=AcceleratorConfig(fsdp_config=FSDPConfig(cpu_offload_extra_parallel_names=["ep"])),
+        )
+
+
 @pytest.mark.parametrize(
     "kwargs,compile_enabled",
     [
diff --git a/tests/checkpoints/test_async_dcp_recovery.py b/tests/checkpoints/test_async_dcp_recovery.py
new file mode 100644
index 0000000000..fdb70d0aa7
--- /dev/null
+++ b/tests/checkpoints/test_async_dcp_recovery.py
@@ -0,0 +1,82 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exercise real asynchronous DCP without MultiOptimizer."""
+
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from scripts.ci.accelerator.launch_plan import LaunchSpec
+from tests.tools.launch_utils import spawn_with_rendezvous
+from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
+
+
+def _build():
+    torch.manual_seed(7)
+    model = torch.nn.Linear(2, 2)
+    return model, torch.optim.AdamW(model.parameters(), lr=0.01)
+
+
+def _step(model, optimizer, inputs):
+    loss = model(inputs).square().sum()
+    loss.backward()
+    grads = {name: param.grad.detach().clone() for name, param in model.named_parameters()}
+    optimizer.step()
+    optimizer.zero_grad()
+    return loss.detach(), grads
+
+
+def _snapshot(model, optimizer):
+    return {
+        name: {
+            "param": param.detach().clone(),
+            **{key: value.detach().clone() for key, value in optimizer.state[param].items() if torch.is_tensor(value)},
+        }
+        for name, param in model.named_parameters()
+    }
+
+
+def _assert_equal(left, right):
+    assert left.keys() == right.keys()
+    for name, fields in left.items():
+        assert fields.keys() == right[name].keys()
+        for field, value in fields.items():
+            assert torch.equal(value, right[name][field]), f"{name}.{field} differs"
+
+
+def _async_resume_worker(rank, checkpoint, endpoint):
+    dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=45))
+    try:
+        parallel = SimpleNamespace(dp_mode="ddp")
+        model, optimizer = _build()
+        _step(model, optimizer, torch.tensor([[1.0, 2.0]]))
+        before = _snapshot(model, optimizer)
+        DistributedCheckpointer.save(
+            checkpoint, {"model": model, "optimizer": optimizer}, save_async=True, parallel_state=parallel
+        )
+        DistributedCheckpointer.wait_for_pending_save()
+        resumed, restored_optimizer = _build()
+        with torch.no_grad():
+            for param in resumed.parameters():
+                param.fill_(-1)
+        DistributedCheckpointer.load(
+            checkpoint, {"model": resumed, "optimizer": restored_optimizer}, parallel_state=parallel
+        )
+        _assert_equal(before, _snapshot(resumed, restored_optimizer))
+        loss, grads = _step(model, optimizer, torch.tensor([[3.0, 4.0]]))
+        restored_loss, restored_grads = _step(resumed, restored_optimizer, torch.tensor([[3.0, 4.0]]))
+        assert torch.equal(loss, restored_loss)
+        assert grads.keys() == restored_grads.keys()
+        assert all(torch.equal(value, restored_grads[name]) for name, value in grads.items())
+        _assert_equal(_snapshot(model, optimizer), _snapshot(resumed, restored_optimizer))
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2))
+def test_async_adamw_recovery_preserves_next_update(tmp_path):
+    spawn_with_rendezvous(_async_resume_worker, args=(str(tmp_path / "checkpoint"),), nprocs=2)
diff --git a/tests/checkpoints/test_checkpoint_callback.py b/tests/checkpoints/test_checkpoint_callback.py
index a3832515de..70e4e40ecd 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)
@@ -699,8 +721,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_a_pre_split_job_cursor_is_not_read(self, mock_dist, tmp_path):
         """Only ``extra_state/rank_{R}.pt`` carries the cursor now.
diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py
index fd8e45cd0c..8da950392b 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 (
@@ -39,9 +40,91 @@
 _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], "xpu", mesh_dim_names=("ep_fsdp",))
+        full_mesh = DeviceMesh.from_group(
+            [singleton_groups[rank], dist.group.WORLD],
+            "xpu",
+            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
+    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 +158,23 @@ 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, "_veomni_scheduler_object": None},
         )
         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, "_veomni_scheduler_object": None},
         )
 
-        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 +245,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 +265,24 @@ 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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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 +303,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 +318,18 @@ 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, "_veomni_scheduler_object": None},
+            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())
@@ -619,7 +705,8 @@ def test_a_marker_missing_on_this_rank_alone_stops_every_rank(self, mock_gps, tm
         model, optimizer, _ = self._build(seed=1)
 
         with (
-            patch("veomni.utils.dist_utils.any_rank_failed", return_value=True) as agreed,
+            patch("veomni.utils.dist_utils.any_rank_failed", side_effect=lambda failed, group=None: failed) as agreed,
+            patch("veomni.checkpoint.dcp_checkpointer._reject_unproven_extra_parallel_optimizer"),
             patch.object(DistributedCheckpointer, "_create_storage_reader"),
             patch("veomni.checkpoint.dcp_checkpointer.dcp") as mock_dcp,
         ):
@@ -631,7 +718,7 @@ def test_a_marker_missing_on_this_rank_alone_stops_every_rank(self, mock_gps, tm
                     state={"model": model, "optimizer": optimizer, "extra_state": {}},
                 )
 
-        agreed.assert_called_once()
+        assert agreed.call_args.args[0] is True
         mock_dcp.load.assert_not_called()
 
 
@@ -657,9 +744,9 @@ def test_load_passes_partial_planner_to_dcp(self, mock_dcp, mock_gps, tmp_path):
             Path(directory).mkdir(parents=True)
             (Path(directory) / layout.DCP_MARKER_FILENAME).write_text("complete")
 
-        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": {}}
 
@@ -744,7 +831,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,
         )
@@ -786,7 +873,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",
         )
@@ -855,7 +942,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,
         )
@@ -1095,7 +1182,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",
@@ -2525,6 +2612,8 @@ 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_expert_noep_checkpoint.py b/tests/checkpoints/test_expert_noep_checkpoint.py
new file mode 100644
index 0000000000..03367b79be
--- /dev/null
+++ b/tests/checkpoints/test_expert_noep_checkpoint.py
@@ -0,0 +1,214 @@
+# 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_world(geometry):
+    return 4 if geometry == "hsdp" else 2
+
+
+def _noep_worker(rank, directory, geometry, optimizer_kind, phase, endpoint):
+    world = _noep_world(geometry)
+    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=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 = _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"
+    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_world(mode):
+    return 4 if mode == "tp" else 2
+
+
+def _unproven_source_worker(rank, directory, mode, endpoint):
+    world = _unproven_world(mode)
+    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 layouts[key]["expert_order"] == "unproven"
+        target = {"model": {"layer.experts.weight": torch.zeros(8, 4)}}
+        with pytest.raises(dcp.CheckpointException, match="noncanonical 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=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=_unproven_world(mode))
+
+
+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 d342d66003..f760b45d2e 100644
--- a/tests/checkpoints/test_expert_shard_layout.py
+++ b/tests/checkpoints/test_expert_shard_layout.py
@@ -3,29 +3,47 @@
 
 """CPU/Gloo regressions for canonical EP-first DCP expert coordinates."""
 
+import os
 from datetime import timedelta
+from pathlib import Path
+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 DTensor, Replicate, Shard
+from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+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,
+    DistributedCheckpointer,
+    ModelState,
+    _expert_layout,
+    _extra_parallel_first_mesh,
+    _extra_parallel_layout_key,
     _ExtraParallelSavePlanner,
+    _get_sharding_plan,
     _ModelStrictLoadPlanner,
+    _process_shard,
+    dcp_to_torch_state_dict,
     drop_extra_parallel_dim,
     restore_extra_parallel_dim,
 )
 
 
+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))
+    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)
@@ -36,6 +54,25 @@ def _layout_worker(rank, world, directory, rendezvous):
         global_weight = torch.arange(8 * 16 * 4, dtype=torch.float32).reshape(8, 16, 4)
         ep_weight = global_weight.chunk(2, 0)[ep_rank]
 
+        # EP-only rank-one values must ignore the unused default FSDP axis.
+        vector = torch.arange(4, dtype=torch.float32) + ep_rank
+        plain = restore_extra_parallel_dim(vector, mesh, submesh)
+        torch.testing.assert_close(plain.to_local(), vector, rtol=0, atol=0)
+        expected_vector = torch.cat((torch.arange(4), torch.arange(4) + 1)).float()
+        torch.testing.assert_close(plain.full_tensor(), expected_vector, rtol=0, atol=0)
+        vector_path = directory + "ep_only_vector"
+        dcp.save({"model": {"experts.vector": plain}}, checkpoint_id=vector_path, planner=_ExtraParallelSavePlanner())
+        target = restore_extra_parallel_dim(torch.zeros_like(vector), mesh, submesh)
+        dcp.load(
+            {"model": {"experts.vector": target}},
+            checkpoint_id=vector_path,
+            planner=_ModelStrictLoadPlanner(strict_model=True),
+        )
+        torch.testing.assert_close(target.to_local(), vector, rtol=0, atol=0)
+        fsdp_vector = DTensor.from_local(vector, submesh, [Shard(0)] if world == 4 else [Replicate(), Shard(0)])
+        with pytest.raises(RuntimeError, match="ExtraParallel shard dimension is outside"):
+            restore_extra_parallel_dim(fsdp_vector, mesh, submesh)
+
         def wrap(dim, *, empty=False):
             local = ep_weight.chunk(2, dim)[fsdp_rank].contiguous()
             if empty:
@@ -72,15 +109,98 @@ def wrap(dim, *, empty=False):
             local_dt = drop_extra_parallel_dim(loaded, submesh, 1 - source_dim)
             torch.testing.assert_close(local_dt.full_tensor(), ep_weight, rtol=0, atol=0)
 
-        # Old row-sharded files are identifiable from the saved chunks, and
-        # can load into either target dimension without a source version tag.
+        # An unversioned row-sharded source can have identical chunks to a
+        # misordered composite TP+EP source. Require source layout evidence.
         path = directory + "legacy_rows"
         dcp.save({"model": {"weight": wrap(1)}}, checkpoint_id=path)
         for target_dim in (0, 1):
             target = {"model": {"weight": wrap(target_dim, empty=True)}}
-            dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=True))
+            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)
+        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(RuntimeError, 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"
+            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,
+                    str(repo_root / "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,
+                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)
+            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()
         placements = [Shard(0), Shard(0)] if world == 4 else [Replicate(), Shard(0), Shard(0)]
@@ -94,19 +214,44 @@ def wrap(dim, *, empty=False):
                     dcp.load(target, checkpoint_id=path, planner=_ModelStrictLoadPlanner(strict_model=strict))
                 assert target[key]["weight"].to_local().count_nonzero() == 0
 
+        # Offline export targets plain CPU tensors: an unversioned stacked
+        # expert saved only in same-axis chunks cannot prove expert order.
+        path = directory + "legacy_ambiguous_export"
+        dcp.save({"model": {"layer.experts.weight": legacy}}, checkpoint_id=path)
+        if rank == 0:
+            with pytest.raises(RuntimeError, match="Cannot determine source expert order"):
+                _get_sharding_plan(path)
+            with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"):
+                _process_shard({"layer.experts.weight": "model.layer.experts.weight"}, 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(RuntimeError, 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()
 
@@ -119,6 +264,145 @@ 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)])
+        # Emulate a legacy checkpoint written before per-tensor provenance existed.
+        dcp.save({"model": {"layer.experts.bias": source}}, checkpoint_id=root, planner=DefaultSavePlanner())
+        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}}
+
+        class OldPartialPlanner(_ExtraParallelSavePlanner):
+            def create_global_plan(self, all_plans):
+                plans, metadata = super().create_global_plan(all_plans)
+                layouts = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"]
+                metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"] = {
+                    key: value for key, value in layouts.items() if value["expert_order"] == "canonical"
+                }
+                return plans, metadata
+
+        dcp.save(source, checkpoint_id=root, planner=OldPartialPlanner())
+        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(RuntimeError, 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:
+        from torch.distributed.tensor.placement_types import _StridedShard
+
+        mesh = init_device_mesh("cpu", (2, 2, 2), mesh_dim_names=("tp", "ep_fsdp", "ep"))
+        tp, fsdp, ep = mesh.get_coordinate()
+        expected = torch.arange(8 * 16 * 4, dtype=torch.float32).reshape(8, 16, 4)
+        local = expected.chunk(2, 0)[ep].chunk(2, 0)[fsdp].chunk(2, 1)[tp].contiguous()
+        old = DTensor.from_local(local, mesh, [Shard(1), Shard(0), Shard(0)])
+        dcp.save({"model": {"layer.experts.weight": old}}, checkpoint_id=root)
+        target = DTensor.from_local(
+            torch.zeros_like(local), mesh, [Shard(1), _StridedShard(0, split_factor=2), Shard(0)]
+        )
+        with pytest.raises(dcp.CheckpointException, match="Cannot determine source expert order"):
+            dcp.load(
+                {"model": {"layer.experts.weight": target}}, checkpoint_id=root, planner=_ModelStrictLoadPlanner(True)
+            )
+        assert target.to_local().count_nonzero() == 0
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=8))
+def test_unversioned_composite_expert_layout_is_rejected(tmp_path):
+    spawn_with_rendezvous(_unversioned_composite_expert_worker, args=(str(tmp_path / "checkpoint"),), nprocs=8)
+
+
 def _lora_model(config=None):
     from veomni.lora import VeOmniLoraConfig, VeOmniLoraModel
 
@@ -134,7 +418,88 @@ def _lora_model(config=None):
     return VeOmniLoraModel(model, config)
 
 
-def _lora_export_worker(rank, world, shard_dim, directory, endpoint):
+@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 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
     from veomni.lora.state_dict import get_lora_state_dict, load_adapter_state_dict, strip_adapter_name
@@ -158,7 +523,15 @@ def _lora_export_worker(rank, world, shard_dim, directory, endpoint):
             sharded = DTensor.from_local(local, mesh["ep_fsdp"], [Shard(shard_dim)])
             module_name, param_name = name.rsplit(".", 1)
             setattr(model.get_submodule(module_name), param_name, torch.nn.Parameter(sharded))
-            model._fqn2spec_info[name] = SpecInfo("ep", Shard(0), name, mesh, fsdp_shard_dim=shard_dim)
+            model._fqn2spec_info[name] = SpecInfo(
+                "ep",
+                Shard(ep_dim),
+                name,
+                mesh,
+                fsdp_shard_dim=shard_dim,
+                global_shape=tuple(value.shape),
+                checkpoint_layout_v2=canonical,
+            )
 
         save_lora_adapter_with_dcp(model, directory)
         exported = load_adapter_state_dict(directory)
@@ -177,6 +550,1068 @@ def _lora_export_worker(rank, world, shard_dim, directory, endpoint):
 
 @pytest.mark.launch_plan(provider=from_parameter("world", device="cpu"))
 @pytest.mark.parametrize("world", [4])
-@pytest.mark.parametrize("shard_dim", [0, 1])
-def test_lora_adapter_export_preserves_expert_shard_layout(world, shard_dim, tmp_path):
-    spawn_with_rendezvous(_lora_export_worker, args=(world, shard_dim, str(tmp_path / "adapter")), nprocs=world)
+@pytest.mark.parametrize("shard_dim,ep_dim", [(0, 0), (1, 0), (0, -3)])
+@pytest.mark.parametrize("canonical", [False, True])
+def test_lora_adapter_export_preserves_expert_shard_layout(world, shard_dim, ep_dim, canonical, tmp_path):
+    spawn_with_rendezvous(
+        _lora_export_worker, args=(world, shard_dim, str(tmp_path / "adapter"), ep_dim, canonical), nprocs=world
+    )
+
+
+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:
+        _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(parallel_state=parallel),
+        )
+        dist.barrier()
+        result = [None]
+        if rank == 0:
+            try:
+                metadata = dcp.FileSystemReader(root).read_metadata()
+                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:
+                result[0] = repr(exc)
+        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"
+        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:
+        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", [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")
+
+
+@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
+
+    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]
+
+        # 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()
+
+
+@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 == "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)}}
+        )
+        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 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")
+            if mode == "v2_plan_conflict":
+                with pytest.raises(dcp.CheckpointException, match="unproven expert layout"):
+                    dcp.save(state, checkpoint_id=root, planner=planner)
+                assert not (Path(root) / ".metadata").exists()
+                completed = torch.tensor(1)
+                dist.all_reduce(completed)
+                assert completed.item() == 2
+                return
+            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 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
+                )
+            }
+            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)
+            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", "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)
+
+
+def _unproven_v2_high_level_save_worker(rank, root, endpoint):
+    from torch.distributed.tensor import distribute_tensor
+
+    from veomni.checkpoint.layout import weights_dir
+    from veomni.distributed.parallel_plan import ParallelPlan
+    from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state, 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.bias = torch.nn.Parameter(torch.ones(4, 4))
+
+        def get_parallel_plan(self):
+            return ParallelPlan({"ep": {"experts.weight": Shard(0)}})
+
+    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",),
+            extra_parallel_sizes=(2,),
+            extra_parallel_placement_innermost=(False,),
+        )
+        parallel = get_parallel_state()
+        model = Model()
+        model.get_parallel_plan().apply(model, parallel.extra_parallel_fsdp_device_mesh)
+        mesh = parallel.extra_parallel_fsdp_device_mesh["ep"]
+        fsdp_mesh = mesh[mesh.mesh_dim_names[:-1]]
+        for name, param in list(model.named_parameters()):
+            module, attr = name.rsplit(".", 1)
+            value = distribute_tensor(param.detach().clone(), fsdp_mesh, [Shard(1)])
+            setattr(model.get_submodule(module), attr, torch.nn.Parameter(value))
+        model._fqn2spec_info["experts.weight"].checkpoint_layout_v2 = True
+
+        with pytest.raises(dcp.CheckpointException, match="unproven expert layout"):
+            DistributedCheckpointer.save(root, {"model": model}, parallel_state=parallel)
+        assert not (Path(weights_dir(root)) / ".metadata").exists()
+        completed = torch.tensor(1)
+        dist.all_reduce(completed)
+        assert completed.item() == 2
+    finally:
+        clear_parallel_state()
+        dist.destroy_process_group()
+
+
+@pytest.mark.launch_plan(provider=from_parameter("world", device="cpu"))
+@pytest.mark.parametrize("world", [2])
+def test_unproven_v2_expert_layout_fails_before_checkpoint_publication(world, tmp_path):
+    spawn_with_rendezvous(_unproven_v2_high_level_save_worker, args=(str(tmp_path / "checkpoint"),), nprocs=world)
+
+
+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=world, timeout=timedelta(seconds=30)
+    )
+    try:
+        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)[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(
+            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, 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, 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="Raw DCP MultiOptimizer"):
+            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_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):
+    import veomni.checkpoint.dcp_checkpointer as checkpoint_module
+    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)
+        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"
+            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 coordinated == ["validating ExtraParallel optimizer checkpoint"]
+        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
+
+    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):
+    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="Raw DCP MultiOptimizer"):
+                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_raw_multi_optimizer_is_rejected_with_subgroup(world):
+    spawn_with_rendezvous(_subgroup_sparse_optimizer_failure_worker, args=(), 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 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(
+            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)}}
+        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:
+                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)
+        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)
+
+
+def test_export_preflight_preserves_existing_hf_and_allows_filtered_lora(tmp_path):
+    import json
+
+    from safetensors.torch import load_file
+
+    from scripts.merge_dcp_to_hf import save_lora_adapter_weights, save_model_weights
+    from veomni.checkpoint.dcp_checkpointer import _get_sharding_plan
+
+    old = {"model.aa.weight": torch.full((2, 2), 1.0), "model.zz.experts.weight": torch.full((2, 2), 2.0)}
+    new = {"model.aa.weight": torch.full((2, 2), 7.0), "model.zz.experts.weight": torch.full((2, 2), 8.0)}
+    old_path, new_path, output = tmp_path / "old", tmp_path / "new", tmp_path / "hf"
+    dcp.save(old, checkpoint_id=old_path)
+    save_model_weights(str(output), str(old_path), save_dtype=None, shard_size=16)
+    index_path = output / "model.safetensors.index.json"
+    original_index = index_path.read_bytes()
+    mapping = json.loads(original_index)["weight_map"]
+    original = {name: load_file(str(output / filename))[name].clone() for name, filename in mapping.items()}
+
+    dcp.save(new, checkpoint_id=new_path, planner=_ExtraParallelSavePlanner())
+    with pytest.raises(RuntimeError, match="noncanonical expert order"):
+        save_model_weights(str(output), str(new_path), save_dtype=None, shard_size=16)
+    assert index_path.read_bytes() == original_index
+    for name, filename in mapping.items():
+        torch.testing.assert_close(load_file(str(output / filename))[name], original[name], rtol=0, atol=0)
+    filtered, _, _ = _get_sharding_plan(new_path, drop_hf_keys={"zz.experts.weight"})
+    assert "aa.weight" in filtered and "zz.experts.weight" not in filtered
+
+    lora_key = "model.base_model.model.proj.lora_A.default.weight"
+    lora_path, lora_output = tmp_path / "lora-dcp", tmp_path / "lora-hf"
+    dcp.save({**new, lora_key: torch.ones(2, 2)}, checkpoint_id=lora_path, planner=_ExtraParallelSavePlanner())
+    save_lora_adapter_weights(str(lora_output), str(lora_path), save_dtype=None)
+    lora = load_file(str(lora_output / "adapter_model.safetensors"))
+    torch.testing.assert_close(lora["base_model.model.proj.lora_A.default.weight"], torch.ones(2, 2), rtol=0, atol=0)
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..ce8e6088c6
--- /dev/null
+++ b/tests/checkpoints/test_extra_parallel_layout_v2.py
@@ -0,0 +1,815 @@
+# 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 DTensor, Replicate, Shard, distribute_tensor
+
+from scripts.ci.accelerator.launch_plan import LaunchSpec
+from veomni.checkpoint.dcp_checkpointer import (
+    _EXPERT_LAYOUTS_KEY,
+    _EXTRA_PARALLEL_LAYOUT_MARKER,
+    DistributedCheckpointer,
+    ModelState,
+    OptimizerState,
+    _apply_extra_parallel_dim,
+    _checkpoint_extra_parallel_signature,
+    _export_extra_parallel_layout_key,
+    _extra_parallel_first_mesh,
+    _extra_parallel_layout_key,
+    _ExtraParallelSavePlanner,
+    _from_local_preserve_device,
+    _get_sharding_plan,
+    _ModelStrictLoadPlanner,
+    _needs_extra_parallel_layout_marker,
+    _process_shard,
+    drop_extra_parallel_dim,
+    drop_extra_parallel_dim_v2,
+    restore_extra_parallel_dim,
+    restore_extra_parallel_dim_v2,
+)
+from veomni.distributed.parallel_plan import SpecInfo
+
+
+pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+
+
+@pytest.mark.parametrize("requires_grad", [False, True])
+@pytest.mark.parametrize("noncontiguous", [False, True])
+@pytest.mark.parametrize("grad_enabled", [False, True])
+def test_legacy_cpu_shards_keep_stride_and_negative_axis(requires_grad, noncontiguous, grad_enabled):
+    class FakeMesh:
+        def __init__(self, shape):
+            self.device_type = "xpu"
+            self.mesh = torch.zeros(shape, dtype=torch.int64)
+            self.ndim = len(shape)
+            self.mesh_dim_names = tuple(f"d{i}" for i in range(self.ndim))
+
+        def size(self, dim=None):
+            return self.mesh.numel() if dim is None else self.mesh.size(dim)
+
+        def get_coordinate(self):
+            return [0] * self.ndim
+
+    full, sub = FakeMesh((2, 2)), FakeMesh((2,))
+    local = torch.arange(24, dtype=torch.float32).reshape(3, 8)[:, ::2]
+    if not noncontiguous:
+        local = local.contiguous()
+    local.requires_grad_(requires_grad)
+    with torch.set_grad_enabled(grad_enabled):
+        original = _from_local_preserve_device(local, sub, [Shard(0)])
+        restored = restore_extra_parallel_dim(original, full, sub, ep_shard_dim=-2, fsdp_shard_dim=0)
+        dropped = drop_extra_parallel_dim(restored, sub, fsdp_shard_dim=-2)
+    assert dropped.placements == (Shard(0),)
+    for value in (original, restored, dropped):
+        assert value.to_local().data_ptr() == local.data_ptr()
+        assert value.to_local().stride() == local.stride()
+        assert value.stride() == local.stride()
+        assert value.requires_grad == requires_grad
+        assert value._local_tensor.requires_grad == requires_grad
+
+
+@pytest.mark.parametrize("grad_enabled", [False, True])
+def test_matching_cpu_mesh_matches_native_dtensor_grad_state(grad_enabled):
+    class FakeMesh:
+        device_type = "cpu"
+        ndim = 1
+        mesh = torch.zeros(2, dtype=torch.int64)
+        mesh_dim_names = ("d0",)
+
+        def size(self, dim=None):
+            return 2
+
+        def get_coordinate(self):
+            return [0]
+
+    local = torch.arange(24, dtype=torch.float32).reshape(3, 8)[:, ::2].detach().requires_grad_(True)
+    mesh = FakeMesh()
+    with torch.set_grad_enabled(grad_enabled):
+        reference = DTensor.from_local(local, device_mesh=mesh, placements=[Shard(0)], run_check=False)
+        actual = _from_local_preserve_device(local, mesh, [Shard(0)])
+    assert actual.requires_grad == reference.requires_grad
+    assert actual._local_tensor.requires_grad == reference._local_tensor.requires_grad
+    assert actual._local_tensor.data_ptr() == local.data_ptr()
+    assert actual.stride() == reference.stride()
+
+
+def _roundtrip(rank, root, ep_shard_dim, fsdp_shard_dim, rows=5):
+    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(rows * 7, dtype=torch.float32).reshape(rows, 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)
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=4))
+def test_v2_strided_layout_with_uneven_fsdp_chunks(tmp_path):
+    torch.multiprocessing.spawn(_roundtrip, args=(str(tmp_path), 0, 0, 6), nprocs=4, join=True)
+
+
+def _six_rank_roundtrip(rank, root, rows):
+    dist.init_process_group("gloo", init_method=f"file://{root}/init-ep2-fsdp3", rank=rank, world_size=6)
+    try:
+        mesh = init_device_mesh("cpu", (3, 2), mesh_dim_names=("fsdp", "ep"))
+        fsdp_mesh = mesh["fsdp"]
+        full = torch.arange(rows * 6, dtype=torch.float32).reshape(rows, 6)
+        local_ep = full.chunk(2, dim=0)[mesh.get_local_rank("ep")]
+        local = distribute_tensor(local_ep.contiguous(), fsdp_mesh, placements=[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}/ep2-fsdp3")
+        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}/ep2-fsdp3")
+        restored = drop_extra_parallel_dim_v2(target, fsdp_mesh, 0, extra_parallel_name="ep")
+        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=6))
+@pytest.mark.parametrize("rows", [7, 10])
+def test_v2_six_rank_uneven_roundtrip(tmp_path, rows):
+    torch.multiprocessing.spawn(_six_rank_roundtrip, args=(str(tmp_path), rows), nprocs=6, join=True)
+
+
+def _ep_one_roundtrip(rank, root):
+    dist.init_process_group("gloo", init_method=f"file://{root}/init-ep1", rank=rank, world_size=4)
+    try:
+        mesh = init_device_mesh("cpu", (4, 1), mesh_dim_names=("fsdp", "ep"))
+        fsdp_mesh = mesh["fsdp"]
+        full = torch.arange(35, dtype=torch.float32).reshape(5, 7)
+        local = distribute_tensor(full, fsdp_mesh, placements=[Shard(0)])
+        combined = restore_extra_parallel_dim_v2(local, mesh, fsdp_mesh, 0, 0, global_shape=tuple(full.shape))
+        dcp.save({"expert": combined}, checkpoint_id=f"{root}/ep1")
+        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}/ep1")
+        restored = drop_extra_parallel_dim_v2(target, fsdp_mesh, 0, extra_parallel_name="ep")
+        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))
+def test_v2_single_ep_rank_with_uneven_fsdp_chunks(tmp_path):
+    torch.multiprocessing.spawn(_ep_one_roundtrip, args=(str(tmp_path),), 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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    assert not list(tmp_path.iterdir())
+
+
+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)
+
+
+@pytest.mark.parametrize("mutation", ["version", "replace_map"])
+@pytest.mark.parametrize("operation", ["state_dict", "load_state_dict"])
+def test_v2_wrapper_rejects_layout_mutation_before_io(monkeypatch, mutation, operation):
+    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
+    )
+    adapter = ModelState(model, parallel_state=SimpleNamespace(dp_mode="fsdp2"))
+    saved = adapter.state_dict()
+    assert _extra_parallel_layout_key(model._fqn2spec_info) in saved
+    before = model.weight.detach().clone()
+    if mutation == "version":
+        spec.checkpoint_layout_v2 = False
+    else:
+        model._fqn2spec_info = dict(model._fqn2spec_info)
+    with pytest.raises(RuntimeError, match="layout changed while checkpoint adapter was active"):
+        getattr(adapter, operation)(saved) if operation == "load_state_dict" else getattr(adapter, operation)()
+    torch.testing.assert_close(model.weight, before, 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_v2_offline_hf_export_keeps_marker_and_tensor_values(tmp_path):
+    from veomni.checkpoint.dcp_checkpointer import dcp_to_torch_state_dict
+
+    marker = _extra_parallel_layout_key({})
+    expected = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+    mask = torch.tensor([True, False])
+    bf16_weight = expected.to(torch.bfloat16)
+    dcp.save(
+        {
+            "model": {
+                marker + "__model.foo.weight": expected,
+                marker + "__model.foo.mask": mask,
+                marker + "__model.foo.bf16_weight": bf16_weight,
+                marker: "ep-outer-v1",
+            }
+        },
+        checkpoint_id=tmp_path,
+    )
+    result = dcp_to_torch_state_dict(tmp_path)
+    assert set(result) == {"model.foo.weight", "model.foo.mask", "model.foo.bf16_weight"}
+    torch.testing.assert_close(result["model.foo.weight"], expected, rtol=0, atol=0)
+    torch.testing.assert_close(result["model.foo.mask"], mask, rtol=0, atol=0)
+    torch.testing.assert_close(result["model.foo.bf16_weight"], bf16_weight, 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})
+
+
+@pytest.mark.parametrize("shard_size", [None, 16])
+def test_v2_offline_hf_export_preserves_names_and_values(tmp_path, shard_size):
+    from veomni.distributed.parallel_state import ParallelState
+
+    marker = _extra_parallel_layout_key({})
+    source = {
+        "model.layer.experts.weight": torch.arange(16, dtype=torch.float32).reshape(4, 4),
+        "model.layer.other.weight": torch.arange(4, dtype=torch.bfloat16),
+    }
+    dcp.save(
+        {"model": {marker: "ep-outer-v1", **{marker + "__" + k: v for k, v in source.items()}}},
+        checkpoint_id=tmp_path,
+        planner=_ExtraParallelSavePlanner(parallel_state=ParallelState(device_type="cpu")),
+    )
+    plan, _, _ = _get_sharding_plan(str(tmp_path), shard_size=shard_size)
+    actual = {}
+    for shard in plan if isinstance(plan, list) else [plan]:
+        actual.update(_process_shard(shard, str(tmp_path)))
+    assert actual.keys() == source.keys()
+    for key, value in source.items():
+        torch.testing.assert_close(actual[key], value, rtol=0, atol=0)
+
+
+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
+
+    from veomni.checkpoint.layout import weights_dir
+
+    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}
+        )
+        for source_axis, target_axis in ((0, 1), (1, 0)):
+            source = model_for(source_axis, ep_local)
+            path = f"{root}/high-{source_axis}"
+            DistributedCheckpointer.save(path, {"model": source}, parallel_state=parallel)
+            metadata = dcp.FileSystemReader(weights_dir(path)).read_metadata()
+            layouts = metadata.planner_data[_EXPERT_LAYOUTS_KEY]["tensors"]
+            assert layouts and all(layout["expert_order"] == "canonical" for layout in layouts.values())
+            target = model_for(target_axis, torch.zeros_like(ep_local))
+            DistributedCheckpointer.load(path, {"model": target}, parallel_state=parallel)
+            expected = distribute_tensor(ep_local, fsdp_mesh, [Shard(target_axis)])
+            torch.testing.assert_close(target.weight.to_local(), expected.to_local(), rtol=0, atol=0)
+
+        # A rank-inconsistent version must not invalidate an existing DCP.
+        import hashlib
+        from pathlib import Path
+
+        path = f"{root}/high-0"
+
+        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()
+        source = model_for(0, ep_local)
+        source._fqn2spec_info["weight"].checkpoint_layout_v2 = rank % 2 == 0
+        with pytest.raises(RuntimeError, match="layout version differs across ranks"):
+            DistributedCheckpointer.save(path, {"model": source}, parallel_state=parallel)
+        assert fingerprint() == before
+        for mismatch in ("axis", "trainable_only", "freeze"):
+            source = model_for(0, ep_local)
+            if mismatch == "axis" and rank % 2:
+                source._fqn2spec_info["weight"].fsdp_shard_dim = 1
+            if mismatch == "freeze" and rank % 2:
+                source.weight.requires_grad_(False)
+            with pytest.raises(RuntimeError, match="layout or optimizer policy differs across ranks"):
+                DistributedCheckpointer.save(
+                    path,
+                    {"model": source},
+                    trainable_only=mismatch == "freeze" or (mismatch == "trainable_only" and rank % 2 == 0),
+                    parallel_state=parallel,
+                )
+            assert fingerprint() == before
+        target = model_for(0, torch.zeros_like(ep_local))
+        DistributedCheckpointer.load(path, {"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)
+
+        # V2 split-layout checkpoints created before per-tensor layout metadata
+        # used the namespace marker and DCP's default planner; these still load.
+        source = model_for(0, ep_local)
+        path = f"{root}/namespaced-before-planner"
+        old_state = ModelState(source, parallel_state=parallel).state_dict()
+        marker = _extra_parallel_layout_key(source._fqn2spec_info)
+        old_state[marker + "__weight"] = DTensor.from_local(
+            source.weight.to_local(),
+            _extra_parallel_first_mesh(mesh),
+            [Shard(0), Shard(0)],
+            run_check=False,
+            shape=full.shape,
+            stride=full.stride(),
+        )
+        dcp.save({"model": old_state}, checkpoint_id=weights_dir(path))
+        for target_axis in (0, 1):
+            target = model_for(target_axis, torch.zeros_like(ep_local))
+            DistributedCheckpointer.load(path, {"model": target}, parallel_state=parallel)
+            expected = distribute_tensor(ep_local, fsdp_mesh, [Shard(target_axis)])
+            torch.testing.assert_close(target.weight.to_local(), expected.to_local(), rtol=0, atol=0)
+
+        # The upstream pre-split layout is intentionally unsupported, even if
+        # its tensors have a versioned marker. Refuse before changing the target.
+        pre_split = f"{root}/pre-split-namespaced"
+        dcp.save({"model": old_state}, checkpoint_id=pre_split)
+        target = model_for(0, torch.zeros_like(ep_local))
+        before = target.weight.to_local().clone()
+        with pytest.raises(FileNotFoundError, match="not resumable"):
+            DistributedCheckpointer.load(pre_split, {"model": target}, parallel_state=parallel)
+        torch.testing.assert_close(target.weight.to_local(), before, 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_logical_signature_normalizes_negative_axes_but_preserves_singleton_axes():
+    mesh = SimpleNamespace(shape=(1, 1), mesh_dim_names=("ep_fsdp", "ep"))
+    first = SpecInfo("ep", Shard(0), "weight", mesh, 0, None, (8, 8), True)
+    second = SpecInfo("ep", Shard(-1), "weight", mesh, -1, None, (8, 8), True)
+    equivalent = SpecInfo("ep", Shard(1), "weight", mesh, 1, None, (8, 8), True)
+
+    def signature(spec):
+        return _checkpoint_extra_parallel_signature({"weight": spec}, optimizer_present=False, trainable_only=False)
+
+    assert signature(first) != signature(second)
+    assert signature(equivalent) == signature(second)
+    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(
+        {"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
+    )
+
+
+@pytest.mark.parametrize("adapter_name", ["model", "optimizer"])
+@pytest.mark.parametrize("operation", ["state_dict", "load_state_dict"])
+def test_stateful_adapter_rejects_layout_mutation_before_io(adapter_name, operation):
+    model = torch.nn.Linear(2, 2, bias=False)
+    mesh = SimpleNamespace(shape=(1, 1))
+    spec = SpecInfo("ep", Shard(0), "weight", mesh, 0, None, (2, 2), True)
+    model._fqn2spec_info = {"weight": spec}
+    parallel = SimpleNamespace(dp_mode="ddp")
+    adapter = (
+        ModelState(model, parallel_state=parallel)
+        if adapter_name == "model"
+        else OptimizerState(model, torch.optim.AdamW(model.parameters()), parallel_state=parallel)
+    )
+    spec.checkpoint_layout_v2 = False
+    with pytest.raises(RuntimeError, match="layout changed while checkpoint adapter was active"):
+        getattr(adapter, operation)({}) if operation == "load_state_dict" else getattr(adapter, operation)()
+
+
+def _singleton_axis_guard_worker(rank, root, singleton):
+    import hashlib
+    from datetime import timedelta
+    from pathlib import Path
+
+    dist.init_process_group(
+        "gloo", init_method=f"file://{root}/init-singleton", rank=rank, world_size=4, timeout=timedelta(seconds=45)
+    )
+    try:
+        shape = (2, 1, 2) if singleton == "fsdp" else (2, 2, 1)
+        mesh = init_device_mesh("cpu", shape, mesh_dim_names=("ep_replicate", "ep_fsdp", "ep"))
+        fsdp_mesh = mesh[("ep_replicate", "ep_fsdp")]
+        ep_local = (
+            torch.arange(64, dtype=torch.float32).reshape(8, 8).chunk(shape[-1], dim=0)[mesh.get_local_rank("ep")]
+        )
+        parallel = SimpleNamespace(
+            dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh}
+        )
+
+        def model_for(axis, ep_axis=0):
+            model = torch.nn.Module()
+            value = distribute_tensor(ep_local.contiguous(), fsdp_mesh, [Replicate(), Shard(axis)])
+            model.register_parameter("weight", torch.nn.Parameter(value))
+            model._fqn2spec_info = {"weight": SpecInfo("ep", Shard(ep_axis), "weight", mesh, axis, None, (8, 8), True)}
+            return model
+
+        path = f"{root}/singleton-checkpoint"
+        DistributedCheckpointer.save(path, {"model": model_for(0)}, 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()
+        with pytest.raises(RuntimeError, match="layout or optimizer policy differs across ranks"):
+            model = model_for(rank % 2 if singleton == "fsdp" else 0, rank % 2 if singleton == "ep" else 0)
+            DistributedCheckpointer.save(path, {"model": model}, parallel_state=parallel)
+        assert fingerprint() == before
+        target = model_for(0)
+        target.weight.data.zero_()
+        DistributedCheckpointer.load(path, {"model": target}, parallel_state=parallel)
+        expected = distribute_tensor(ep_local.contiguous(), fsdp_mesh, [Replicate(), 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))
+@pytest.mark.parametrize("singleton", ["fsdp", "ep"])
+def test_singleton_axis_disagreement_preserves_checkpoint(tmp_path, singleton):
+    torch.multiprocessing.spawn(_singleton_axis_guard_worker, args=(str(tmp_path), singleton), 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/checkpoints/test_extra_parallel_optimizer_keys.py b/tests/checkpoints/test_extra_parallel_optimizer_keys.py
new file mode 100644
index 0000000000..9029b6016f
--- /dev/null
+++ b/tests/checkpoints/test_extra_parallel_optimizer_keys.py
@@ -0,0 +1,225 @@
+"""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",
+        )
+
+
+def test_plain_optimizer_state_with_multiple_fsdp_ranks_is_rejected(monkeypatch):
+    from types import SimpleNamespace
+
+    from torch.distributed.tensor import Shard
+
+    class FakeMesh:
+        mesh_dim_names = ("ep_fsdp", "ep")
+
+        def __getitem__(self, _):
+            return SimpleNamespace(size=lambda: 2)
+
+    monkeypatch.setattr(dcp_checkpointer, "_validate_extra_parallel_meshes", lambda _: None)
+    spec = SimpleNamespace(placement=Shard(0), para_fsdp_mesh=FakeMesh())
+    state = {"state.experts.weight.exp_avg": torch.ones(2, 3)}
+    with pytest.raises(RuntimeError, match="no proven replica layout"):
+        dcp_checkpointer._apply_extra_parallel_dim(
+            state, {"experts.weight": spec}, None, "restore", key_match="substring"
+        )
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("noncontiguous", [False, True])
+@pytest.mark.parametrize("requires_grad", [False, True])
+@pytest.mark.parametrize("grad_enabled", [False, True])
+def test_same_axis_offloaded_shards_keep_cpu_storage_and_strided_layout(
+    dtype, noncontiguous, requires_grad, grad_enabled
+):
+    from torch.distributed.tensor import Shard
+    from torch.distributed.tensor.placement_types import _StridedShard
+
+    class FakeMesh:
+        def __init__(self, shape):
+            self.device_type = "xpu"
+            self.mesh = torch.zeros(shape, dtype=torch.int64)
+            self.ndim = len(shape)
+            self.mesh_dim_names = tuple(f"d{i}" for i in range(self.ndim))
+
+        def size(self, dim=None):
+            return self.mesh.numel() if dim is None else self.mesh.size(dim)
+
+        def get_coordinate(self):
+            return [0] * self.ndim
+
+    full, sub = FakeMesh((2, 2)), FakeMesh((2,))
+    if noncontiguous:
+        local = torch.arange(24, dtype=dtype).reshape(3, 8)[:, ::2]
+    else:
+        local = torch.arange(12, dtype=dtype).reshape(3, 4)
+    local.requires_grad_(requires_grad)
+    with torch.set_grad_enabled(grad_enabled):
+        original = dcp_checkpointer._from_local_preserve_device(local, sub, [Shard(0)])
+        restored = dcp_checkpointer.restore_extra_parallel_dim(original, full, sub, ep_shard_dim=-2, fsdp_shard_dim=0)
+        dropped = dcp_checkpointer.drop_extra_parallel_dim(restored, sub, fsdp_shard_dim=-2)
+    assert restored.shape == (12, 4)
+    assert restored.placements == (_StridedShard(0, split_factor=2), Shard(0))
+    assert dropped.shape == (6, 4)
+    assert dropped.placements == (Shard(0),)
+    for value in (original, restored, dropped):
+        assert value.dtype == dtype
+        assert value.to_local().device.type == "cpu"
+        assert value.to_local().data_ptr() == local.data_ptr()
+        assert value.to_local().stride() == local.stride()
+        assert value.stride() == local.stride()
+        assert value.requires_grad == requires_grad
+        assert value._local_tensor.requires_grad == requires_grad
+
+
+def _plain_extra_parallel_roundtrip(rank, root):
+    from datetime import timedelta
+
+    import torch.distributed as dist
+    import torch.distributed.checkpoint as dcp
+    from torch.distributed.device_mesh import init_device_mesh
+
+    dist.init_process_group(
+        "gloo", init_method=f"file://{root}/init", rank=rank, world_size=2, timeout=timedelta(seconds=30)
+    )
+    try:
+        mesh = init_device_mesh("cpu", (1, 2), mesh_dim_names=("ep_fsdp", "ep"))
+        submesh = mesh["ep_fsdp"]
+        local = torch.full((2, 3), float(rank + 1))
+        saved = dcp_checkpointer.restore_extra_parallel_dim(local, mesh, submesh, ep_shard_dim=0)
+        expected = torch.cat([torch.ones((2, 3)), torch.full((2, 3), 2.0)])
+        torch.testing.assert_close(saved.full_tensor(), expected, rtol=0, atol=0)
+        dcp.save({"expert": saved}, checkpoint_id=f"{root}/checkpoint")
+        restored = dcp_checkpointer.restore_extra_parallel_dim(torch.zeros_like(local), mesh, submesh, ep_shard_dim=0)
+        dcp.load({"expert": restored}, checkpoint_id=f"{root}/checkpoint")
+        dropped = dcp_checkpointer.drop_extra_parallel_dim(restored, submesh)
+        assert not isinstance(dropped, torch.distributed.tensor.DTensor)
+        torch.testing.assert_close(dropped, local, rtol=0, atol=0)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2))
+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):
+    from datetime import timedelta
+    from pathlib import Path
+    from types import SimpleNamespace
+
+    import torch.distributed as dist
+    import torch.distributed.checkpoint as dcp
+    from torch.distributed.tensor import Shard
+
+    class FakeMesh:
+        mesh_dim_names = ("ep_fsdp", "ep")
+
+        def __getitem__(self, _):
+            return SimpleNamespace(size=lambda: 2)
+
+    class FakeOptimizer:
+        _is_multi_optimizer = True
+
+        def __init__(self):
+            self.state_dict_calls = 0
+
+        def state_dict(self):
+            self.state_dict_calls += 1
+            return {"state.experts.weight.exp_avg": torch.ones(2, 3)} if rank == 0 else {}
+
+    dist.init_process_group(
+        "gloo", init_method=f"file://{root}/init", rank=rank, world_size=2, timeout=timedelta(seconds=30)
+    )
+    try:
+        dcp_checkpointer._validate_extra_parallel_meshes = lambda _: None
+        spec = SimpleNamespace(placement=Shard(0), para_fsdp_mesh=FakeMesh())
+        model = SimpleNamespace(_fqn2spec_info={"experts.weight": spec})
+        parallel = SimpleNamespace(dp_mode="fsdp2")
+        leaf = FakeOptimizer()
+        optimizer = dcp_checkpointer.OptimizerState(model, leaf, parallel_state=parallel)
+        checkpoint = Path(root) / "rejected"
+        with pytest.raises(RuntimeError, match="Raw DCP MultiOptimizer does not prove parameter groups"):
+            dcp.save({"optimizer": optimizer}, checkpoint_id=checkpoint)
+        # The owner-aware codec rejects raw MultiOptimizer DCP before either
+        # rank can inspect an unproven, potentially asymmetric shard layout.
+        assert leaf.state_dict_calls == 0
+        assert not (checkpoint / ".metadata").exists()
+        value = torch.ones(())
+        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_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/tests/checkpoints/test_legacy_strided_input.py b/tests/checkpoints/test_legacy_strided_input.py
new file mode 100644
index 0000000000..3b04f5dfa5
--- /dev/null
+++ b/tests/checkpoints/test_legacy_strided_input.py
@@ -0,0 +1,42 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: Apache-2.0
+
+"""Legacy FSDP-first DTensor input remains supported by the drop helper."""
+
+from datetime import timedelta
+
+import pytest
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import DTensor, Shard, distribute_tensor
+from torch.distributed.tensor.placement_types import _StridedShard
+
+from scripts.ci.accelerator.launch_plan import LaunchSpec
+from tests.tools.launch_utils import spawn_with_rendezvous
+from veomni.checkpoint.dcp_checkpointer import drop_extra_parallel_dim_v2
+
+
+def _worker(rank, endpoint):
+    dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=4, timeout=timedelta(seconds=60))
+    try:
+        for geometry, rows in (((2, 2), 6), ((4, 1), 5)):
+            mesh = init_device_mesh("cpu", geometry, mesh_dim_names=("ep_fsdp", "ep"))
+            fsdp = mesh["ep_fsdp"]
+            full = torch.arange(rows * 7, dtype=torch.float32).reshape(rows, 7)
+            ep_slice = full.chunk(geometry[1], dim=0)[mesh.get_local_rank("ep")]
+            local = distribute_tensor(ep_slice.contiguous(), fsdp, [Shard(0)])
+            first = _StridedShard(0, split_factor=2) if geometry[1] == 2 else Shard(0)
+            saved = DTensor.from_local(
+                local.to_local(), mesh, [first, Shard(0)], run_check=False, shape=full.shape, stride=full.stride()
+            )
+            restored = drop_extra_parallel_dim_v2(saved, fsdp, 0, extra_parallel_name="ep")
+            assert tuple(restored.shape) == tuple(local.shape)
+            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(nproc=4, device="cpu"))
+def test_legacy_fsdp_first_strided_and_ep1_drop():
+    spawn_with_rendezvous(_worker, args=(), nprocs=4)
diff --git a/tests/checkpoints/test_multi_export_receipt.py b/tests/checkpoints/test_multi_export_receipt.py
new file mode 100644
index 0000000000..442360a0eb
--- /dev/null
+++ b/tests/checkpoints/test_multi_export_receipt.py
@@ -0,0 +1,490 @@
+"""Same-step HF and epoch saves reuse the model owners' DCP transaction."""
+
+import gc
+import hashlib
+import weakref
+from datetime import timedelta
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock, 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 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, _export_wrote_dcp
+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
+
+
+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)
+    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 _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
+
+    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)
+
+        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):
+                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 = ("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()
+            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 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 = {}
+        trainer.channel_loss_callback = None
+        trainer.offline_cache_writer = None
+        if scope in ("omni", "mixed"):
+            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)
+        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 any(generations) and (scope != "omni" or all(generations))
+        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
+        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)
+        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
+        if rank == 0:
+            assert snapshot == {
+                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)
+            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_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 in ("omni", "mixed"):
+            modules = {runtime.module_name: runtime for runtime in restored}
+            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()
+        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)
+            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:
+                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", "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_ordinary_owner_restore_failure_reaches_every_rank(tmp_path):
+    from tests.tools.launch_utils import spawn_with_rendezvous
+
+    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
new file mode 100644
index 0000000000..447cec432c
--- /dev/null
+++ b/tests/checkpoints/test_multi_optimizer_proven_resume.py
@@ -0,0 +1,2796 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# 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, runner_capabilities
+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:
+        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_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 _adapter_layout_mutation_worker(rank, endpoint):
+    from datetime import timedelta
+
+    import torch.distributed as dist
+
+    from veomni.checkpoint.dcp_checkpointer import OptimizerState, _ProvenModelState, _ProvenMultiOptimizerState
+
+    class OptimizerProbe:
+        def __init__(self):
+            self.entered = False
+
+        def _sparse_state_dict(self, **kwargs):
+            self.entered = True
+            dist.all_reduce(torch.ones(1))
+            return {}
+
+    dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=30))
+    try:
+        for kind in ("model", "optimizer"):
+            for operation in ("state_dict", "load_state_dict"):
+                model = torch.nn.Linear(1, 1)
+                optimizer = OptimizerProbe()
+                parallel = SimpleNamespace(dp_mode="ddp")
+                if kind == "model":
+                    adapter = _ProvenModelState(model, optimizer=optimizer, parallel_state=parallel, group=None)
+                else:
+                    adapter = _ProvenMultiOptimizerState(
+                        model, optimizer, parallel_state=parallel, generation="g", expected_layout="l", group=None
+                    )
+                if rank == 0:
+                    model._fqn2spec_info = {}
+                with pytest.raises(RuntimeError, match="layout changed|failed on another rank"):
+                    getattr(adapter, operation)({}) if operation == "load_state_dict" else getattr(
+                        adapter, operation
+                    )()
+                assert not optimizer.entered
+                tensor = torch.ones(1)
+                dist.all_reduce(tensor)
+                assert tensor.item() == 2
+
+        model = torch.nn.Linear(1, 1)
+        model._fqn2spec_info = {}
+        adapter = OptimizerState(
+            model, torch.optim.AdamW(model.parameters()), parallel_state=SimpleNamespace(dp_mode="fsdp2")
+        )
+        if rank == 0:
+            model._fqn2spec_info = {}
+        with pytest.raises(RuntimeError, match="layout changed|ExtraParallel|failed on another rank"):
+            adapter.state_dict()
+        tensor = torch.ones(1)
+        dist.all_reduce(tensor)
+        assert tensor.item() == 2
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu", nproc=2))
+def test_adapter_layout_mutation_is_coordinated_before_collectives():
+    from tests.tools.launch_utils import spawn_with_rendezvous
+
+    spawn_with_rendezvous(_adapter_layout_mutation_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
+
+    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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+        target, restored = _model_and_optimizer()
+        original_model = checkpoint_module._ProvenModelState
+        original_optimizer = checkpoint_module._ProvenMultiOptimizerState
+        original_reader = DistributedCheckpointer._create_storage_reader
+        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)
+
+        with (
+            patch.object(checkpoint_module, "_ProvenModelState", side_effect=model_state),
+            patch.object(checkpoint_module, "_ProvenMultiOptimizerState", side_effect=optimizer_state),
+            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, "_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)
+        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("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
+
+    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"):
+        DistributedCheckpointer.save(
+            str(tmp_path),
+            {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+            save_async=True,
+            parallel_state=SimpleNamespace(dp_mode="ddp"),
+        )
+    assert not list(tmp_path.iterdir())
+
+
+def test_failed_final_load_barrier_does_not_certify_multi_optimizer():
+    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.config = SimpleNamespace(dcp_evict_source_cache=False)
+    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
+    leaf.checkpoint_abort_file_backed_state.assert_called_once()
+
+
+@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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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, "_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,
+    )
+    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 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, "_veomni_scheduler_object": None},
+            stage_dir=str(stage),
+            parallel_state=SimpleNamespace(dp_mode="ddp"),
+        )
+    assert stale_sidecar.exists()
+    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
+
+    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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+            stage_dir=stage,
+        )
+        target, restored = _model_and_optimizer()
+        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)
+            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()
+
+
+@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
+
+    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, "_veomni_scheduler_object": None},
+                parallel_state=parallel,
+            )
+        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("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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+                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.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
+
+    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.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
+
+    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()
+
+
+@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
+
+    spawn_with_rendezvous(_inventory_hash_failure_worker, args=(), 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, "_veomni_scheduler_object": None},
+                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_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, endpoint):
+    import contextlib
+    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:
+
+        class TestRuntime(runtime_module.VeOmniModelRuntime):
+            def setup(self):
+                pass
+
+            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
+
+            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"):
+            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.launch_plan(provider=LaunchSpec(nproc=2, device="cpu"))
+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=(), 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, "_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, "_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, "_veomni_scheduler_object": None},
+                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, "_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, "_veomni_scheduler_object": None},
+            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, weights_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({"model": model.state_dict()}, checkpoint_id=weights_dir(str(tmp_path)))
+    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, "_veomni_scheduler_object": None},
+            parallel_state=SimpleNamespace(dp_mode="ddp"),
+        )
+
+
+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, "_veomni_scheduler_object": None},
+        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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    torch.testing.assert_close(target.weight, original, rtol=0, atol=0)
+
+
+@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(param_dtype)
+        leaf = AnyPrecisionAdamW(
+            model.parameters(),
+            momentum_dtype=moment_dtype,
+            variance_dtype=moment_dtype,
+            compensation_buffer_dtype=moment_dtype,
+        )
+        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")
+    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
+    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, "_veomni_scheduler_object": None},
+        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):
+    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, "_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,
+    )
+    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, "_veomni_scheduler_object": None},
+            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, "_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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+            trainable_only=True,
+        )
+    torch.testing.assert_close(target.weight, original, rtol=0, atol=0)
+
+
+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, "_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, "_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)
+    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_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, "_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)
+    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)
+
+
+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_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)},
+                "_veomni_scheduler_object": None,
+            },
+            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": {}, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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, "_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,
+    )
+    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, "_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, "_veomni_scheduler_object": None},
+            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, "_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, "_veomni_scheduler_object": None},
+                parallel_state=parallel,
+            )
+        return
+    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)
+    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, "_veomni_scheduler_object": None},
+            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, "_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, "_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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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, "_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)
+    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, "_veomni_scheduler_object": None},
+            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, "_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)
+    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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    torch.testing.assert_close(target.weight, before, rtol=0, atol=0)
+
+
+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"}, "_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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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, "_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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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, _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):
+    model, optimizer = _model_and_optimizer()
+    parallel = SimpleNamespace(dp_mode="ddp")
+    generation = DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target, restored = _model_and_optimizer()
+    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
+
+
+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, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    DistributedCheckpointer.save(
+        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, "_veomni_scheduler_object": None}
+    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, "_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, "_veomni_scheduler_object": None},
+            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, "_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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    torch.testing.assert_close(target.weight, before, rtol=0, atol=0)
+
+
+@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, "_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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    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 _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
+
+    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
+    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
+
+    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, "_veomni_scheduler_object": None},
+            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, "_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_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)
+
+
+@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
+
+
+@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
+
+    def unordered(step):
+        return sum(float(value) for value in {"0.125", "1e16", "-1e16"})
+
+    def unordered_constant(step):
+        return step in {1, 2, 3}
+
+    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
+
+    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)
+    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_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")
+    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()
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..7fbea95aeb
--- /dev/null
+++ b/tests/checkpoints/test_multi_optimizer_rope_cache.py
@@ -0,0 +1,699 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: Apache-2.0
+
+"""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
+
+
+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
+
+    model, optimizer = _build()
+    runtime = SimpleNamespace(
+        model=model,
+        optimizer=optimizer,
+        model_name="policy",
+        train_args=SimpleNamespace(checkpoint=SimpleNamespace(manager="dcp", save_async=False)),
+    )
+    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):
+    def __init__(self, theta=10000, rope_class=Qwen3RotaryEmbedding):
+        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": theta},
+        )
+        self.rope = rope_class(config)
+
+
+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)
+    optimizer.step()
+    optimizer.zero_grad()
+
+
+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(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()}
+    trace.update({f"grad.{name}": value.grad.detach().clone() for name, value in model.named_parameters()})
+    optimizer.step()
+    optimizer.zero_grad()
+    return trace
+
+
+def _rope_transitions(model, optimizer):
+    result = {}
+    for index, length in enumerate((8, 16, 17, 16, 15, 32, 24, 4)):
+        trace = _rope_update(model, optimizer, length)
+        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, 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},
+        parallel_state=SimpleNamespace(dp_mode="ddp"),
+    )
+    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"):
+        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_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
+
+    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_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))
+    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, "_veomni_scheduler_object": None},
+            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")
+    with monkeypatch.context() as old:
+        old.setattr(dynamic_rope, "dynamic_rope_checkpoint_contract", lambda _model: [])
+        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, "_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, "_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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    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_nonfinite_canonical_rope_is_rejected(tmp_path, monkeypatch):
+    from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+
+    model, optimizer = _build()
+    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):
+        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()
+
+
+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:
+        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, "_veomni_scheduler_object": None},
+                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_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)
+
+
+@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:
+        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="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_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)
+    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)
+
+
+@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, "_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, "_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)
+
+
+@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, "_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, "_veomni_scheduler_object": None},
+            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")
+    with pytest.raises(RuntimeError, match="Cannot prove dynamic RoPE checkpoint cache"):
+        DistributedCheckpointer.save(
+            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_multi_optimizer_trainer_step.py b/tests/checkpoints/test_multi_optimizer_trainer_step.py
new file mode 100644
index 0000000000..8d920039c0
--- /dev/null
+++ b/tests/checkpoints/test_multi_optimizer_trainer_step.py
@@ -0,0 +1,338 @@
+# 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
+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.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"])
+@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
+    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(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
+
+    with (
+        patch("veomni.models.checkpoint.manager.build_checkpointer", return_value=DistributedCheckpointer),
+        patch("torch.distributed.barrier"),
+    ):
+        source = build(instrumented in ("source", "both"))
+        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(instrumented in ("target", "both"))
+        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
+        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 == (4 if mode == "matching" else 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)
+        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
+            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"
+                },
+                "train_dataloader": {"cursor": 2},
+            }
+        )
+        with (
+            patch.object(callback, "_read_current", return_value=state),
+            patch("veomni.trainer.callbacks.global_state_callback.set_device_rng_state"),
+        ):
+            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,
+                    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
+        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",
+        "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)
+
+
+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
+
+    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"))
+        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()
+        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", "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/tests/checkpoints/test_optimizer_checkpoint_adapter.py b/tests/checkpoints/test_optimizer_checkpoint_adapter.py
new file mode 100644
index 0000000000..4b825c1a3c
--- /dev/null
+++ b/tests/checkpoints/test_optimizer_checkpoint_adapter.py
@@ -0,0 +1,208 @@
+# 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 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})
+        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, "_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,
+    )
+    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, "_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, "_veomni_scheduler_object": None},
+            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, "_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, "_veomni_scheduler_object": None},
+            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, "_veomni_scheduler_object": None},
+            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()), "_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
new file mode 100644
index 0000000000..06e8e9f2c3
--- /dev/null
+++ b/tests/checkpoints/test_qwen_omni_nested_checkpoint_hooks.py
@@ -0,0 +1,96 @@
+# 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."""
+
+from types import SimpleNamespace
+
+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
+
+
+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_omni_hook_allows_peft_dcp_roundtrip(tmp_path, model_class):
+    peft = pytest.importorskip("peft")
+
+    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, "_veomni_scheduler_object": None}, parallel_state=parallel
+    )
+    target, restored = build()
+    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:
+        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
new file mode 100644
index 0000000000..40c06c93ae
--- /dev/null
+++ b/tests/checkpoints/test_rank_local_optimizer.py
@@ -0,0 +1,440 @@
+# 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 scripts.ci.accelerator.launch_plan import LaunchSpec
+from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
+from veomni.checkpoint.rank_local_optimizer import (
+    local_optimizer_state,
+    recover_frozen_optimizer_ownership,
+    restore_optimizer_state,
+)
+from veomni.distributed.parallel_plan import SpecInfo
+from veomni.optim.muon import DistributedMuon
+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)
+    )
+    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.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)
+
+
+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, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+        target, restored = build()
+        DistributedCheckpointer.load(
+            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, {})
+        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.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)
+
+
+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
+        if kind == "swap_bf16":
+            values = values.to(torch.bfloat16)
+        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")
+        elif kind == "anyprecision":
+            from veomni.optim.optimizer import AnyPrecisionAdamW
+
+            leaf = AnyPrecisionAdamW([model.weight], lr=0.01)
+        elif kind in ("swap", "swap_bf16"):
+            from veomni.optim.swap_adamw import SwapAdamW
+
+            if inactive_owner and owner == 1:
+                model.weight.requires_grad_(False)
+            leaf = SwapAdamW([model.weight], lr=0.01, pin_memory=False)
+            if inactive_owner and owner == 1:
+                model.weight.requires_grad_(True)
+        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, {})
+            if kind in ("swap", "swap_bf16") and state:
+                state = {"step": torch.tensor(state["step"]), **leaf._host_states[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": {
+                    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, "_veomni_scheduler_object": None},
+                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, "_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)
+            advance(0.2)
+            torch.testing.assert_close(snapshot(), reference["after"], rtol=0, atol=0)
+    finally:
+        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", "swap"])
+@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,
+    )
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu"))
+@pytest.mark.parametrize("kind", ["swap", "swap_bf16"])
+@pytest.mark.parametrize("same_axis", [False, True])
+def test_swap_rank_owned_optimizer_fresh_process_resume(tmp_path, kind, same_axis):
+    for phase in ("save", "load"):
+        mp.spawn(
+            _cold_optimizer_worker,
+            args=(str(tmp_path), True, phase, kind, True, same_axis),
+            nprocs=4,
+            join=True,
+        )
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(nproc=4, device="cpu"))
+@pytest.mark.parametrize("canonical", [False, True])
+def test_same_axis_muon_fresh_process_resume(tmp_path, canonical):
+    mp.spawn(_cold_optimizer_worker, args=(str(tmp_path), False, "save", "muon", canonical, True), nprocs=4, join=True)
+    mp.spawn(_cold_optimizer_worker, args=(str(tmp_path), False, "load", "muon", canonical, 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 _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, "_veomni_scheduler_object": None},
+                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, "_veomni_scheduler_object": None},
+                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()
+
+
+@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)
+
+
+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, "_veomni_scheduler_object": None},
+            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, "_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)
+        assert completed.item() == 2
+    finally:
+        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
new file mode 100644
index 0000000000..5478fdf93f
--- /dev/null
+++ b/tests/checkpoints/test_rank_local_ple.py
@@ -0,0 +1,82 @@
+# 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.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)
+    )
+    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, "_veomni_scheduler_object": None},
+                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, "_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)
+            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()
+
+
+@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)
diff --git a/tests/checkpoints/test_ranklocal_provenance_integration.py b/tests/checkpoints/test_ranklocal_provenance_integration.py
new file mode 100644
index 0000000000..3ba877599b
--- /dev/null
+++ b/tests/checkpoints/test_ranklocal_provenance_integration.py
@@ -0,0 +1,60 @@
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+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.tools.launch_utils import spawn_with_rendezvous
+from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
+from veomni.distributed.parallel_plan import SpecInfo
+from veomni.optim.optimizer import MultiOptimizer
+
+
+def _worker(rank, path, endpoint):
+    dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=45))
+    try:
+        mesh = init_device_mesh("cpu", (1, 2), mesh_dim_names=("ep_fsdp", "ep"))
+        parallel = SimpleNamespace(
+            dp_mode="fsdp2", extra_parallel_names=("ep",), extra_parallel_fsdp_device_mesh={"ep": mesh}
+        )
+
+        def build(seed):
+            torch.manual_seed(seed)
+            model = torch.nn.Module()
+            model.experts = torch.nn.Module()
+            values = torch.arange(16, dtype=torch.float32).reshape(4, 4) / 16
+            model.experts.weight = torch.nn.Parameter(distribute_tensor(values, mesh, [Shard(1), Shard(0)]))
+            model._fqn2spec_info = {"experts.weight": SpecInfo("ep", Shard(0), "experts.weight", mesh, 1, 1)}
+            leaf = torch.optim.AdamW(model.parameters(), lr=0.01)
+            return model, MultiOptimizer(model, {"leaf": leaf}, ["leaf"])
+
+        source, optimizer = build(1)
+        for param in source.parameters():
+            param.grad = torch.ones_like(param)
+        optimizer.step()
+        optimizer.zero_grad()
+        state = {"model": source, "optimizer": optimizer, "_veomni_scheduler_object": None}
+        DistributedCheckpointer.save(path, state, parallel_state=parallel)
+        target, restored = build(2)
+        recovered = {"model": target, "optimizer": restored, "_veomni_scheduler_object": None}
+        DistributedCheckpointer.load(path, recovered, parallel_state=parallel)
+        for left, right in zip(source.parameters(), target.parameters(), strict=True):
+            torch.testing.assert_close(left.to_local(), right.to_local(), rtol=0, atol=0)
+            left.grad = torch.ones_like(left)
+            right.grad = torch.ones_like(right)
+        optimizer.step()
+        restored.step()
+        for left, right in zip(source.parameters(), target.parameters(), strict=True):
+            torch.testing.assert_close(left.to_local(), right.to_local(), rtol=0, atol=0)
+        dist.barrier()
+    finally:
+        dist.destroy_process_group()
+
+
+@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)
diff --git a/tests/checkpoints/test_swap_proven_multi_resume.py b/tests/checkpoints/test_swap_proven_multi_resume.py
new file mode 100644
index 0000000000..4d274a94db
--- /dev/null
+++ b/tests/checkpoints/test_swap_proven_multi_resume.py
@@ -0,0 +1,644 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Licensed under the Apache License, Version 2.0.
+"""CPU round-trip checks for host-offloaded state in proven MultiOptimizer DCP."""
+
+from types import SimpleNamespace
+
+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
+from veomni.optim.swap_adamw import SwapAdamW
+
+
+pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+
+
+def _build(*, amsgrad=False, dtype=torch.float32):
+    model = torch.nn.Linear(3, 2).to(dtype)
+    leaf = SwapAdamW(model.parameters(), lr=0.01, betas=(0.9, 0.95), amsgrad=amsgrad, pin_memory=False)
+    return model, leaf, MultiOptimizer(model, {"swap": leaf}, ["swap"])
+
+
+def _step(model, optimizer):
+    for index, parameter in enumerate(model.parameters()):
+        if parameter.requires_grad:
+            parameter.grad = torch.full_like(parameter, (index + 1) * 0.125)
+    optimizer.step()
+    optimizer.zero_grad()
+
+
+@pytest.mark.parametrize("amsgrad", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_swap_proven_resume_matches_next_update(tmp_path, amsgrad, dtype):
+    torch.manual_seed(18)
+    model, source, optimizer = _build(amsgrad=amsgrad, dtype=dtype)
+    _step(model, optimizer)
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target, restored, resumed = _build(amsgrad=amsgrad, dtype=dtype)
+    DistributedCheckpointer.load(
+        str(tmp_path),
+        {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    for left, right in zip(model.parameters(), target.parameters(), strict=True):
+        torch.testing.assert_close(left, right, rtol=0, atol=0)
+        assert source.state[left]["step"] == restored.state[right]["step"]
+        for field in source._host_states[left]:
+            torch.testing.assert_close(
+                source._host_states[left][field], restored._host_states[right][field], rtol=0, atol=0
+            )
+    _step(model, optimizer)
+    _step(target, resumed)
+    for left, right in zip(model.parameters(), target.parameters(), strict=True):
+        torch.testing.assert_close(left, right, rtol=0, atol=0)
+        assert source.state[left]["step"] == restored.state[right]["step"]
+        for field in source._host_states[left]:
+            torch.testing.assert_close(
+                source._host_states[left][field], restored._host_states[right][field], rtol=0, atol=0
+            )
+
+
+@pytest.mark.parametrize("history", [False, True])
+def test_swap_preserves_frozen_or_absent_parameter_history(tmp_path, history):
+    model = torch.nn.Linear(3, 2)
+    if not history:
+        model.bias.requires_grad_(False)
+    source = SwapAdamW(model.parameters(), lr=0.01, pin_memory=False)
+    optimizer = MultiOptimizer(model, {"swap": source}, ["swap"])
+    if not history:
+        model.bias.requires_grad_(True)
+        model.weight.grad = torch.full_like(model.weight, 0.125)
+        optimizer.step()
+        optimizer.zero_grad()
+        assert model.bias not in source.state
+    else:
+        _step(model, optimizer)
+        model.bias.requires_grad_(False)
+
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target = torch.nn.Linear(3, 2)
+    target.bias.requires_grad_(not history)
+    restored = SwapAdamW(target.parameters(), lr=0.01, pin_memory=False)
+    resumed = MultiOptimizer(target, {"swap": restored}, ["swap"])
+    DistributedCheckpointer.load(
+        str(tmp_path),
+        {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    assert (model.bias in source.state) is (target.bias in restored.state)
+    assert target.bias.requires_grad is (not history)
+    if history:
+        assert source.state[model.bias]["step"] == restored.state[target.bias]["step"]
+        for field in source._host_states[model.bias]:
+            torch.testing.assert_close(
+                source._host_states[model.bias][field], restored._host_states[target.bias][field], rtol=0, atol=0
+            )
+        model.bias.requires_grad_(True)
+        target.bias.requires_grad_(True)
+    _step(model, optimizer)
+    _step(target, resumed)
+    for left, right in zip(model.parameters(), target.parameters(), strict=True):
+        torch.testing.assert_close(left, right, rtol=0, atol=0)
+
+
+def test_swap_empty_leaf_roundtrip_has_no_synthetic_state(tmp_path):
+    model = torch.nn.Linear(3, 2)
+    for param in model.parameters():
+        param.requires_grad_(False)
+    source = SwapAdamW(model.parameters(), lr=0.01, pin_memory=False)
+    optimizer = MultiOptimizer(model, {"swap": source}, ["swap"])
+    assert not source.state
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target = torch.nn.Linear(3, 2)
+    for param in target.parameters():
+        param.requires_grad_(False)
+    restored = SwapAdamW(target.parameters(), lr=0.01, pin_memory=False)
+    resumed = MultiOptimizer(target, {"swap": restored}, ["swap"])
+    DistributedCheckpointer.load(
+        str(tmp_path),
+        {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    assert not source.state and not restored.state
+    for left, right in zip(model.parameters(), target.parameters(), strict=True):
+        left.requires_grad_(True)
+        right.requires_grad_(True)
+    _step(model, optimizer)
+    _step(target, resumed)
+    for left, right in zip(model.parameters(), target.parameters(), strict=True):
+        torch.testing.assert_close(left, right, rtol=0, atol=0)
+
+
+def test_swap_all_frozen_history_recovers_without_synthetic_step(tmp_path):
+    model, source, optimizer = _build()
+    _step(model, optimizer)
+    for param in model.parameters():
+        param.requires_grad_(False)
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target = torch.nn.Linear(3, 2)
+    for param in target.parameters():
+        param.requires_grad_(False)
+    restored = SwapAdamW(target.parameters(), lr=0.01, pin_memory=False)
+    resumed = MultiOptimizer(target, {"swap": restored}, ["swap"])
+    assert not restored.state
+    DistributedCheckpointer.load(
+        str(tmp_path),
+        {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    for left, right in zip(model.parameters(), target.parameters(), strict=True):
+        assert not right.requires_grad
+        assert source.state[left]["step"] == restored.state[right]["step"] == 1
+        for field in source._state_keys:
+            torch.testing.assert_close(
+                source._host_states[left][field], restored._host_states[right][field], rtol=0, atol=0
+            )
+        left.requires_grad_(True)
+        right.requires_grad_(True)
+    _step(model, optimizer)
+    _step(target, resumed)
+    for left, right in zip(model.parameters(), target.parameters(), strict=True):
+        torch.testing.assert_close(left, right, rtol=0, atol=0)
+
+
+def test_swap_restore_hooks_observe_live_state(tmp_path):
+    model, source, optimizer = _build()
+    _step(model, optimizer)
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target, leaf, resumed = _build()
+    live_state = leaf.state
+    seen = []
+
+    def before(current, state_dict):
+        assert current.state is live_state
+        seen.append("before")
+
+    def after(current):
+        assert current.state is live_state
+        source_param = next(model.parameters())
+        target_param = next(target.parameters())
+        torch.testing.assert_close(
+            current._host_states[target_param]["exp_avg"], source._host_states[source_param]["exp_avg"], rtol=0, atol=0
+        )
+        seen.append("after")
+
+    leaf.register_load_state_dict_pre_hook(before)
+    leaf.register_load_state_dict_post_hook(after)
+    DistributedCheckpointer.load(
+        str(tmp_path),
+        {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    assert seen == ["before", "after"]
+
+
+def test_swap_checkpoint_state_hooks_observe_live_state(tmp_path):
+    model, _, optimizer = _build()
+    _step(model, optimizer)
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target, leaf, resumed = _build()
+    seen = []
+
+    def before(opt):
+        assert all(isinstance(state["step"], int) for state in opt.state.values())
+        seen.append("before")
+
+    def after(opt, state_dict):
+        assert all(isinstance(state["step"], int) for state in opt.state.values())
+        assert all(isinstance(state["step"], int) for state in state_dict["state"].values() if "step" in state)
+        seen.append("after")
+        return state_dict
+
+    leaf.register_state_dict_pre_hook(before)
+    leaf.register_state_dict_post_hook(after)
+    DistributedCheckpointer.load(
+        str(tmp_path),
+        {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    assert seen and seen.count("before") == seen.count("after")
+
+
+def test_swap_and_fp32_master_leaf_preserve_both_updates(tmp_path):
+    from tests.checkpoints.test_optimizer_checkpoint_adapter import ToyAdamW
+
+    def build():
+        model = torch.nn.Linear(3, 2)
+        swap = SwapAdamW([model.weight], lr=0.01, pin_memory=False)
+        master = ToyAdamW([model.bias])
+        return model, swap, master, MultiOptimizer(model, {"swap": swap, "master": master}, ["swap", "master"])
+
+    source, swap, master, 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_swap, restored_master, resumed = build()
+    DistributedCheckpointer.load(
+        str(tmp_path),
+        {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    for field in ("master_weight", "exp_avg", "exp_avg_sq"):
+        assert restored_master.state[target.bias][field].dtype == torch.float32
+        torch.testing.assert_close(
+            master.state[source.bias][field], restored_master.state[target.bias][field], rtol=0, atol=0
+        )
+    for field in swap._state_keys:
+        torch.testing.assert_close(
+            swap._host_states[source.weight][field], restored_swap._host_states[target.weight][field], rtol=0, atol=0
+        )
+    _step(source, optimizer)
+    _step(target, resumed)
+    for left, right in zip(source.parameters(), target.parameters(), strict=True):
+        torch.testing.assert_close(left, right, rtol=0, atol=0)
+    torch.testing.assert_close(
+        master.state[source.bias]["master_weight"], restored_master.state[target.bias]["master_weight"], rtol=0, atol=0
+    )
+
+
+def test_swap_rejects_nonfresh_host_state_before_loading(tmp_path):
+    model, _, optimizer = _build()
+    _step(model, optimizer)
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target, leaf, resumed = _build()
+    first = next(target.parameters())
+    leaf._host_states[first]["exp_avg"].fill_(0.25)
+    before = {name: param.detach().clone() for name, param in target.named_parameters()}
+    with pytest.raises(RuntimeError, match="pristine host state"):
+        DistributedCheckpointer.load(
+            str(tmp_path),
+            {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    for name, param in target.named_parameters():
+        torch.testing.assert_close(param, before[name], rtol=0, atol=0)
+    assert torch.all(leaf._host_states[first]["exp_avg"] == 0.25)
+    assert not getattr(leaf, "_veomni_dcp_restore_failed", False)
+
+
+def test_swap_failed_apply_keeps_host_buffers_and_blocks_update(tmp_path, monkeypatch):
+    model, _, optimizer = _build()
+    _step(model, optimizer)
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target, leaf, resumed = _build()
+    original = {
+        param: {key: value.clone() for key, value in states.items()} for param, states in leaf._host_states.items()
+    }
+
+    def fail_before_apply(current, state_dict):
+        for param, states in current._veomni_dcp_staged_state.items():
+            for field in current._state_keys:
+                assert states[field].data_ptr() != current._host_states[param][field].data_ptr()
+                states[field].fill_(17)
+        raise RuntimeError("injected staged restore failure")
+
+    monkeypatch.setattr(SwapAdamW, "load_state_dict", fail_before_apply)
+    with pytest.raises(RuntimeError, match="injected staged restore failure"):
+        DistributedCheckpointer.load(
+            str(tmp_path),
+            {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    assert "_veomni_dcp_staged_state" not in vars(leaf)
+    for param, states in original.items():
+        for field, value in states.items():
+            torch.testing.assert_close(leaf._host_states[param][field], value, rtol=0, atol=0)
+    with pytest.raises(RuntimeError, match="failed checkpoint restore"):
+        leaf.step()
+    with pytest.raises(RuntimeError, match="failed checkpoint restore"):
+        resumed.step()
+
+
+def test_failed_swap_restore_blocks_all_multi_optimizer_leaves(tmp_path, monkeypatch):
+    def build():
+        model = torch.nn.Linear(3, 2)
+        adam = torch.optim.AdamW([model.weight], lr=0.01)
+        swap = SwapAdamW([model.bias], lr=0.01, pin_memory=False)
+        return model, MultiOptimizer(model, {"adam": adam, "swap": swap}, ["adam", "swap"])
+
+    model, optimizer = build()
+    _step(model, optimizer)
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(
+        str(tmp_path),
+        {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    target, resumed = build()
+
+    def fail(*_args, **_kwargs):
+        raise RuntimeError("injected later optimizer leaf failure")
+
+    monkeypatch.setattr(SwapAdamW, "load_state_dict", fail)
+    with pytest.raises(RuntimeError, match="later optimizer leaf failure"):
+        DistributedCheckpointer.load(
+            str(tmp_path),
+            {"model": target, "optimizer": resumed, "_veomni_scheduler_object": None},
+            parallel_state=parallel,
+        )
+    for param in target.parameters():
+        param.grad = torch.ones_like(param)
+    before = target.weight.detach().clone()
+    adam_step = resumed.optimizers_dict["adam"].state[target.weight]["step"].clone()
+    with pytest.raises(RuntimeError, match="failed checkpoint restore"):
+        resumed.step()
+    torch.testing.assert_close(target.weight, before, rtol=0, atol=0)
+    torch.testing.assert_close(resumed.optimizers_dict["adam"].state[target.weight]["step"], adam_step, rtol=0, atol=0)
+
+
+def test_swap_rejects_inexact_checkpoint_step_before_saving(tmp_path):
+    model, leaf, optimizer = _build()
+    first = next(model.parameters())
+    leaf.state[first]["step"] = 2**24 + 1
+    with pytest.raises(RuntimeError, match="cannot be represented exactly"):
+        DistributedCheckpointer.save(
+            str(tmp_path),
+            {"model": model, "optimizer": optimizer, "_veomni_scheduler_object": None},
+            parallel_state=SimpleNamespace(dp_mode="ddp"),
+        )
+    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)
+    _step(model, source)
+    model.bias.requires_grad_(False)
+    parallel = SimpleNamespace(dp_mode="ddp")
+    DistributedCheckpointer.save(str(tmp_path), {"model": model, "optimizer": source}, parallel_state=parallel)
+
+    target = torch.nn.Linear(3, 2)
+    target.bias.requires_grad_(False)
+    restored = SwapAdamW(target.parameters(), lr=0.01, pin_memory=False)
+    DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": restored}, parallel_state=parallel)
+    assert not target.bias.requires_grad
+    assert restored.state[target.bias]["step"] == source.state[model.bias]["step"]
+    for field in source._state_keys:
+        torch.testing.assert_close(
+            source._host_states[model.bias][field], restored._host_states[target.bias][field], rtol=0, atol=0
+        )
+    model.bias.requires_grad_(True)
+    target.bias.requires_grad_(True)
+    _step(model, source)
+    _step(target, restored)
+    torch.testing.assert_close(model.bias, target.bias, rtol=0, atol=0)
+
+
+def test_single_swap_absent_frozen_history_stays_absent(tmp_path):
+    model = torch.nn.Linear(3, 2)
+    model.bias.requires_grad_(False)
+    optimizer = SwapAdamW(model.parameters(), pin_memory=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(target.parameters(), pin_memory=False)
+    DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": resumed}, parallel_state=parallel)
+    assert not target.bias.requires_grad
+    assert target.bias not in resumed.state
+    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)
+    _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(target.parameters(), pin_memory=False)
+    seen = []
+
+    def model_hook(module, _incompatible_keys):
+        assert not module.bias.requires_grad
+        seen.append(True)
+
+    target.register_load_state_dict_post_hook(model_hook)
+    DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": resumed}, parallel_state=parallel)
+    assert seen
+    assert not target.bias.requires_grad
+    assert resumed.state[target.bias]["step"] == optimizer.state[model.bias]["step"]
+
+
+def test_single_swap_frozen_history_rejects_optimizer_load_hook_before_loading(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(target.parameters(), pin_memory=False)
+    resumed.register_load_state_dict_pre_hook(lambda *_: None)
+    before = target.weight.clone()
+    with pytest.raises(RuntimeError, match="does not support optimizer load-state hooks"):
+        DistributedCheckpointer.load(str(tmp_path), {"model": target, "optimizer": resumed}, parallel_state=parallel)
+    assert not target.bias.requires_grad
+    assert target.bias not in resumed.state
+    torch.testing.assert_close(target.weight, before, rtol=0, atol=0)
+
+
+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/tests/checkpoints/test_tp_expert_save_preflight.py b/tests/checkpoints/test_tp_expert_save_preflight.py
new file mode 100644
index 0000000000..6f1bccc634
--- /dev/null
+++ b/tests/checkpoints/test_tp_expert_save_preflight.py
@@ -0,0 +1,134 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: Apache-2.0
+"""Unsupported TP expert saves must fail before checkpoint writes."""
+
+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 veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
+from veomni.distributed.parallel_plan import SpecInfo
+from veomni.distributed.parallel_state import ParallelState
+from veomni.utils.dist_utils import raise_if_any_rank_failed
+
+
+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")),
+        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(
+                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 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()
+
+
+@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, 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")
+    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)
+
+
+@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 _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/checkpoints/test_tp_fsdp_provenance.py b/tests/checkpoints/test_tp_fsdp_provenance.py
new file mode 100644
index 0000000000..1dee4ddb9b
--- /dev/null
+++ b/tests/checkpoints/test_tp_fsdp_provenance.py
@@ -0,0 +1,242 @@
+# 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 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):
+    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")
+    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/tests/ci/_npu_moe_fp32_cold_worker.py b/tests/ci/_npu_moe_fp32_cold_worker.py
new file mode 100644
index 0000000000..1cbbe34c8a
--- /dev/null
+++ b/tests/ci/_npu_moe_fp32_cold_worker.py
@@ -0,0 +1,233 @@
+# 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
+from torch import nn
+from torch.distributed._tensor import DTensor, Shard
+
+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, 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):
+    def __init__(self):
+        super().__init__()
+        config = SimpleNamespace(
+            num_experts=4,
+            hidden_size=256,
+            moe_intermediate_size=32,
+            hidden_act="silu",
+            ops_config=SimpleNamespace(moe_implementation="fused_npu"),
+        )
+        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__()
+        # 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):
+        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 _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 _assert_selective_placement(model, optimizer, resident):
+    expected = "npu" if resident else "cpu"
+    for param in model.layer.experts.parameters():
+        assert _local(param).device.type == expected
+        state = _optimizer_state(optimizer, param)
+        assert all(_local(state[key]).device.type == expected for key in ("exp_avg", "exp_avg_sq"))
+
+
+def _snapshot(model, optimizer):
+    result = {}
+    for name, param in model.named_parameters():
+        state = _optimizer_state(optimizer, param)
+        parameter_device = _local(param).device.type
+        assert all(_local(state[part]).device.type == parameter_device for part in ("exp_avg", "exp_avg_sq"))
+        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()
+    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
+        assert _local(param.grad).device.type == _local(param).device.type
+        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, *, selective=False, resident=False):
+    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)
+        _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,),
+        )
+        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,
+            cpu_offload_extra_parallel_names=["ep"] if selective else [],
+            cpu_offload_exclude_module_patterns=["layer.experts"] if resident else [],
+        )
+        assert model.layer.experts._veomni_fp32_expert_wgrad
+        if selective:
+            expected = "npu" if resident else "cpu"
+            assert all(_local(param).device.type == expected for param in model.layer.experts.parameters())
+        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)
+            # 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, "_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, "_veomni_scheduler_object": None},
+                parallel_state=get_parallel_state(),
+            )
+        if selective:
+            _assert_selective_placement(model, optimizer, resident)
+        before = _snapshot(model, optimizer)
+        after = _step(model, optimizer, rank, ep, 2)
+        if selective:
+            _assert_selective_placement(model, optimizer, resident)
+        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()
+
+
+def run_selective(rank, phase, output_dir, resident, endpoint):
+    run(rank, phase, output_dir, 2, endpoint, selective=True, resident=resident)
+
+
+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..b738c80c60
--- /dev/null
+++ b/tests/ci/test_npu_moe_fp32_cold.py
@@ -0,0 +1,208 @@
+# 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
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch.distributed.checkpoint import FileSystemReader
+
+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
+
+
+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 left.dtype == right.dtype and left.shape == right.shape
+        assert torch.equal(left, right)
+    else:
+        assert left == right
+
+
+def _assert_close_across_devices(left, right, path="root"):
+    if isinstance(left, dict):
+        assert left.keys() == right.keys()
+        for key in left:
+            _assert_close_across_devices(left[key], right[key], f"{path}.{key}")
+    elif isinstance(left, torch.Tensor):
+        assert left.shape == right.shape and left.dtype == right.dtype
+        if not left.is_floating_point():
+            assert torch.equal(left, right)
+            return
+        assert torch.isfinite(left).all() and torch.isfinite(right).all()
+        difference = left.double() - right.double()
+        denominator = max(
+            torch.linalg.vector_norm(left.double()).item(), torch.linalg.vector_norm(right.double()).item()
+        )
+        relative_l2 = torch.linalg.vector_norm(difference).item() / denominator if denominator else 0.0
+        max_abs = difference.abs().max().item() if difference.numel() else 0.0
+        print(f"{path}: max_abs={max_abs:.9g}, relative_l2={relative_l2:.9g}", flush=True)
+        torch.testing.assert_close(left, right, rtol=1e-5, atol=1e-12)
+        assert relative_l2 <= 1e-5
+    else:
+        assert left == right
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+def test_cross_device_comparison_rejects_precision_or_dtype_regressions():
+    reference = {"tensor": torch.tensor([1.0, 2.0], dtype=torch.float32)}
+    _assert_close_across_devices(reference, {"tensor": reference["tensor"].clone()})
+    with pytest.raises(AssertionError):
+        _assert_close_across_devices(reference, {"tensor": torch.tensor([1.0, 2.001])})
+    with pytest.raises(AssertionError):
+        _assert_close_across_devices(reference, {"tensor": reference["tensor"].double()})
+
+
+def _assert_ep2_checkpoint_layout(checkpoint_dir):
+    weights = FileSystemReader(weights_dir(str(checkpoint_dir))).read_metadata().state_dict_metadata
+    for name, shape in (
+        ("layer.experts.gate_up_proj", (4, 64, 256)),
+        ("layer.experts.down_proj", (4, 256, 32)),
+    ):
+        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
+
+
+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):
+    _require_two_native_npus()
+    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)
+    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")
+
+
+@pytest.mark.launch_plan(provider=_native_npu_launch)
+def test_native_fp32_expert_selective_offload_cold_resume(tmp_path):
+    _require_two_native_npus()
+    from tests.tools.launch_utils import spawn_with_rendezvous
+
+    results = {}
+    for resident in (False, True):
+        output = tmp_path / ("resident" if resident else "offloaded")
+        output.mkdir()
+        for phase in ("fresh", "resume"):
+            spawn_with_rendezvous(run_selective, args=(phase, str(output), resident), nprocs=2)
+        results[resident] = []
+        for rank in range(2):
+            fresh = torch.load(output / f"ep2-fresh-rank{rank}.pt", map_location="cpu", weights_only=True)
+            resumed = torch.load(output / f"ep2-resume-rank{rank}.pt", map_location="cpu", weights_only=True)
+            _assert_equal(fresh, resumed)
+            results[resident].append(fresh)
+        _assert_ep2_checkpoint_layout(output / "ep2-dcp")
+    for offloaded, resident in zip(results[False], results[True], strict=True):
+        # AdamW executes on CPU for offloaded shards and on NPU for resident
+        # shards. Their floating-point arithmetic need not be bitwise identical.
+        _assert_close_across_devices(offloaded, resident)
diff --git a/tests/ci/test_npu_moe_precision.py b/tests/ci/test_npu_moe_precision.py
new file mode 100644
index 0000000000..375446088a
--- /dev/null
+++ b/tests/ci/test_npu_moe_precision.py
@@ -0,0 +1,231 @@
+# 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("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",
+    [
+        ({"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_fp8_policy_is_ignored_for_npu_experts_without_skipping_fp32_validation(monkeypatch):
+    from veomni.distributed import torch_parallelize
+
+    state = SimpleNamespace(dp_mode="fsdp2", tp_size=1, pp_size=1, tp_enabled=False, fsdp_enabled=True)
+    monkeypatch.setattr(torch_parallelize, "get_parallel_state", lambda: state)
+    calls = []
+    monkeypatch.setattr(
+        torch_parallelize, "parallelize_model_fsdp2", lambda **kwargs: calls.append(kwargs) or kwargs["model"]
+    )
+    model = _model_with_experts()
+    mixed = MixedPrecisionConfig(npu_expert_fp32_wgrad=True)
+    assert (
+        torch_parallelize.build_parallelize_model(
+            model, mixed_precision=mixed, enable_gradient_checkpointing=False, expert_weight_fp8_all_gather=True
+        )
+        is model
+    )
+    assert calls[0]["expert_weight_fp8_all_gather"] is False
+    assert calls[0]["mixed_precision"].npu_expert_fp32_wgrad is True
+
+    calls.clear()
+    with pytest.raises(ValueError, match="FP32 wgrad requires"):
+        torch_parallelize.build_parallelize_model(
+            model,
+            mixed_precision=MixedPrecisionConfig(npu_expert_fp32_wgrad=True, reduce_dtype="bfloat16"),
+            enable_gradient_checkpointing=False,
+            expert_weight_fp8_all_gather=True,
+        )
+    assert not calls
+
+
+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
+    from veomni.distributed.torch_parallelize import _fsdp_target_classes, _is_fsdp_wrap_target
+    from veomni.models.configuration_base import as_veomni_config
+    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"
+    as_veomni_config(config, validate=False)._set_ops_implementation(eager_ops_config())
+    with 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/models/test_optimizer_checkpoint_preflight.py b/tests/models/test_optimizer_checkpoint_preflight.py
new file mode 100644
index 0000000000..5f44d3faf6
--- /dev/null
+++ b/tests/models/test_optimizer_checkpoint_preflight.py
@@ -0,0 +1,227 @@
+# Copyright 2025 Bytedance Ltd. and/or its affiliates
+# Licensed under the Apache License, Version 2.0.
+
+"""MultiOptimizer checkpoint compatibility is checked before any training step."""
+
+import os
+from contextlib import nullcontext
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+
+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"))
+
+
+def _runtime(multi, save_async, manager="dcp"):
+    class StubRuntime(model_runtime.VeOmniModelRuntime):
+        def setup(self):
+            pass
+
+        def _build_model(self):
+            self.model = torch.nn.Linear(2, 2)
+
+        def _freeze_model_module(self):
+            pass
+
+        def _build_parallelized_model(self):
+            pass
+
+        def _build_optimizer(self):
+            if multi == "fail":
+                raise RuntimeError("injected optimizer build failure")
+            self.optimizer = None if multi is None else SimpleNamespace(_is_multi_optimizer=multi)
+
+        def _build_model_assets(self):
+            pass
+
+        def build_checkpoint(self):
+            pass
+
+    args = SimpleNamespace(optimizer=SimpleNamespace(type="adamw"))
+    train = SimpleNamespace(checkpoint=SimpleNamespace(load_path=None, save_async=save_async, manager=manager))
+    return StubRuntime(args, train=train)
+
+
+@pytest.mark.parametrize("multi,save_async", [(None, True), (False, False), (False, True), (True, False)])
+def test_compatible_optimizer_checkpoint_modes(monkeypatch, multi, save_async):
+    monkeypatch.setattr(model_runtime, "use_parallel_state", lambda *_: nullcontext())
+    optimizer = _runtime(multi, save_async).optimizer
+    assert (None if optimizer is None else optimizer._is_multi_optimizer) is multi
+
+
+def test_async_multi_optimizer_fails_at_runtime_build(monkeypatch):
+    monkeypatch.setattr(model_runtime, "use_parallel_state", lambda *_: nullcontext())
+    with pytest.raises(RuntimeError, match="synchronous"):
+        _runtime(True, True)
+
+
+def test_other_checkpoint_managers_keep_their_own_async_policy(monkeypatch):
+    monkeypatch.setattr(model_runtime, "use_parallel_state", lambda *_: nullcontext())
+    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
+
+    runtime = ModuleRuntime.__new__(ModuleRuntime)
+    runtime.model_name = "llm"
+    runtime.model = torch.nn.Linear(2, 2)
+    runtime.train_args = SimpleNamespace(checkpoint=SimpleNamespace(save_async=True))
+    monkeypatch.setattr(runtime, "_scoped", nullcontext)
+    monkeypatch.setattr(runtime, "_apply_recompute_policy", lambda: None)
+    monkeypatch.setattr(
+        runtime, "_build_optimizer", lambda: setattr(runtime, "optimizer", SimpleNamespace(_is_multi_optimizer=True))
+    )
+    reached_checkpoint = []
+    monkeypatch.setattr(runtime, "build_checkpoint", lambda: reached_checkpoint.append(True))
+    with pytest.raises(RuntimeError, match="synchronous"):
+        runtime.build_after_omni_model_wrap()
+    assert not reached_checkpoint
+
+
+@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)
+    monkeypatch.setattr(ModuleRuntime, "_scoped", lambda self: nullcontext())
+    monkeypatch.setattr(ModuleRuntime, "_build_model", lambda self: setattr(self, "model", torch.nn.Linear(2, 2)))
+    monkeypatch.setattr(ModuleRuntime, "_build_model_assets", lambda self: None)
+    monkeypatch.setattr(ModuleRuntime, "_freeze_model_module", lambda self: None)
+    monkeypatch.setattr(ModuleRuntime, "_build_parallelized_model", lambda self: None)
+    monkeypatch.setattr(
+        ModuleRuntime,
+        "_build_optimizer",
+        lambda self: setattr(self, "optimizer", None if multi is None else SimpleNamespace(_is_multi_optimizer=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", 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:
+        with pytest.raises(RuntimeError, match="synchronous"):
+            ModuleRuntime(SimpleNamespace(accelerator=accelerator), "llm", **kwargs)
+        assert not reached_checkpoint
+    else:
+        ModuleRuntime(SimpleNamespace(accelerator=accelerator), "llm", **kwargs)
+        assert reached_checkpoint == [True]
+
+
+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, 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)
+        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_optimizer_checkpoint_failures_reach_all_ranks():
+    from tests.tools.launch_utils import spawn_with_rendezvous
+
+    spawn_with_rendezvous(_preflight_failure_worker, args=(), nprocs=2)
+
+
+def _accelerator_worker(rank, endpoint):
+    from veomni.utils.device import get_device_type, get_dist_comm_backend, get_torch_device
+
+    device = get_device_type()
+    get_torch_device().set_device(f"{device}:{rank}")
+    dist.init_process_group(
+        get_dist_comm_backend(), store=endpoint.connect(), rank=rank, world_size=2, timeout=timedelta(seconds=60)
+    )
+    try:
+        model_runtime.use_parallel_state = lambda *_: nullcontext()
+        _runtime(True, False)
+        with pytest.raises(RuntimeError, match="synchronous|checking MultiOptimizer checkpoint mode"):
+            _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)
+        dist.all_reduce(done)
+        assert done.item() == 2
+    finally:
+        dist.destroy_process_group()
+
+
+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.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)
diff --git a/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py b/tests/ops/moe_experts/test_deepgemm_fp8_distributed.py
index 791b630b3f..84f4d18a08 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,33 +313,46 @@ 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, "_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):
+                    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, "_veomni_scheduler_object": None}
                 )
                 for (group, param, key), expected_value in saved_optim.items():
                     actual_value = opt.optimizers_dict[group].state[param][key]
@@ -347,18 +367,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 +383,45 @@ 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, "_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)
+        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/ops/moe_experts/test_npu_fp32_wgrad.py b/tests/ops/moe_experts/test_npu_fp32_wgrad.py
new file mode 100644
index 0000000000..d58f3374ce
--- /dev/null
+++ b/tests/ops/moe_experts/test_npu_fp32_wgrad.py
@@ -0,0 +1,174 @@
+# 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 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
+
+
+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),
+    )
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+@pytest.mark.parametrize("enabled", [False, True])
+def test_npu_wrapper_accepts_quantization_options_and_forwards_fp32_wgrad(monkeypatch, enabled):
+    from veomni.ops.kernels.moe_experts.standard import npu
+
+    captured = {}
+
+    def forward(*args, **kwargs):
+        captured.update(kwargs)
+        return args[3]
+
+    monkeypatch.setattr(npu, "npu_fused_moe_forward", forward)
+    tensor = torch.ones(1)
+    result = npu.wrapper(
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        tensor,
+        num_experts=1,
+        block_size=64,
+        expert_dtype="fp4",
+        fp32_wgrad=enabled,
+    )
+    assert result is tensor
+    assert captured["fp32_wgrad"] is enabled
+
+
+@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]
+    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())
+
+
+@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]
+    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.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()
+    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.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)
+    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))
+
+
+@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"):
+        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/tests/optim/test_multi_optimizer_sparse_lr.py b/tests/optim/test_multi_optimizer_sparse_lr.py
new file mode 100644
index 0000000000..747afd15ed
--- /dev/null
+++ b/tests/optim/test_multi_optimizer_sparse_lr.py
@@ -0,0 +1,253 @@
+"""An unused optimizer leaf still has a scheduler-controlled learning rate."""
+
+from types import SimpleNamespace
+
+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
+
+
+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)
+    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_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
+
+
+@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(), 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"])
+    result = optimizer._sparse_state_dict()
+    assert calls == [True]
+    assert result["param_groups.weight.lr"] == 0.0025
+    assert not leaf.state
+
+
+@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(), 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()]
+    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, "_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()
+        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, "_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):
+        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_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, "_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, "_veomni_scheduler_object": None},
+        parallel_state=parallel,
+    )
+    assert restored_bias.param_groups[0]["lr"] == 0.0025
+
+
+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"])
+
+        from veomni.optim.lr_scheduler import build_lr_scheduler
+
+        model, first, second, optimizer = build(42)
+        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()
+        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,
+                "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)
+
+        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)
+        )
+        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)
+    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)
diff --git a/tests/optim/test_muon_unit.py b/tests/optim/test_muon_unit.py
index fdc6a3ca19..626862ba5a 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,53 @@ 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, "_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)
+
+
+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 +386,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 +403,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 +415,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 +422,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 +1136,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 +1160,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:
diff --git a/tests/seed_omni/model/test_hf_layout.py b/tests/seed_omni/model/test_hf_layout.py
index 007ca86c7d..e06d88eb93 100644
--- a/tests/seed_omni/model/test_hf_layout.py
+++ b/tests/seed_omni/model/test_hf_layout.py
@@ -492,7 +492,32 @@ def test_an_offline_embedding_run_skips_the_export_check(tmp_path, monkeypatch):
     assert len(checked) == 1
 
 
-def test_custom_layout_subclass_overrides_hooks_and_per_module_dtypes(tmp_path):
+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])
+
+
+@pytest.mark.parametrize("namespaced", [False, True])
+def test_custom_layout_subclass_overrides_hooks_and_per_module_dtypes(tmp_path, monkeypatch, namespaced):
     from veomni.arguments.omni_arguments_types import _check_hf_source_exportable
 
     custom_type = "custom_named_shards_omni"
@@ -594,6 +619,19 @@ def finalize_exported_shards(self, source_path: str, save_path: str, weight_file
         assert split_model.modules_dict[FAKE_A].proj.weight.dtype == torch.bfloat16
         assert split_model.modules_dict[FAKE_B].proj.weight.dtype == torch.float32
 
+        if namespaced:
+            from veomni.checkpoint.dcp_checkpointer import (
+                ModelState,
+                _extra_parallel_layout_key,
+                _namespace_extra_parallel_layout_state,
+            )
+
+            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: (model.modules_dict[FAKE_A], None)}, str(out))
         assert not (out / "model.safetensors.index.json").exists()
diff --git a/tests/seed_omni/runtime/test_module_runtime.py b/tests/seed_omni/runtime/test_module_runtime.py
index 737389356c..7b63092261 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:
@@ -62,6 +66,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
@@ -141,7 +233,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 +369,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 +684,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 +764,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 bbc2138f4b..f2807453d7 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
 
@@ -97,10 +97,12 @@ def test_parallelize_forwards_module_skip_decision(
                 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,
+                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,
diff --git a/tests/seed_omni/runtime/test_omni_model_runtime.py b/tests/seed_omni/runtime/test_omni_model_runtime.py
index 06d946bc0c..d43f95c5e4 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,
@@ -414,6 +415,29 @@ 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 test_composed_wrap_leaves_cache_reader_modules_on_meta(monkeypatch):
     """Under ``fsdp_scope='model'``, cache readers must not go through weight materialisation."""
     model = _model()
@@ -442,10 +466,12 @@ 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,
+                cpu_offload_extra_parallel_names=[],
+                cpu_offload_exclude_module_patterns=[],
                 max_load_broadcast_size=20.0,
                 expert_weight_fp8_all_gather=True,
                 expert_fsdp_shard_dim=0,
diff --git a/tests/utils/test_extra_parallel_clip_grad_norm.py b/tests/utils/test_extra_parallel_clip_grad_norm.py
index dac6bc2fd0..4ce9fab767 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
@@ -39,6 +40,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 +131,183 @@ 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(fused, checkpoint_dir):
+    from torch.utils._pytree import tree_map
+
+    from veomni.checkpoint.dcp_checkpointer import DistributedCheckpointer
+    from veomni.optim.lr_scheduler import build_lr_scheduler
+
+    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
+
+    def optimizer_snapshot(optimizer):
+        return tree_map(
+            lambda tensor: local(tensor).detach().cpu().clone() if torch.is_tensor(tensor) else tensor,
+            optimizer.state_dict(),
+        )
+
+    def optimizer_lrs(optimizer):
+        return {name: [group["lr"] for group in leaf.param_groups] for name, leaf in optimizer.optimizers_dict.items()}
+
+    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=fused, optimizer_type="adamw")
+                for model in (reference, candidate)
+            ]
+            schedulers = [
+                build_lr_scheduler(optimizer, train_steps=20, lr=0.125, lr_decay_style="linear")
+                for optimizer in optimizers
+            ]
+            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=1e-5, atol=1e-12)
+                    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=1e-5, atol=1e-12)
+                before_clip = [local(parameter.grad).detach().cpu().clone() for parameter in candidate.parameters()]
+                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)
+                assert norms[1].item() > 0.5
+                coefficient = torch.clamp(0.5 / (norms[1].detach().cpu() + 1e-6), max=1.0)
+                for left, right, before in zip(
+                    reference.parameters(), candidate.parameters(), before_clip, strict=True
+                ):
+                    clipped = local(right.grad).cpu()
+                    torch.testing.assert_close(local(left.grad).cpu(), clipped, rtol=1e-5, atol=1e-12)
+                    torch.testing.assert_close(clipped, before * coefficient, rtol=1e-5, atol=1e-12)
+                for optimizer, scheduler in zip(optimizers, schedulers, strict=True):
+                    optimizer.step()
+                    scheduler.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)
+            checkpoint = str(Path(checkpoint_dir) / f"resident-{len(resident)}")
+            assert any(lr < 0.125 for rates in optimizer_lrs(optimizers[1]).values() for lr in rates)
+            for leaf in optimizers[1].optimizers_dict.values():
+                for parameter, state in leaf.state.items():
+                    if local(parameter).device.type == "cpu":
+                        for value in (parameter, state["exp_avg"], state["exp_avg_sq"]):
+                            assert local(value).device.type == "cpu"
+                            assert value.dtype == torch.float32
+            saved = {
+                "model": candidate,
+                "optimizer": optimizers[1],
+                "extra_state": {"lr_scheduler": schedulers[1].state_dict()},
+                "_veomni_scheduler_object": schedulers[1],
+            }
+            DistributedCheckpointer.save(checkpoint, saved, parallel_state=ps)
+            resumed = build(True, resident)
+            resumed_optimizer = build_optimizer(
+                resumed, lr=0.125, weight_decay=0.0, fused=fused, optimizer_type="adamw"
+            )
+            resumed_scheduler = build_lr_scheduler(
+                resumed_optimizer, train_steps=20, lr=0.125, lr_decay_style="linear"
+            )
+            restored = {
+                "model": resumed,
+                "optimizer": resumed_optimizer,
+                "extra_state": {},
+                "_veomni_scheduler_object": resumed_scheduler,
+            }
+            DistributedCheckpointer.load(checkpoint, restored, parallel_state=ps)
+            resumed_scheduler.load_state_dict(restored["extra_state"]["lr_scheduler"])
+
+            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
+            )
+            assert schedulers[1].state_dict() == resumed_scheduler.state_dict()
+            assert optimizer_lrs(optimizers[1]) == optimizer_lrs(resumed_optimizer)
+            for model, optimizer, scheduler in (
+                (candidate, optimizers[1], schedulers[1]),
+                (resumed, resumed_optimizer, resumed_scheduler),
+            ):
+                model(torch.ones(16, device=device)).backward()
+                veomni_clip_grad_norm(model, 0.5)
+                optimizer.step()
+                scheduler.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
+            )
+            assert schedulers[1].state_dict() == resumed_scheduler.state_dict()
+            assert optimizer_lrs(optimizers[1]) == optimizer_lrs(resumed_optimizer)
+            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"])
+
+
 # (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,16 +514,22 @@ 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")
+    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}")
     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 +541,8 @@ def main():
         scenarios = list(reversed(scenarios))
 
     try:
+        if cli_args.suite == "selective_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,
@@ -374,7 +561,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",
@@ -385,6 +574,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
 
@@ -395,6 +588,12 @@ 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)
+@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)
+
+
 # ---------------------------------------------------------------------------
 # 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..033d800f8d 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,47 @@ 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)
+
+
+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
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=[],
diff --git a/veomni/arguments/arguments_types.py b/veomni/arguments/arguments_types.py
index f8e64541be..d46a241fb4 100644
--- a/veomni/arguments/arguments_types.py
+++ b/veomni/arguments/arguments_types.py
@@ -96,7 +96,8 @@ class OptimizerConfig:
                 "Optimizer-managed state offload. 'none' disables it (training state may still be "
                 "offloaded by model.accelerator.fsdp_config.offload); 'cpu' keeps the moment states "
                 "in host buffers and streams them to the device in batches, lowering peak device "
-                "memory (type='adamw' only)."
+                "memory (type='adamw' only). This cannot be combined with selective ExtraParallel "
+                "CPU offload."
             )
         },
     )
@@ -281,6 +282,10 @@ 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)
         # Annotations are not enforced for YAML values (see constraints), so the
         # literal is validated explicitly; otherwise a typo would silently fall
         # back to the stock path in the builder.
@@ -621,8 +626,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.")
@@ -726,6 +744,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."},
+    )
     expert_fsdp_shard_dim: Optional[int] = field(
         default=None,
         metadata={
@@ -761,6 +787,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; "
@@ -1093,6 +1126,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.
@@ -1101,6 +1140,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 "
@@ -1998,6 +2043,10 @@ def __post_init__(self):
                 "model.optimizer.state_offload='cpu' cannot be combined with "
                 "model.accelerator.fsdp_config.offload; both move training state to host."
             )
+        if self.optimizer.state_offload == "cpu" and self.accelerator.fsdp_config.cpu_offload_extra_parallel_names:
+            raise ValueError(
+                "model.optimizer.state_offload='cpu' cannot be combined with selective ExtraParallel CPU offload."
+            )
         # ep_sharded_stream_load only runs on the every-rank-reads path, so it is
         # mutually exclusive with broadcast_model_weights_from_rank0. Fail early
         # instead of silently ignoring the flag.
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 8aee5244cf..414a454848 100644
--- a/veomni/checkpoint/dcp_checkpointer.py
+++ b/veomni/checkpoint/dcp_checkpointer.py
@@ -15,10 +15,15 @@
 
 import gc
 import hashlib
+import json
 import os
+import re
 import shutil
+import uuid
+import weakref
 from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
 from dataclasses import replace
 from datetime import timedelta
 from typing import Any, Dict, Optional, Set, Union
@@ -26,14 +31,19 @@
 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,
     load,
 )
 from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
-from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE, Metadata
+from torch.distributed.checkpoint.metadata import (
+    STATE_DICT_TYPE,
+    BytesStorageMetadata,
+    Metadata,
+    TensorStorageMetadata,
+)
 from torch.distributed.checkpoint.state_dict import (
     StateDictOptions,
     get_model_state_dict,
@@ -44,9 +54,9 @@
 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.device import empty_cache, get_device_type, synchronize
 from ..utils.dist_utils import any_rank_failed, raise_if_any_rank_failed
 from ..utils.dtensor_utils import rewrap_dtensor_local
 from .checkpointer import CheckpointerBase
@@ -55,11 +65,31 @@
     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,
+    _clear_swap_checkpoint_staging,
+    _dcp_parameter_name_map,
+    _inventory_digests,
+    _iter_optimizer_leaves,
+    _multi_extra_snapshot,
+    _preflight_multi_optimizer_layout,
+    _prepare_multi_optimizer_saved_states,
+    _validate_fresh_swap_optimizer,
+    _validate_source_optimizer_states,
+)
+from .optimizer_resume_layout import (
+    MULTI_LAYOUT_MARKER,
+    capture_multi_layout,
+    finish_multi_layout,
+    frozen_parameter_digest,
+)
 
 
 logger = logging.get_logger(__name__)
@@ -69,42 +99,494 @@
 _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"
+_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):
-    """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:
+    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
-    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 isinstance(ep, _StridedShard):
+        return {"expert_order": "unproven"}
     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)
+    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, _StridedShard)) 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"
+    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,
+        "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 _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
+    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
+    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, _StridedShard)) and tensor.device_mesh.size(axis) > 1
+    ]
+    return len(dims) != len(set(dims))
+
+
+def _verified_rank_local_keys(adapter, *, loading):
+    """Only the proven optimizer adapter can supply rank-local inventory keys."""
+    if (
+        type(adapter) is not _ProvenMultiOptimizerState
+        or adapter._load is not loading
+        or not adapter.should_extra_parallel_aware
+    ):
+        return frozenset()
+    return adapter._verified_rank_local_keys
+
+
+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,
+        rank_local_adapter=None,
+        allow_staged_copies=False,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.parallel_state = parallel_state
+        self.rank_local_adapter = rank_local_adapter
+        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)
+            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 = {}
+        rank_local_keys = _verified_rank_local_keys(self.rank_local_adapter, loading=False)
+        for key, value in self.state_dict.items():
+            if key in rank_local_keys:
+                continue
+            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(
+                    logical_key, value, self.parallel_state, self.ordinary_fqns, self.shared_lora_fqns
+                )
+                or _tp_only_source_layout(
+                    logical_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
+            elif _unidentified_same_axis_shards(value) or (
+                "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})
 
     def create_global_plan(self, all_plans):
@@ -113,7 +595,24 @@ 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
+        has_v2_model = any(
+            key.startswith("model.") and _is_extra_parallel_layout_marker(key.split(".")[-1])
+            for key in metadata.state_dict_metadata
+        )
+        if has_v2_model:
+            unproven = sorted(
+                key
+                for key, value in layouts.items()
+                if key.startswith("model.") and value.get("expert_order") != "canonical"
+            )
+            if unproven:
+                raise RuntimeError(
+                    f"Cannot save canonical ExtraParallel checkpoint with unproven expert layout: {unproven[0]}"
+                )
         metadata.planner_data = {
             **(metadata.planner_data or {}),
             _EXPERT_LAYOUTS_KEY: {"version": 1, "tensors": layouts},
@@ -124,11 +623,51 @@ 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):
+    def __init__(
+        self, strict_model: bool, expected_save_id=None, legacy_expert_source_layout=None, *, rank_local_adapter=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
+        self.rank_local_adapter = rank_local_adapter
 
     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(str(path[-1]).split(".")[-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"
+                    )
+        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")
         self._validate_expert_layouts()
         plan = super().create_local_plan()
         if not self.strict_model:
@@ -153,37 +692,67 @@ def create_local_plan(self):
     def _validate_expert_layouts(self):
         planner_data = self.metadata.planner_data or {}
         source = planner_data.get(_EXPERT_LAYOUTS_KEY)
+        rank_local_keys = _verified_rank_local_keys(self.rank_local_adapter, loading=True)
+        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 key in rank_local_keys:
+                    continue
+                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.
-        if planner_data.get("__veomni_ep_first_layout__") == 1:
+        # 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({})
+        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():
+            if key in source_tensors or key in rank_local_keys:
+                continue
             layout = _expert_layout(value)
             saved = self.metadata.state_dict_metadata.get(key)
-            if layout is None or saved is None:
+            if saved is None or not torch.is_tensor(value) or not hasattr(saved, "size"):
                 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 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)
-                for chunk in saved.chunks
+            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 (
+                is_expert
+                and bool(saved.size)
+                and (has_v2_marker or (source is not None and self.legacy_expert_source_layout is None))
             ):
+                raise RuntimeError(
+                    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
+            # 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 not is_expert:
+                    continue
+            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. "
@@ -220,6 +789,647 @@ 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 _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, saving=False
+):
+    error, has_v2, logical_signature = None, False, None
+    try:
+        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)
+        has_v2 = _needs_extra_parallel_layout_marker(specs)
+        if has_v2:
+            _extra_parallel_layout_key(specs)
+        if state.get("optimizer") is not None and getattr(ps, "dp_mode", None) == "fsdp2":
+            if has_v2 and not allow_proven_optimizer:
+                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 named_params))
+                    if trainable_only
+                    else ()
+                ),
+                runtime_ndims={
+                    name: tensor.ndim
+                    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
+    raise_if_any_rank_failed(error, "checking canonical ExtraParallel optimizer state", group=group)
+    _check_extra_parallel_signature_across_ranks(has_v2, logical_signature, group=group)
+    error = None
+    try:
+        if has_v2:
+            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 _checkpoint_extra_parallel_signature(
+    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 = []
+    for name, spec in sorted((specs or {}).items()):
+        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
+        fsdp_axis = getattr(spec, "fsdp_shard_dim", None)
+        persistent_axis = getattr(spec, "persistent_fsdp_shard_dim", None)
+        layout_v2 = bool(getattr(spec, "checkpoint_layout_v2", False))
+        if persistent_axis is not None:
+            # Persistent shards do not use the ordinary FSDP axis or v2 marker.
+            fsdp_axis, layout_v2 = None, None
+        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
+        mesh_names = getattr(mesh, "mesh_dim_names", None)
+        entries.append(
+            (
+                name,
+                spec.para_name,
+                type(placement).__name__,
+                axis,
+                fsdp_axis,
+                persistent_axis,
+                tuple(shape) if shape is not None else None,
+                layout_v2,
+                tuple(mesh_names) if mesh_names is not None else None,
+                tuple(mesh.shape) if mesh is not None else None,
+            )
+        )
+    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):
+    """Reject rank-inconsistent checkpoint policy before collectives or file mutation."""
+    if not dist.is_initialized():
+        return
+    device = "cpu" if dist.get_backend(group) == "gloo" else get_device_type()
+    presence = torch.tensor([int(has_v2), int(not has_v2)], dtype=torch.int32, device=device)
+    dist.all_reduce(presence, op=dist.ReduceOp.MAX, group=group)
+    if bool(presence[0]) and bool(presence[1]):
+        raise RuntimeError("checkpoint ExtraParallel layout version differs across ranks")
+    if not bool(presence[0]):
+        return
+    digest = hashlib.sha256(repr(signature).encode("utf-8")).digest()
+    low = torch.tensor(list(digest), dtype=torch.int32, device=device)
+    high = low.clone()
+    dist.all_reduce(low, op=dist.ReduceOp.MIN, group=group)
+    dist.all_reduce(high, op=dist.ReduceOp.MAX, group=group)
+    if not torch.equal(low, high):
+        raise RuntimeError("checkpoint ExtraParallel layout or optimizer policy differs across ranks")
+
+
+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 _assert_extra_parallel_layout_unchanged(owner):
+    specs = owner.extra_parallel_fqn2spec_info
+    if specs is not getattr(owner.model, "_fqn2spec_info", None) or (
+        owner._layout_signature != _runtime_extra_parallel_layout_signature(specs)
+    ):
+        raise RuntimeError("ExtraParallel layout changed while checkpoint adapter was active")
+
+
+def _coordinate_extra_parallel_layout_check(owner, group):
+    error = None
+    try:
+        _assert_extra_parallel_layout_unchanged(owner)
+    except BaseException as exc:
+        error = exc
+    raise_if_any_rank_failed(error, "checking checkpoint adapter layout", group=group)
+
+
+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 CPU optimizer shards without moving or changing their strides."""
+    if local_tensor.device.type != "cpu" or device_mesh.device_type == "cpu":
+        return DTensor.from_local(
+            local_tensor, device_mesh=device_mesh, placements=placements, run_check=False, shape=shape, stride=stride
+        )
+    wrapped = rewrap_dtensor_local(
+        local_tensor,
+        mesh=device_mesh,
+        placements=placements,
+        shape=shape,
+        stride=stride,
+    )
+    if wrapped.requires_grad != local_tensor.requires_grad:
+        wrapped = DTensor(local_tensor, wrapped._spec, requires_grad=local_tensor.requires_grad)
+    return wrapped
+
+
+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 or ep_size == 1 or isinstance(loaded_tensor.placements[fsdp_dim], _StridedShard):
+            # 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 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],
@@ -236,8 +1446,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.<FQN>.``, e.g.
       ``"state.model.layers.0.mlp.experts.gate_proj.exp_avg"``.
 
     Non-tensor values and 0-D tensors are skipped unconditionally — they
@@ -254,12 +1463,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}"
@@ -275,12 +1490,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
@@ -288,6 +1497,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).
@@ -295,12 +1514,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
@@ -326,14 +1562,28 @@ 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):
+        _assert_extra_parallel_layout_unchanged(self)
         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"
@@ -341,7 +1591,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()
@@ -351,7 +1602,13 @@ 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
+        _assert_extra_parallel_layout_unchanged(self)
+        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")
 
@@ -368,6 +1625,45 @@ 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
+    )
+
+
+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.
 
@@ -409,11 +1705,18 @@ class instead ends an absent param at ``step`` 1 (the checkpoint
             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 = (
+            _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"
         )
@@ -421,28 +1724,51 @@ def __init__(self, model, optimizer, parallel_state=None, *, load: bool = False)
 
     def state_dict(self):
         if self.should_extra_parallel_aware:
-            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()
-            optim_sd_with_extra_parallel_dim = self.get_state_dict_with_extra_parallel_dim_preprocess(
-                vanilla_optim_sd, "restore"
-            )
-            return optim_sd_with_extra_parallel_dim
-
+            # Reject unsupported raw DCP on every rank before any optimizer
+            # can enter its own collectives, even if a local check fails.
+            error = None
+            try:
+                _assert_extra_parallel_layout_unchanged(self)
+                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, "validating ExtraParallel optimizer checkpoint", group=self.process_group)
+            raise RuntimeError("ExtraParallel optimizer checkpoint preflight unexpectedly succeeded")
+
+        _assert_extra_parallel_layout_unchanged(self)
+        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):
         """Restore optimizer state from the loaded DCP state dict."""
-        optim_state_from_dcp_load = state_dict
+        _assert_extra_parallel_layout_unchanged(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")
+        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(
@@ -481,6 +1807,147 @@ 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
+
+    def state_dict(self):
+        _coordinate_extra_parallel_layout_check(self, self.group)
+        result, error = None, None
+        try:
+            result = super().state_dict()
+            logical_keys = {key for raw in result if (key := _export_extra_parallel_layout_key(raw)) is not None}
+            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)
+        return result
+
+    def load_state_dict(self, state_dict):
+        _coordinate_extra_parallel_layout_check(self, self.group)
+        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,
+        expected_layout,
+        group,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.generation = generation
+        self.expected_layout = expected_layout
+        self.group = group
+        self._verified_rank_local_keys = frozenset()
+
+    def state_dict(self):
+        _coordinate_extra_parallel_layout_check(self, self.group)
+        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())
+                self._verified_rank_local_keys = frozenset(inventory)
+                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
+        result[MULTI_LAYOUT_MARKER] = json.dumps(
+            {
+                "generation": self.generation,
+                "layout": self.expected_layout,
+                "inventory_digests": _inventory_digests(inventory, group=self.group),
+            },
+            sort_keys=True,
+        )
+        return result
+
+    def load_state_dict(self, state_dict):
+        _coordinate_extra_parallel_layout_check(self, self.group)
+        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)
+        digests = _inventory_digests(inventory, group=self.group)
+        error = (
+            RuntimeError("MultiOptimizer inventory identity changed during tensor loading")
+            if marker["inventory_digests"] != digests
+            else None
+        )
+        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)
+        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.
@@ -491,7 +1958,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
@@ -499,19 +1967,21 @@ def drop_extra_parallel_dim(loaded_tensor: torch.Tensor, device_mesh: DeviceMesh
     """
 
     num_placements = len(loaded_tensor.placements)
-    if num_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 (num_placements in (2, 3) and not has_fsdp_shard):
         tensor_to_put = loaded_tensor.to_local()
     elif num_placements == 2:
-        tensor_to_put = rewrap_dtensor_local(
-            loaded_tensor._local_tensor,
-            mesh=device_mesh,
-            placements=[Shard(fsdp_shard_dim)],
+        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 = rewrap_dtensor_local(
-            loaded_tensor._local_tensor,
-            mesh=device_mesh,
-            placements=[Replicate(), Shard(fsdp_shard_dim)],
+        tensor_to_put = _from_local_preserve_device(
+            loaded_tensor._local_tensor, device_mesh=device_mesh, placements=[Replicate(), Shard(fsdp_shard_dim)]
         )
     else:
         raise RuntimeError(
@@ -557,6 +2027,11 @@ def restore_extra_parallel_dim(
     assert fsdp_mesh.ndim in (2, 3), f"global_mesh.ndim must be 2 or 3, got {fsdp_mesh.ndim}"
 
     if isinstance(orgin_tensor, DTensor):
+        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
         # 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().
@@ -569,21 +2044,18 @@ def restore_extra_parallel_dim(
             placements = [Replicate(), fsdp_placement, Shard(ep_shard_dim)]
         else:
             placements = [fsdp_placement, Shard(ep_shard_dim)]
-        dtensor = rewrap_dtensor_local(
-            orgin_tensor._local_tensor,
-            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. Use the helper here too:
-        # ``DTensor.from_local`` copies a host-resident local onto the mesh
-        # device, which would materialize the swap optimizer's host state on the
-        # accelerator during a save.
-        dtensor = rewrap_dtensor_local(
-            orgin_tensor,
-            mesh=extra_parallel_fsdp_mesh,
-            placements=[Shard(0)],
-        )
+        # 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")
+        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 = _from_local_preserve_device(orgin_tensor, device_mesh=fsdp_mesh, placements=placements)
     else:
         raise RuntimeError(f"origin_tensor - {orgin_tensor} is not a tensor!")
 
@@ -684,8 +2156,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
@@ -856,6 +2331,140 @@ def copy_markers() -> None:
         raise RuntimeError("checkpoint promotion failed on another rank; no completion marker was written")
 
 
+@contextmanager
+def _single_swap_frozen_resume(model, optimizer, reader, *, group=None, trainable_only=False):
+    """Include saved frozen SwapAdamW owners in the legacy DCP template."""
+    from ..optim.swap_adamw import SwapAdamW
+
+    if type(optimizer) is not SwapAdamW:
+        yield ()
+        return
+    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:
+        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)
+            }
+            present = {key for key, value in fields.items() if value is not None}
+            if not present:
+                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:
+                raise RuntimeError("Frozen SwapAdamW history does not support optimizer load-state hooks")
+            for field, value in fields.items():
+                expected_shape = torch.Size([]) if field == "step" else param.shape
+                expected_dtype = torch.float32 if field == "step" else param.dtype
+                if (
+                    not isinstance(value, TensorStorageMetadata)
+                    or value.size != expected_shape
+                    or value.properties.dtype != expected_dtype
+                ):
+                    raise RuntimeError(f"SwapAdamW frozen state {name}.{field} has incompatible shape or dtype")
+            if param not in optimizer._host_states:
+                allocated.append(param)
+                optimizer._ensure_param_state(param, bool(optimizer._group_by_param[param]["amsgrad"]))
+            param.requires_grad_(True)
+            changed.append(param)
+    except BaseException as exc:
+        error = exc
+    try:
+        raise_if_any_rank_failed(error, "preparing frozen SwapAdamW checkpoint state", group=group)
+    except BaseException:
+        for param in changed:
+            param.requires_grad_(False)
+        for param in allocated:
+            optimizer._host_states.pop(param, None)
+            optimizer._device_states.pop(param, None)
+            optimizer.state.pop(param, None)
+        raise
+    try:
+        yield changed
+    except BaseException:
+        optimizer._veomni_dcp_restore_failed = True
+        raise
+    finally:
+        for param in changed:
+            param.requires_grad_(False)
+
+
+@contextmanager
+def _single_swap_frozen_model_view(changed):
+    for param in changed:
+        param.requires_grad_(False)
+    try:
+        yield
+        if any(param.requires_grad for param in changed):
+            raise RuntimeError("Model load-state hooks changed frozen SwapAdamW parameter trainability")
+    finally:
+        for param in changed:
+            param.requires_grad_(True)
+
+
 class DistributedCheckpointer(CheckpointerBase):
     """
     Distributed checkpointer for torch.distributed.checkpoint
@@ -951,19 +2560,93 @@ def save(
         return:
             None
         """
-        if "model" not in state:
-            raise ValueError("Model must be provided to save a distributed checkpoint.")
-
-        # 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, 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")
+        _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
+        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 "_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)
+                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)
+            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:
+                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:
+            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 = 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
         # multi-module job calls this once per module, and a single key would have
@@ -976,54 +2659,106 @@ 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)
-
-        # The per-rank extra state 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 per-rank pickles
-        # 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
+        raise_if_any_rank_failed(error, "preparing checkpoint files")
+        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")
+            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
+            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,
+                parallel_state=parallel_state if parallel_state is not None else get_parallel_state(),
+                source_model=state["model"],
             )
 
             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,
+                            expected_layout=initial_layout,
+                            group=None,
                         )
-                    },
-                    storage_writer=cls._create_storage_writer(os.path.join(write_root, OPTIMIZER_DIRNAME)),
+                        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)),
+                    ),
+                    "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,
                     timeout_seconds=save_timeout_seconds,
+                    parallel_state=parallel_state if parallel_state is not None else get_parallel_state(),
+                    source_model=state["model"],
                 )
         except BaseException:
             if stage_path is not None and _local_rank() == 0:
@@ -1034,6 +2769,7 @@ def save(
             _promote_staged_checkpoint(stage_path, model_root, 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:
@@ -1074,6 +2810,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
@@ -1088,6 +2825,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
@@ -1097,11 +2837,13 @@ 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:
-            raise ValueError("Model must be provided to load a distributed checkpoint.")
+        if isinstance(state, dict):
+            state.pop("_veomni_optimizer_identity_verified", None)
+            state.pop("_veomni_optimizer_checkpoint_generation", None)
+        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
+        )
 
         def _model_state() -> "ModelState":
             return ModelState(state["model"], trainable_only=trainable_only, parallel_state=parallel_state)
@@ -1112,10 +2854,20 @@ 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
 
+        error = 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")
+        except BaseException as exc:
+            error = exc
+        raise_if_any_rank_failed(error, "inspecting checkpoint layout", group=process_group)
+
         model_root = model_dir(path, module)
         # Every directory this load will read has to be complete, or ``dcp.load``
         # fails with a bare "metadata is None" that names neither the path nor
@@ -1166,27 +2918,199 @@ def _optimizer_state() -> "OptimizerState":
         raise_if_any_rank_failed(error, "validating the checkpoint to load", group=process_group)
         if error is not None:
             raise error
-        dcp.load(
-            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),
-        )
-        if wants_optimizer:
+        specs = getattr(state["model"], "_fqn2spec_info", None)
+        if _needs_extra_parallel_layout_marker(specs):
+            marker_error = None
+            try:
+                _read_extra_parallel_marker(
+                    cls._create_storage_reader(weights_dir(path, module)), _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 multi:
+            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)
+            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, 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={"optimizer": _optimizer_state()},
-                storage_reader=cls._create_storage_reader(optimizer_dir(path, module)),
+                state_dict={"model": model_state},
+                storage_reader=reader,
                 process_group=process_group,
-                # The strict check looks for missing ``model`` keys; there are none
-                # in an optimizer-only load, so strictness has nothing to say here.
-                planner=_ModelStrictLoadPlanner(strict_model=False),
+                planner=planner,
             )
-
-        cls._load_extra_state(checkpoint_dir=model_root, state=state)
-
-        logger.info_rank0(f"Loaded checkpoint from {model_root}")
-
-        return state
+            try:
+                _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, reader, planner, error = None, None, None, None
+                try:
+                    optimizer_state = _ProvenMultiOptimizerState(
+                        model=state["model"],
+                        optimizer=optimizer,
+                        parallel_state=parallel_state,
+                        load=True,
+                        generation=generation,
+                        expected_layout=expected_layout,
+                        group=process_group,
+                    )
+                    reader = cls._create_storage_reader(optimizer_path)
+                    planner = _ModelStrictLoadPlanner(
+                        strict_model=False,
+                        expected_save_id=optimizer_id,
+                        rank_local_adapter=optimizer_state if optimizer_state.should_extra_parallel_aware else None,
+                    )
+                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=reader,
+                    process_group=process_group,
+                    planner=planner,
+                )
+            except BaseException:
+                from ..optim.swap_adamw import SwapAdamW
+
+                optimizer._veomni_dcp_restore_failed = True
+                for leaf in _iter_optimizer_leaves(optimizer):
+                    if type(leaf) is SwapAdamW:
+                        leaf._veomni_dcp_restore_failed = True
+                raise
+            finally:
+                _clear_swap_checkpoint_staging(optimizer)
+            if extra is not None:
+                state["extra_state"].update(extra)
+            state["_veomni_optimizer_identity_verified"] = True
+            state["_veomni_optimizer_checkpoint_generation"] = generation
+            return state
+        model_state, optimizer_state, model_reader, reader, error = None, None, None, None, None
+        try:
+            model_state = _model_state()
+            optimizer_state = _optimizer_state() if wants_optimizer else None
+            model_reader = cls._create_storage_reader(weights_dir(path, module))
+            if wants_optimizer:
+                reader = cls._create_storage_reader(optimizer_dir(path, module))
+        except BaseException as exc:
+            error = exc
+        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 = (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 = any_rank_failed(bool(frozen_swap_owners), group=process_group)
+            error = None
+            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)
+            model_root = model_dir(path, module)
+            error = None
+            try:
+                with _single_swap_frozen_model_view(frozen_swap_owners):
+                    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 = 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(model_root)
+            logger.info_rank0(f"Loaded checkpoint from {model_root}")
+            return state
 
     @classmethod
     def wait_for_pending_save(cls) -> None:
@@ -1288,6 +3212,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.
 
@@ -1300,7 +3226,21 @@ 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:
+            optimizer_adapter = save_state.get("optimizer")
+            planner = _ExtraParallelSavePlanner(
+                dedup_save_to_lowest_rank=save_to_lowest_rank,
+                parallel_state=parallel_state,
+                source_model=source_model,
+                rank_local_adapter=optimizer_adapter
+                if type(optimizer_adapter) is _ProvenMultiOptimizerState
+                else None,
+                allow_staged_copies=save_async,
+            )
+        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
@@ -1326,9 +3266,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
@@ -1397,6 +3342,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.*
@@ -1421,11 +3370,60 @@ def _normalize_key(key: str) -> Optional[str]:
         return key[6:]
 
 
+def _validate_expert_export_provenance(metadata: Metadata, model_keys, legacy_expert_source_layout=None) -> None:
+    """Reject expert layouts that a plain CPU export cannot prove safe."""
+    planner_data = metadata.planner_data or {}
+    recorded = planner_data.get(_EXPERT_LAYOUTS_KEY)
+    layouts = {}
+    if recorded is not None:
+        if recorded["version"] != 1:
+            raise RuntimeError(f"Unsupported checkpoint expert layout version: {recorded['version']}")
+        layouts = recorded["tensors"]
+        for key in model_keys & layouts.keys():
+            if layouts[key]["expert_order"] != "canonical":
+                raise RuntimeError(
+                    f"Checkpoint tensor {key!r} has noncanonical expert order; migrate it before export"
+                )
+
+    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.
+    # Plain CPU export targets do not carry the EP mesh needed to distinguish it.
+    for key in model_keys:
+        tensor_metadata = metadata.state_dict_metadata[key]
+        shape = getattr(tensor_metadata, "size", ())
+        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 has_v2_marker or (recorded is not None and legacy_expert_source_layout is None):
+            raise RuntimeError(
+                f"Cannot determine source expert order for {key!r}: checkpoint metadata omits its layout. "
+                "Verify and migrate the original training layout before exporting."
+            )
+        if legacy_expert_source_layout is None and not all(
+            tuple(chunk.sizes) == tuple(shape) for chunk in tensor_metadata.chunks
+        ):
+            raise RuntimeError(
+                f"Cannot determine source expert order for {key!r} in this unversioned checkpoint. "
+                "Verify and migrate the original training layout before exporting."
+            )
+
+
 def _get_sharding_plan(
     checkpoint_path: Union[str, os.PathLike],
     shard_size: int = None,
     save_dtype: Optional[Union[str, torch.dtype]] = None,
     drop_hf_keys: Optional[Set[str]] = None,
+    *,
+    legacy_expert_source_layout=None,
+    include_hf_key=None,
 ):
     """
     Create sharding plan from checkpoint metadata without loading weights.
@@ -1435,23 +3433,41 @@ def _get_sharding_plan(
         total_size: Total size in bytes
         all_dcp_keys: All valid DCP model keys
     """
+    _validate_legacy_expert_source_layout(legacy_expert_source_layout)
     reader = FileSystemReader(checkpoint_path)
     metadata = reader.read_metadata()
 
     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():
         hf_key = _normalize_key(key)
         if hf_key:
+            if include_hf_key is not None and not include_hf_key(hf_key):
+                continue
             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"):
@@ -1474,6 +3490,8 @@ def _get_sharding_plan(
             tensor_infos.append({"dcp_key": key, "hf_key": hf_key, "size": byte_size, "metadata": tensor_meta})
             all_dcp_keys.append(key)
 
+    _validate_expert_export_provenance(metadata, set(all_dcp_keys), legacy_expert_source_layout)
+
     if dropped:
         logger.info(f"Excluded {dropped} tensor(s) named by `drop_hf_keys` from the plan")
 
@@ -1502,6 +3520,8 @@ def _get_sharding_plan(
     if current_shard:
         shards.append(current_shard)
     if shard_size is None:
+        if not shards and include_hf_key is not None:
+            return {}, total_size, all_dcp_keys
         assert len(shards) == 1, "Shard size None should result in a single shard"
         shards = shards[0]
     return shards, total_size, all_dcp_keys
@@ -1511,13 +3531,23 @@ 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()
 
     state_dict = OrderedDict()
     dcp_keys_to_load = list(shard_keys.values())
 
+    # A partial offline shard must carry the checked layout marker too, so
+    # the same planner can validate the source namespace before copying data.
+    marker = "model." + _extra_parallel_layout_key({})
+    if marker in metadata.state_dict_metadata:
+        _read_extra_parallel_marker(reader, marker.removeprefix("model."))
+        state_dict[marker] = ""
+
     for dcp_key in dcp_keys_to_load:
         tensor_metadata = metadata.state_dict_metadata[dcp_key]
         if not hasattr(tensor_metadata.properties, "dtype"):
@@ -1534,7 +3564,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,
     )
 
@@ -1567,19 +3597,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)
+    shard, _, _ = _get_sharding_plan(save_checkpoint_path, legacy_expert_source_layout=legacy_expert_source_layout)
 
-    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
diff --git a/veomni/checkpoint/dynamic_rope.py b/veomni/checkpoint/dynamic_rope.py
new file mode 100644
index 0000000000..bcaa7457cb
--- /dev/null
+++ b/veomni/checkpoint/dynamic_rope.py
@@ -0,0 +1,146 @@
+# 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 importlib
+import json
+import math
+
+import torch
+
+from .tensor_bytes import tensor_byte_chunks
+
+
+def _is_native_qwen3_dynamic(module):
+    """Recognize upstream and VeOmni-generated Qwen3 rotary implementations."""
+    cls = type(module)
+    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
+    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 (
+        _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):
+    value = value.detach().cpu()
+    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):
+        digest.update(chunk)
+    return digest.hexdigest()
+
+
+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")
+                # 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](
+                    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()
+                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 tuple(inv.shape) == tuple(initial_inv.shape)
+                    and inv_digest == initial_digest
+                    and matches_canonical
+                    and math.isfinite(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 '<root>'}") from exc
+            if not valid:
+                raise RuntimeError(
+                    f"Dynamic RoPE cache for {name or '<root>'} depends on prior sequence lengths; "
+                    "its nonpersistent state cannot yet be restored exactly"
+                )
+            # 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/layout.py b/veomni/checkpoint/layout.py
index 23eb56779e..f2dc9ee211 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/checkpoint/multi_optimizer_checkpoint.py b/veomni/checkpoint/multi_optimizer_checkpoint.py
new file mode 100644
index 0000000000..a2d8fb234a
--- /dev/null
+++ b/veomni/checkpoint/multi_optimizer_checkpoint.py
@@ -0,0 +1,640 @@
+# 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 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,
+    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 _iter_optimizer_leaves(optimizer):
+    yield from optimizer.optimizers_dict.values()
+
+
+def _validate_fresh_swap_optimizer(leaf):
+    """Verify SwapAdamW's eagerly allocated buffers still have fresh state."""
+    if "_veomni_dcp_staged_state" in vars(leaf) or getattr(leaf, "_veomni_dcp_restore_failed", False):
+        raise RuntimeError("SwapAdamW checkpoint resume requires a fresh optimizer without state")
+    params = {param for group in leaf.param_groups for param in group["params"]}
+    if set(leaf.state) != set(leaf._host_states) or set(leaf.state) != set(leaf._device_states):
+        raise RuntimeError("SwapAdamW checkpoint resume requires a fresh optimizer without state")
+    if not set(leaf.state).issubset(params):
+        raise RuntimeError("SwapAdamW checkpoint resume found unowned state")
+    for param, state in leaf.state.items():
+        host = leaf._host_states[param]
+        if (
+            type(state.get("step")) is not int
+            or state["step"] != 0
+            or set(state) != {"step", *leaf._state_keys}
+            or set(host) != set(leaf._state_keys)
+        ):
+            raise RuntimeError("SwapAdamW checkpoint resume requires a fresh optimizer without state")
+        for value in host.values():
+            if value.device.type != "cpu" or bool(torch.count_nonzero(value).item()):
+                raise RuntimeError("SwapAdamW checkpoint resume requires pristine host state")
+
+
+def _clear_swap_checkpoint_staging(optimizer):
+    from ..optim.swap_adamw import SwapAdamW
+
+    for leaf in _iter_optimizer_leaves(optimizer):
+        if type(leaf) is SwapAdamW:
+            leaf.__dict__.pop("_veomni_dcp_staged_state", None)
+
+
+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
+    from ..optim.swap_adamw import SwapAdamW
+
+    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) in (torch.optim.AdamW, SwapAdamW):
+        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
+    from ..optim.swap_adamw import SwapAdamW
+
+    for leaf in _iter_optimizer_leaves(optimizer):
+        if type(leaf) is SwapAdamW:
+            if getattr(leaf, "_veomni_dcp_restore_failed", False) or "_veomni_dcp_staged_state" in vars(leaf):
+                raise RuntimeError("SwapAdamW cannot save an incomplete checkpoint restore")
+            if set(leaf.state) != set(leaf._host_states) or set(leaf.state) != set(leaf._device_states):
+                raise RuntimeError("SwapAdamW checkpoint state ownership is inconsistent")
+        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")
+                if type(leaf) is SwapAdamW:
+                    local_param = param.to_local() if isinstance(param, DTensor) else param
+                    host = leaf._host_states[param]
+                    if set(host) != set(leaf._state_keys) or any(
+                        value.device.type != "cpu" or value.shape != local_param.shape or value.dtype != param.dtype
+                        for value in host.values()
+                    ):
+                        raise RuntimeError("SwapAdamW checkpoint host buffers have an incompatible layout")
+                for field, value in state.items():
+                    if field == "step" and type(leaf) is SwapAdamW:
+                        # SwapAdamW serializes this integer as float32 for AdamW
+                        # compatibility; larger values would silently round.
+                        if type(value) is not int or not 0 <= value <= 2**24:
+                            raise RuntimeError("SwapAdamW step cannot be represented exactly in its checkpoint")
+                        continue
+                    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")
+                        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)
+                            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 value.dtype != expected_dtype:
+                            raise RuntimeError(f"MultiOptimizer state dtype cannot be safely restored for {field}")
+
+
+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 _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 = 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:
+        if getattr(optimizer, "_veomni_dcp_restore_failed", False):
+            raise RuntimeError("MultiOptimizer checkpoint resume requires a newly constructed optimizer")
+        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
+
+        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):
+            from ..optim.swap_adamw import SwapAdamW
+
+            if type(leaf) is SwapAdamW:
+                _validate_fresh_swap_optimizer(leaf)
+            elif 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")
+        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
+
+        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 = getattr(optimizer, "_is_multi_optimizer", False) is True
+    if not multi:
+        return
+    error = None
+    try:
+        from ..optim.checkpoint_adapter import rank_local_adamw_adapter
+        from ..optim.optimizer import AnyPrecisionAdamW
+        from ..optim.swap_adamw import SwapAdamW
+
+        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):
+            staged = {} if type(leaf) is SwapAdamW and not validate_only else None
+            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")
+                            if type(leaf) is SwapAdamW and step_dtype != torch.float32:
+                                raise RuntimeError("SwapAdamW checkpoint step must be a float32 scalar")
+                            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 type(leaf) is not SwapAdamW
+                                and (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]
+                            ]
+                        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 type(leaf) is SwapAdamW or (
+                                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:
+                        if staged is not None:
+                            staged[param] = state
+                        else:
+                            leaf.state[param] = state
+            if staged is not None:
+                if "_veomni_dcp_staged_state" in vars(leaf):
+                    raise RuntimeError("SwapAdamW checkpoint staging is already active")
+                leaf._veomni_dcp_staged_state = staged
+    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..69ce22ac59
--- /dev/null
+++ b/veomni/checkpoint/optimizer_resume_layout.py
@@ -0,0 +1,602 @@
+# 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 bisect
+import dis
+import hashlib
+import json
+import marshal
+import math
+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__"
+
+_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):
+    if value is None or type(value) in (bool, int, float, str):
+        return value
+    if type(value) is 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 _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
+    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_functions(scheduler):
+    """Describe LambdaLR callables using the existing checkpoint identity format."""
+
+    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.
+        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",
+                # Hash-randomized set iteration can change floating-point LR
+                # results across processes without changing the recipe hash.
+                "BUILD_SET",
+                "SET_ADD",
+                "SET_UPDATE",
+            ):
+                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")
+                if instruction.argval not in _SCHEDULER_MATH_MEMBERS:
+                    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:
+                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) 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")
+
+        functions.append(
+            {
+                "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 ()],
+            }
+        )
+    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(
+    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
+    from ..optim.optimizer import AnyPrecisionAdamW, MultiOptimizer
+    from ..optim.swap_adamw import SwapAdamW
+
+    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}")
+    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")
+    from .dynamic_rope import _is_native_qwen3_dynamic_reset_alias, dynamic_rope_checkpoint_contract
+
+    model_aliases, model_tensors = {}, {}
+    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)
+            )
+            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):
+            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(), {}, {}
+    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,
+            SwapAdamW,
+        ):
+            raise RuntimeError("MultiOptimizer checkpoint requires an explicitly supported optimizer leaf")
+        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")
+        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,
+        # 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),
+        "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_v2"] = 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()
+
+
+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 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
+    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)
+    except BaseException as exc:
+        error = exc
+    raise_if_any_rank_failed(error, "capturing MultiOptimizer checkpoint layout", group=group)
+    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,
+    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 = None
+    error = None
+    try:
+        local = _local_layout(
+            model,
+            optimizer,
+            scheduler_snapshot,
+            scheduler_object,
+            extra_digest,
+            writing=writing,
+            frozen_digest=frozen_digest,
+        )
+    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
+        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]
+    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..313416d8d4
--- /dev/null
+++ b/veomni/checkpoint/rank_local_optimizer.py
@@ -0,0 +1,216 @@
+# 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 io
+import json
+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)
+
+
+@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/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/npu_moe_precision.py b/veomni/distributed/npu_moe_precision.py
new file mode 100644
index 0000000000..ddce0fac9e
--- /dev/null
+++ b/veomni/distributed/npu_moe_precision.py
@@ -0,0 +1,59 @@
+# 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 the optional NPU expert weight-gradient precision."""
+
+import torch
+
+
+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 (
+        not mixed.enable
+        or mixed.param_dtype != "bfloat16"
+        or mixed.reduce_dtype != "float32"
+        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"
+        )
+
+
+def validate_qwen4_experts(model):
+    """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"
+    ]
+    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(
+        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(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/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/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index 5e88a734c0..fca6089ebd 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -410,6 +410,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,
     *,
@@ -462,6 +470,32 @@ 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 validate_expert_precision_config
+
+        validate_expert_precision_config(
+            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")
+
+    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
@@ -500,7 +534,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()
     if expert_weight_fp8_all_gather and compile_config.enable:
         raise ValueError("expert_weight_fp8_all_gather does not support torch.compile.")
 
@@ -662,9 +695,34 @@ 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)
+    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:
@@ -714,6 +772,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
@@ -808,7 +870,24 @@ 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])
+                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.pop("offload_policy")
+                    matched_resident_patterns.update(resident_matches)
+                    resident_module_fqns.append(module_fqns[id(_para_mod)])
+                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)
+                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)
@@ -829,6 +908,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():
@@ -851,6 +938,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)
@@ -860,6 +953,8 @@ def parallelize_model_fsdp2(
     # `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
@@ -924,13 +1019,22 @@ def parallelize_model_fsdp2(
         module_adapter_path=kwargs.pop("module_adapter_path", None),
     )
 
+    # 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 expert_weight_fp8_all_gather:
         from .fsdp2.fp8_all_gather import register_deepgemm_fp8_all_gather
 
         count = register_deepgemm_fp8_all_gather(model)
         logger.info_rank0(f"Enabled DeepGEMM FP8 all-gather for {count} expert parameters.")
 
-    if materialize_device == "cpu":
+    if materialize_device == "cpu" or cpu_offloaded_modules:
         _move_buffers_to_device(model, get_device_type())
 
     # Register grad norm clipping method for FSDP2
@@ -1059,6 +1163,20 @@ def build_parallelize_model(
             Ignored when the model contains no matching experts.
     """
     parallel_state = get_parallel_state()
+    compile_config = compile_config or CompileConfig()
+    if mixed_precision.npu_expert_fp32_wgrad:
+        from .npu_moe_precision import validate_expert_precision_config
+
+        validate_expert_precision_config(
+            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")
     validate_expert_fsdp_shard_dim(expert_fsdp_shard_dim, fsdp_mode=parallel_state.dp_mode)
     if expert_weight_fp8_all_gather is not False:
         from .fsdp2.fp8_all_gather import resolve_fp8_all_gather
@@ -1075,7 +1193,6 @@ def build_parallelize_model(
         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 expert_weight_fp8_all_gather and compile_config.enable:
         raise ValueError("expert_weight_fp8_all_gather does not support torch.compile.")
 
diff --git a/veomni/models/checkpoint/manager.py b/veomni/models/checkpoint/manager.py
index b8bde5c051..fdf0d4e691 100644
--- a/veomni/models/checkpoint/manager.py
+++ b/veomni/models/checkpoint/manager.py
@@ -33,9 +33,88 @@
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import torch.distributed as dist
+from torch.optim.lr_scheduler import LambdaLR
 
 from ...checkpoint import CheckpointerBase, build_checkpointer, layout
 from ...utils import helper
+from ...utils.dist_utils import raise_if_any_rank_failed
+
+
+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:
+        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) is not LambdaLR or scheduler.optimizer is not optimizer:
+        raise RuntimeError("MultiOptimizer scheduler implementation or binding is unsupported")
+    return {"type": "LambdaLR", "functions": _scheduler_functions(scheduler)}
+
+
+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:
@@ -86,6 +165,10 @@ 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._loaded_checkpoint_path = None
         self.checkpointer: CheckpointerBase = build_checkpointer(
             ckpt_manager=self.config.manager,
             dist_backend=runtime.args.accelerator.fsdp_config.fsdp_mode,
@@ -186,6 +269,15 @@ 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:
+            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:
@@ -198,6 +290,11 @@ 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()
         if load_dir is None:
             return
@@ -207,14 +304,73 @@ 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"])
-        dist.barrier()
+        multi = getattr(self.runtime.optimizer, "_is_multi_optimizer", False) is True
+        error = None
+        try:
+            self.checkpointer.load(load_dir, state, **self._checkpointer_kwargs())
+        except BaseException as exc:
+            error = exc
+        raise_if_any_rank_failed(error, "loading checkpoint")
+        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
+        raise_if_any_rank_failed(error, "restoring trainer step and scheduler")
+        if multi:
+            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")
+        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
+        self._loaded_checkpoint_path = load_dir
         logger.info_rank0(f"Load distributed checkpoint from {load_dir} successfully!")
 
     def save_dcp(self, state: "TrainerState") -> None:
@@ -223,14 +379,42 @@ 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
+        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)
+        except BaseException as exc:
+            error = exc
+        if multi:
+            raise_if_any_rank_failed(error, "preparing MultiOptimizer model extra state")
+        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,
@@ -239,8 +423,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!")
 
@@ -294,6 +487,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/models/model_runtime.py b/veomni/models/model_runtime.py
index 1768d26cbf..8355ed98a2 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__)
 
 
@@ -143,10 +160,68 @@ def __init__(
             self._build_model()
             self._freeze_model_module()
             self._build_parallelized_model()
-            self._build_optimizer()
+            from ..utils.dist_utils import raise_if_any_rank_failed
+
+            canonical = getattr(self.train_args.checkpoint, "extra_parallel_canonical_layout", False)
+            layout_error = None
+            try:
+                if canonical:
+                    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")
+            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",))
+            )
+            self._build_optimizer_checked(checkpoint_path if recover_ownership else None)
         self._build_model_assets()
         self.build_checkpoint()
 
+    def _build_optimizer_checked(self, checkpoint_path: Optional[str] = None) -> None:
+        """Coordinate local build failures before entering checkpoint preflight."""
+        from ..utils.dist_utils import raise_if_any_rank_failed
+
+        error = None
+        try:
+            if checkpoint_path is not None:
+                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()
+        except BaseException as exc:
+            error = exc
+        raise_if_any_rank_failed(error, "building model optimizer")
+        self._validate_optimizer_checkpoint_mode()
+
+    def _validate_optimizer_checkpoint_mode(self) -> None:
+        """Reject unsupported checkpoint modes on all ranks before training."""
+        from ..utils.dist_utils import raise_if_any_rank_failed
+
+        checkpoint = getattr(self.train_args, "checkpoint", None)
+        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")
+                if getattr(self.optimizer, "_is_multi_optimizer", False) is True
+                else None
+            )
+            raise_if_any_rank_failed(error, "checking MultiOptimizer checkpoint mode")
+
     def __getattr__(self, name: str) -> Any:
         """Forward unshadowed :class:`torch.nn.Module` APIs to the wrapped model."""
         model = object.__getattribute__(self, "model")
@@ -419,6 +494,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
@@ -477,6 +554,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,
@@ -664,7 +743,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 {}
 
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 a66b86cba8..6930c69f81 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,18 @@ 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
+
+            with use_parallel_state("base"):
+                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
         # materialisation path: ``skip_hf_weight_load`` still calls ``to_empty``.
@@ -317,6 +329,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/models/seed_omni/accelerated/omni_module/omni_module_runtime.py b/veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py
index ac4ddb8362..90e0286d06 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
@@ -139,6 +139,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()
@@ -174,9 +177,27 @@ 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()
+                    self._build_optimizer_checked()
                     self.build_checkpoint()
 
     @property
@@ -458,7 +479,7 @@ def build_after_omni_model_wrap(self) -> None:
         """
         with self._scoped():
             self._apply_recompute_policy()
-            self._build_optimizer()
+            self._build_optimizer_checked()
             self.build_checkpoint()
 
     def _apply_recompute_policy(self) -> None:
diff --git a/veomni/models/seed_omni/utils/hf_layout.py b/veomni/models/seed_omni/utils/hf_layout.py
index b46f30b231..e92123a359 100644
--- a/veomni/models/seed_omni/utils/hf_layout.py
+++ b/veomni/models/seed_omni/utils/hf_layout.py
@@ -246,10 +246,17 @@ def extract_module_save_state(
         parallel_state: Any = None,
     ) -> Mapping[str, torch.Tensor]:
         """``{module_key: tensor}`` to export for live ``model`` of ``module_name``."""
-        from ....checkpoint.dcp_checkpointer import ModelState
+        from ....checkpoint.dcp_checkpointer import ModelState, _export_extra_parallel_layout_key
 
         del module_name
-        return ModelState(model, parallel_state=parallel_state).state_dict()
+        state = {}
+        for key, tensor in ModelState(model, parallel_state=parallel_state).state_dict().items():
+            key = _export_extra_parallel_layout_key(key)
+            if key is not None:
+                if key in state:
+                    raise RuntimeError(f"Duplicate exported ExtraParallel tensor key: {key}")
+                state[key] = tensor
+        return state
 
     def build_fqn_to_index_mapping(self, weight_files: Mapping[str, str]) -> dict[str, int]:
         """Map each source key in ``weight_files`` to its 1-based ``HuggingFaceStorageWriter`` shard index."""
@@ -750,7 +757,12 @@ def save_hf_source_checkpoint(
     """
     from torch.distributed.checkpoint import HuggingFaceStorageWriter
 
+    from ....checkpoint.dcp_checkpointer import (
+        _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
@@ -759,31 +771,38 @@ def save_hf_source_checkpoint(
 
     save_state: dict[str, torch.Tensor] = {}
     for name, (model, parallel_state) in modules.items():
-        module_state = layout.extract_module_save_state(name, model, parallel_state)
-        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 = layout.extract_module_save_state(name, model, parallel_state)
+            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/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 1671b4124a..6aa1e46f63 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
@@ -1371,7 +1371,7 @@
          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(
@@ -1420,7 +1420,7 @@
 +            unused,
 +            self.down_proj,
 +            self.gate_up_proj,
-+            num_experts=self.num_experts,
++            **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_gpu.py b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_gpu.py
index c84c79c381..4c9a782345 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
@@ -1426,7 +1426,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(
@@ -1456,7 +1456,7 @@ def forward(
             unused,
             self.down_proj,
             self.gate_up_proj,
-            num_experts=self.num_experts,
+            **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 92d460dfbe..0c78820abc 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
@@ -1371,7 +1371,7 @@
          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(
@@ -1420,7 +1420,7 @@
 +            unused,
 +            self.down_proj,
 +            self.gate_up_proj,
-+            num_experts=self.num_experts,
++            **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.py b/veomni/models/transformers/qwen4_exp/generated/patched_modeling_qwen4_exp_npu.py
index f19348a144..e2804b7084 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
@@ -1426,7 +1426,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(
@@ -1456,7 +1456,7 @@ def forward(
             unused,
             self.down_proj,
             self.gate_up_proj,
-            num_experts=self.num_experts,
+            **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 28f5987ed9..4e02cd4795 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
@@ -1056,7 +1056,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(
@@ -1086,7 +1086,7 @@ def forward(
             unused,
             self.down_proj,
             self.gate_up_proj,
-            num_experts=self.num_experts,
+            **dict(num_experts=self.num_experts, **({"fp32_wgrad": True} if self._veomni_fp32_expert_wgrad else {})),
         )
         # --- Patch.1 ---
 
diff --git a/veomni/ops/kernels/moe_experts/standard/npu.py b/veomni/ops/kernels/moe_experts/standard/npu.py
index b3180fa4f5..12e95f0c4a 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
 
@@ -382,8 +399,11 @@ def wrapper(
     assume_distinct_experts: bool = False,
     block_size: int = 128,
     expert_dtype: str = "fp8",
+    fp32_wgrad: bool = False,
 ) -> Tensor:
     """Call the NPU fused MoE path. Empty weights are ``None``."""
+    # Quantization is performed by the QAT row, not the NPU fused row.
+    del block_size, expert_dtype
     # Empty tensor = unused layout. Registry args stay tensors.
     return npu_fused_moe_forward(
         num_experts,
@@ -396,4 +416,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/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..bd887abd0b 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,13 @@
 
 
 logger = logging.get_logger(__name__)
+_CHECKPOINT_SCHEDULER_BINDINGS = 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()
 
 
 class MultiLRScheduler(dict):
@@ -43,9 +51,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 +92,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 +108,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 +119,10 @@ 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)
+    return scheduler
 
 
 def get_constant_schedule_with_warmup(
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 b294039298..5a54f493a8 100644
--- a/veomni/optim/optimizer.py
+++ b/veomni/optim/optimizer.py
@@ -15,6 +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
@@ -22,6 +26,9 @@
 from torch.distributed._tensor import DTensor
 from torch.distributed.checkpoint.state_dict import (
     StateDictOptions,
+    _flatten_optim_state_dict,
+    _get_fqns,
+    _state_dict_fn,
     get_optimizer_state_dict,
     set_optimizer_state_dict,
 )
@@ -150,6 +157,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__(
@@ -176,6 +235,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):
         """
@@ -255,6 +398,116 @@ def step(self, closure=None):
                     p.data.addcdiv_(exp_avg, centered_variance, value=-step_size)
 
 
+def _empty_optimizer_state_dict(model, opt):
+    packed_groups = []
+    id_to_name = {}
+
+    def capture_params(leaf):
+        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:
+            # 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:
+        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})
+
+
+@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)
+
+
+@contextmanager
+def _dcp_swap_load_template(optimizer):
+    from .swap_adamw import SwapAdamW
+
+    staged = getattr(optimizer, "_veomni_dcp_staged_state", None)
+    if type(optimizer) is not SwapAdamW or staged is None:
+        yield
+        return
+    if "_veomni_dcp_live_state" in vars(optimizer):
+        raise RuntimeError("SwapAdamW checkpoint load is already active")
+    live = optimizer.state
+    optimizer._veomni_dcp_live_state = live
+    # PyTorch probes absent parameters through ``state[param]``. Keep those
+    # empty probe entries out of the proven source-owner inventory.
+    optimizer.state = defaultdict(dict, staged)
+    try:
+        yield
+    finally:
+        optimizer.state = live
+        del optimizer._veomni_dcp_live_state
+
+
 class MultiOptimizer(Optimizer, Stateful):
     """
     A container that handles multiple optimizers (for extra_parallel and non-extra_parallel parameters when extra_parallel+fsdp2 is enabled)
@@ -276,73 +529,117 @@ def __init__(
         self._is_multi_optimizer: bool = True
         self.key_names = key_names
 
-    def step(self) -> None:
+    def step(self, closure=None):
+        if getattr(self, "_veomni_dcp_restore_failed", False):
+            raise RuntimeError("MultiOptimizer cannot train after a failed checkpoint restore; rebuild the optimizer")
+        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))}"
-                )
-            else:
-                logger.info_rank0("No clashes when merging MultiOptimizer state dicts")
-            merged.update(sd)
+        # A snapshot must not let a raw DCP load overwrite live optimizer
+        # tensors before ``load_state_dict`` can reject unproven provenance.
+        import copy
 
-        return merged
+        from .checkpoint_adapter import rank_local_adamw_adapter
+        from .muon import DistributedMuon
+        from .swap_adamw import SwapAdamW
+
+        for optimizer in self.optimizers_dict.values():
+            if type(optimizer) not in (torch.optim.AdamW, AnyPrecisionAdamW, DistributedMuon, SwapAdamW) and (
+                rank_local_adamw_adapter(optimizer) is None
+            ):
+                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()
+        }
 
     def _sparse_state_dict(
         self,
+        *,
+        coordinate=False,
+        group=None,
     ) -> 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))
-            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_context
+
+                with rank_local_adamw_context(opt), _dcp_swap_load_template(opt):
+                    if not opt.state and not getattr(opt, "_veomni_dcp_staged_state", None):
+                        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) & 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:
         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_swap_load_template(optimizer),
+                    _dcp_reuse_owned_anyprecision_state(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()]
@@ -492,7 +789,14 @@ def build_optimizer(
     selects the fused kernel from the backend (``FUSED_ADAMW_DEVICE_TYPES``) and
     takes the elementwise path elsewhere.
     """
+    if isinstance(betas, list):
+        betas = tuple(betas)
     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:
@@ -829,6 +1133,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
diff --git a/veomni/optim/swap_adamw.py b/veomni/optim/swap_adamw.py
index 5608262867..aae576b895 100644
--- a/veomni/optim/swap_adamw.py
+++ b/veomni/optim/swap_adamw.py
@@ -30,10 +30,9 @@
   supported: under ExtraParallel the factory builds one ``SwapAdamW`` per
   parameter group inside a ``MultiOptimizer``. Validated on Ascend NPU + FSDP2.
 * Checkpoint save/load reads and writes the host buffers directly (see
-  ``state_dict`` / ``load_state_dict``). Resuming across a changed FSDP/EP
-  mesh (resharding) works through DCP's byte-level planning — an EP=4 → EP=2
-  resume at world 8 replayed steps 11-20 with matching printed loss and
-  grad_norm in manual validation — but is not covered by automated tests.
+  ``state_dict`` / ``load_state_dict``). Proven MultiOptimizer recovery
+  requires the same optimizer ownership and FSDP/EP mesh; it does not support
+  elastic resharding.
 * The update kernel is ``torch._fused_adamw_`` on CUDA and Ascend NPU, issued
   once per ``(device, dtype)`` bucket inside a param group, within each
   swapped-in batch (torch's own fused path groups the same way). Other backends
@@ -321,12 +320,68 @@ 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).
         """
-        self._bind_host_states()
-        state_dict = super().state_dict()
+        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()
+            state_dict = super().state_dict()
+        else:
+            # DCP loads into independent CPU buffers until its provenance and
+            # inventory checks have completed. Never expose the live host
+            # buffers as the load template. User hooks still observe the live
+            # optimizer, as they do outside DCP; only the serialized values
+            # come from the staged mapping.
+            live = getattr(self, "_veomni_dcp_live_state", self.state)
+            current = self.state
+            self.state = live
+            try:
+                for hook in self._optimizer_state_dict_pre_hooks.values():
+                    hook(self)
+            finally:
+                self.state = current
+            pre_hooks = self._optimizer_state_dict_pre_hooks
+            post_hooks = self._optimizer_state_dict_post_hooks
+            self._optimizer_state_dict_pre_hooks = type(pre_hooks)()
+            self._optimizer_state_dict_post_hooks = type(post_hooks)()
+            self.state = staged
+            try:
+                state_dict = super().state_dict()
+            finally:
+                self.state = current
+                self._optimizer_state_dict_pre_hooks = pre_hooks
+                self._optimizer_state_dict_post_hooks = post_hooks
+            # A normal SwapAdamW state-dict post hook sees Python integer
+            # steps. Keep that hook contract even though DCP's scratch uses
+            # tensors as its eventual destination.
+            state_dict["state"] = {
+                key: ({**value, "step": int(value["step"].item())} if "step" in value else value)
+                for key, value in state_dict["state"].items()
+            }
+            self.state = live
+            try:
+                for hook in post_hooks.values():
+                    result = hook(self, state_dict)
+                    if result is not None:
+                        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
@@ -349,11 +404,42 @@ def load_state_dict(self, state_dict) -> None:
         values for keys the planner does not overwrite (partial loads), so resume
         into a fresh optimizer to start omitted parameters from zero.
         """
+        state_dict = state_dict.copy()
+        live_state = getattr(self, "_veomni_dcp_live_state", None)
+        for hook in self._optimizer_load_state_dict_pre_hooks.values():
+            if live_state is not None and not getattr(hook, "_veomni_internal_dcp_filter", False):
+                staged_state = self.state
+                self.state = live_state
+                try:
+                    result = hook(self, state_dict)
+                finally:
+                    self.state = staged_state
+            else:
+                result = hook(self, state_dict)
+            if result is not None:
+                state_dict = result
+
+        if live_state is None:
+            self._apply_state_dict(state_dict)
+            for hook in self._optimizer_load_state_dict_post_hooks.values():
+                hook(self)
+        else:
+            staged_state = self.state
+            self.state = live_state
+            try:
+                self._apply_state_dict(state_dict)
+                for hook in self._optimizer_load_state_dict_post_hooks.values():
+                    hook(self)
+            finally:
+                self.state = staged_state
+
+    def _apply_state_dict(self, state_dict) -> None:
         groups = self.param_groups
         saved_groups = state_dict["param_groups"]
         if len(groups) != len(saved_groups):
             raise ValueError("loaded state dict has a different number of parameter groups")
         id_map = {}
+        group_updates = []
         for group, saved_group in zip(groups, saved_groups):
             if len(group["params"]) != len(saved_group["params"]):
                 raise ValueError(
@@ -369,12 +455,16 @@ def load_state_dict(self, state_dict) -> None:
             id_map.update(zip(saved_group["params"], group["params"]))
             # ``amsgrad`` is structural (it decides whether a third moment buffer
             # exists); keep the constructor's value rather than the checkpoint's.
-            group.update({k: v for k, v in saved_group.items() if k not in ("params", "amsgrad")})
+            group_updates.append({k: v for k, v in saved_group.items() if k not in ("params", "amsgrad")})
 
         loaded = set()
+        pending = []
+        staged = getattr(self, "_veomni_dcp_staged_state", None)
         for key, value in state_dict["state"].items():
             param = id_map.get(key)
             if param is None:
+                if staged is not None:
+                    raise RuntimeError("SwapAdamW checkpoint contains an unowned parameter state")
                 continue
             # A parameter absent from the checkpoint still gets an entry from the
             # flattened loader, carrying only empty dicts (``{"exp_avg": {}, ...}``),
@@ -382,6 +472,30 @@ def load_state_dict(self, state_dict) -> None:
             # existing or the entry being truthy.
             if any(isinstance(v, (torch.Tensor, int, float)) for v in value.values()):
                 loaded.add(param)
+            pending.append((param, value))
+
+        if staged is not None:
+            if loaded != set(staged):
+                raise RuntimeError("SwapAdamW loaded state does not match the proven checkpoint inventory")
+            for param, value in pending:
+                if param not in loaded:
+                    continue
+                if set(value) != {"step", *self._state_keys}:
+                    raise RuntimeError("SwapAdamW checkpoint parameter state is incomplete")
+                step = value["step"]
+                if not torch.is_tensor(step) or step.ndim != 0 or step.dtype != torch.float32:
+                    raise RuntimeError("SwapAdamW checkpoint step must be a float32 scalar")
+                step_value = step.item()
+                if not math.isfinite(step_value) or not 0 <= step_value <= 2**24 or not step_value.is_integer():
+                    raise RuntimeError("SwapAdamW checkpoint step is not an exact non-negative integer")
+                for field in self._state_keys:
+                    moment = value[field]
+                    if not torch.is_tensor(moment) or moment.shape != param.shape or moment.dtype != param.dtype:
+                        raise RuntimeError(f"SwapAdamW checkpoint state {field} has an incompatible shape or dtype")
+
+        for group, update in zip(groups, group_updates):
+            group.update(update)
+        for param, value in pending:
             self._load_param_state(param, value)
         # ``torch.optim.Optimizer.load_state_dict`` installs a fresh state mapping,
         # so an unseen parameter carries no history there; reset ours to match.
@@ -391,6 +505,11 @@ def load_state_dict(self, state_dict) -> None:
             for host_tensor in host_state.values():
                 host_tensor.zero_()
             self.state[param]["step"] = 0
+        if staged is not None:
+            for param in set(self._host_states) - loaded:
+                self._host_states.pop(param)
+                self._device_states.pop(param)
+                self.state.pop(param, None)
 
     def _load_param_state(self, p: Tensor, value: dict) -> None:
         # Reached from a direct ``load_state_dict`` call. The DCP resume path goes
@@ -658,6 +777,8 @@ def _update_chunk(self, params: Sequence[Tensor], group: dict[str, Any]) -> None
 
     @torch.no_grad()
     def step(self, closure=None):
+        if getattr(self, "_veomni_dcp_restore_failed", False):
+            raise RuntimeError("SwapAdamW cannot train after a failed checkpoint restore; rebuild the optimizer")
         loss = None
         if closure is not None:
             with torch.enable_grad():
diff --git a/veomni/trainer/callbacks/checkpoint_callback.py b/veomni/trainer/callbacks/checkpoint_callback.py
index bb05980d89..4faecf140a 100644
--- a/veomni/trainer/callbacks/checkpoint_callback.py
+++ b/veomni/trainer/callbacks/checkpoint_callback.py
@@ -30,9 +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
 
 
@@ -43,6 +47,52 @@
 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)
+
+
+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(
+                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")
+
+
 class CheckpointCallback(Callback):
     """Schedule DCP / HF / LoRA I/O and the one-shot tokenizer/config export."""
 
@@ -57,6 +107,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._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()
@@ -64,6 +128,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")
@@ -72,21 +137,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 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)."
@@ -101,12 +178,29 @@ def on_epoch_end(self, state: TrainerState, **kwargs):
                 )
 
     def _save_dcp(self, state: TrainerState):
-        self.trainer.save_dcp(state)
+        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):
+                self._async_save_failed = True
+                self._last_dcp_step = -1
+            raise
         self._last_dcp_step = state.global_step
 
     def _save_hf(self, state: TrainerState, stage: str = "step_end"):
-        self.trainer.save_hf_or_lora(state, stage=stage)
+        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):
+                self._async_save_failed = True
+                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"]
diff --git a/veomni/trainer/callbacks/global_state_callback.py b/veomni/trainer/callbacks/global_state_callback.py
index daf24df15a..bf70cb3874 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
@@ -44,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
 
@@ -59,6 +60,116 @@
 _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
+
+
+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.
 
@@ -160,10 +271,13 @@ 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)
-        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}
+        step_root, path_error = None, None
+        try:
+            step_root = layout.step_dir(args.train.checkpoint.save_path, state.global_step)
+        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
@@ -181,11 +295,31 @@ 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)
+            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
+        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 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:
+            extra_payload["_veomni_optimizer_generations"] = generations
         write_error: Optional[Exception] = None
         try:
             for path, blob in (
@@ -247,6 +381,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
@@ -263,44 +404,134 @@ def load_global_state(self) -> Optional[Dict[str, Any]]:
 
         found = current_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
-        self.trainer.state.global_step = global_state["global_step"]
-        self._restore_position(global_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")
+        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"))
-        if self.trainer.start_step == 0 and self.trainer.train_dataloader is not None:
-            iter(self.trainer.train_dataloader)
+        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")
+        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:
+                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}, "
diff --git a/veomni/utils/save_safetensor_utils.py b/veomni/utils/save_safetensor_utils.py
index 405391dd92..687bdda8b5 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
@@ -279,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, 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 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.
@@ -322,15 +337,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():
@@ -343,6 +366,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.
 
@@ -357,6 +381,9 @@ 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.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
 
@@ -366,17 +393,57 @@ 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,
+        _ExtraParallelSavePlanner,
+        _needs_extra_parallel_layout_marker,
+        _validate_extra_parallel_tensor_v2,
+    )
+    from .dist_utils import raise_if_any_rank_failed
+
+    fqn2spec_info = getattr(model, "_fqn2spec_info", None)
+    error = None
+    try:
+        if _needs_extra_parallel_layout_marker(fqn2spec_info):
+            _extra_parallel_layout_key(fqn2spec_info)
+    except BaseException as exc:
+        error = exc
+    raise_if_any_rank_failed(error, "checking LoRA ExtraParallel 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 +455,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,14 +471,29 @@ 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"],
-                    ep_shard_dim=spec.placement.dim,
-                    fsdp_shard_dim=spec.fsdp_shard_dim,
-                )
+            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"],
+                        ep_shard_dim=spec.placement.dim,
+                        fsdp_shard_dim=spec.fsdp_shard_dim,
+                    )
 
     # 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.
@@ -424,36 +505,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=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()