From 11ec209a8ac7434324c4c742580985fa7d5f4c6f Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 20 Aug 2026 12:52:26 -0700
Subject: [PATCH 01/20] perf(checkpoint): stream Nemotron TE expert weights

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../checkpoint/_backports/hf_storage.py       |  14 ++
 .../components/checkpoint/checkpointing.py    | 132 ++++++++++++-
 .../checkpoint/state_dict_adapter.py          |  60 ++++++
 .../models/nemotron_v3/state_dict_adapter.py  | 181 +++++++++++++++++-
 .../checkpoint/test_checkpointing.py          | 134 ++++++++++++-
 .../test_nemotron_v3_state_dict_adapter.py    |  58 ++++++
 6 files changed, 574 insertions(+), 5 deletions(-)

diff --git a/nemo_automodel/components/checkpoint/_backports/hf_storage.py b/nemo_automodel/components/checkpoint/_backports/hf_storage.py
index 3ba517f72e..5759c99e06 100644
--- a/nemo_automodel/components/checkpoint/_backports/hf_storage.py
+++ b/nemo_automodel/components/checkpoint/_backports/hf_storage.py
@@ -276,6 +276,14 @@ def __init__(self, path: str, token: str | None = None, key_mapping: dict[str, s
             super().__init__(path=path)
 
         self.key_mapping = key_mapping
+        self._metadata_cache: Metadata | None = None
+
+    def reset(self, checkpoint_id: str | os.PathLike | None = None) -> None:
+        """Reset reader state and invalidate cached metadata when the checkpoint changes."""
+        previous_path = str(self.path)
+        super().reset(checkpoint_id)
+        if checkpoint_id is not None and str(self.path) != previous_path:
+            self._metadata_cache = None
 
     def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
         per_file: dict[str, list[ReadItem]] = {}
@@ -372,6 +380,11 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
         return fut
 
     def read_metadata(self) -> Metadata:
+        if self._metadata_cache is not None:
+            if getattr(self._metadata_cache, "storage_meta", None) is not None:
+                self._metadata_cache.storage_meta.load_id = self.load_id
+            return self._metadata_cache
+
         state_dict_metadata: dict[str, TensorStorageMetadata] = {}
         storage_data: dict[MetadataIndex, _HFStorageInfo] = {}
 
@@ -443,6 +456,7 @@ def read_metadata(self) -> Metadata:
             metadata.storage_meta = StorageMeta()
         metadata.storage_meta.load_id = self.load_id  # type: ignore[union-attr]
 
+        self._metadata_cache = metadata
         return metadata
 
 
diff --git a/nemo_automodel/components/checkpoint/checkpointing.py b/nemo_automodel/components/checkpoint/checkpointing.py
index 7b9920dfe7..47638510f0 100644
--- a/nemo_automodel/components/checkpoint/checkpointing.py
+++ b/nemo_automodel/components/checkpoint/checkpointing.py
@@ -70,7 +70,7 @@
     requires_tensor_merging,
 )
 from nemo_automodel.components.checkpoint.lifecycle import CheckpointLifecycle
-from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
+from nemo_automodel.components.checkpoint.state_dict_adapter import CheckpointLoadGroup, StateDictAdapter
 from nemo_automodel.components.checkpoint.stateful_wrappers import ModelState, OptimizerState
 from nemo_automodel.components.checkpoint.utils import (
     ensure_tied_lm_head,
@@ -760,6 +760,112 @@ def load_optimizer(
         self._do_load(state_dict, os.path.join(weights_path, "optim"))
         optimizer_state.load_state_dict(state_dict)
 
+    def _load_model_in_checkpoint_groups(
+        self,
+        model_state: ModelState,
+        adapter: StateDictAdapter,
+        model_path: str,
+        storage_reader: StorageReader,
+    ) -> None:
+        """Load one adapter-owned dependency group at a time.
+
+        Args:
+            model_state: Wrapper around the single model part whose final parameter storage is populated.
+            adapter: Model-owned adapter that maps checkpoint tensors into dependency-complete groups.
+            model_path: Hugging Face safetensors checkpoint directory.
+            storage_reader: Reader already bound to ``model_path``. Its metadata may be reused across groups.
+
+        Raises:
+            RuntimeError: If a group requests checkpoint keys that are missing or if the adapter yields no groups.
+            TypeError: If the adapter yields an object that is not a :class:`CheckpointLoadGroup`.
+            ValueError: If groups are empty or request the same checkpoint key more than once.
+        """
+        model_part = _unwrap_ddp_model(model_state.model[0])
+        ep_mesh_dims = [dim for dim in self.moe_mesh.mesh_dim_names if dim != "pp"] if self.moe_mesh is not None else []
+        ep_mesh = self.moe_mesh[tuple(ep_mesh_dims)] if ep_mesh_dims else self.moe_mesh
+        started = time.monotonic()
+        checkpoint_keys = _get_checkpoint_metadata_keys(model_path, storage_reader)
+        metadata_seconds = time.monotonic() - started
+
+        storage_seconds = 0.0
+        install_seconds = 0.0
+        requested_checkpoint_keys: set[str] = set()
+        loaded_native_keys: set[str] = set()
+        requested_bytes = 0
+        max_group_bytes = 0
+        group_count = 0
+
+        process_group = getattr(self, "process_group", None)
+        process_group_kwargs = {"process_group": process_group} if process_group is not None else {}
+
+        for group in adapter.iter_checkpoint_load_groups(model_part, device_mesh=ep_mesh):
+            if not isinstance(group, CheckpointLoadGroup):
+                raise TypeError(
+                    f"{type(adapter).__name__}.iter_checkpoint_load_groups yielded {type(group).__name__}, "
+                    "expected CheckpointLoadGroup"
+                )
+            if not group.destinations:
+                raise ValueError(f"{type(adapter).__name__} yielded an empty checkpoint load group")
+
+            group_keys = set(group.destinations)
+            duplicate_keys = sorted(group_keys & requested_checkpoint_keys)
+            if duplicate_keys:
+                raise ValueError(
+                    f"{type(adapter).__name__} requested {len(duplicate_keys)} checkpoint keys in multiple groups "
+                    f"(examples={duplicate_keys[:5]})"
+                )
+            missing_keys = sorted(group_keys - checkpoint_keys)
+            if missing_keys:
+                raise RuntimeError(
+                    f"Checkpoint {model_path} is missing {len(missing_keys)} keys required by streaming group "
+                    f"{group_count + 1} (examples={missing_keys[:5]})"
+                )
+
+            group_bytes = sum(estimate_tensor_bytes(tensor) for tensor in group.destinations.values())
+            requested_bytes += group_bytes
+            max_group_bytes = max(max_group_bytes, group_bytes)
+            requested_checkpoint_keys |= group_keys
+
+            storage_started = time.monotonic()
+            # The reader is already bound to ``model_path``. Omitting checkpoint_id prevents DCP from resetting it
+            # for every group, which lets the HF reader reuse parsed safetensors metadata.
+            dcp.load(group.destinations, storage_reader=storage_reader, **process_group_kwargs)
+            storage_seconds += time.monotonic() - storage_started
+
+            install_started = time.monotonic()
+            group.install()
+            install_seconds += time.monotonic() - install_started
+            loaded_native_keys |= set(group.native_keys)
+            group_count += 1
+
+            del group
+
+        if group_count == 0:
+            raise RuntimeError(
+                f"{type(adapter).__name__} opted into streaming checkpoint loading but yielded no groups"
+            )
+
+        if model_state.uses_tied_lm_head and not model_state.is_peft:
+            ensure_tied_lm_head(model_part)
+
+        total_seconds = time.monotonic() - started
+        requested_gb = requested_bytes / (1 << 30)
+        max_group_gb = max_group_bytes / (1 << 30)
+        logger.info(
+            "load_model: streamed %.2f GB in %d groups over %.2fs "
+            "(%.2f GB/s overall | max group %.2f GB, metadata %.2fs, storage read %.2fs, install %.2fs, "
+            "native keys %d)",
+            requested_gb,
+            group_count,
+            total_seconds,
+            requested_gb / max(total_seconds, 1e-9),
+            max_group_gb,
+            metadata_seconds,
+            storage_seconds,
+            install_seconds,
+            len(loaded_native_keys),
+        )
+
     @torch.no_grad()
     def load_model(
         self,
@@ -849,6 +955,30 @@ def load_model(
             )
             and not self.config.dequantize_base_checkpoint
         )
+        supports_streaming_checkpoint_load = (
+            isinstance(state_dict_adapter, StateDictAdapter)
+            and state_dict_adapter.supports_streaming_checkpoint_load
+            and not self.config.dequantize_base_checkpoint
+        )
+        if (
+            is_init_step
+            and is_safetensors
+            and len(model_state.model) == 1
+            and world_size == 1
+            and supports_streaming_checkpoint_load
+            and not allow_checkpoint_key_subset
+        ):
+            storage_reader = self._get_storage_reader(
+                model_path,
+                key_mapping=None,
+                is_init_step=True,
+                is_safetensors=True,
+            )
+            if storage_reader is None:
+                raise RuntimeError(f"No safetensors storage reader is available for streaming load from {model_path}")
+            self._load_model_in_checkpoint_groups(model_state, state_dict_adapter, model_path, storage_reader)
+            return
+
         single_device_custom_safetensors = (
             is_safetensors and is_custom_model and world_size == 1 and not can_load_without_full_copy
         )
diff --git a/nemo_automodel/components/checkpoint/state_dict_adapter.py b/nemo_automodel/components/checkpoint/state_dict_adapter.py
index 11371af940..c87ab8a058 100644
--- a/nemo_automodel/components/checkpoint/state_dict_adapter.py
+++ b/nemo_automodel/components/checkpoint/state_dict_adapter.py
@@ -13,12 +13,41 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional
 
+import torch
+
 if TYPE_CHECKING:
     from torch.distributed.device_mesh import DeviceMesh
 
 
+@dataclass
+class CheckpointLoadGroup:
+    """One dependency-complete unit of a bounded checkpoint load.
+
+    Attributes:
+        destinations: Mapping from checkpoint FQNs to destination tensors. Each tensor has the exact shape and axis
+            order recorded for that FQN in checkpoint metadata. A destination may alias final model storage or may be
+            bounded temporary storage consumed by :meth:`install`.
+        native_keys: Native model state-dict keys completed by this group. Their tensors may have arbitrary ranks and
+            model-owned layouts; every key must be fully installed before :meth:`install` returns.
+    """
+
+    destinations: dict[str, torch.Tensor]
+    native_keys: frozenset[str]
+
+    def install(self) -> None:
+        """Install loaded destinations into final model storage.
+
+        The default implementation is a no-op for destinations that already alias final model storage. Allocating
+        conversion groups override this method, complete their model-owned transformation, and release references to
+        temporary storage before returning.
+        """
+        return None
+
+
 class StateDictAdapter(ABC):
     """Abstract base class for state dict transformations.
 
@@ -28,6 +57,7 @@ class StateDictAdapter(ABC):
 
     _supports_write_through_checkpoint_load: bool = False
     _supports_checkpoint_load_without_full_copy: bool = False
+    _supports_streaming_checkpoint_load: bool = False
 
     @property
     def supports_write_through_checkpoint_load(self) -> bool:
@@ -48,6 +78,36 @@ def supports_checkpoint_load_without_full_copy(self) -> bool:
         """
         return self._supports_checkpoint_load_without_full_copy
 
+    @property
+    def supports_streaming_checkpoint_load(self) -> bool:
+        """Whether the adapter can load a base checkpoint as bounded dependency groups.
+
+        Adapters should opt in only when :meth:`iter_checkpoint_load_groups` covers every required pretrained model
+        tensor and each yielded group writes through to final storage or installs bounded temporary storage before the
+        next group is requested.
+        """
+        return self._supports_streaming_checkpoint_load
+
+    def iter_checkpoint_load_groups(
+        self,
+        model_part: torch.nn.Module,
+        device_mesh: Optional["DeviceMesh"] = None,
+    ) -> Iterator[CheckpointLoadGroup]:
+        """Yield dependency-complete groups for bounded base-checkpoint loading.
+
+        Args:
+            model_part: Model part that owns the final parameter and buffer storage populated by yielded groups.
+            device_mesh: Optional device mesh describing the final distributed tensor placements.
+
+        Returns:
+            Iterator of groups. Each group's destination tensors have checkpoint-native shapes and axis order; the
+            iterator must not retain an installed group's temporary tensors after advancing.
+
+        Raises:
+            NotImplementedError: If the adapter does not implement streaming checkpoint loading.
+        """
+        raise NotImplementedError(f"{type(self).__name__} does not implement streaming checkpoint loading")
+
     @abstractmethod
     def to_hf(self, state_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
         """Convert from native model state dict to HuggingFace format.
diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
index 4a71be7ef9..c60db2e356 100644
--- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
@@ -14,20 +14,36 @@
 
 import logging
 import re
+from collections.abc import Iterator
+from dataclasses import dataclass
 from typing import Any, Optional
 
 import torch
 from torch.distributed.device_mesh import DeviceMesh
 
-from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
+from nemo_automodel.components.checkpoint.state_dict_adapter import CheckpointLoadGroup, StateDictAdapter
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.moe.config import MoEConfig
 from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin
+from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
 
 logger = logging.getLogger(__name__)
 
 _MAMBA_FP32_PARAMS_TO_BARE = re.compile(r"(\.mixer)\._fp32_params\.")
 _MAMBA_FP32_PARAM_NAMES = ("A_log", "dt_bias", "D")
+_STREAMABLE_EXPERT_WEIGHT = re.compile(
+    r"^(?P<root>(?:.+\.)?layers\.\d+\.mixer\.experts)\."
+    r"(?P<projection>gate_and_up_projs|down_projs)$"
+)
+
+
+@dataclass
+class _NemotronExpertGroupBuilder:
+    """Mutable builder for one dependency-complete expert layer group."""
+
+    destinations: dict[str, torch.Tensor]
+    expert_ids_by_projection: dict[str, set[int]]
+    native_keys: set[str]
 
 
 def _strip_mamba_fp32_holder_key(key: str) -> str:
@@ -113,6 +129,169 @@ def __init__(
             "model.layers.{}.mixer.experts.{}.down_proj.weight": "model.layers.{}.mixer.experts.down_projs",
         }
 
+    @property
+    def supports_streaming_checkpoint_load(self) -> bool:
+        """Whether this adapter can stream split experts into final single-device storage.
+
+        TE/DeepEP is configured for distributed runs, but the MoE layer deliberately falls back to ``GroupedExperts``
+        when world size is one. The adapter configuration still says ``experts="te"``, so the legacy converter assumes
+        the grouped tensors are TE stack copies and rebuilds them. For non-gated ReLU-squared experts, the runtime
+        grouped tensors instead provide per-expert transposed views that write through to final storage. Expert bias
+        and other activations retain the allocating fallback until their layouts have concrete coverage.
+        """
+        return (
+            self.moe_config is not None
+            and self.backend.experts == "te"
+            and self.moe_config.expert_activation == "relu2"
+            and not self.moe_config.expert_bias
+        )
+
+    def iter_checkpoint_load_groups(
+        self,
+        model_part: torch.nn.Module,
+        device_mesh: Optional["DeviceMesh"] = None,
+    ) -> Iterator[CheckpointLoadGroup]:
+        """Yield final-storage destinations for the single-device TE fallback.
+
+        Args:
+            model_part: Nemotron V3 model whose ordinary parameters and grouped expert parameters are final load
+                destinations. Native input projections have shape [experts, hidden, expert_hidden], and native down
+                projections have shape [experts, expert_hidden, hidden].
+            device_mesh: Must be ``None``. Distributed expert parameters require a rank-symmetric group plan.
+
+        Returns:
+            Iterator whose first group contains ordinary parameter aliases and whose remaining groups contain one
+            complete MoE layer each. Expert destinations use HF shapes [expert_hidden, hidden] for ``up_proj`` and
+            [hidden, expert_hidden] for ``down_proj`` and are transposed views of final grouped parameter storage.
+
+        Raises:
+            RuntimeError: If this adapter configuration did not opt into streaming checkpoint loading.
+            ValueError: If a distributed mesh, an allocating ordinary conversion, an expert bias, or an incomplete
+                expert group is encountered.
+        """
+        if not self.supports_streaming_checkpoint_load:
+            raise RuntimeError(f"{type(self).__name__} does not support streaming for backend {self.backend.experts}")
+        if device_mesh is not None:
+            raise ValueError("Nemotron V3 streaming checkpoint loading currently requires a single-device model")
+        moe_config = self.moe_config
+        if moe_config is None:
+            raise RuntimeError("Nemotron V3 streaming checkpoint loading requires an MoE configuration")
+
+        ordinary_destinations: dict[str, torch.Tensor] = {}
+        ordinary_native_keys: set[str] = set()
+        expert_groups: dict[str, _NemotronExpertGroupBuilder] = {}
+        expected_expert_ids = set(range(moe_config.n_routed_experts))
+
+        for parameter_name, parameter in model_part.named_parameters():
+            native_name = canonical_parameter_fqn(parameter_name)
+            # PEFT is applied before the base checkpoint is loaded. Its adapter parameters are intentionally absent
+            # from the pretrained checkpoint and retain the initialization performed by initialize_model_weights().
+            if "lora" in native_name:
+                continue
+
+            expert_match = _STREAMABLE_EXPERT_WEIGHT.match(native_name)
+            if expert_match is not None:
+                expert_root = expert_match.group("root")
+                native_projection_name = expert_match.group("projection")
+                if native_projection_name == "gate_and_up_projs":
+                    projection_name = "up_proj"
+                    expected_shape = (
+                        moe_config.n_routed_experts,
+                        moe_config.expert_dim,
+                        moe_config.moe_inter_dim,
+                    )
+                else:
+                    projection_name = "down_proj"
+                    expected_shape = (
+                        moe_config.n_routed_experts,
+                        moe_config.moe_inter_dim,
+                        moe_config.expert_dim,
+                    )
+
+                if tuple(parameter.shape) != expected_shape:
+                    raise ValueError(
+                        f"Grouped expert parameter {native_name} has shape {tuple(parameter.shape)}, "
+                        f"expected {expected_shape}"
+                    )
+
+                builder = expert_groups.setdefault(
+                    expert_root,
+                    _NemotronExpertGroupBuilder(destinations={}, expert_ids_by_projection={}, native_keys=set()),
+                )
+                split_weights = self._split_experts_weights(parameter.detach(), moe_config.n_routed_experts)
+                expert_ids = list(self._last_expert_ids)
+                for expert_id, expert_weight in zip(expert_ids, split_weights, strict=True):
+                    checkpoint_name = self._native_key_to_hf(f"{expert_root}.{expert_id}.{projection_name}.weight")
+                    if checkpoint_name in builder.destinations:
+                        raise ValueError(f"Duplicate expert checkpoint destination: {checkpoint_name}")
+                    checkpoint_destination = expert_weight.transpose(0, 1)
+                    expected_checkpoint_shape = (
+                        (moe_config.moe_inter_dim, moe_config.expert_dim)
+                        if projection_name == "up_proj"
+                        else (moe_config.expert_dim, moe_config.moe_inter_dim)
+                    )
+                    if tuple(checkpoint_destination.shape) != expected_checkpoint_shape:
+                        raise ValueError(
+                            f"Expert checkpoint destination {checkpoint_name} has shape "
+                            f"{tuple(checkpoint_destination.shape)}, expected {expected_checkpoint_shape}"
+                        )
+                    builder.destinations[checkpoint_name] = checkpoint_destination
+                builder.expert_ids_by_projection[native_projection_name] = set(expert_ids)
+                builder.native_keys.add(native_name)
+                del split_weights
+                continue
+
+            destination = parameter.detach()
+            converted = self.convert_single_tensor_to_hf(native_name, destination, quantization=False)
+            if len(converted) != 1:
+                raise ValueError(
+                    f"Ordinary Nemotron parameter {native_name} produced {len(converted)} checkpoint destinations"
+                )
+            checkpoint_name, checkpoint_destination = converted[0]
+            if not isinstance(checkpoint_destination, torch.Tensor):
+                raise ValueError(
+                    f"Ordinary Nemotron destination {checkpoint_name} is {type(checkpoint_destination).__name__}, "
+                    "expected Tensor"
+                )
+            aliases_parameter = (
+                checkpoint_destination.device == destination.device
+                and checkpoint_destination.untyped_storage().data_ptr() == destination.untyped_storage().data_ptr()
+            )
+            if not aliases_parameter:
+                raise ValueError(
+                    f"Ordinary Nemotron destination {checkpoint_name} does not alias final parameter {native_name}"
+                )
+            if checkpoint_name in ordinary_destinations:
+                raise ValueError(f"Duplicate ordinary Nemotron checkpoint destination: {checkpoint_name}")
+            ordinary_destinations[checkpoint_name] = checkpoint_destination
+            ordinary_native_keys.add(native_name)
+
+        for group_key, builder in expert_groups.items():
+            for projection_name in ("gate_and_up_projs", "down_projs"):
+                expert_ids = builder.expert_ids_by_projection.get(projection_name, set())
+                if expert_ids != expected_expert_ids:
+                    missing_experts = sorted(expected_expert_ids - expert_ids)
+                    raise ValueError(
+                        f"Incomplete expert group {group_key}.{projection_name}: "
+                        f"missing expert ids {missing_experts[:10]}"
+                    )
+
+        if not expert_groups:
+            raise ValueError("Nemotron V3 streaming load plan found no grouped expert layers")
+
+        if ordinary_destinations:
+            yield CheckpointLoadGroup(
+                destinations=ordinary_destinations,
+                native_keys=frozenset(ordinary_native_keys),
+            )
+
+        for group_key in sorted(expert_groups):
+            builder = expert_groups[group_key]
+            yield CheckpointLoadGroup(
+                destinations=builder.destinations,
+                native_keys=frozenset(builder.native_keys),
+            )
+
     @property
     def _hf_prefix(self) -> str:
         """Return the source checkpoint's public Nemotron-H model prefix."""
diff --git a/tests/unit_tests/checkpoint/test_checkpointing.py b/tests/unit_tests/checkpoint/test_checkpointing.py
index 36af0c99e4..6ca15d5a66 100644
--- a/tests/unit_tests/checkpoint/test_checkpointing.py
+++ b/tests/unit_tests/checkpoint/test_checkpointing.py
@@ -34,6 +34,7 @@
 from nemo_automodel.components.checkpoint._backports.hf_storage import (
     _DIFFUSERS_INDEX_FN,
     _extract_file_index_with_status,
+    _HuggingFaceStorageReader,
     get_fqn_to_dtype_mapping,
     get_fqn_to_file_index_mapping,
 )
@@ -63,7 +64,7 @@
     is_cloud_path,
     save_config,
 )
-from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
+from nemo_automodel.components.checkpoint.state_dict_adapter import CheckpointLoadGroup, StateDictAdapter
 from nemo_automodel.components.checkpoint.stateful_wrappers import (
     ModelState,
     OptimizerState,
@@ -1696,7 +1697,7 @@ def test_single_device_custom_model_uses_fast_path(self, mock_load_full, mock_lo
     @patch("nemo_automodel.components.checkpoint.checkpointing._is_safetensors_checkpoint", return_value=True)
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype")
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model")
-    @pytest.mark.parametrize("load_capability", ["write_through", "without_full_copy"])
+    @pytest.mark.parametrize("load_capability", ["write_through", "without_full_copy", "streaming"])
     @pytest.mark.parametrize("dequantize_base_checkpoint", [False, True])
     def test_single_device_adapter_without_full_copy_routes_by_quantization(
         self,
@@ -1715,6 +1716,7 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
         model.state_dict_adapter = MagicMock(spec=StateDictAdapter)
         model.state_dict_adapter.supports_write_through_checkpoint_load = load_capability == "write_through"
         model.state_dict_adapter.supports_checkpoint_load_without_full_copy = load_capability == "without_full_copy"
+        model.state_dict_adapter.supports_streaming_checkpoint_load = load_capability == "streaming"
         mock_state_dict = {"layer.weight": torch.randn(4, 4), "layer.bias": torch.randn(4)}
         mock_load_hf.return_value = mock_state_dict
 
@@ -1734,7 +1736,8 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
                 "nemo_automodel.components.checkpoint.checkpointing._maybe_adapt_state_dict_from_hf",
                 side_effect=lambda model_part, state_dict, **kwargs: state_dict,
             ),
-            patch.object(checkpointer, "_get_storage_reader", return_value=None),
+            patch.object(checkpointer, "_get_storage_reader", return_value=MagicMock()),
+            patch.object(checkpointer, "_load_model_in_checkpoint_groups") as mock_streaming_load,
             patch.object(checkpointer, "_do_load", return_value=mock_state_dict) as mock_dcp_load,
         ):
             mock_model_state = mock_model_state_cls.return_value
@@ -1746,14 +1749,139 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
         if dequantize_base_checkpoint:
             mock_load_full.assert_called_once()
             mock_load_hf.assert_called_once()
+            mock_streaming_load.assert_not_called()
+            mock_dcp_load.assert_not_called()
+        elif load_capability == "streaming":
+            mock_load_full.assert_not_called()
+            mock_load_hf.assert_not_called()
+            mock_streaming_load.assert_called_once()
             mock_dcp_load.assert_not_called()
         else:
             mock_load_full.assert_not_called()
             mock_load_hf.assert_not_called()
+            mock_streaming_load.assert_not_called()
             mock_dcp_load.assert_called_once()
             assert "load_model:" in caplog.text
 
 
+def test_checkpoint_group_executor_installs_before_requesting_next_group(caplog):
+    checkpointer = TestLoadModelCustomModelGuard()._make_checkpointer()
+    model_state = SimpleNamespace(
+        model=[torch.nn.Linear(2, 2)],
+        uses_tied_lm_head=False,
+        is_peft=False,
+    )
+    adapter = MagicMock(spec=StateDictAdapter)
+    reader = MagicMock()
+    reader.read_metadata.return_value = SimpleNamespace(
+        state_dict_metadata={"checkpoint.0": object(), "checkpoint.1": object()}
+    )
+
+    active_groups = 0
+    peak_active_groups = 0
+    installed: dict[int, torch.Tensor] = {}
+
+    def iter_groups(model_part, device_mesh=None):
+        """Yield groups after verifying the prior group's destination was installed.
+
+        Args:
+            model_part: Model whose parameters have arbitrary shapes; this test only verifies object identity.
+            device_mesh: Optional distributed mesh, expected to be ``None`` for this single-process test.
+
+        Returns:
+            Iterator of groups containing destination tensors of shape [elements].
+        """
+        nonlocal active_groups, peak_active_groups
+        assert model_part is model_state.model[0]
+        assert device_mesh is None
+        for group_id in range(2):
+            assert active_groups == 0
+            active_groups += 1
+            peak_active_groups = max(peak_active_groups, active_groups)
+            destination = torch.empty(3)
+            group = CheckpointLoadGroup(
+                destinations={f"checkpoint.{group_id}": destination},
+                native_keys=frozenset({f"native.{group_id}"}),
+            )
+
+            def install(group_id=group_id, destination=destination):
+                """Record one loaded destination tensor.
+
+                Args:
+                    group_id: Sequential group identifier.
+                    destination: Loaded tensor of shape [elements].
+                """
+                nonlocal active_groups
+                installed[group_id] = destination.clone()
+                active_groups -= 1
+
+            group.install = install
+            yield group
+            assert group_id in installed
+
+    def load_group(destinations, **kwargs):
+        """Fill one checkpoint destination mapping as DCP would.
+
+        Args:
+            destinations: Mapping containing one tensor of shape [elements].
+            **kwargs: DCP reader arguments.
+        """
+        assert kwargs["storage_reader"] is reader
+        key = next(iter(destinations))
+        destinations[key].fill_(int(key.rsplit(".", 1)[1]) + 1)
+
+    adapter.iter_checkpoint_load_groups.side_effect = iter_groups
+    caplog.set_level(logging.INFO)
+    with patch("nemo_automodel.components.checkpoint.checkpointing.dcp.load", side_effect=load_group) as mock_load:
+        checkpointer._load_model_in_checkpoint_groups(model_state, adapter, "/checkpoint", reader)
+
+    assert peak_active_groups == 1
+    assert active_groups == 0
+    assert torch.equal(installed[0], torch.ones(3))
+    assert torch.equal(installed[1], torch.full((3,), 2.0))
+    assert mock_load.call_count == 2
+    reader.read_metadata.assert_called_once()
+    assert "streamed" in caplog.text
+    assert "2 groups" in caplog.text
+
+
+def test_checkpoint_group_executor_reuses_hf_metadata_and_loads_exact_values(tmp_path):
+    checkpoint = {
+        "checkpoint.0": torch.arange(6, dtype=torch.float32).reshape(2, 3),
+        "checkpoint.1": torch.arange(4, dtype=torch.bfloat16).reshape(2, 2),
+    }
+    save_file(checkpoint, tmp_path / "model.safetensors")
+
+    backing_tensors = {
+        key: torch.empty(value.shape[1], value.shape[0], dtype=value.dtype) for key, value in checkpoint.items()
+    }
+    destinations = {key: backing.transpose(0, 1) for key, backing in backing_tensors.items()}
+    assert all(not destination.is_contiguous() for destination in destinations.values())
+    groups = [
+        CheckpointLoadGroup(
+            destinations={key: destinations[key]},
+            native_keys=frozenset({f"native.{index}"}),
+        )
+        for index, key in enumerate(checkpoint)
+    ]
+    adapter = MagicMock(spec=StateDictAdapter)
+    adapter.iter_checkpoint_load_groups.return_value = iter(groups)
+    model_state = SimpleNamespace(
+        model=[torch.nn.Linear(2, 2)],
+        uses_tied_lm_head=False,
+        is_peft=False,
+    )
+    reader = _HuggingFaceStorageReader(str(tmp_path))
+    checkpointer = TestLoadModelCustomModelGuard()._make_checkpointer()
+
+    with patch.object(reader.fs, "ls", wraps=reader.fs.ls) as mock_list_files:
+        checkpointer._load_model_in_checkpoint_groups(model_state, adapter, str(tmp_path), reader)
+
+    assert mock_list_files.call_count == 1
+    for key, expected in checkpoint.items():
+        torch.testing.assert_close(destinations[key], expected)
+
+
 class TestLoadModelCheckpointKeySubset:
     """Test allow_checkpoint_key_subset support for torch_save exports."""
 
diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
index e5b133a821..6652ae3b64 100644
--- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
+++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
@@ -126,6 +126,64 @@ def test_write_through_load_capability_requires_aliasing_expert_backend(self, co
         adapter.backend.dispatcher = "mok"
         assert adapter.supports_write_through_checkpoint_load is False
 
+    def test_te_single_device_fallback_streams_through_grouped_parameter_views(self, config, moe_config, backend):
+        backend.experts = "te"
+        adapter = NemotronV3StateDictAdapter(config, moe_config, backend)
+
+        class _FakeSingleDeviceFallbackExperts(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.gate_and_up_projs = torch.nn.Parameter(
+                    torch.zeros(moe_config.n_routed_experts, moe_config.expert_dim, moe_config.moe_inter_dim)
+                )
+                self.down_projs = torch.nn.Parameter(
+                    torch.zeros(moe_config.n_routed_experts, moe_config.moe_inter_dim, moe_config.expert_dim)
+                )
+
+        model = torch.nn.Module()
+        model.model = torch.nn.Module()
+        model.model.embed_tokens = torch.nn.Embedding(8, moe_config.expert_dim)
+        model.model.embed_tokens.lora_A = torch.nn.Linear(moe_config.expert_dim, 2, bias=False)
+        lora_initial = model.model.embed_tokens.lora_A.weight.detach().clone()
+        layer = torch.nn.Module()
+        layer.mixer = torch.nn.Module()
+        layer.mixer.experts = _FakeSingleDeviceFallbackExperts()
+        experts = layer.mixer.experts
+        model.model.layers = torch.nn.ModuleList([layer])
+
+        groups = list(adapter.iter_checkpoint_load_groups(model))
+        assert adapter.supports_streaming_checkpoint_load is True
+        assert len(groups) == 2
+
+        all_destinations = {key: value for group in groups for key, value in group.destinations.items()}
+        assert set(all_destinations) == {
+            "backbone.embeddings.weight",
+            *{
+                f"backbone.layers.0.mixer.experts.{expert_id}.{projection}.weight"
+                for expert_id in range(moe_config.n_routed_experts)
+                for projection in ("up_proj", "down_proj")
+            },
+        }
+        for expert_id in range(moe_config.n_routed_experts):
+            up_destination = all_destinations[f"backbone.layers.0.mixer.experts.{expert_id}.up_proj.weight"]
+            down_destination = all_destinations[f"backbone.layers.0.mixer.experts.{expert_id}.down_proj.weight"]
+            assert up_destination.untyped_storage().data_ptr() == experts.gate_and_up_projs.untyped_storage().data_ptr()
+            assert down_destination.untyped_storage().data_ptr() == experts.down_projs.untyped_storage().data_ptr()
+            assert not up_destination.is_contiguous()
+            assert not down_destination.is_contiguous()
+
+        assert not any("lora" in key for key in all_destinations)
+        for group in groups[1:]:
+            for destination in group.destinations.values():
+                destination.fill_(3)
+            group.install()
+        assert torch.count_nonzero(experts.gate_and_up_projs) == experts.gate_and_up_projs.numel()
+        assert torch.count_nonzero(experts.down_projs) == experts.down_projs.numel()
+        torch.testing.assert_close(model.model.embed_tokens.lora_A.weight, lora_initial)
+
+        moe_config.expert_bias = True
+        assert adapter.supports_streaming_checkpoint_load is False
+
     def test_from_hf_map_structure(self, config, moe_config, backend):
         """Test from_hf_map structure."""
         adapter = NemotronV3StateDictAdapter(config, moe_config, backend)

From 4eef4718766b528817e084ba482f57b8fe9ec2bc Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 20 Aug 2026 14:48:25 -0700
Subject: [PATCH 02/20] perf(checkpoint): stream Ling checkpoint transforms

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../models/ling_v2/state_dict_adapter.py      | 340 +++++++++++++++++-
 .../test_ling_v2_state_dict_adapter.py        | 100 +++++-
 2 files changed, 434 insertions(+), 6 deletions(-)

diff --git a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
index 9e1391c5b4..d650c5e9cb 100644
--- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
@@ -41,16 +41,19 @@
 """
 
 import re
+from collections.abc import Iterator
+from dataclasses import dataclass
 from typing import Any, Optional
 
 import torch
 from torch.distributed.device_mesh import DeviceMesh
 
-from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
+from nemo_automodel.components.checkpoint.state_dict_adapter import CheckpointLoadGroup, StateDictAdapter
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.ling_v2.config import BailingMoeV2Config
 from nemo_automodel.components.moe.config import MoEConfig
 from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin
+from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
 
 # Map of single-key renames applied in both directions.  Each tuple is
 # (HF substring, native substring); replacement is whole-substring and
@@ -64,6 +67,78 @@
 )
 
 _LAYER_QKV_RE = re.compile(r"^(?P<prefix>(?:.*\.)?layers\.\d+)\.attention\.query_key_value\.weight$")
+_NATIVE_LAYER_QKV_RE = re.compile(r"^(?P<prefix>(?:.*\.)?layers\.\d+)\.self_attn\.(?P<projection>[qkv])_proj\.weight$")
+_STREAMABLE_EXPERT_WEIGHT = re.compile(
+    r"^(?P<root>(?:.*\.)?layers\.\d+\.mlp\.experts)\."
+    r"(?P<projection>gate_and_up_projs|down_projs)$"
+)
+
+
+@dataclass
+class _LingQKVGroupBuilder:
+    """Mutable builder for one fused-QKV checkpoint dependency group.
+
+    Attributes:
+        projections: Mapping from ``q``, ``k``, and ``v`` to a native parameter name and its final tensor. Query
+            tensors have shape [q_hidden, hidden]; key and value tensors have shape [kv_hidden, hidden].
+    """
+
+    projections: dict[str, tuple[str, torch.Tensor]]
+
+
+@dataclass
+class _LingExpertGroupBuilder:
+    """Mutable builder for one split-expert checkpoint dependency group.
+
+    Attributes:
+        destinations: Checkpoint gate/up views of shape [expert_hidden, hidden] and down views of shape
+            [hidden, expert_hidden]. Every view aliases a final grouped expert parameter.
+        native_keys: Names of the final grouped expert parameters completed by ``destinations``.
+        projections: Native grouped projection names already added to the builder.
+    """
+
+    destinations: dict[str, torch.Tensor]
+    native_keys: set[str]
+    projections: set[str]
+
+
+@dataclass
+class _LingFusedQKVLoadGroup(CheckpointLoadGroup):
+    """Bounded fused-QKV destination and its final split parameter storage.
+
+    Attributes:
+        destinations: One checkpoint tensor of shape [q_heads * head_dim + 2 * kv_heads * head_dim, hidden]. The
+            tensor is temporary storage released before the next checkpoint group is requested.
+        native_keys: Names of the three final Q, K, and V parameters completed by this group.
+        q_proj: Final query parameter of shape [q_heads * head_dim, hidden], mutated by :meth:`install`.
+        k_proj: Final key parameter of shape [kv_heads * head_dim, hidden], mutated by :meth:`install`.
+        v_proj: Final value parameter of shape [kv_heads * head_dim, hidden], mutated by :meth:`install`.
+    """
+
+    q_proj: torch.Tensor
+    k_proj: torch.Tensor
+    v_proj: torch.Tensor
+
+    @torch.no_grad()
+    def install(self) -> None:
+        """Split the loaded fused tensor and copy its Q, K, and V slices into final parameter storage."""
+        if len(self.destinations) != 1:
+            raise ValueError(f"Ling fused-QKV group expected one checkpoint destination, got {len(self.destinations)}")
+        fused_qkv = next(iter(self.destinations.values()))
+        expected_shape = (self.q_proj.shape[0] + self.k_proj.shape[0] + self.v_proj.shape[0], self.q_proj.shape[1])
+        if tuple(fused_qkv.shape) != expected_shape:
+            raise ValueError(
+                f"Ling fused-QKV destination has shape {tuple(fused_qkv.shape)}, expected {expected_shape}"
+            )
+
+        q_value, k_value, v_value = torch.split(
+            fused_qkv,
+            [self.q_proj.shape[0], self.k_proj.shape[0], self.v_proj.shape[0]],
+            dim=0,
+        )
+        self.q_proj.copy_(q_value)
+        self.k_proj.copy_(k_value)
+        self.v_proj.copy_(v_value)
 
 
 def _rename_hf_to_native(key: str) -> str:
@@ -99,6 +174,261 @@ def __init__(
         self.dtype = dtype
         self._uses_model_prefix = True
 
+    @property
+    def supports_streaming_checkpoint_load(self) -> bool:
+        """Whether Ling can bound its single-device fused-QKV and expert checkpoint transformations."""
+        return (
+            self.backend.experts in {"torch", "torch_mm"}
+            and self.backend.dispatcher == "torch"
+            and not self.moe_config.expert_bias
+        )
+
+    def _add_streaming_expert_destinations(
+        self,
+        builder: _LingExpertGroupBuilder,
+        *,
+        expert_root: str,
+        native_name: str,
+        projection: str,
+        grouped_weight: torch.Tensor,
+    ) -> None:
+        """Add split HF views of one final grouped expert parameter.
+
+        Args:
+            builder: Mutable dependency group that retains the produced final-storage views.
+            expert_root: Checkpoint prefix ending in ``layers.{layer}.mlp.experts``.
+            native_name: Native name completed by ``grouped_weight``.
+            projection: Either ``gate_and_up_projs`` or ``down_projs``.
+            grouped_weight: Final gate/up tensor of shape [experts, hidden, 2 * expert_hidden] or final down tensor of
+                shape [experts, expert_hidden, hidden]. The emitted checkpoint destinations are non-contiguous views
+                that alias this storage.
+
+        Raises:
+            ValueError: If the projection is duplicated or its tensor shape is incompatible with the Ling layout.
+        """
+        if projection in builder.projections:
+            raise ValueError(f"Duplicate Ling expert projection {expert_root}.{projection}")
+
+        n_experts = self.moe_config.n_routed_experts
+        expert_hidden = self.moe_config.moe_inter_dim
+        if projection == "gate_and_up_projs":
+            expected_shape = (n_experts, self.moe_config.dim, 2 * expert_hidden)
+        else:
+            expected_shape = (n_experts, expert_hidden, self.moe_config.dim)
+        if tuple(grouped_weight.shape) != expected_shape:
+            raise ValueError(
+                f"Ling grouped expert parameter {native_name} has shape {tuple(grouped_weight.shape)}, "
+                f"expected {expected_shape}"
+            )
+
+        for expert_id, expert_weight in enumerate(grouped_weight.unbind(0)):
+            if projection == "gate_and_up_projs":
+                checkpoint_tensors = {
+                    "gate_proj": expert_weight[:, :expert_hidden].transpose(0, 1),
+                    "up_proj": expert_weight[:, expert_hidden:].transpose(0, 1),
+                }
+            else:
+                checkpoint_tensors = {"down_proj": expert_weight.transpose(0, 1)}
+            for checkpoint_projection, checkpoint_destination in checkpoint_tensors.items():
+                checkpoint_name = f"{expert_root}.{expert_id}.{checkpoint_projection}.weight"
+                if checkpoint_name in builder.destinations:
+                    raise ValueError(f"Duplicate Ling expert checkpoint destination {checkpoint_name}")
+                builder.destinations[checkpoint_name] = checkpoint_destination
+
+        builder.native_keys.add(native_name)
+        builder.projections.add(projection)
+
+    @staticmethod
+    def _build_qkv_load_group(prefix: str, builder: _LingQKVGroupBuilder) -> _LingFusedQKVLoadGroup:
+        """Allocate one fused checkpoint tensor backed by three final projection parameters.
+
+        Args:
+            prefix: Native layer prefix ending in ``layers.{layer}``.
+            builder: Q/K/V parameters with shapes [q_hidden, hidden], [kv_hidden, hidden], and [kv_hidden, hidden].
+
+        Returns:
+            Load group with one temporary tensor of shape [q_hidden + 2 * kv_hidden, hidden]. Its installation copies
+            the three row ranges into the final Q/K/V parameters.
+
+        Raises:
+            ValueError: If a projection is missing or the final parameters disagree in rank, shape, device, or dtype.
+        """
+        missing_projections = {"q", "k", "v"} - builder.projections.keys()
+        if missing_projections:
+            raise ValueError(f"Incomplete Ling fused-QKV group {prefix}: missing {sorted(missing_projections)}")
+        q_name, q_proj = builder.projections["q"]
+        k_name, k_proj = builder.projections["k"]
+        v_name, v_proj = builder.projections["v"]
+        if q_proj.ndim != 2 or k_proj.ndim != 2 or v_proj.ndim != 2:
+            raise ValueError(
+                f"Ling fused-QKV group {prefix} requires rank-2 parameters, got "
+                f"Q={tuple(q_proj.shape)}, K={tuple(k_proj.shape)}, V={tuple(v_proj.shape)}"
+            )
+        if k_proj.shape != v_proj.shape or q_proj.shape[1] != k_proj.shape[1]:
+            raise ValueError(
+                f"Ling fused-QKV group {prefix} has incompatible shapes "
+                f"Q={tuple(q_proj.shape)}, K={tuple(k_proj.shape)}, V={tuple(v_proj.shape)}"
+            )
+        if q_proj.device != k_proj.device or q_proj.device != v_proj.device:
+            raise ValueError(f"Ling fused-QKV group {prefix} spans multiple devices")
+        if q_proj.dtype != k_proj.dtype or q_proj.dtype != v_proj.dtype:
+            raise ValueError(f"Ling fused-QKV group {prefix} spans multiple dtypes")
+
+        checkpoint_name = f"{prefix}.attention.query_key_value.weight"
+        fused_destination = torch.empty(
+            (q_proj.shape[0] + k_proj.shape[0] + v_proj.shape[0], q_proj.shape[1]),
+            dtype=q_proj.dtype,
+            device=q_proj.device,
+        )
+        return _LingFusedQKVLoadGroup(
+            destinations={checkpoint_name: fused_destination},
+            native_keys=frozenset({q_name, k_name, v_name}),
+            q_proj=q_proj,
+            k_proj=k_proj,
+            v_proj=v_proj,
+        )
+
+    @staticmethod
+    def _build_expert_load_group(
+        expert_root: str,
+        builder: _LingExpertGroupBuilder,
+        n_experts: int,
+    ) -> CheckpointLoadGroup:
+        """Validate and freeze one complete layer of final-storage expert destinations.
+
+        Args:
+            expert_root: Checkpoint prefix ending in ``layers.{layer}.mlp.experts``.
+            builder: Per-expert gate, up, and down checkpoint views. Gate/up views have shape [expert_hidden, hidden],
+                down views have shape [hidden, expert_hidden], and all views alias final grouped parameters.
+            n_experts: Number of checkpoint experts required for the layer.
+
+        Returns:
+            No-op installation group containing exactly three destinations per expert.
+
+        Raises:
+            ValueError: If either grouped projection or any split expert destination is missing.
+        """
+        missing_projections = {"gate_and_up_projs", "down_projs"} - builder.projections
+        if missing_projections:
+            raise ValueError(f"Incomplete Ling expert group {expert_root}: missing {sorted(missing_projections)}")
+        expected_destinations = 3 * n_experts
+        if len(builder.destinations) != expected_destinations:
+            raise ValueError(
+                f"Ling expert group {expert_root} has {len(builder.destinations)} checkpoint destinations, "
+                f"expected {expected_destinations}"
+            )
+        return CheckpointLoadGroup(
+            destinations=builder.destinations,
+            native_keys=frozenset(builder.native_keys),
+        )
+
+    def iter_checkpoint_load_groups(
+        self,
+        model_part: torch.nn.Module,
+        device_mesh: DeviceMesh | None = None,
+    ) -> Iterator[CheckpointLoadGroup]:
+        """Yield final-storage expert groups and bounded fused-QKV groups for Ling.
+
+        Args:
+            model_part: Single-device Ling model. Attention Q/K/V parameters have shapes [q_hidden, hidden],
+                [kv_hidden, hidden], and [kv_hidden, hidden]. Grouped gate/up experts have shape
+                [experts, hidden, 2 * expert_hidden], and grouped down experts have shape
+                [experts, expert_hidden, hidden].
+            device_mesh: Must be ``None``. Distributed Ling loads continue to use the rank-sharded DCP path.
+
+        Returns:
+            Iterator whose ordinary and expert destinations alias final model storage. Each attention group owns one
+            temporary checkpoint tensor of shape [q_hidden + 2 * kv_hidden, hidden] and installs it before advancing.
+
+        Raises:
+            RuntimeError: If the configured expert backend did not opt into streaming.
+            ValueError: If a distributed mesh, unsupported tensor layout, incomplete dependency group, or allocating
+                ordinary destination is encountered.
+        """
+        if not self.supports_streaming_checkpoint_load:
+            raise RuntimeError(
+                f"{type(self).__name__} does not support streaming for experts={self.backend.experts}, "
+                f"dispatcher={self.backend.dispatcher}"
+            )
+        if device_mesh is not None:
+            raise ValueError("Ling streaming checkpoint loading currently requires a single-device model")
+
+        ordinary_destinations: dict[str, torch.Tensor] = {}
+        ordinary_native_keys: set[str] = set()
+        qkv_groups: dict[str, _LingQKVGroupBuilder] = {}
+        expert_groups: dict[str, _LingExpertGroupBuilder] = {}
+        n_experts = self.moe_config.n_routed_experts
+
+        for parameter_name, parameter in model_part.named_parameters():
+            native_name = canonical_parameter_fqn(parameter_name)
+            if "lora" in native_name:
+                continue
+
+            qkv_match = _NATIVE_LAYER_QKV_RE.match(native_name)
+            if qkv_match is not None:
+                prefix = qkv_match.group("prefix")
+                projection = qkv_match.group("projection")
+                builder = qkv_groups.setdefault(prefix, _LingQKVGroupBuilder(projections={}))
+                if projection in builder.projections:
+                    raise ValueError(f"Duplicate Ling {projection.upper()} projection for {prefix}")
+                builder.projections[projection] = (native_name, parameter.detach())
+                continue
+
+            expert_match = _STREAMABLE_EXPERT_WEIGHT.match(native_name)
+            if expert_match is not None:
+                expert_root = expert_match.group("root")
+                projection = expert_match.group("projection")
+                builder = expert_groups.setdefault(
+                    expert_root,
+                    _LingExpertGroupBuilder(destinations={}, native_keys=set(), projections=set()),
+                )
+                self._add_streaming_expert_destinations(
+                    builder,
+                    expert_root=expert_root,
+                    native_name=native_name,
+                    projection=projection,
+                    grouped_weight=parameter.detach(),
+                )
+                continue
+
+            destination = parameter.detach()
+            converted = self.convert_single_tensor_to_hf(native_name, destination, quantization=False)
+            if len(converted) != 1:
+                raise ValueError(f"Ordinary Ling parameter {native_name} produced {len(converted)} destinations")
+            checkpoint_name, checkpoint_destination = converted[0]
+            if not isinstance(checkpoint_destination, torch.Tensor):
+                raise ValueError(
+                    f"Ordinary Ling destination {checkpoint_name} is {type(checkpoint_destination).__name__}, "
+                    "expected Tensor"
+                )
+            aliases_parameter = (
+                checkpoint_destination.device == destination.device
+                and checkpoint_destination.untyped_storage().data_ptr() == destination.untyped_storage().data_ptr()
+            )
+            if not aliases_parameter:
+                raise ValueError(f"Ordinary Ling destination {checkpoint_name} does not alias {native_name}")
+            if checkpoint_name in ordinary_destinations:
+                raise ValueError(f"Duplicate ordinary Ling checkpoint destination {checkpoint_name}")
+            ordinary_destinations[checkpoint_name] = checkpoint_destination
+            ordinary_native_keys.add(native_name)
+
+        if not qkv_groups:
+            raise ValueError("Ling streaming load plan found no fused-QKV layers")
+        if not expert_groups:
+            raise ValueError("Ling streaming load plan found no grouped expert layers")
+
+        if ordinary_destinations:
+            yield CheckpointLoadGroup(
+                destinations=ordinary_destinations,
+                native_keys=frozenset(ordinary_native_keys),
+            )
+
+        for prefix, builder in qkv_groups.items():
+            yield self._build_qkv_load_group(prefix, builder)
+
+        for expert_root, builder in expert_groups.items():
+            yield self._build_expert_load_group(expert_root, builder, n_experts)
+
     # ---- HF -> native ----------------------------------------------------
 
     def from_hf(
@@ -167,9 +497,9 @@ def to_hf(
                     hf_state_dict[k] = v
                 continue
 
-            m = re.match(r"^(?P<prefix>(?:.*\.)?layers\.\d+)\.self_attn\.(?P<proj>[qkv])_proj\.weight$", fqn)
+            m = _NATIVE_LAYER_QKV_RE.match(fqn)
             if m:
-                pending_qkv.setdefault(m.group("prefix"), {})[m.group("proj")] = tensor
+                pending_qkv.setdefault(m.group("prefix"), {})[m.group("projection")] = tensor
                 continue
 
             hf_state_dict[_rename_native_to_hf(fqn)] = tensor
@@ -201,8 +531,8 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
         if converted is not None:
             return converted
 
-        m = re.match(r"^(?P<prefix>(?:.*\.)?layers\.\d+)\.self_attn\.(?P<proj>[qkv])_proj\.weight$", fqn)
+        m = _NATIVE_LAYER_QKV_RE.match(fqn)
         if m:
-            return [(f"{m.group('prefix')}.attention.{m.group('proj')}_proj.weight", tensor)]
+            return [(f"{m.group('prefix')}.attention.{m.group('projection')}_proj.weight", tensor)]
 
         return [(_rename_native_to_hf(fqn), tensor)]
diff --git a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
index a25ac29d56..d3e68106af 100644
--- a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
+++ b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
@@ -12,13 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import weakref
+
 import pytest
 import torch
 
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.ling_v2.config import BailingMoeV2Config
+from nemo_automodel.components.models.ling_v2.model import BailingMoeV2ForCausalLM
 from nemo_automodel.components.models.ling_v2.state_dict_adapter import BailingMoeV2StateDictAdapter
 from nemo_automodel.components.moe.config import MoEConfig
+from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
 
 
 @pytest.fixture
@@ -78,7 +82,7 @@ def backend_config():
         attn="sdpa",
         linear="torch",
         rms_norm="torch",
-        experts="torch",
+        experts="torch_mm",
         dispatcher="torch",
         enable_hf_state_dict_adapter=False,
     )
@@ -222,3 +226,97 @@ def test_hf_to_native_to_hf_is_lossless_for_qkv_path(self, adapter, config):
                 hf_sd[f"model.layers.{L}.attention.dense.weight"],
             )
         torch.testing.assert_close(roundtripped["model.word_embeddings.weight"], hf_sd["model.word_embeddings.weight"])
+
+
+class TestStreamingLoadGroups:
+    def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
+        self,
+        adapter,
+        config,
+        moe_config,
+        backend_config,
+    ):
+        model = BailingMoeV2ForCausalLM(config=config, moe_config=moe_config, backend=backend_config)
+        parameters = {canonical_parameter_fqn(name): parameter for name, parameter in model.named_parameters()}
+        groups = adapter.iter_checkpoint_load_groups(model)
+        loaded_native_keys = set()
+
+        ordinary_group = next(groups)
+        loaded_native_keys.update(ordinary_group.native_keys)
+        assert "model.word_embeddings.weight" in ordinary_group.destinations
+        assert ordinary_group.destinations["model.word_embeddings.weight"].untyped_storage().data_ptr() == (
+            parameters["model.embed_tokens.weight"].untyped_storage().data_ptr()
+        )
+
+        pending_group = None
+        for layer_id in range(config.num_hidden_layers):
+            qkv_group = next(groups) if pending_group is None else pending_group
+            pending_group = None
+            checkpoint_name = f"model.layers.{layer_id}.attention.query_key_value.weight"
+            assert set(qkv_group.destinations) == {checkpoint_name}
+            loaded_native_keys.update(qkv_group.native_keys)
+
+            fused_qkv = qkv_group.destinations[checkpoint_name]
+            fused_values = torch.arange(fused_qkv.numel(), dtype=torch.float32).reshape(fused_qkv.shape)
+            fused_values = fused_values.to(fused_qkv.dtype)
+            fused_qkv.copy_(fused_values)
+            temporary_ref = weakref.ref(fused_qkv)
+            qkv_group.install()
+
+            q_size = config.num_attention_heads * config.head_dim
+            kv_size = config.num_key_value_heads * config.head_dim
+            torch.testing.assert_close(
+                parameters[f"model.layers.{layer_id}.self_attn.q_proj.weight"], fused_values[:q_size]
+            )
+            torch.testing.assert_close(
+                parameters[f"model.layers.{layer_id}.self_attn.k_proj.weight"],
+                fused_values[q_size : q_size + kv_size],
+            )
+            torch.testing.assert_close(
+                parameters[f"model.layers.{layer_id}.self_attn.v_proj.weight"], fused_values[q_size + kv_size :]
+            )
+
+            del fused_qkv, qkv_group
+            pending_group = next(groups)
+            assert temporary_ref() is None
+
+        expert_group = pending_group
+        assert expert_group is not None
+        loaded_native_keys.update(expert_group.native_keys)
+        gate_key = "model.layers.1.mlp.experts.0.gate_proj.weight"
+        gate_destination = expert_group.destinations[gate_key]
+        gate_destination.fill_(7.0)
+        expert_group.install()
+        grouped_gate_up = parameters["model.layers.1.mlp.experts.gate_and_up_projs"]
+        torch.testing.assert_close(
+            grouped_gate_up[0, :, : config.moe_intermediate_size],
+            torch.full_like(grouped_gate_up[0, :, : config.moe_intermediate_size], 7.0),
+        )
+
+        with pytest.raises(StopIteration):
+            next(groups)
+
+        assert loaded_native_keys == set(parameters)
+
+    def test_streaming_guards_backend_and_distributed_mesh(self, adapter, config, moe_config):
+        assert adapter.supports_streaming_checkpoint_load is True
+        with pytest.raises(ValueError, match="single-device"):
+            next(adapter.iter_checkpoint_load_groups(torch.nn.Linear(2, 2), device_mesh=object()))
+
+        unsupported_backend = BackendConfig(
+            attn="sdpa",
+            linear="torch",
+            rms_norm="torch",
+            experts="te",
+            dispatcher="deepep",
+            enable_hf_state_dict_adapter=False,
+        )
+        unsupported_adapter = BailingMoeV2StateDictAdapter(
+            config=config,
+            moe_config=moe_config,
+            backend=unsupported_backend,
+            dtype=torch.float32,
+        )
+        assert unsupported_adapter.supports_streaming_checkpoint_load is False
+        with pytest.raises(RuntimeError, match="does not support streaming"):
+            next(unsupported_adapter.iter_checkpoint_load_groups(torch.nn.Linear(2, 2)))

From a4bebd0a0a36e6ca01caf6ac10a64069a8aa15e8 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 20 Aug 2026 14:59:03 -0700
Subject: [PATCH 03/20] fix(checkpoint): include Ling routing buffers

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../models/ling_v2/state_dict_adapter.py      | 20 ++++++++---------
 .../test_ling_v2_state_dict_adapter.py        | 22 +++++++++++++------
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
index d650c5e9cb..40d261f620 100644
--- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
@@ -330,10 +330,10 @@ def iter_checkpoint_load_groups(
         """Yield final-storage expert groups and bounded fused-QKV groups for Ling.
 
         Args:
-            model_part: Single-device Ling model. Attention Q/K/V parameters have shapes [q_hidden, hidden],
-                [kv_hidden, hidden], and [kv_hidden, hidden]. Grouped gate/up experts have shape
-                [experts, hidden, 2 * expert_hidden], and grouped down experts have shape
-                [experts, expert_hidden, hidden].
+            model_part: Single-device Ling model whose parameters and persistent buffers are final destinations.
+                Attention Q/K/V parameters have shapes [q_hidden, hidden], [kv_hidden, hidden], and
+                [kv_hidden, hidden]. Grouped gate/up experts have shape [experts, hidden, 2 * expert_hidden], and
+                grouped down experts have shape [experts, expert_hidden, hidden].
             device_mesh: Must be ``None``. Distributed Ling loads continue to use the rank-sharded DCP path.
 
         Returns:
@@ -359,8 +359,8 @@ def iter_checkpoint_load_groups(
         expert_groups: dict[str, _LingExpertGroupBuilder] = {}
         n_experts = self.moe_config.n_routed_experts
 
-        for parameter_name, parameter in model_part.named_parameters():
-            native_name = canonical_parameter_fqn(parameter_name)
+        for state_name, state_tensor in model_part.state_dict(keep_vars=True).items():
+            native_name = canonical_parameter_fqn(state_name)
             if "lora" in native_name:
                 continue
 
@@ -371,7 +371,7 @@ def iter_checkpoint_load_groups(
                 builder = qkv_groups.setdefault(prefix, _LingQKVGroupBuilder(projections={}))
                 if projection in builder.projections:
                     raise ValueError(f"Duplicate Ling {projection.upper()} projection for {prefix}")
-                builder.projections[projection] = (native_name, parameter.detach())
+                builder.projections[projection] = (native_name, state_tensor.detach())
                 continue
 
             expert_match = _STREAMABLE_EXPERT_WEIGHT.match(native_name)
@@ -387,14 +387,14 @@ def iter_checkpoint_load_groups(
                     expert_root=expert_root,
                     native_name=native_name,
                     projection=projection,
-                    grouped_weight=parameter.detach(),
+                    grouped_weight=state_tensor.detach(),
                 )
                 continue
 
-            destination = parameter.detach()
+            destination = state_tensor.detach()
             converted = self.convert_single_tensor_to_hf(native_name, destination, quantization=False)
             if len(converted) != 1:
-                raise ValueError(f"Ordinary Ling parameter {native_name} produced {len(converted)} destinations")
+                raise ValueError(f"Ordinary Ling state tensor {native_name} produced {len(converted)} destinations")
             checkpoint_name, checkpoint_destination = converted[0]
             if not isinstance(checkpoint_destination, torch.Tensor):
                 raise ValueError(
diff --git a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
index d3e68106af..7d56186532 100644
--- a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
+++ b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
@@ -237,7 +237,9 @@ def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
         backend_config,
     ):
         model = BailingMoeV2ForCausalLM(config=config, moe_config=moe_config, backend=backend_config)
-        parameters = {canonical_parameter_fqn(name): parameter for name, parameter in model.named_parameters()}
+        model_state = {
+            canonical_parameter_fqn(name): tensor for name, tensor in model.state_dict(keep_vars=True).items()
+        }
         groups = adapter.iter_checkpoint_load_groups(model)
         loaded_native_keys = set()
 
@@ -245,7 +247,13 @@ def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
         loaded_native_keys.update(ordinary_group.native_keys)
         assert "model.word_embeddings.weight" in ordinary_group.destinations
         assert ordinary_group.destinations["model.word_embeddings.weight"].untyped_storage().data_ptr() == (
-            parameters["model.embed_tokens.weight"].untyped_storage().data_ptr()
+            model_state["model.embed_tokens.weight"].untyped_storage().data_ptr()
+        )
+        correction_bias = ordinary_group.destinations["model.layers.1.mlp.gate.expert_bias"]
+        correction_bias.fill_(3.0)
+        torch.testing.assert_close(
+            model_state["model.layers.1.mlp.gate.e_score_correction_bias"],
+            torch.full_like(model_state["model.layers.1.mlp.gate.e_score_correction_bias"], 3.0),
         )
 
         pending_group = None
@@ -266,14 +274,14 @@ def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
             q_size = config.num_attention_heads * config.head_dim
             kv_size = config.num_key_value_heads * config.head_dim
             torch.testing.assert_close(
-                parameters[f"model.layers.{layer_id}.self_attn.q_proj.weight"], fused_values[:q_size]
+                model_state[f"model.layers.{layer_id}.self_attn.q_proj.weight"], fused_values[:q_size]
             )
             torch.testing.assert_close(
-                parameters[f"model.layers.{layer_id}.self_attn.k_proj.weight"],
+                model_state[f"model.layers.{layer_id}.self_attn.k_proj.weight"],
                 fused_values[q_size : q_size + kv_size],
             )
             torch.testing.assert_close(
-                parameters[f"model.layers.{layer_id}.self_attn.v_proj.weight"], fused_values[q_size + kv_size :]
+                model_state[f"model.layers.{layer_id}.self_attn.v_proj.weight"], fused_values[q_size + kv_size :]
             )
 
             del fused_qkv, qkv_group
@@ -287,7 +295,7 @@ def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
         gate_destination = expert_group.destinations[gate_key]
         gate_destination.fill_(7.0)
         expert_group.install()
-        grouped_gate_up = parameters["model.layers.1.mlp.experts.gate_and_up_projs"]
+        grouped_gate_up = model_state["model.layers.1.mlp.experts.gate_and_up_projs"]
         torch.testing.assert_close(
             grouped_gate_up[0, :, : config.moe_intermediate_size],
             torch.full_like(grouped_gate_up[0, :, : config.moe_intermediate_size], 7.0),
@@ -296,7 +304,7 @@ def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
         with pytest.raises(StopIteration):
             next(groups)
 
-        assert loaded_native_keys == set(parameters)
+        assert loaded_native_keys == set(model_state)
 
     def test_streaming_guards_backend_and_distributed_mesh(self, adapter, config, moe_config):
         assert adapter.supports_streaming_checkpoint_load is True

From e66ce5512fbc8db88f9ba58ceb4d887c6c9c643e Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 20 Aug 2026 15:03:30 -0700
Subject: [PATCH 04/20] fix(checkpoint): include Nemotron routing buffers

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../models/nemotron_v3/state_dict_adapter.py   | 18 +++++++++---------
 .../test_nemotron_v3_state_dict_adapter.py     | 12 ++++++++++++
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
index c60db2e356..e72b7c7391 100644
--- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
@@ -154,9 +154,9 @@ def iter_checkpoint_load_groups(
         """Yield final-storage destinations for the single-device TE fallback.
 
         Args:
-            model_part: Nemotron V3 model whose ordinary parameters and grouped expert parameters are final load
-                destinations. Native input projections have shape [experts, hidden, expert_hidden], and native down
-                projections have shape [experts, expert_hidden, hidden].
+            model_part: Nemotron V3 model whose parameters and persistent buffers are final load destinations. Native
+                input projections have shape [experts, hidden, expert_hidden], and native down projections have shape
+                [experts, expert_hidden, hidden].
             device_mesh: Must be ``None``. Distributed expert parameters require a rank-symmetric group plan.
 
         Returns:
@@ -182,8 +182,8 @@ def iter_checkpoint_load_groups(
         expert_groups: dict[str, _NemotronExpertGroupBuilder] = {}
         expected_expert_ids = set(range(moe_config.n_routed_experts))
 
-        for parameter_name, parameter in model_part.named_parameters():
-            native_name = canonical_parameter_fqn(parameter_name)
+        for state_name, state_tensor in model_part.state_dict(keep_vars=True).items():
+            native_name = canonical_parameter_fqn(state_name)
             # PEFT is applied before the base checkpoint is loaded. Its adapter parameters are intentionally absent
             # from the pretrained checkpoint and retain the initialization performed by initialize_model_weights().
             if "lora" in native_name:
@@ -208,9 +208,9 @@ def iter_checkpoint_load_groups(
                         moe_config.expert_dim,
                     )
 
-                if tuple(parameter.shape) != expected_shape:
+                if tuple(state_tensor.shape) != expected_shape:
                     raise ValueError(
-                        f"Grouped expert parameter {native_name} has shape {tuple(parameter.shape)}, "
+                        f"Grouped expert parameter {native_name} has shape {tuple(state_tensor.shape)}, "
                         f"expected {expected_shape}"
                     )
 
@@ -218,7 +218,7 @@ def iter_checkpoint_load_groups(
                     expert_root,
                     _NemotronExpertGroupBuilder(destinations={}, expert_ids_by_projection={}, native_keys=set()),
                 )
-                split_weights = self._split_experts_weights(parameter.detach(), moe_config.n_routed_experts)
+                split_weights = self._split_experts_weights(state_tensor.detach(), moe_config.n_routed_experts)
                 expert_ids = list(self._last_expert_ids)
                 for expert_id, expert_weight in zip(expert_ids, split_weights, strict=True):
                     checkpoint_name = self._native_key_to_hf(f"{expert_root}.{expert_id}.{projection_name}.weight")
@@ -241,7 +241,7 @@ def iter_checkpoint_load_groups(
                 del split_weights
                 continue
 
-            destination = parameter.detach()
+            destination = state_tensor.detach()
             converted = self.convert_single_tensor_to_hf(native_name, destination, quantization=False)
             if len(converted) != 1:
                 raise ValueError(
diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
index 6652ae3b64..277a2e7314 100644
--- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
+++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
@@ -147,6 +147,11 @@ def __init__(self):
         lora_initial = model.model.embed_tokens.lora_A.weight.detach().clone()
         layer = torch.nn.Module()
         layer.mixer = torch.nn.Module()
+        layer.mixer.gate = torch.nn.Module()
+        layer.mixer.gate.register_buffer(
+            "e_score_correction_bias",
+            torch.zeros(moe_config.n_routed_experts, dtype=torch.float32),
+        )
         layer.mixer.experts = _FakeSingleDeviceFallbackExperts()
         experts = layer.mixer.experts
         model.model.layers = torch.nn.ModuleList([layer])
@@ -158,6 +163,7 @@ def __init__(self):
         all_destinations = {key: value for group in groups for key, value in group.destinations.items()}
         assert set(all_destinations) == {
             "backbone.embeddings.weight",
+            "backbone.layers.0.mixer.gate.e_score_correction_bias",
             *{
                 f"backbone.layers.0.mixer.experts.{expert_id}.{projection}.weight"
                 for expert_id in range(moe_config.n_routed_experts)
@@ -173,6 +179,12 @@ def __init__(self):
             assert not down_destination.is_contiguous()
 
         assert not any("lora" in key for key in all_destinations)
+        correction_bias = all_destinations["backbone.layers.0.mixer.gate.e_score_correction_bias"]
+        correction_bias.fill_(2)
+        torch.testing.assert_close(
+            layer.mixer.gate.e_score_correction_bias,
+            torch.full_like(layer.mixer.gate.e_score_correction_bias, 2),
+        )
         for group in groups[1:]:
             for destination in group.destinations.values():
                 destination.fill_(3)

From 7c46005c191b03840b80e5010a9daf5147406e2e Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 20 Aug 2026 15:07:00 -0700
Subject: [PATCH 05/20] fix(checkpoint): skip runtime-only load state

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../components/models/ling_v2/state_dict_adapter.py        | 4 +++-
 .../components/models/nemotron_v3/state_dict_adapter.py    | 7 ++++++-
 .../nemotron_v3/test_nemotron_v3_state_dict_adapter.py     | 7 +++++++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
index 40d261f620..5d1219f8c8 100644
--- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
@@ -361,8 +361,10 @@ def iter_checkpoint_load_groups(
 
         for state_name, state_tensor in model_part.state_dict(keep_vars=True).items():
             native_name = canonical_parameter_fqn(state_name)
-            if "lora" in native_name:
+            if "lora" in native_name or native_name.endswith("._extra_state"):
                 continue
+            if not isinstance(state_tensor, torch.Tensor):
+                raise ValueError(f"Ling state entry {native_name} is {type(state_tensor).__name__}, expected Tensor")
 
             qkv_match = _NATIVE_LAYER_QKV_RE.match(native_name)
             if qkv_match is not None:
diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
index e72b7c7391..163c8b6a3a 100644
--- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
@@ -186,8 +186,13 @@ def iter_checkpoint_load_groups(
             native_name = canonical_parameter_fqn(state_name)
             # PEFT is applied before the base checkpoint is loaded. Its adapter parameters are intentionally absent
             # from the pretrained checkpoint and retain the initialization performed by initialize_model_weights().
-            if "lora" in native_name:
+            # Transformer Engine's runtime-only ``_extra_state`` records are likewise absent from HF checkpoints.
+            if "lora" in native_name or native_name.endswith("._extra_state"):
                 continue
+            if not isinstance(state_tensor, torch.Tensor):
+                raise ValueError(
+                    f"Nemotron V3 state entry {native_name} is {type(state_tensor).__name__}, expected Tensor"
+                )
 
             expert_match = _STREAMABLE_EXPERT_WEIGHT.match(native_name)
             if expert_match is not None:
diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
index 277a2e7314..ac14bfde56 100644
--- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
+++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
@@ -140,6 +140,12 @@ def __init__(self):
                     torch.zeros(moe_config.n_routed_experts, moe_config.moe_inter_dim, moe_config.expert_dim)
                 )
 
+            def get_extra_state(self) -> dict[str, bool]:
+                return {"runtime_only": True}
+
+            def set_extra_state(self, state: object) -> None:
+                del state
+
         model = torch.nn.Module()
         model.model = torch.nn.Module()
         model.model.embed_tokens = torch.nn.Embedding(8, moe_config.expert_dim)
@@ -179,6 +185,7 @@ def __init__(self):
             assert not down_destination.is_contiguous()
 
         assert not any("lora" in key for key in all_destinations)
+        assert not any("_extra_state" in key for key in all_destinations)
         correction_bias = all_destinations["backbone.layers.0.mixer.gate.e_score_correction_bias"]
         correction_bias.fill_(2)
         torch.testing.assert_close(

From f66dbe84ff3206cbef010a8dd8630550f7578876 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 20 Aug 2026 19:39:34 -0700
Subject: [PATCH 06/20] docs(checkpoint): clarify sequential load groups

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../checkpoint/state_dict_adapter.py          | 29 +++++++++----------
 .../models/ling_v2/state_dict_adapter.py      | 22 +++++++-------
 .../models/nemotron_v3/state_dict_adapter.py  | 16 +++++-----
 3 files changed, 32 insertions(+), 35 deletions(-)

diff --git a/nemo_automodel/components/checkpoint/state_dict_adapter.py b/nemo_automodel/components/checkpoint/state_dict_adapter.py
index c87ab8a058..642437a9c7 100644
--- a/nemo_automodel/components/checkpoint/state_dict_adapter.py
+++ b/nemo_automodel/components/checkpoint/state_dict_adapter.py
@@ -25,25 +25,23 @@
 
 @dataclass
 class CheckpointLoadGroup:
-    """One dependency-complete unit of a bounded checkpoint load.
+    """One set of checkpoint tensors that must be loaded and installed together.
 
     Attributes:
-        destinations: Mapping from checkpoint FQNs to destination tensors. Each tensor has the exact shape and axis
-            order recorded for that FQN in checkpoint metadata. A destination may alias final model storage or may be
-            bounded temporary storage consumed by :meth:`install`.
-        native_keys: Native model state-dict keys completed by this group. Their tensors may have arbitrary ranks and
-            model-owned layouts; every key must be fully installed before :meth:`install` returns.
+        destinations: Checkpoint keys and the tensors DCP should fill. A tensor can use model weight memory directly
+            or be a temporary value needed for conversion.
+        native_keys: Model state-dict keys completed by this group. Every key must be ready in the model before
+            :meth:`install` returns.
     """
 
     destinations: dict[str, torch.Tensor]
     native_keys: frozenset[str]
 
     def install(self) -> None:
-        """Install loaded destinations into final model storage.
+        """Finish converting this group's loaded values into model weights.
 
-        The default implementation is a no-op for destinations that already alias final model storage. Allocating
-        conversion groups override this method, complete their model-owned transformation, and release references to
-        temporary storage before returning.
+        The default does nothing because the destinations already use model weight memory. A group that uses temporary
+        tensors overrides this method, performs its conversion, and releases those tensors before the next group.
         """
         return None
 
@@ -80,11 +78,10 @@ def supports_checkpoint_load_without_full_copy(self) -> bool:
 
     @property
     def supports_streaming_checkpoint_load(self) -> bool:
-        """Whether the adapter can load a base checkpoint as bounded dependency groups.
+        """Whether the adapter can load a base checkpoint one dependency group at a time.
 
         Adapters should opt in only when :meth:`iter_checkpoint_load_groups` covers every required pretrained model
-        tensor and each yielded group writes through to final storage or installs bounded temporary storage before the
-        next group is requested.
+        tensor. Each group must finish updating the model and release any temporary tensors before the next group.
         """
         return self._supports_streaming_checkpoint_load
 
@@ -93,15 +90,15 @@ def iter_checkpoint_load_groups(
         model_part: torch.nn.Module,
         device_mesh: Optional["DeviceMesh"] = None,
     ) -> Iterator[CheckpointLoadGroup]:
-        """Yield dependency-complete groups for bounded base-checkpoint loading.
+        """Yield the groups to load sequentially from a base checkpoint.
 
         Args:
             model_part: Model part that owns the final parameter and buffer storage populated by yielded groups.
             device_mesh: Optional device mesh describing the final distributed tensor placements.
 
         Returns:
-            Iterator of groups. Each group's destination tensors have checkpoint-native shapes and axis order; the
-            iterator must not retain an installed group's temporary tensors after advancing.
+            Groups whose destination tensors match the checkpoint shapes and dimension order. After a group is
+            installed, the iterator must not keep its temporary tensors alive.
 
         Raises:
             NotImplementedError: If the adapter does not implement streaming checkpoint loading.
diff --git a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
index 5d1219f8c8..3e86643dfb 100644
--- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
@@ -92,7 +92,7 @@ class _LingExpertGroupBuilder:
 
     Attributes:
         destinations: Checkpoint gate/up views of shape [expert_hidden, hidden] and down views of shape
-            [hidden, expert_hidden]. Every view aliases a final grouped expert parameter.
+            [hidden, expert_hidden]. Every view uses the grouped model weights as its memory.
         native_keys: Names of the final grouped expert parameters completed by ``destinations``.
         projections: Native grouped projection names already added to the builder.
     """
@@ -104,7 +104,7 @@ class _LingExpertGroupBuilder:
 
 @dataclass
 class _LingFusedQKVLoadGroup(CheckpointLoadGroup):
-    """Bounded fused-QKV destination and its final split parameter storage.
+    """One fused QKV checkpoint tensor and the three model weights it fills.
 
     Attributes:
         destinations: One checkpoint tensor of shape [q_heads * head_dim + 2 * kv_heads * head_dim, hidden]. The
@@ -176,7 +176,7 @@ def __init__(
 
     @property
     def supports_streaming_checkpoint_load(self) -> bool:
-        """Whether Ling can bound its single-device fused-QKV and expert checkpoint transformations."""
+        """Whether Ling can load single-device QKV and expert transformations one layer at a time."""
         return (
             self.backend.experts in {"torch", "torch_mm"}
             and self.backend.dispatcher == "torch"
@@ -201,7 +201,7 @@ def _add_streaming_expert_destinations(
             projection: Either ``gate_and_up_projs`` or ``down_projs``.
             grouped_weight: Final gate/up tensor of shape [experts, hidden, 2 * expert_hidden] or final down tensor of
                 shape [experts, expert_hidden, hidden]. The emitted checkpoint destinations are non-contiguous views
-                that alias this storage.
+                that use this model weight memory.
 
         Raises:
             ValueError: If the projection is duplicated or its tensor shape is incompatible with the Ling layout.
@@ -299,7 +299,7 @@ def _build_expert_load_group(
         Args:
             expert_root: Checkpoint prefix ending in ``layers.{layer}.mlp.experts``.
             builder: Per-expert gate, up, and down checkpoint views. Gate/up views have shape [expert_hidden, hidden],
-                down views have shape [hidden, expert_hidden], and all views alias final grouped parameters.
+                down views have shape [hidden, expert_hidden], and all views use the grouped model weights as memory.
             n_experts: Number of checkpoint experts required for the layer.
 
         Returns:
@@ -327,7 +327,7 @@ def iter_checkpoint_load_groups(
         model_part: torch.nn.Module,
         device_mesh: DeviceMesh | None = None,
     ) -> Iterator[CheckpointLoadGroup]:
-        """Yield final-storage expert groups and bounded fused-QKV groups for Ling.
+        """Yield Ling checkpoint tensors in groups that are installed before the next group is read.
 
         Args:
             model_part: Single-device Ling model whose parameters and persistent buffers are final destinations.
@@ -337,8 +337,8 @@ def iter_checkpoint_load_groups(
             device_mesh: Must be ``None``. Distributed Ling loads continue to use the rank-sharded DCP path.
 
         Returns:
-            Iterator whose ordinary and expert destinations alias final model storage. Each attention group owns one
-            temporary checkpoint tensor of shape [q_hidden + 2 * kv_hidden, hidden] and installs it before advancing.
+            Iterator whose ordinary and expert destinations use model weight memory directly. Each attention group
+            uses one temporary fused-QKV tensor and splits it into Q, K, and V before the iterator advances.
 
         Raises:
             RuntimeError: If the configured expert backend did not opt into streaming.
@@ -403,12 +403,12 @@ def iter_checkpoint_load_groups(
                     f"Ordinary Ling destination {checkpoint_name} is {type(checkpoint_destination).__name__}, "
                     "expected Tensor"
                 )
-            aliases_parameter = (
+            uses_parameter_memory = (
                 checkpoint_destination.device == destination.device
                 and checkpoint_destination.untyped_storage().data_ptr() == destination.untyped_storage().data_ptr()
             )
-            if not aliases_parameter:
-                raise ValueError(f"Ordinary Ling destination {checkpoint_name} does not alias {native_name}")
+            if not uses_parameter_memory:
+                raise ValueError(f"Ordinary Ling destination {checkpoint_name} does not use {native_name} memory")
             if checkpoint_name in ordinary_destinations:
                 raise ValueError(f"Duplicate ordinary Ling checkpoint destination {checkpoint_name}")
             ordinary_destinations[checkpoint_name] = checkpoint_destination
diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
index 163c8b6a3a..84b7f0748b 100644
--- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
@@ -136,8 +136,8 @@ def supports_streaming_checkpoint_load(self) -> bool:
         TE/DeepEP is configured for distributed runs, but the MoE layer deliberately falls back to ``GroupedExperts``
         when world size is one. The adapter configuration still says ``experts="te"``, so the legacy converter assumes
         the grouped tensors are TE stack copies and rebuilds them. For non-gated ReLU-squared experts, the runtime
-        grouped tensors instead provide per-expert transposed views that write through to final storage. Expert bias
-        and other activations retain the allocating fallback until their layouts have concrete coverage.
+        grouped tensors instead provide per-expert transposed views of the model's weight memory. Expert bias and
+        other activations keep the full-checkpoint fallback until their layouts have concrete coverage.
         """
         return (
             self.moe_config is not None
@@ -160,9 +160,9 @@ def iter_checkpoint_load_groups(
             device_mesh: Must be ``None``. Distributed expert parameters require a rank-symmetric group plan.
 
         Returns:
-            Iterator whose first group contains ordinary parameter aliases and whose remaining groups contain one
-            complete MoE layer each. Expert destinations use HF shapes [expert_hidden, hidden] for ``up_proj`` and
-            [hidden, expert_hidden] for ``down_proj`` and are transposed views of final grouped parameter storage.
+            Iterator whose first group contains ordinary parameters and whose remaining groups contain one complete
+            MoE layer each. Expert destinations use HF shapes [expert_hidden, hidden] for ``up_proj`` and
+            [hidden, expert_hidden] for ``down_proj`` while using the grouped model weights as their memory.
 
         Raises:
             RuntimeError: If this adapter configuration did not opt into streaming checkpoint loading.
@@ -258,13 +258,13 @@ def iter_checkpoint_load_groups(
                     f"Ordinary Nemotron destination {checkpoint_name} is {type(checkpoint_destination).__name__}, "
                     "expected Tensor"
                 )
-            aliases_parameter = (
+            uses_parameter_memory = (
                 checkpoint_destination.device == destination.device
                 and checkpoint_destination.untyped_storage().data_ptr() == destination.untyped_storage().data_ptr()
             )
-            if not aliases_parameter:
+            if not uses_parameter_memory:
                 raise ValueError(
-                    f"Ordinary Nemotron destination {checkpoint_name} does not alias final parameter {native_name}"
+                    f"Ordinary Nemotron destination {checkpoint_name} does not use model parameter {native_name} memory"
                 )
             if checkpoint_name in ordinary_destinations:
                 raise ValueError(f"Duplicate ordinary Nemotron checkpoint destination: {checkpoint_name}")

From e5c27b48bae7b3fa9bd2013f458d3bfe73bfd7c1 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 20 Aug 2026 19:49:21 -0700
Subject: [PATCH 07/20] fix(checkpoint): discover adapter keys from tensor
 shapes

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../checkpoint/state_dict_adapter.py          |  8 +++--
 .../models/gemma4_moe/state_dict_adapter.py   | 18 -----------
 .../test_state_dict_adapter_capabilities.py   | 31 +++++++++++++++++++
 3 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/nemo_automodel/components/checkpoint/state_dict_adapter.py b/nemo_automodel/components/checkpoint/state_dict_adapter.py
index 642437a9c7..1ccccc7036 100644
--- a/nemo_automodel/components/checkpoint/state_dict_adapter.py
+++ b/nemo_automodel/components/checkpoint/state_dict_adapter.py
@@ -152,7 +152,7 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
         pass
 
     def get_hf_state_dict_keys(self, state_dict: dict[str, Any]) -> list[str]:
-        """Return the Hugging Face keys produced by ``to_hf``.
+        """Return the Hugging Face keys produced by ``to_hf`` without converting real weights.
 
         Args:
             state_dict: Native model state mapping. Tensor values may have
@@ -162,7 +162,11 @@ def get_hf_state_dict_keys(self, state_dict: dict[str, Any]) -> list[str]:
         Returns:
             Hugging Face state-dict keys in adapter iteration order.
         """
-        return list(self.to_hf(state_dict, exclude_key_regex=r".*_extra_state.*", quantization=False))
+        shape_only_state = {
+            key: torch.empty_like(value, device="meta") if isinstance(value, torch.Tensor) else value
+            for key, value in state_dict.items()
+        }
+        return list(self.to_hf(shape_only_state, exclude_key_regex=r".*_extra_state.*", quantization=False))
 
     def map_peft_target_module_to_hf(self, name: str) -> str:
         """Translate a PEFT target-module name to the HuggingFace layout.
diff --git a/nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py b/nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py
index 1d60dabaea..121f16eecb 100644
--- a/nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py
+++ b/nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py
@@ -382,24 +382,6 @@ def _supports_ep_load_destination(tensor: Any, n_experts: int) -> bool:
             for mesh_dim, placement in enumerate(tensor.placements)
         )
 
-    def get_hf_state_dict_keys(self, state_dict: dict[str, Any]) -> list[str]:
-        """Return converted keys without gathering real expert weights.
-
-        Args:
-            state_dict: Native Gemma4 state mapping. Expert tensors have shape
-                ``[local_experts, hidden, 2 * expert_hidden]`` for fused gate-up
-                weights or ``[local_experts, expert_hidden, hidden]`` for down
-                weights. Other tensor values retain their model-owned layouts.
-
-        Returns:
-            Hugging Face state-dict keys in adapter iteration order.
-        """
-        meta_state_dict = {
-            key: torch.empty_like(value, device="meta") if isinstance(value, torch.Tensor) else value
-            for key, value in state_dict.items()
-        }
-        return list(self.to_hf(meta_state_dict, exclude_key_regex=r".*_extra_state.*"))
-
     def _gather_expert_tensor(
         self,
         tensor: torch.Tensor,
diff --git a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
index dea2e7f1de..a3bc5cd8e5 100644
--- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
+++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
@@ -50,6 +50,37 @@
 from nemo_automodel.components.models.qwen3_vl_moe.state_dict_adapter import Qwen3VLMoeStateDictAdapter
 
 
+class _AllocatingKeysAdapter(StateDictAdapter):
+    """Test adapter whose normal conversion joins two tensors."""
+
+    def __init__(self) -> None:
+        self.converted_devices: list[torch.device] = []
+
+    def to_hf(self, state_dict, **kwargs):
+        self.converted_devices = [value.device for value in state_dict.values() if isinstance(value, torch.Tensor)]
+        return {"fused.weight": torch.cat((state_dict.pop("q.weight"), state_dict.pop("k.weight")), dim=0)}
+
+    def from_hf(self, hf_state_dict, device_mesh=None, **kwargs):
+        return hf_state_dict
+
+    def convert_single_tensor_to_hf(self, fqn, tensor, **kwargs):
+        return [(fqn, tensor)]
+
+
+def test_get_hf_state_dict_keys_uses_shape_only_tensors() -> None:
+    adapter = _AllocatingKeysAdapter()
+    state_dict = {
+        "q.weight": torch.ones(4, 8),
+        "k.weight": torch.ones(2, 8),
+    }
+
+    keys = adapter.get_hf_state_dict_keys(state_dict)
+
+    assert keys == ["fused.weight"]
+    assert adapter.converted_devices == [torch.device("meta"), torch.device("meta")]
+    assert list(state_dict) == ["q.weight", "k.weight"]
+
+
 def _assert_destinations_write_through(
     adapter: StateDictAdapter,
     state_dict: dict[str, torch.Tensor],

From 14ab539ab1e9366a53eccddada0d9faeaf361159 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 20 Aug 2026 21:41:35 -0700
Subject: [PATCH 08/20] refactor(checkpoint): share direct load groups

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../components/checkpoint/checkpointing.py    |  49 +++--
 .../checkpoint/state_dict_adapter.py          | 165 ++++++++++++++--
 .../models/ling_v2/state_dict_adapter.py      |  78 +++-----
 .../models/nemotron_v3/state_dict_adapter.py  | 186 ++----------------
 .../components/moe/state_dict_mixin.py        |   7 +-
 .../checkpoint/test_checkpointing.py          |  63 ++++--
 .../test_state_dict_adapter_capabilities.py   |  58 ++++++
 .../test_ling_v2_state_dict_adapter.py        |  21 +-
 .../test_nemotron_v3_state_dict_adapter.py    |  16 +-
 9 files changed, 355 insertions(+), 288 deletions(-)

diff --git a/nemo_automodel/components/checkpoint/checkpointing.py b/nemo_automodel/components/checkpoint/checkpointing.py
index 47638510f0..cc75cd2a5c 100644
--- a/nemo_automodel/components/checkpoint/checkpointing.py
+++ b/nemo_automodel/components/checkpoint/checkpointing.py
@@ -767,20 +767,22 @@ def _load_model_in_checkpoint_groups(
         model_path: str,
         storage_reader: StorageReader,
     ) -> None:
-        """Load one adapter-owned dependency group at a time.
+        """Load one adapter-owned dependency group at a time and verify complete model coverage.
 
         Args:
             model_state: Wrapper around the single model part whose final parameter storage is populated.
-            adapter: Model-owned adapter that maps checkpoint tensors into dependency-complete groups.
+            adapter: Model-owned adapter that maps model tensors into dependency-complete checkpoint groups.
             model_path: Hugging Face safetensors checkpoint directory.
             storage_reader: Reader already bound to ``model_path``. Its metadata may be reused across groups.
 
         Raises:
-            RuntimeError: If a group requests checkpoint keys that are missing or if the adapter yields no groups.
+            RuntimeError: If checkpoint keys are missing, if model tensors are not covered, or if no groups are yielded.
             TypeError: If the adapter yields an object that is not a :class:`CheckpointLoadGroup`.
             ValueError: If groups are empty or request the same checkpoint key more than once.
         """
         model_part = _unwrap_ddp_model(model_state.model[0])
+        checkpoint_state = adapter.get_checkpoint_load_state(model_part)
+        expected_native_keys = set(checkpoint_state)
         ep_mesh_dims = [dim for dim in self.moe_mesh.mesh_dim_names if dim != "pp"] if self.moe_mesh is not None else []
         ep_mesh = self.moe_mesh[tuple(ep_mesh_dims)] if ep_mesh_dims else self.moe_mesh
         started = time.monotonic()
@@ -798,7 +800,7 @@ def _load_model_in_checkpoint_groups(
         process_group = getattr(self, "process_group", None)
         process_group_kwargs = {"process_group": process_group} if process_group is not None else {}
 
-        for group in adapter.iter_checkpoint_load_groups(model_part, device_mesh=ep_mesh):
+        for group in adapter.iter_checkpoint_load_groups(checkpoint_state, device_mesh=ep_mesh):
             if not isinstance(group, CheckpointLoadGroup):
                 raise TypeError(
                     f"{type(adapter).__name__}.iter_checkpoint_load_groups yielded {type(group).__name__}, "
@@ -806,6 +808,8 @@ def _load_model_in_checkpoint_groups(
                 )
             if not group.destinations:
                 raise ValueError(f"{type(adapter).__name__} yielded an empty checkpoint load group")
+            if not group.native_keys:
+                raise ValueError(f"{type(adapter).__name__} yielded a checkpoint load group with no model tensors")
 
             group_keys = set(group.destinations)
             duplicate_keys = sorted(group_keys & requested_checkpoint_keys)
@@ -817,10 +821,24 @@ def _load_model_in_checkpoint_groups(
             missing_keys = sorted(group_keys - checkpoint_keys)
             if missing_keys:
                 raise RuntimeError(
-                    f"Checkpoint {model_path} is missing {len(missing_keys)} keys required by streaming group "
+                    f"Checkpoint {model_path} is missing {len(missing_keys)} keys required by checkpoint load group "
                     f"{group_count + 1} (examples={missing_keys[:5]})"
                 )
 
+            group_native_keys = set(group.native_keys)
+            duplicate_native_keys = sorted(group_native_keys & loaded_native_keys)
+            if duplicate_native_keys:
+                raise ValueError(
+                    f"{type(adapter).__name__} completed {len(duplicate_native_keys)} model tensors in multiple "
+                    f"groups (examples={duplicate_native_keys[:5]})"
+                )
+            unexpected_native_keys = sorted(group_native_keys - expected_native_keys)
+            if unexpected_native_keys:
+                raise ValueError(
+                    f"{type(adapter).__name__} reported {len(unexpected_native_keys)} unknown model tensors "
+                    f"(examples={unexpected_native_keys[:5]})"
+                )
+
             group_bytes = sum(estimate_tensor_bytes(tensor) for tensor in group.destinations.values())
             requested_bytes += group_bytes
             max_group_bytes = max(max_group_bytes, group_bytes)
@@ -835,14 +853,19 @@ def _load_model_in_checkpoint_groups(
             install_started = time.monotonic()
             group.install()
             install_seconds += time.monotonic() - install_started
-            loaded_native_keys |= set(group.native_keys)
+            loaded_native_keys |= group_native_keys
             group_count += 1
 
             del group
 
         if group_count == 0:
+            raise RuntimeError(f"{type(adapter).__name__} opted into checkpoint load groups but yielded no groups")
+
+        missing_native_keys = sorted(expected_native_keys - loaded_native_keys)
+        if missing_native_keys:
             raise RuntimeError(
-                f"{type(adapter).__name__} opted into streaming checkpoint loading but yielded no groups"
+                f"{type(adapter).__name__} checkpoint load groups omitted {len(missing_native_keys)} model tensors "
+                f"(examples={missing_native_keys[:5]})"
             )
 
         if model_state.uses_tied_lm_head and not model_state.is_peft:
@@ -852,7 +875,7 @@ def _load_model_in_checkpoint_groups(
         requested_gb = requested_bytes / (1 << 30)
         max_group_gb = max_group_bytes / (1 << 30)
         logger.info(
-            "load_model: streamed %.2f GB in %d groups over %.2fs "
+            "load_model: loaded %.2f GB in %d checkpoint groups over %.2fs "
             "(%.2f GB/s overall | max group %.2f GB, metadata %.2fs, storage read %.2fs, install %.2fs, "
             "native keys %d)",
             requested_gb,
@@ -955,9 +978,9 @@ def load_model(
             )
             and not self.config.dequantize_base_checkpoint
         )
-        supports_streaming_checkpoint_load = (
+        supports_checkpoint_load_groups = (
             isinstance(state_dict_adapter, StateDictAdapter)
-            and state_dict_adapter.supports_streaming_checkpoint_load
+            and state_dict_adapter.supports_checkpoint_load_groups
             and not self.config.dequantize_base_checkpoint
         )
         if (
@@ -965,7 +988,7 @@ def load_model(
             and is_safetensors
             and len(model_state.model) == 1
             and world_size == 1
-            and supports_streaming_checkpoint_load
+            and supports_checkpoint_load_groups
             and not allow_checkpoint_key_subset
         ):
             storage_reader = self._get_storage_reader(
@@ -975,7 +998,9 @@ def load_model(
                 is_safetensors=True,
             )
             if storage_reader is None:
-                raise RuntimeError(f"No safetensors storage reader is available for streaming load from {model_path}")
+                raise RuntimeError(
+                    f"No safetensors storage reader is available for checkpoint load groups from {model_path}"
+                )
             self._load_model_in_checkpoint_groups(model_state, state_dict_adapter, model_path, storage_reader)
             return
 
diff --git a/nemo_automodel/components/checkpoint/state_dict_adapter.py b/nemo_automodel/components/checkpoint/state_dict_adapter.py
index 1ccccc7036..14721d2c56 100644
--- a/nemo_automodel/components/checkpoint/state_dict_adapter.py
+++ b/nemo_automodel/components/checkpoint/state_dict_adapter.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
 from dataclasses import dataclass
@@ -19,6 +21,8 @@
 
 import torch
 
+from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
+
 if TYPE_CHECKING:
     from torch.distributed.device_mesh import DeviceMesh
 
@@ -29,7 +33,7 @@ class CheckpointLoadGroup:
 
     Attributes:
         destinations: Checkpoint keys and the tensors DCP should fill. A tensor can use model weight memory directly
-            or be a temporary value needed for conversion.
+            or be a temporary tensor needed for conversion.
         native_keys: Model state-dict keys completed by this group. Every key must be ready in the model before
             :meth:`install` returns.
     """
@@ -38,7 +42,7 @@ class CheckpointLoadGroup:
     native_keys: frozenset[str]
 
     def install(self) -> None:
-        """Finish converting this group's loaded values into model weights.
+        """Finish converting this group's loaded tensors into model weights.
 
         The default does nothing because the destinations already use model weight memory. A group that uses temporary
         tensors overrides this method, performs its conversion, and releases those tensors before the next group.
@@ -55,7 +59,7 @@ class StateDictAdapter(ABC):
 
     _supports_write_through_checkpoint_load: bool = False
     _supports_checkpoint_load_without_full_copy: bool = False
-    _supports_streaming_checkpoint_load: bool = False
+    _supports_checkpoint_load_groups: bool = False
 
     @property
     def supports_write_through_checkpoint_load(self) -> bool:
@@ -77,33 +81,162 @@ def supports_checkpoint_load_without_full_copy(self) -> bool:
         return self._supports_checkpoint_load_without_full_copy
 
     @property
-    def supports_streaming_checkpoint_load(self) -> bool:
-        """Whether the adapter can load a base checkpoint one dependency group at a time.
+    def supports_checkpoint_load_groups(self) -> bool:
+        """Whether the adapter can load a base checkpoint in one or more memory-safe groups.
+
+        Adapters should opt in only when :meth:`iter_checkpoint_load_groups` covers every tensor returned by
+        :meth:`get_checkpoint_load_state`. Each group must finish updating the model and release any temporary tensors
+        before the next group.
+        """
+        return self._supports_checkpoint_load_groups
+
+    def get_checkpoint_load_state(self, model_part: torch.nn.Module) -> dict[str, torch.Tensor]:
+        """Collect the model tensors that a Hugging Face base checkpoint must populate.
+
+        Args:
+            model_part: Model whose parameters and persistent buffers use their native model layouts. Tensor ranks and
+                axis orders are model-specific. LoRA tensors and runtime-only ``_extra_state`` entries are excluded
+                because they are not stored in the Hugging Face base checkpoint.
 
-        Adapters should opt in only when :meth:`iter_checkpoint_load_groups` covers every required pretrained model
-        tensor. Each group must finish updating the model and release any temporary tensors before the next group.
+        Returns:
+            Mapping from canonical native state-dict names to detached tensors with the same shapes, strides, dtypes,
+            devices, and storage as the model parameters and persistent buffers.
+
+        Raises:
+            ValueError: If a checkpoint-owned state entry is not a tensor or canonical names collide.
         """
-        return self._supports_streaming_checkpoint_load
+        checkpoint_state: dict[str, torch.Tensor] = {}
+        for state_name, state_value in model_part.state_dict(keep_vars=True).items():
+            native_name = canonical_parameter_fqn(state_name)
+            if "lora" in native_name or native_name.rsplit(".", 1)[-1] == "_extra_state":
+                continue
+            if not isinstance(state_value, torch.Tensor):
+                raise ValueError(
+                    f"Checkpoint-owned state entry {native_name} is {type(state_value).__name__}, expected Tensor"
+                )
+            if native_name in checkpoint_state:
+                raise ValueError(f"Multiple model state entries resolve to checkpoint key {native_name}")
+            checkpoint_state[native_name] = state_value.detach()
+        return checkpoint_state
+
+    def _checkpoint_load_views(
+        self,
+        native_name: str,
+        native_tensor: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        """Convert one model tensor into checkpoint-layout views of the same storage.
+
+        Args:
+            native_name: Canonical native state-dict name for ``native_tensor``.
+            native_tensor: Model tensor with arbitrary rank and model-specific axis order. Every returned checkpoint
+                tensor must be a view of this tensor's storage; non-contiguous views are allowed.
+
+        Returns:
+            Mapping from Hugging Face checkpoint names to tensors with checkpoint-specific shapes and axis orders.
+            Every tensor uses ``native_tensor`` storage, so DCP writes update the model directly.
+
+        Raises:
+            ValueError: If conversion returns no tensors, returns duplicate names or non-tensors, or allocates storage.
+        """
+        if native_tensor.is_meta:
+            raise ValueError(f"Checkpoint load destination {native_name} is still on the meta device")
+
+        converted = self.convert_single_tensor_to_hf(
+            native_name,
+            native_tensor,
+            quantization=False,
+            for_checkpoint_load=True,
+            track_inplace_load=False,
+        )
+        if not converted:
+            raise ValueError(f"Checkpoint conversion for {native_name} produced no destinations")
+
+        destinations: dict[str, torch.Tensor] = {}
+        source_storage = native_tensor.untyped_storage().data_ptr()
+        for checkpoint_name, checkpoint_tensor in converted:
+            if not isinstance(checkpoint_tensor, torch.Tensor):
+                raise ValueError(
+                    f"Checkpoint destination {checkpoint_name} is {type(checkpoint_tensor).__name__}, expected Tensor"
+                )
+            if checkpoint_tensor.device != native_tensor.device:
+                raise ValueError(
+                    f"Checkpoint destination {checkpoint_name} is on {checkpoint_tensor.device}, "
+                    f"but model tensor {native_name} is on {native_tensor.device}"
+                )
+            if checkpoint_tensor.untyped_storage().data_ptr() != source_storage:
+                raise ValueError(
+                    f"Checkpoint destination {checkpoint_name} does not use model tensor {native_name} storage"
+                )
+            if checkpoint_name in destinations:
+                raise ValueError(f"Checkpoint conversion for {native_name} produced duplicate key {checkpoint_name}")
+            destinations[checkpoint_name] = checkpoint_tensor
+        return destinations
+
+    def _single_checkpoint_load_view(
+        self,
+        native_name: str,
+        native_tensor: torch.Tensor,
+    ) -> tuple[str, torch.Tensor]:
+        """Return the only checkpoint-layout view produced for one model tensor.
+
+        Args:
+            native_name: Canonical native state-dict name for ``native_tensor``.
+            native_tensor: Model tensor with arbitrary rank and model-specific axis order. The returned tensor has the
+                checkpoint layout and uses the same storage.
+
+        Returns:
+            Hugging Face checkpoint name and its tensor view.
+
+        Raises:
+            ValueError: If conversion produces anything other than one checkpoint tensor.
+        """
+        destinations = self._checkpoint_load_views(native_name, native_tensor)
+        if len(destinations) != 1:
+            raise ValueError(
+                f"Checkpoint conversion for {native_name} produced {len(destinations)} destinations, expected one"
+            )
+        return next(iter(destinations.items()))
 
     def iter_checkpoint_load_groups(
         self,
-        model_part: torch.nn.Module,
-        device_mesh: Optional["DeviceMesh"] = None,
+        checkpoint_state: dict[str, torch.Tensor],
+        device_mesh: DeviceMesh | None = None,
     ) -> Iterator[CheckpointLoadGroup]:
-        """Yield the groups to load sequentially from a base checkpoint.
+        """Build one direct-write group for adapters whose conversions return model-storage views.
 
         Args:
-            model_part: Model part that owns the final parameter and buffer storage populated by yielded groups.
+            checkpoint_state: Canonical native names and model tensors with model-specific shapes and axis orders.
+                Each converted checkpoint tensor must be a view of its corresponding model tensor.
             device_mesh: Optional device mesh describing the final distributed tensor placements.
 
         Returns:
-            Groups whose destination tensors match the checkpoint shapes and dimension order. After a group is
-            installed, the iterator must not keep its temporary tensors alive.
+            One group containing every checkpoint-layout view. Model adapters override this method only when a
+            conversion needs temporary tensors that must be installed and released before later groups are loaded.
 
         Raises:
-            NotImplementedError: If the adapter does not implement streaming checkpoint loading.
+            RuntimeError: If the adapter has not opted into checkpoint load groups.
+            ValueError: If a distributed mesh is provided or two model tensors map to the same checkpoint key.
         """
-        raise NotImplementedError(f"{type(self).__name__} does not implement streaming checkpoint loading")
+        if not self.supports_checkpoint_load_groups:
+            raise RuntimeError(f"{type(self).__name__} does not support checkpoint load groups")
+        if device_mesh is not None:
+            raise ValueError("The default checkpoint load group requires a single-device model")
+
+        destinations: dict[str, torch.Tensor] = {}
+        for native_name, native_tensor in checkpoint_state.items():
+            converted = self._checkpoint_load_views(native_name, native_tensor)
+            duplicate_keys = sorted(converted.keys() & destinations.keys())
+            if duplicate_keys:
+                raise ValueError(
+                    f"Multiple model tensors map to {len(duplicate_keys)} checkpoint keys "
+                    f"(examples={duplicate_keys[:5]})"
+                )
+            destinations.update(converted)
+
+        yield CheckpointLoadGroup(
+            destinations=destinations,
+            native_keys=frozenset(checkpoint_state),
+        )
 
     @abstractmethod
     def to_hf(self, state_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
diff --git a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
index 3e86643dfb..2806acf3a9 100644
--- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
@@ -53,7 +53,6 @@
 from nemo_automodel.components.models.ling_v2.config import BailingMoeV2Config
 from nemo_automodel.components.moe.config import MoEConfig
 from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin
-from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
 
 # Map of single-key renames applied in both directions.  Each tuple is
 # (HF substring, native substring); replacement is whole-substring and
@@ -175,7 +174,7 @@ def __init__(
         self._uses_model_prefix = True
 
     @property
-    def supports_streaming_checkpoint_load(self) -> bool:
+    def supports_checkpoint_load_groups(self) -> bool:
         """Whether Ling can load single-device QKV and expert transformations one layer at a time."""
         return (
             self.backend.experts in {"torch", "torch_mm"}
@@ -183,7 +182,7 @@ def supports_streaming_checkpoint_load(self) -> bool:
             and not self.moe_config.expert_bias
         )
 
-    def _add_streaming_expert_destinations(
+    def _add_expert_checkpoint_views(
         self,
         builder: _LingExpertGroupBuilder,
         *,
@@ -221,19 +220,13 @@ def _add_streaming_expert_destinations(
                 f"expected {expected_shape}"
             )
 
-        for expert_id, expert_weight in enumerate(grouped_weight.unbind(0)):
-            if projection == "gate_and_up_projs":
-                checkpoint_tensors = {
-                    "gate_proj": expert_weight[:, :expert_hidden].transpose(0, 1),
-                    "up_proj": expert_weight[:, expert_hidden:].transpose(0, 1),
-                }
-            else:
-                checkpoint_tensors = {"down_proj": expert_weight.transpose(0, 1)}
-            for checkpoint_projection, checkpoint_destination in checkpoint_tensors.items():
-                checkpoint_name = f"{expert_root}.{expert_id}.{checkpoint_projection}.weight"
-                if checkpoint_name in builder.destinations:
-                    raise ValueError(f"Duplicate Ling expert checkpoint destination {checkpoint_name}")
-                builder.destinations[checkpoint_name] = checkpoint_destination
+        checkpoint_views = self._checkpoint_load_views(native_name, grouped_weight)
+        duplicate_keys = sorted(checkpoint_views.keys() & builder.destinations.keys())
+        if duplicate_keys:
+            raise ValueError(
+                f"Duplicate Ling expert checkpoint destinations for {expert_root} (examples={duplicate_keys[:5]})"
+            )
+        builder.destinations.update(checkpoint_views)
 
         builder.native_keys.add(native_name)
         builder.projections.add(projection)
@@ -324,16 +317,16 @@ def _build_expert_load_group(
 
     def iter_checkpoint_load_groups(
         self,
-        model_part: torch.nn.Module,
+        checkpoint_state: dict[str, torch.Tensor],
         device_mesh: DeviceMesh | None = None,
     ) -> Iterator[CheckpointLoadGroup]:
         """Yield Ling checkpoint tensors in groups that are installed before the next group is read.
 
         Args:
-            model_part: Single-device Ling model whose parameters and persistent buffers are final destinations.
-                Attention Q/K/V parameters have shapes [q_hidden, hidden], [kv_hidden, hidden], and
-                [kv_hidden, hidden]. Grouped gate/up experts have shape [experts, hidden, 2 * expert_hidden], and
-                grouped down experts have shape [experts, expert_hidden, hidden].
+            checkpoint_state: Canonical native names and final Ling model tensors. Attention Q/K/V parameters have
+                shapes [q_hidden, hidden], [kv_hidden, hidden], and [kv_hidden, hidden]. Grouped gate/up experts have
+                shape [experts, hidden, 2 * expert_hidden], and grouped down experts have shape
+                [experts, expert_hidden, hidden].
             device_mesh: Must be ``None``. Distributed Ling loads continue to use the rank-sharded DCP path.
 
         Returns:
@@ -341,17 +334,17 @@ def iter_checkpoint_load_groups(
             uses one temporary fused-QKV tensor and splits it into Q, K, and V before the iterator advances.
 
         Raises:
-            RuntimeError: If the configured expert backend did not opt into streaming.
+            RuntimeError: If the configured expert backend does not support checkpoint load groups.
             ValueError: If a distributed mesh, unsupported tensor layout, incomplete dependency group, or allocating
                 ordinary destination is encountered.
         """
-        if not self.supports_streaming_checkpoint_load:
+        if not self.supports_checkpoint_load_groups:
             raise RuntimeError(
-                f"{type(self).__name__} does not support streaming for experts={self.backend.experts}, "
+                f"{type(self).__name__} does not support checkpoint load groups for experts={self.backend.experts}, "
                 f"dispatcher={self.backend.dispatcher}"
             )
         if device_mesh is not None:
-            raise ValueError("Ling streaming checkpoint loading currently requires a single-device model")
+            raise ValueError("Ling checkpoint load groups currently require a single-device model")
 
         ordinary_destinations: dict[str, torch.Tensor] = {}
         ordinary_native_keys: set[str] = set()
@@ -359,13 +352,7 @@ def iter_checkpoint_load_groups(
         expert_groups: dict[str, _LingExpertGroupBuilder] = {}
         n_experts = self.moe_config.n_routed_experts
 
-        for state_name, state_tensor in model_part.state_dict(keep_vars=True).items():
-            native_name = canonical_parameter_fqn(state_name)
-            if "lora" in native_name or native_name.endswith("._extra_state"):
-                continue
-            if not isinstance(state_tensor, torch.Tensor):
-                raise ValueError(f"Ling state entry {native_name} is {type(state_tensor).__name__}, expected Tensor")
-
+        for native_name, state_tensor in checkpoint_state.items():
             qkv_match = _NATIVE_LAYER_QKV_RE.match(native_name)
             if qkv_match is not None:
                 prefix = qkv_match.group("prefix")
@@ -373,7 +360,7 @@ def iter_checkpoint_load_groups(
                 builder = qkv_groups.setdefault(prefix, _LingQKVGroupBuilder(projections={}))
                 if projection in builder.projections:
                     raise ValueError(f"Duplicate Ling {projection.upper()} projection for {prefix}")
-                builder.projections[projection] = (native_name, state_tensor.detach())
+                builder.projections[projection] = (native_name, state_tensor)
                 continue
 
             expert_match = _STREAMABLE_EXPERT_WEIGHT.match(native_name)
@@ -384,40 +371,25 @@ def iter_checkpoint_load_groups(
                     expert_root,
                     _LingExpertGroupBuilder(destinations={}, native_keys=set(), projections=set()),
                 )
-                self._add_streaming_expert_destinations(
+                self._add_expert_checkpoint_views(
                     builder,
                     expert_root=expert_root,
                     native_name=native_name,
                     projection=projection,
-                    grouped_weight=state_tensor.detach(),
+                    grouped_weight=state_tensor,
                 )
                 continue
 
-            destination = state_tensor.detach()
-            converted = self.convert_single_tensor_to_hf(native_name, destination, quantization=False)
-            if len(converted) != 1:
-                raise ValueError(f"Ordinary Ling state tensor {native_name} produced {len(converted)} destinations")
-            checkpoint_name, checkpoint_destination = converted[0]
-            if not isinstance(checkpoint_destination, torch.Tensor):
-                raise ValueError(
-                    f"Ordinary Ling destination {checkpoint_name} is {type(checkpoint_destination).__name__}, "
-                    "expected Tensor"
-                )
-            uses_parameter_memory = (
-                checkpoint_destination.device == destination.device
-                and checkpoint_destination.untyped_storage().data_ptr() == destination.untyped_storage().data_ptr()
-            )
-            if not uses_parameter_memory:
-                raise ValueError(f"Ordinary Ling destination {checkpoint_name} does not use {native_name} memory")
+            checkpoint_name, checkpoint_destination = self._single_checkpoint_load_view(native_name, state_tensor)
             if checkpoint_name in ordinary_destinations:
                 raise ValueError(f"Duplicate ordinary Ling checkpoint destination {checkpoint_name}")
             ordinary_destinations[checkpoint_name] = checkpoint_destination
             ordinary_native_keys.add(native_name)
 
         if not qkv_groups:
-            raise ValueError("Ling streaming load plan found no fused-QKV layers")
+            raise ValueError("Ling checkpoint load plan found no fused-QKV layers")
         if not expert_groups:
-            raise ValueError("Ling streaming load plan found no grouped expert layers")
+            raise ValueError("Ling checkpoint load plan found no grouped expert layers")
 
         if ordinary_destinations:
             yield CheckpointLoadGroup(
diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
index 84b7f0748b..8d81633e7a 100644
--- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
@@ -14,36 +14,20 @@
 
 import logging
 import re
-from collections.abc import Iterator
-from dataclasses import dataclass
 from typing import Any, Optional
 
 import torch
 from torch.distributed.device_mesh import DeviceMesh
 
-from nemo_automodel.components.checkpoint.state_dict_adapter import CheckpointLoadGroup, StateDictAdapter
+from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.moe.config import MoEConfig
 from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin
-from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
 
 logger = logging.getLogger(__name__)
 
 _MAMBA_FP32_PARAMS_TO_BARE = re.compile(r"(\.mixer)\._fp32_params\.")
 _MAMBA_FP32_PARAM_NAMES = ("A_log", "dt_bias", "D")
-_STREAMABLE_EXPERT_WEIGHT = re.compile(
-    r"^(?P<root>(?:.+\.)?layers\.\d+\.mixer\.experts)\."
-    r"(?P<projection>gate_and_up_projs|down_projs)$"
-)
-
-
-@dataclass
-class _NemotronExpertGroupBuilder:
-    """Mutable builder for one dependency-complete expert layer group."""
-
-    destinations: dict[str, torch.Tensor]
-    expert_ids_by_projection: dict[str, set[int]]
-    native_keys: set[str]
 
 
 def _strip_mamba_fp32_holder_key(key: str) -> str:
@@ -130,14 +114,14 @@ def __init__(
         }
 
     @property
-    def supports_streaming_checkpoint_load(self) -> bool:
-        """Whether this adapter can stream split experts into final single-device storage.
+    def supports_checkpoint_load_groups(self) -> bool:
+        """Whether the shared direct-view group can load this single-device runtime layout.
 
         TE/DeepEP is configured for distributed runs, but the MoE layer deliberately falls back to ``GroupedExperts``
         when world size is one. The adapter configuration still says ``experts="te"``, so the legacy converter assumes
-        the grouped tensors are TE stack copies and rebuilds them. For non-gated ReLU-squared experts, the runtime
-        grouped tensors instead provide per-expert transposed views of the model's weight memory. Expert bias and
-        other activations keep the full-checkpoint fallback until their layouts have concrete coverage.
+        the grouped tensors are TE stack copies and rebuilds them. The shared checkpoint-group implementation instead
+        verifies that every converted tensor is a view of the actual runtime model storage. Expert bias and other
+        activations keep the full-checkpoint fallback until their layouts have concrete coverage.
         """
         return (
             self.moe_config is not None
@@ -146,157 +130,6 @@ def supports_streaming_checkpoint_load(self) -> bool:
             and not self.moe_config.expert_bias
         )
 
-    def iter_checkpoint_load_groups(
-        self,
-        model_part: torch.nn.Module,
-        device_mesh: Optional["DeviceMesh"] = None,
-    ) -> Iterator[CheckpointLoadGroup]:
-        """Yield final-storage destinations for the single-device TE fallback.
-
-        Args:
-            model_part: Nemotron V3 model whose parameters and persistent buffers are final load destinations. Native
-                input projections have shape [experts, hidden, expert_hidden], and native down projections have shape
-                [experts, expert_hidden, hidden].
-            device_mesh: Must be ``None``. Distributed expert parameters require a rank-symmetric group plan.
-
-        Returns:
-            Iterator whose first group contains ordinary parameters and whose remaining groups contain one complete
-            MoE layer each. Expert destinations use HF shapes [expert_hidden, hidden] for ``up_proj`` and
-            [hidden, expert_hidden] for ``down_proj`` while using the grouped model weights as their memory.
-
-        Raises:
-            RuntimeError: If this adapter configuration did not opt into streaming checkpoint loading.
-            ValueError: If a distributed mesh, an allocating ordinary conversion, an expert bias, or an incomplete
-                expert group is encountered.
-        """
-        if not self.supports_streaming_checkpoint_load:
-            raise RuntimeError(f"{type(self).__name__} does not support streaming for backend {self.backend.experts}")
-        if device_mesh is not None:
-            raise ValueError("Nemotron V3 streaming checkpoint loading currently requires a single-device model")
-        moe_config = self.moe_config
-        if moe_config is None:
-            raise RuntimeError("Nemotron V3 streaming checkpoint loading requires an MoE configuration")
-
-        ordinary_destinations: dict[str, torch.Tensor] = {}
-        ordinary_native_keys: set[str] = set()
-        expert_groups: dict[str, _NemotronExpertGroupBuilder] = {}
-        expected_expert_ids = set(range(moe_config.n_routed_experts))
-
-        for state_name, state_tensor in model_part.state_dict(keep_vars=True).items():
-            native_name = canonical_parameter_fqn(state_name)
-            # PEFT is applied before the base checkpoint is loaded. Its adapter parameters are intentionally absent
-            # from the pretrained checkpoint and retain the initialization performed by initialize_model_weights().
-            # Transformer Engine's runtime-only ``_extra_state`` records are likewise absent from HF checkpoints.
-            if "lora" in native_name or native_name.endswith("._extra_state"):
-                continue
-            if not isinstance(state_tensor, torch.Tensor):
-                raise ValueError(
-                    f"Nemotron V3 state entry {native_name} is {type(state_tensor).__name__}, expected Tensor"
-                )
-
-            expert_match = _STREAMABLE_EXPERT_WEIGHT.match(native_name)
-            if expert_match is not None:
-                expert_root = expert_match.group("root")
-                native_projection_name = expert_match.group("projection")
-                if native_projection_name == "gate_and_up_projs":
-                    projection_name = "up_proj"
-                    expected_shape = (
-                        moe_config.n_routed_experts,
-                        moe_config.expert_dim,
-                        moe_config.moe_inter_dim,
-                    )
-                else:
-                    projection_name = "down_proj"
-                    expected_shape = (
-                        moe_config.n_routed_experts,
-                        moe_config.moe_inter_dim,
-                        moe_config.expert_dim,
-                    )
-
-                if tuple(state_tensor.shape) != expected_shape:
-                    raise ValueError(
-                        f"Grouped expert parameter {native_name} has shape {tuple(state_tensor.shape)}, "
-                        f"expected {expected_shape}"
-                    )
-
-                builder = expert_groups.setdefault(
-                    expert_root,
-                    _NemotronExpertGroupBuilder(destinations={}, expert_ids_by_projection={}, native_keys=set()),
-                )
-                split_weights = self._split_experts_weights(state_tensor.detach(), moe_config.n_routed_experts)
-                expert_ids = list(self._last_expert_ids)
-                for expert_id, expert_weight in zip(expert_ids, split_weights, strict=True):
-                    checkpoint_name = self._native_key_to_hf(f"{expert_root}.{expert_id}.{projection_name}.weight")
-                    if checkpoint_name in builder.destinations:
-                        raise ValueError(f"Duplicate expert checkpoint destination: {checkpoint_name}")
-                    checkpoint_destination = expert_weight.transpose(0, 1)
-                    expected_checkpoint_shape = (
-                        (moe_config.moe_inter_dim, moe_config.expert_dim)
-                        if projection_name == "up_proj"
-                        else (moe_config.expert_dim, moe_config.moe_inter_dim)
-                    )
-                    if tuple(checkpoint_destination.shape) != expected_checkpoint_shape:
-                        raise ValueError(
-                            f"Expert checkpoint destination {checkpoint_name} has shape "
-                            f"{tuple(checkpoint_destination.shape)}, expected {expected_checkpoint_shape}"
-                        )
-                    builder.destinations[checkpoint_name] = checkpoint_destination
-                builder.expert_ids_by_projection[native_projection_name] = set(expert_ids)
-                builder.native_keys.add(native_name)
-                del split_weights
-                continue
-
-            destination = state_tensor.detach()
-            converted = self.convert_single_tensor_to_hf(native_name, destination, quantization=False)
-            if len(converted) != 1:
-                raise ValueError(
-                    f"Ordinary Nemotron parameter {native_name} produced {len(converted)} checkpoint destinations"
-                )
-            checkpoint_name, checkpoint_destination = converted[0]
-            if not isinstance(checkpoint_destination, torch.Tensor):
-                raise ValueError(
-                    f"Ordinary Nemotron destination {checkpoint_name} is {type(checkpoint_destination).__name__}, "
-                    "expected Tensor"
-                )
-            uses_parameter_memory = (
-                checkpoint_destination.device == destination.device
-                and checkpoint_destination.untyped_storage().data_ptr() == destination.untyped_storage().data_ptr()
-            )
-            if not uses_parameter_memory:
-                raise ValueError(
-                    f"Ordinary Nemotron destination {checkpoint_name} does not use model parameter {native_name} memory"
-                )
-            if checkpoint_name in ordinary_destinations:
-                raise ValueError(f"Duplicate ordinary Nemotron checkpoint destination: {checkpoint_name}")
-            ordinary_destinations[checkpoint_name] = checkpoint_destination
-            ordinary_native_keys.add(native_name)
-
-        for group_key, builder in expert_groups.items():
-            for projection_name in ("gate_and_up_projs", "down_projs"):
-                expert_ids = builder.expert_ids_by_projection.get(projection_name, set())
-                if expert_ids != expected_expert_ids:
-                    missing_experts = sorted(expected_expert_ids - expert_ids)
-                    raise ValueError(
-                        f"Incomplete expert group {group_key}.{projection_name}: "
-                        f"missing expert ids {missing_experts[:10]}"
-                    )
-
-        if not expert_groups:
-            raise ValueError("Nemotron V3 streaming load plan found no grouped expert layers")
-
-        if ordinary_destinations:
-            yield CheckpointLoadGroup(
-                destinations=ordinary_destinations,
-                native_keys=frozenset(ordinary_native_keys),
-            )
-
-        for group_key in sorted(expert_groups):
-            builder = expert_groups[group_key]
-            yield CheckpointLoadGroup(
-                destinations=builder.destinations,
-                native_keys=frozenset(builder.native_keys),
-            )
-
     @property
     def _hf_prefix(self) -> str:
         """Return the source checkpoint's public Nemotron-H model prefix."""
@@ -486,7 +319,12 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
         # emitted HF keys stay under ``mtp.`` instead of ``backbone.``.
         if fqn.startswith("mtp."):
             fqn = _strip_mamba_fp32_holder_key(fqn)
-            expert_split = self._convert_single_merged_expert_to_hf_split_experts(fqn, tensor, prefix_override="mtp.")
+            expert_split = self._convert_single_merged_expert_to_hf_split_experts(
+                fqn,
+                tensor,
+                prefix_override="mtp.",
+                **kwargs,
+            )
             result = expert_split if expert_split is not None else [(fqn, tensor)]
             result = [(key, _upcast_mamba_fp32_state_tensor(key, value)) for key, value in result]
             if exclude_key_regex:
diff --git a/nemo_automodel/components/moe/state_dict_mixin.py b/nemo_automodel/components/moe/state_dict_mixin.py
index 766522b9ab..5a28467436 100644
--- a/nemo_automodel/components/moe/state_dict_mixin.py
+++ b/nemo_automodel/components/moe/state_dict_mixin.py
@@ -829,6 +829,7 @@ def _convert_single_merged_expert_to_hf_split_experts(
         *,
         prefix_override: str | None = None,
         for_checkpoint_load: bool = False,
+        track_inplace_load: bool = True,
         **kwargs,
     ) -> list[tuple[str, torch.Tensor]]:
         """Convert one grouped expert tensor to Hugging Face's per-expert layout.
@@ -845,6 +846,8 @@ def _convert_single_merged_expert_to_hf_split_experts(
                 outside the main backbone, e.g. ``"mtp."`` for the MTP head.
             for_checkpoint_load: Return views that DCP will completely overwrite. Save/export callers leave this
                 disabled so converted tensors preserve their current values in contiguous storage.
+            track_inplace_load: Record direct-write views for a later ``from_hf`` call. Checkpoint load groups disable
+                this because they write into the model and finish the load without calling ``from_hf``.
             **kwargs: Absorbed for forward-compatibility with base callers
                 that forward arbitrary state-dict kwargs (e.g. ``exclude_key_regex``).
 
@@ -903,7 +906,7 @@ def checkpoint_load_destination(view: torch.Tensor, source: torch.Tensor) -> tor
                 and not quantization
                 and gate_up_storage_is_model_weight
             )
-            if inplace_ok:
+            if inplace_ok and track_inplace_load:
                 self._register_inplace_loaded_key(fqn, prefix_override)
 
             result = []
@@ -956,7 +959,7 @@ def checkpoint_load_destination(view: torch.Tensor, source: torch.Tensor) -> tor
                 and not quantization
                 and down_storage_is_model_weight
             )
-            if inplace_ok:
+            if inplace_ok and track_inplace_load:
                 self._register_inplace_loaded_key(fqn, prefix_override)
 
             result = []
diff --git a/tests/unit_tests/checkpoint/test_checkpointing.py b/tests/unit_tests/checkpoint/test_checkpointing.py
index 6ca15d5a66..4e87ae8d9d 100644
--- a/tests/unit_tests/checkpoint/test_checkpointing.py
+++ b/tests/unit_tests/checkpoint/test_checkpointing.py
@@ -1697,7 +1697,7 @@ def test_single_device_custom_model_uses_fast_path(self, mock_load_full, mock_lo
     @patch("nemo_automodel.components.checkpoint.checkpointing._is_safetensors_checkpoint", return_value=True)
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype")
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model")
-    @pytest.mark.parametrize("load_capability", ["write_through", "without_full_copy", "streaming"])
+    @pytest.mark.parametrize("load_capability", ["write_through", "without_full_copy", "groups"])
     @pytest.mark.parametrize("dequantize_base_checkpoint", [False, True])
     def test_single_device_adapter_without_full_copy_routes_by_quantization(
         self,
@@ -1716,7 +1716,7 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
         model.state_dict_adapter = MagicMock(spec=StateDictAdapter)
         model.state_dict_adapter.supports_write_through_checkpoint_load = load_capability == "write_through"
         model.state_dict_adapter.supports_checkpoint_load_without_full_copy = load_capability == "without_full_copy"
-        model.state_dict_adapter.supports_streaming_checkpoint_load = load_capability == "streaming"
+        model.state_dict_adapter.supports_checkpoint_load_groups = load_capability == "groups"
         mock_state_dict = {"layer.weight": torch.randn(4, 4), "layer.bias": torch.randn(4)}
         mock_load_hf.return_value = mock_state_dict
 
@@ -1737,7 +1737,7 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
                 side_effect=lambda model_part, state_dict, **kwargs: state_dict,
             ),
             patch.object(checkpointer, "_get_storage_reader", return_value=MagicMock()),
-            patch.object(checkpointer, "_load_model_in_checkpoint_groups") as mock_streaming_load,
+            patch.object(checkpointer, "_load_model_in_checkpoint_groups") as mock_grouped_load,
             patch.object(checkpointer, "_do_load", return_value=mock_state_dict) as mock_dcp_load,
         ):
             mock_model_state = mock_model_state_cls.return_value
@@ -1749,17 +1749,17 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
         if dequantize_base_checkpoint:
             mock_load_full.assert_called_once()
             mock_load_hf.assert_called_once()
-            mock_streaming_load.assert_not_called()
+            mock_grouped_load.assert_not_called()
             mock_dcp_load.assert_not_called()
-        elif load_capability == "streaming":
+        elif load_capability == "groups":
             mock_load_full.assert_not_called()
             mock_load_hf.assert_not_called()
-            mock_streaming_load.assert_called_once()
+            mock_grouped_load.assert_called_once()
             mock_dcp_load.assert_not_called()
         else:
             mock_load_full.assert_not_called()
             mock_load_hf.assert_not_called()
-            mock_streaming_load.assert_not_called()
+            mock_grouped_load.assert_not_called()
             mock_dcp_load.assert_called_once()
             assert "load_model:" in caplog.text
 
@@ -1781,18 +1781,24 @@ def test_checkpoint_group_executor_installs_before_requesting_next_group(caplog)
     peak_active_groups = 0
     installed: dict[int, torch.Tensor] = {}
 
-    def iter_groups(model_part, device_mesh=None):
+    checkpoint_state = {
+        "native.0": torch.empty(3),
+        "native.1": torch.empty(3),
+    }
+    adapter.get_checkpoint_load_state.return_value = checkpoint_state
+
+    def iter_groups(load_state, device_mesh=None):
         """Yield groups after verifying the prior group's destination was installed.
 
         Args:
-            model_part: Model whose parameters have arbitrary shapes; this test only verifies object identity.
+            load_state: Native model tensors of shape [elements] that the groups must cover.
             device_mesh: Optional distributed mesh, expected to be ``None`` for this single-process test.
 
         Returns:
             Iterator of groups containing destination tensors of shape [elements].
         """
         nonlocal active_groups, peak_active_groups
-        assert model_part is model_state.model[0]
+        assert load_state is checkpoint_state
         assert device_mesh is None
         for group_id in range(2):
             assert active_groups == 0
@@ -1841,8 +1847,8 @@ def load_group(destinations, **kwargs):
     assert torch.equal(installed[1], torch.full((3,), 2.0))
     assert mock_load.call_count == 2
     reader.read_metadata.assert_called_once()
-    assert "streamed" in caplog.text
-    assert "2 groups" in caplog.text
+    assert "loaded" in caplog.text
+    assert "2 checkpoint groups" in caplog.text
 
 
 def test_checkpoint_group_executor_reuses_hf_metadata_and_loads_exact_values(tmp_path):
@@ -1865,6 +1871,9 @@ def test_checkpoint_group_executor_reuses_hf_metadata_and_loads_exact_values(tmp
         for index, key in enumerate(checkpoint)
     ]
     adapter = MagicMock(spec=StateDictAdapter)
+    adapter.get_checkpoint_load_state.return_value = {
+        f"native.{index}": backing_tensors[key] for index, key in enumerate(checkpoint)
+    }
     adapter.iter_checkpoint_load_groups.return_value = iter(groups)
     model_state = SimpleNamespace(
         model=[torch.nn.Linear(2, 2)],
@@ -1882,6 +1891,36 @@ def test_checkpoint_group_executor_reuses_hf_metadata_and_loads_exact_values(tmp
         torch.testing.assert_close(destinations[key], expected)
 
 
+def test_checkpoint_group_executor_rejects_incomplete_model_coverage():
+    checkpointer = TestLoadModelCustomModelGuard()._make_checkpointer()
+    model_state = SimpleNamespace(
+        model=[torch.nn.Linear(2, 2)],
+        uses_tied_lm_head=False,
+        is_peft=False,
+    )
+    adapter = MagicMock(spec=StateDictAdapter)
+    adapter.get_checkpoint_load_state.return_value = {
+        "native.loaded": torch.empty(2),
+        "native.omitted": torch.empty(2),
+    }
+    adapter.iter_checkpoint_load_groups.return_value = iter(
+        [
+            CheckpointLoadGroup(
+                destinations={"checkpoint.loaded": torch.empty(2)},
+                native_keys=frozenset({"native.loaded"}),
+            )
+        ]
+    )
+    reader = MagicMock()
+    reader.read_metadata.return_value = SimpleNamespace(state_dict_metadata={"checkpoint.loaded": object()})
+
+    with (
+        patch("nemo_automodel.components.checkpoint.checkpointing.dcp.load"),
+        pytest.raises(RuntimeError, match="omitted 1 model tensors.*native.omitted"),
+    ):
+        checkpointer._load_model_in_checkpoint_groups(model_state, adapter, "/checkpoint", reader)
+
+
 class TestLoadModelCheckpointKeySubset:
     """Test allow_checkpoint_key_subset support for torch_save exports."""
 
diff --git a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
index a3bc5cd8e5..e642609199 100644
--- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
+++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
@@ -67,6 +67,23 @@ def convert_single_tensor_to_hf(self, fqn, tensor, **kwargs):
         return [(fqn, tensor)]
 
 
+class _TransposeLoadAdapter(_AllocatingKeysAdapter):
+    """Test adapter whose checkpoint layout transposes each model tensor."""
+
+    _supports_checkpoint_load_groups = True
+
+    def convert_single_tensor_to_hf(self, fqn, tensor, **kwargs):
+        self.checkpoint_conversion_kwargs = kwargs
+        return [(f"checkpoint.{fqn}", tensor.transpose(0, 1))]
+
+
+class _CopyingLoadAdapter(_TransposeLoadAdapter):
+    """Invalid direct-load adapter that allocates a new checkpoint tensor."""
+
+    def convert_single_tensor_to_hf(self, fqn, tensor, **kwargs):
+        return [(f"checkpoint.{fqn}", tensor.transpose(0, 1).contiguous())]
+
+
 def test_get_hf_state_dict_keys_uses_shape_only_tensors() -> None:
     adapter = _AllocatingKeysAdapter()
     state_dict = {
@@ -81,6 +98,47 @@ def test_get_hf_state_dict_keys_uses_shape_only_tensors() -> None:
     assert list(state_dict) == ["q.weight", "k.weight"]
 
 
+def test_shared_checkpoint_group_uses_model_views_and_excludes_non_base_state() -> None:
+    class _ModelWithRuntimeState(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.zeros(2, 3))
+            self.lora_A = torch.nn.Parameter(torch.ones(1, 3))
+            self.register_buffer("running_scale", torch.zeros(2, 1))
+
+        def get_extra_state(self) -> dict[str, bool]:
+            return {"runtime_only": True}
+
+        def set_extra_state(self, state: object) -> None:
+            del state
+
+    model = _ModelWithRuntimeState()
+    adapter = _TransposeLoadAdapter()
+
+    checkpoint_state = adapter.get_checkpoint_load_state(model)
+    groups = list(adapter.iter_checkpoint_load_groups(checkpoint_state))
+
+    assert set(checkpoint_state) == {"weight", "running_scale"}
+    assert len(groups) == 1
+    assert groups[0].native_keys == frozenset(checkpoint_state)
+    assert adapter.checkpoint_conversion_kwargs["for_checkpoint_load"] is True
+    assert adapter.checkpoint_conversion_kwargs["track_inplace_load"] is False
+    weight_destination = groups[0].destinations["checkpoint.weight"]
+    assert tuple(weight_destination.shape) == (3, 2)
+    assert not weight_destination.is_contiguous()
+    assert weight_destination.untyped_storage().data_ptr() == model.weight.untyped_storage().data_ptr()
+    weight_destination.fill_(4)
+    torch.testing.assert_close(model.weight, torch.full_like(model.weight, 4))
+
+
+def test_shared_checkpoint_group_rejects_allocating_conversion() -> None:
+    source = torch.zeros(2, 3)
+    adapter = _CopyingLoadAdapter()
+
+    with pytest.raises(ValueError, match="does not use model tensor weight storage"):
+        list(adapter.iter_checkpoint_load_groups({"weight": source}))
+
+
 def _assert_destinations_write_through(
     adapter: StateDictAdapter,
     state_dict: dict[str, torch.Tensor],
diff --git a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
index 7d56186532..5e81cdc0f5 100644
--- a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
+++ b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
@@ -22,7 +22,6 @@
 from nemo_automodel.components.models.ling_v2.model import BailingMoeV2ForCausalLM
 from nemo_automodel.components.models.ling_v2.state_dict_adapter import BailingMoeV2StateDictAdapter
 from nemo_automodel.components.moe.config import MoEConfig
-from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
 
 
 @pytest.fixture
@@ -228,7 +227,7 @@ def test_hf_to_native_to_hf_is_lossless_for_qkv_path(self, adapter, config):
         torch.testing.assert_close(roundtripped["model.word_embeddings.weight"], hf_sd["model.word_embeddings.weight"])
 
 
-class TestStreamingLoadGroups:
+class TestCheckpointLoadGroups:
     def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
         self,
         adapter,
@@ -237,10 +236,8 @@ def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
         backend_config,
     ):
         model = BailingMoeV2ForCausalLM(config=config, moe_config=moe_config, backend=backend_config)
-        model_state = {
-            canonical_parameter_fqn(name): tensor for name, tensor in model.state_dict(keep_vars=True).items()
-        }
-        groups = adapter.iter_checkpoint_load_groups(model)
+        model_state = adapter.get_checkpoint_load_state(model)
+        groups = adapter.iter_checkpoint_load_groups(model_state)
         loaded_native_keys = set()
 
         ordinary_group = next(groups)
@@ -306,10 +303,10 @@ def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
 
         assert loaded_native_keys == set(model_state)
 
-    def test_streaming_guards_backend_and_distributed_mesh(self, adapter, config, moe_config):
-        assert adapter.supports_streaming_checkpoint_load is True
+    def test_checkpoint_load_groups_guard_backend_and_distributed_mesh(self, adapter, config, moe_config):
+        assert adapter.supports_checkpoint_load_groups is True
         with pytest.raises(ValueError, match="single-device"):
-            next(adapter.iter_checkpoint_load_groups(torch.nn.Linear(2, 2), device_mesh=object()))
+            next(adapter.iter_checkpoint_load_groups({}, device_mesh=object()))
 
         unsupported_backend = BackendConfig(
             attn="sdpa",
@@ -325,6 +322,6 @@ def test_streaming_guards_backend_and_distributed_mesh(self, adapter, config, mo
             backend=unsupported_backend,
             dtype=torch.float32,
         )
-        assert unsupported_adapter.supports_streaming_checkpoint_load is False
-        with pytest.raises(RuntimeError, match="does not support streaming"):
-            next(unsupported_adapter.iter_checkpoint_load_groups(torch.nn.Linear(2, 2)))
+        assert unsupported_adapter.supports_checkpoint_load_groups is False
+        with pytest.raises(RuntimeError, match="does not support checkpoint load groups"):
+            next(unsupported_adapter.iter_checkpoint_load_groups({}))
diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
index ac14bfde56..5d68a35a8f 100644
--- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
+++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
@@ -162,9 +162,11 @@ def set_extra_state(self, state: object) -> None:
         experts = layer.mixer.experts
         model.model.layers = torch.nn.ModuleList([layer])
 
-        groups = list(adapter.iter_checkpoint_load_groups(model))
-        assert adapter.supports_streaming_checkpoint_load is True
-        assert len(groups) == 2
+        checkpoint_state = adapter.get_checkpoint_load_state(model)
+        groups = list(adapter.iter_checkpoint_load_groups(checkpoint_state))
+        assert adapter.supports_checkpoint_load_groups is True
+        assert len(groups) == 1
+        assert groups[0].native_keys == frozenset(checkpoint_state)
 
         all_destinations = {key: value for group in groups for key, value in group.destinations.items()}
         assert set(all_destinations) == {
@@ -192,16 +194,16 @@ def set_extra_state(self, state: object) -> None:
             layer.mixer.gate.e_score_correction_bias,
             torch.full_like(layer.mixer.gate.e_score_correction_bias, 2),
         )
-        for group in groups[1:]:
-            for destination in group.destinations.values():
+        for checkpoint_name, destination in all_destinations.items():
+            if ".experts." in checkpoint_name:
                 destination.fill_(3)
-            group.install()
+        groups[0].install()
         assert torch.count_nonzero(experts.gate_and_up_projs) == experts.gate_and_up_projs.numel()
         assert torch.count_nonzero(experts.down_projs) == experts.down_projs.numel()
         torch.testing.assert_close(model.model.embed_tokens.lora_A.weight, lora_initial)
 
         moe_config.expert_bias = True
-        assert adapter.supports_streaming_checkpoint_load is False
+        assert adapter.supports_checkpoint_load_groups is False
 
     def test_from_hf_map_structure(self, config, moe_config, backend):
         """Test from_hf_map structure."""

From 42e51fe9f490feb557f730f2eb3092662c37080f Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Fri, 21 Aug 2026 09:25:59 -0700
Subject: [PATCH 09/20] refactor(checkpoint): use standard DCP for single-GPU
 MoE

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../checkpoint/_backports/hf_storage.py       |  14 -
 .../components/checkpoint/checkpointing.py    | 157 +--------
 .../checkpoint/state_dict_adapter.py          | 188 -----------
 .../models/ling_v2/state_dict_adapter.py      | 306 +-----------------
 .../models/nemotron_v3/state_dict_adapter.py  |  17 -
 .../components/moe/state_dict_mixin.py        |  28 +-
 .../checkpoint/test_checkpointing.py          | 171 +---------
 .../test_state_dict_adapter_capabilities.py   |  76 +----
 .../test_ling_v2_state_dict_adapter.py        | 145 +++------
 .../test_nemotron_v3_state_dict_adapter.py    | 108 +++----
 tests/unit_tests/moe/test_state_dict_mixin.py |  31 +-
 11 files changed, 159 insertions(+), 1082 deletions(-)

diff --git a/nemo_automodel/components/checkpoint/_backports/hf_storage.py b/nemo_automodel/components/checkpoint/_backports/hf_storage.py
index 5759c99e06..3ba517f72e 100644
--- a/nemo_automodel/components/checkpoint/_backports/hf_storage.py
+++ b/nemo_automodel/components/checkpoint/_backports/hf_storage.py
@@ -276,14 +276,6 @@ def __init__(self, path: str, token: str | None = None, key_mapping: dict[str, s
             super().__init__(path=path)
 
         self.key_mapping = key_mapping
-        self._metadata_cache: Metadata | None = None
-
-    def reset(self, checkpoint_id: str | os.PathLike | None = None) -> None:
-        """Reset reader state and invalidate cached metadata when the checkpoint changes."""
-        previous_path = str(self.path)
-        super().reset(checkpoint_id)
-        if checkpoint_id is not None and str(self.path) != previous_path:
-            self._metadata_cache = None
 
     def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
         per_file: dict[str, list[ReadItem]] = {}
@@ -380,11 +372,6 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
         return fut
 
     def read_metadata(self) -> Metadata:
-        if self._metadata_cache is not None:
-            if getattr(self._metadata_cache, "storage_meta", None) is not None:
-                self._metadata_cache.storage_meta.load_id = self.load_id
-            return self._metadata_cache
-
         state_dict_metadata: dict[str, TensorStorageMetadata] = {}
         storage_data: dict[MetadataIndex, _HFStorageInfo] = {}
 
@@ -456,7 +443,6 @@ def read_metadata(self) -> Metadata:
             metadata.storage_meta = StorageMeta()
         metadata.storage_meta.load_id = self.load_id  # type: ignore[union-attr]
 
-        self._metadata_cache = metadata
         return metadata
 
 
diff --git a/nemo_automodel/components/checkpoint/checkpointing.py b/nemo_automodel/components/checkpoint/checkpointing.py
index cc75cd2a5c..7b9920dfe7 100644
--- a/nemo_automodel/components/checkpoint/checkpointing.py
+++ b/nemo_automodel/components/checkpoint/checkpointing.py
@@ -70,7 +70,7 @@
     requires_tensor_merging,
 )
 from nemo_automodel.components.checkpoint.lifecycle import CheckpointLifecycle
-from nemo_automodel.components.checkpoint.state_dict_adapter import CheckpointLoadGroup, StateDictAdapter
+from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
 from nemo_automodel.components.checkpoint.stateful_wrappers import ModelState, OptimizerState
 from nemo_automodel.components.checkpoint.utils import (
     ensure_tied_lm_head,
@@ -760,135 +760,6 @@ def load_optimizer(
         self._do_load(state_dict, os.path.join(weights_path, "optim"))
         optimizer_state.load_state_dict(state_dict)
 
-    def _load_model_in_checkpoint_groups(
-        self,
-        model_state: ModelState,
-        adapter: StateDictAdapter,
-        model_path: str,
-        storage_reader: StorageReader,
-    ) -> None:
-        """Load one adapter-owned dependency group at a time and verify complete model coverage.
-
-        Args:
-            model_state: Wrapper around the single model part whose final parameter storage is populated.
-            adapter: Model-owned adapter that maps model tensors into dependency-complete checkpoint groups.
-            model_path: Hugging Face safetensors checkpoint directory.
-            storage_reader: Reader already bound to ``model_path``. Its metadata may be reused across groups.
-
-        Raises:
-            RuntimeError: If checkpoint keys are missing, if model tensors are not covered, or if no groups are yielded.
-            TypeError: If the adapter yields an object that is not a :class:`CheckpointLoadGroup`.
-            ValueError: If groups are empty or request the same checkpoint key more than once.
-        """
-        model_part = _unwrap_ddp_model(model_state.model[0])
-        checkpoint_state = adapter.get_checkpoint_load_state(model_part)
-        expected_native_keys = set(checkpoint_state)
-        ep_mesh_dims = [dim for dim in self.moe_mesh.mesh_dim_names if dim != "pp"] if self.moe_mesh is not None else []
-        ep_mesh = self.moe_mesh[tuple(ep_mesh_dims)] if ep_mesh_dims else self.moe_mesh
-        started = time.monotonic()
-        checkpoint_keys = _get_checkpoint_metadata_keys(model_path, storage_reader)
-        metadata_seconds = time.monotonic() - started
-
-        storage_seconds = 0.0
-        install_seconds = 0.0
-        requested_checkpoint_keys: set[str] = set()
-        loaded_native_keys: set[str] = set()
-        requested_bytes = 0
-        max_group_bytes = 0
-        group_count = 0
-
-        process_group = getattr(self, "process_group", None)
-        process_group_kwargs = {"process_group": process_group} if process_group is not None else {}
-
-        for group in adapter.iter_checkpoint_load_groups(checkpoint_state, device_mesh=ep_mesh):
-            if not isinstance(group, CheckpointLoadGroup):
-                raise TypeError(
-                    f"{type(adapter).__name__}.iter_checkpoint_load_groups yielded {type(group).__name__}, "
-                    "expected CheckpointLoadGroup"
-                )
-            if not group.destinations:
-                raise ValueError(f"{type(adapter).__name__} yielded an empty checkpoint load group")
-            if not group.native_keys:
-                raise ValueError(f"{type(adapter).__name__} yielded a checkpoint load group with no model tensors")
-
-            group_keys = set(group.destinations)
-            duplicate_keys = sorted(group_keys & requested_checkpoint_keys)
-            if duplicate_keys:
-                raise ValueError(
-                    f"{type(adapter).__name__} requested {len(duplicate_keys)} checkpoint keys in multiple groups "
-                    f"(examples={duplicate_keys[:5]})"
-                )
-            missing_keys = sorted(group_keys - checkpoint_keys)
-            if missing_keys:
-                raise RuntimeError(
-                    f"Checkpoint {model_path} is missing {len(missing_keys)} keys required by checkpoint load group "
-                    f"{group_count + 1} (examples={missing_keys[:5]})"
-                )
-
-            group_native_keys = set(group.native_keys)
-            duplicate_native_keys = sorted(group_native_keys & loaded_native_keys)
-            if duplicate_native_keys:
-                raise ValueError(
-                    f"{type(adapter).__name__} completed {len(duplicate_native_keys)} model tensors in multiple "
-                    f"groups (examples={duplicate_native_keys[:5]})"
-                )
-            unexpected_native_keys = sorted(group_native_keys - expected_native_keys)
-            if unexpected_native_keys:
-                raise ValueError(
-                    f"{type(adapter).__name__} reported {len(unexpected_native_keys)} unknown model tensors "
-                    f"(examples={unexpected_native_keys[:5]})"
-                )
-
-            group_bytes = sum(estimate_tensor_bytes(tensor) for tensor in group.destinations.values())
-            requested_bytes += group_bytes
-            max_group_bytes = max(max_group_bytes, group_bytes)
-            requested_checkpoint_keys |= group_keys
-
-            storage_started = time.monotonic()
-            # The reader is already bound to ``model_path``. Omitting checkpoint_id prevents DCP from resetting it
-            # for every group, which lets the HF reader reuse parsed safetensors metadata.
-            dcp.load(group.destinations, storage_reader=storage_reader, **process_group_kwargs)
-            storage_seconds += time.monotonic() - storage_started
-
-            install_started = time.monotonic()
-            group.install()
-            install_seconds += time.monotonic() - install_started
-            loaded_native_keys |= group_native_keys
-            group_count += 1
-
-            del group
-
-        if group_count == 0:
-            raise RuntimeError(f"{type(adapter).__name__} opted into checkpoint load groups but yielded no groups")
-
-        missing_native_keys = sorted(expected_native_keys - loaded_native_keys)
-        if missing_native_keys:
-            raise RuntimeError(
-                f"{type(adapter).__name__} checkpoint load groups omitted {len(missing_native_keys)} model tensors "
-                f"(examples={missing_native_keys[:5]})"
-            )
-
-        if model_state.uses_tied_lm_head and not model_state.is_peft:
-            ensure_tied_lm_head(model_part)
-
-        total_seconds = time.monotonic() - started
-        requested_gb = requested_bytes / (1 << 30)
-        max_group_gb = max_group_bytes / (1 << 30)
-        logger.info(
-            "load_model: loaded %.2f GB in %d checkpoint groups over %.2fs "
-            "(%.2f GB/s overall | max group %.2f GB, metadata %.2fs, storage read %.2fs, install %.2fs, "
-            "native keys %d)",
-            requested_gb,
-            group_count,
-            total_seconds,
-            requested_gb / max(total_seconds, 1e-9),
-            max_group_gb,
-            metadata_seconds,
-            storage_seconds,
-            install_seconds,
-            len(loaded_native_keys),
-        )
-
     @torch.no_grad()
     def load_model(
         self,
@@ -978,32 +849,6 @@ def load_model(
             )
             and not self.config.dequantize_base_checkpoint
         )
-        supports_checkpoint_load_groups = (
-            isinstance(state_dict_adapter, StateDictAdapter)
-            and state_dict_adapter.supports_checkpoint_load_groups
-            and not self.config.dequantize_base_checkpoint
-        )
-        if (
-            is_init_step
-            and is_safetensors
-            and len(model_state.model) == 1
-            and world_size == 1
-            and supports_checkpoint_load_groups
-            and not allow_checkpoint_key_subset
-        ):
-            storage_reader = self._get_storage_reader(
-                model_path,
-                key_mapping=None,
-                is_init_step=True,
-                is_safetensors=True,
-            )
-            if storage_reader is None:
-                raise RuntimeError(
-                    f"No safetensors storage reader is available for checkpoint load groups from {model_path}"
-                )
-            self._load_model_in_checkpoint_groups(model_state, state_dict_adapter, model_path, storage_reader)
-            return
-
         single_device_custom_safetensors = (
             is_safetensors and is_custom_model and world_size == 1 and not can_load_without_full_copy
         )
diff --git a/nemo_automodel/components/checkpoint/state_dict_adapter.py b/nemo_automodel/components/checkpoint/state_dict_adapter.py
index 14721d2c56..a1bbb4e90f 100644
--- a/nemo_automodel/components/checkpoint/state_dict_adapter.py
+++ b/nemo_automodel/components/checkpoint/state_dict_adapter.py
@@ -12,44 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 from abc import ABC, abstractmethod
-from collections.abc import Iterator
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 
-from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
-
 if TYPE_CHECKING:
     from torch.distributed.device_mesh import DeviceMesh
 
 
-@dataclass
-class CheckpointLoadGroup:
-    """One set of checkpoint tensors that must be loaded and installed together.
-
-    Attributes:
-        destinations: Checkpoint keys and the tensors DCP should fill. A tensor can use model weight memory directly
-            or be a temporary tensor needed for conversion.
-        native_keys: Model state-dict keys completed by this group. Every key must be ready in the model before
-            :meth:`install` returns.
-    """
-
-    destinations: dict[str, torch.Tensor]
-    native_keys: frozenset[str]
-
-    def install(self) -> None:
-        """Finish converting this group's loaded tensors into model weights.
-
-        The default does nothing because the destinations already use model weight memory. A group that uses temporary
-        tensors overrides this method, performs its conversion, and releases those tensors before the next group.
-        """
-        return None
-
-
 class StateDictAdapter(ABC):
     """Abstract base class for state dict transformations.
 
@@ -59,7 +30,6 @@ class StateDictAdapter(ABC):
 
     _supports_write_through_checkpoint_load: bool = False
     _supports_checkpoint_load_without_full_copy: bool = False
-    _supports_checkpoint_load_groups: bool = False
 
     @property
     def supports_write_through_checkpoint_load(self) -> bool:
@@ -80,164 +50,6 @@ def supports_checkpoint_load_without_full_copy(self) -> bool:
         """
         return self._supports_checkpoint_load_without_full_copy
 
-    @property
-    def supports_checkpoint_load_groups(self) -> bool:
-        """Whether the adapter can load a base checkpoint in one or more memory-safe groups.
-
-        Adapters should opt in only when :meth:`iter_checkpoint_load_groups` covers every tensor returned by
-        :meth:`get_checkpoint_load_state`. Each group must finish updating the model and release any temporary tensors
-        before the next group.
-        """
-        return self._supports_checkpoint_load_groups
-
-    def get_checkpoint_load_state(self, model_part: torch.nn.Module) -> dict[str, torch.Tensor]:
-        """Collect the model tensors that a Hugging Face base checkpoint must populate.
-
-        Args:
-            model_part: Model whose parameters and persistent buffers use their native model layouts. Tensor ranks and
-                axis orders are model-specific. LoRA tensors and runtime-only ``_extra_state`` entries are excluded
-                because they are not stored in the Hugging Face base checkpoint.
-
-        Returns:
-            Mapping from canonical native state-dict names to detached tensors with the same shapes, strides, dtypes,
-            devices, and storage as the model parameters and persistent buffers.
-
-        Raises:
-            ValueError: If a checkpoint-owned state entry is not a tensor or canonical names collide.
-        """
-        checkpoint_state: dict[str, torch.Tensor] = {}
-        for state_name, state_value in model_part.state_dict(keep_vars=True).items():
-            native_name = canonical_parameter_fqn(state_name)
-            if "lora" in native_name or native_name.rsplit(".", 1)[-1] == "_extra_state":
-                continue
-            if not isinstance(state_value, torch.Tensor):
-                raise ValueError(
-                    f"Checkpoint-owned state entry {native_name} is {type(state_value).__name__}, expected Tensor"
-                )
-            if native_name in checkpoint_state:
-                raise ValueError(f"Multiple model state entries resolve to checkpoint key {native_name}")
-            checkpoint_state[native_name] = state_value.detach()
-        return checkpoint_state
-
-    def _checkpoint_load_views(
-        self,
-        native_name: str,
-        native_tensor: torch.Tensor,
-    ) -> dict[str, torch.Tensor]:
-        """Convert one model tensor into checkpoint-layout views of the same storage.
-
-        Args:
-            native_name: Canonical native state-dict name for ``native_tensor``.
-            native_tensor: Model tensor with arbitrary rank and model-specific axis order. Every returned checkpoint
-                tensor must be a view of this tensor's storage; non-contiguous views are allowed.
-
-        Returns:
-            Mapping from Hugging Face checkpoint names to tensors with checkpoint-specific shapes and axis orders.
-            Every tensor uses ``native_tensor`` storage, so DCP writes update the model directly.
-
-        Raises:
-            ValueError: If conversion returns no tensors, returns duplicate names or non-tensors, or allocates storage.
-        """
-        if native_tensor.is_meta:
-            raise ValueError(f"Checkpoint load destination {native_name} is still on the meta device")
-
-        converted = self.convert_single_tensor_to_hf(
-            native_name,
-            native_tensor,
-            quantization=False,
-            for_checkpoint_load=True,
-            track_inplace_load=False,
-        )
-        if not converted:
-            raise ValueError(f"Checkpoint conversion for {native_name} produced no destinations")
-
-        destinations: dict[str, torch.Tensor] = {}
-        source_storage = native_tensor.untyped_storage().data_ptr()
-        for checkpoint_name, checkpoint_tensor in converted:
-            if not isinstance(checkpoint_tensor, torch.Tensor):
-                raise ValueError(
-                    f"Checkpoint destination {checkpoint_name} is {type(checkpoint_tensor).__name__}, expected Tensor"
-                )
-            if checkpoint_tensor.device != native_tensor.device:
-                raise ValueError(
-                    f"Checkpoint destination {checkpoint_name} is on {checkpoint_tensor.device}, "
-                    f"but model tensor {native_name} is on {native_tensor.device}"
-                )
-            if checkpoint_tensor.untyped_storage().data_ptr() != source_storage:
-                raise ValueError(
-                    f"Checkpoint destination {checkpoint_name} does not use model tensor {native_name} storage"
-                )
-            if checkpoint_name in destinations:
-                raise ValueError(f"Checkpoint conversion for {native_name} produced duplicate key {checkpoint_name}")
-            destinations[checkpoint_name] = checkpoint_tensor
-        return destinations
-
-    def _single_checkpoint_load_view(
-        self,
-        native_name: str,
-        native_tensor: torch.Tensor,
-    ) -> tuple[str, torch.Tensor]:
-        """Return the only checkpoint-layout view produced for one model tensor.
-
-        Args:
-            native_name: Canonical native state-dict name for ``native_tensor``.
-            native_tensor: Model tensor with arbitrary rank and model-specific axis order. The returned tensor has the
-                checkpoint layout and uses the same storage.
-
-        Returns:
-            Hugging Face checkpoint name and its tensor view.
-
-        Raises:
-            ValueError: If conversion produces anything other than one checkpoint tensor.
-        """
-        destinations = self._checkpoint_load_views(native_name, native_tensor)
-        if len(destinations) != 1:
-            raise ValueError(
-                f"Checkpoint conversion for {native_name} produced {len(destinations)} destinations, expected one"
-            )
-        return next(iter(destinations.items()))
-
-    def iter_checkpoint_load_groups(
-        self,
-        checkpoint_state: dict[str, torch.Tensor],
-        device_mesh: DeviceMesh | None = None,
-    ) -> Iterator[CheckpointLoadGroup]:
-        """Build one direct-write group for adapters whose conversions return model-storage views.
-
-        Args:
-            checkpoint_state: Canonical native names and model tensors with model-specific shapes and axis orders.
-                Each converted checkpoint tensor must be a view of its corresponding model tensor.
-            device_mesh: Optional device mesh describing the final distributed tensor placements.
-
-        Returns:
-            One group containing every checkpoint-layout view. Model adapters override this method only when a
-            conversion needs temporary tensors that must be installed and released before later groups are loaded.
-
-        Raises:
-            RuntimeError: If the adapter has not opted into checkpoint load groups.
-            ValueError: If a distributed mesh is provided or two model tensors map to the same checkpoint key.
-        """
-        if not self.supports_checkpoint_load_groups:
-            raise RuntimeError(f"{type(self).__name__} does not support checkpoint load groups")
-        if device_mesh is not None:
-            raise ValueError("The default checkpoint load group requires a single-device model")
-
-        destinations: dict[str, torch.Tensor] = {}
-        for native_name, native_tensor in checkpoint_state.items():
-            converted = self._checkpoint_load_views(native_name, native_tensor)
-            duplicate_keys = sorted(converted.keys() & destinations.keys())
-            if duplicate_keys:
-                raise ValueError(
-                    f"Multiple model tensors map to {len(duplicate_keys)} checkpoint keys "
-                    f"(examples={duplicate_keys[:5]})"
-                )
-            destinations.update(converted)
-
-        yield CheckpointLoadGroup(
-            destinations=destinations,
-            native_keys=frozenset(checkpoint_state),
-        )
-
     @abstractmethod
     def to_hf(self, state_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
         """Convert from native model state dict to HuggingFace format.
diff --git a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
index 2806acf3a9..cfd647cd2f 100644
--- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
@@ -41,14 +41,12 @@
 """
 
 import re
-from collections.abc import Iterator
-from dataclasses import dataclass
 from typing import Any, Optional
 
 import torch
 from torch.distributed.device_mesh import DeviceMesh
 
-from nemo_automodel.components.checkpoint.state_dict_adapter import CheckpointLoadGroup, StateDictAdapter
+from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.ling_v2.config import BailingMoeV2Config
 from nemo_automodel.components.moe.config import MoEConfig
@@ -67,77 +65,6 @@
 
 _LAYER_QKV_RE = re.compile(r"^(?P<prefix>(?:.*\.)?layers\.\d+)\.attention\.query_key_value\.weight$")
 _NATIVE_LAYER_QKV_RE = re.compile(r"^(?P<prefix>(?:.*\.)?layers\.\d+)\.self_attn\.(?P<projection>[qkv])_proj\.weight$")
-_STREAMABLE_EXPERT_WEIGHT = re.compile(
-    r"^(?P<root>(?:.*\.)?layers\.\d+\.mlp\.experts)\."
-    r"(?P<projection>gate_and_up_projs|down_projs)$"
-)
-
-
-@dataclass
-class _LingQKVGroupBuilder:
-    """Mutable builder for one fused-QKV checkpoint dependency group.
-
-    Attributes:
-        projections: Mapping from ``q``, ``k``, and ``v`` to a native parameter name and its final tensor. Query
-            tensors have shape [q_hidden, hidden]; key and value tensors have shape [kv_hidden, hidden].
-    """
-
-    projections: dict[str, tuple[str, torch.Tensor]]
-
-
-@dataclass
-class _LingExpertGroupBuilder:
-    """Mutable builder for one split-expert checkpoint dependency group.
-
-    Attributes:
-        destinations: Checkpoint gate/up views of shape [expert_hidden, hidden] and down views of shape
-            [hidden, expert_hidden]. Every view uses the grouped model weights as its memory.
-        native_keys: Names of the final grouped expert parameters completed by ``destinations``.
-        projections: Native grouped projection names already added to the builder.
-    """
-
-    destinations: dict[str, torch.Tensor]
-    native_keys: set[str]
-    projections: set[str]
-
-
-@dataclass
-class _LingFusedQKVLoadGroup(CheckpointLoadGroup):
-    """One fused QKV checkpoint tensor and the three model weights it fills.
-
-    Attributes:
-        destinations: One checkpoint tensor of shape [q_heads * head_dim + 2 * kv_heads * head_dim, hidden]. The
-            tensor is temporary storage released before the next checkpoint group is requested.
-        native_keys: Names of the three final Q, K, and V parameters completed by this group.
-        q_proj: Final query parameter of shape [q_heads * head_dim, hidden], mutated by :meth:`install`.
-        k_proj: Final key parameter of shape [kv_heads * head_dim, hidden], mutated by :meth:`install`.
-        v_proj: Final value parameter of shape [kv_heads * head_dim, hidden], mutated by :meth:`install`.
-    """
-
-    q_proj: torch.Tensor
-    k_proj: torch.Tensor
-    v_proj: torch.Tensor
-
-    @torch.no_grad()
-    def install(self) -> None:
-        """Split the loaded fused tensor and copy its Q, K, and V slices into final parameter storage."""
-        if len(self.destinations) != 1:
-            raise ValueError(f"Ling fused-QKV group expected one checkpoint destination, got {len(self.destinations)}")
-        fused_qkv = next(iter(self.destinations.values()))
-        expected_shape = (self.q_proj.shape[0] + self.k_proj.shape[0] + self.v_proj.shape[0], self.q_proj.shape[1])
-        if tuple(fused_qkv.shape) != expected_shape:
-            raise ValueError(
-                f"Ling fused-QKV destination has shape {tuple(fused_qkv.shape)}, expected {expected_shape}"
-            )
-
-        q_value, k_value, v_value = torch.split(
-            fused_qkv,
-            [self.q_proj.shape[0], self.k_proj.shape[0], self.v_proj.shape[0]],
-            dim=0,
-        )
-        self.q_proj.copy_(q_value)
-        self.k_proj.copy_(k_value)
-        self.v_proj.copy_(v_value)
 
 
 def _rename_hf_to_native(key: str) -> str:
@@ -174,234 +101,9 @@ def __init__(
         self._uses_model_prefix = True
 
     @property
-    def supports_checkpoint_load_groups(self) -> bool:
-        """Whether Ling can load single-device QKV and expert transformations one layer at a time."""
-        return (
-            self.backend.experts in {"torch", "torch_mm"}
-            and self.backend.dispatcher == "torch"
-            and not self.moe_config.expert_bias
-        )
-
-    def _add_expert_checkpoint_views(
-        self,
-        builder: _LingExpertGroupBuilder,
-        *,
-        expert_root: str,
-        native_name: str,
-        projection: str,
-        grouped_weight: torch.Tensor,
-    ) -> None:
-        """Add split HF views of one final grouped expert parameter.
-
-        Args:
-            builder: Mutable dependency group that retains the produced final-storage views.
-            expert_root: Checkpoint prefix ending in ``layers.{layer}.mlp.experts``.
-            native_name: Native name completed by ``grouped_weight``.
-            projection: Either ``gate_and_up_projs`` or ``down_projs``.
-            grouped_weight: Final gate/up tensor of shape [experts, hidden, 2 * expert_hidden] or final down tensor of
-                shape [experts, expert_hidden, hidden]. The emitted checkpoint destinations are non-contiguous views
-                that use this model weight memory.
-
-        Raises:
-            ValueError: If the projection is duplicated or its tensor shape is incompatible with the Ling layout.
-        """
-        if projection in builder.projections:
-            raise ValueError(f"Duplicate Ling expert projection {expert_root}.{projection}")
-
-        n_experts = self.moe_config.n_routed_experts
-        expert_hidden = self.moe_config.moe_inter_dim
-        if projection == "gate_and_up_projs":
-            expected_shape = (n_experts, self.moe_config.dim, 2 * expert_hidden)
-        else:
-            expected_shape = (n_experts, expert_hidden, self.moe_config.dim)
-        if tuple(grouped_weight.shape) != expected_shape:
-            raise ValueError(
-                f"Ling grouped expert parameter {native_name} has shape {tuple(grouped_weight.shape)}, "
-                f"expected {expected_shape}"
-            )
-
-        checkpoint_views = self._checkpoint_load_views(native_name, grouped_weight)
-        duplicate_keys = sorted(checkpoint_views.keys() & builder.destinations.keys())
-        if duplicate_keys:
-            raise ValueError(
-                f"Duplicate Ling expert checkpoint destinations for {expert_root} (examples={duplicate_keys[:5]})"
-            )
-        builder.destinations.update(checkpoint_views)
-
-        builder.native_keys.add(native_name)
-        builder.projections.add(projection)
-
-    @staticmethod
-    def _build_qkv_load_group(prefix: str, builder: _LingQKVGroupBuilder) -> _LingFusedQKVLoadGroup:
-        """Allocate one fused checkpoint tensor backed by three final projection parameters.
-
-        Args:
-            prefix: Native layer prefix ending in ``layers.{layer}``.
-            builder: Q/K/V parameters with shapes [q_hidden, hidden], [kv_hidden, hidden], and [kv_hidden, hidden].
-
-        Returns:
-            Load group with one temporary tensor of shape [q_hidden + 2 * kv_hidden, hidden]. Its installation copies
-            the three row ranges into the final Q/K/V parameters.
-
-        Raises:
-            ValueError: If a projection is missing or the final parameters disagree in rank, shape, device, or dtype.
-        """
-        missing_projections = {"q", "k", "v"} - builder.projections.keys()
-        if missing_projections:
-            raise ValueError(f"Incomplete Ling fused-QKV group {prefix}: missing {sorted(missing_projections)}")
-        q_name, q_proj = builder.projections["q"]
-        k_name, k_proj = builder.projections["k"]
-        v_name, v_proj = builder.projections["v"]
-        if q_proj.ndim != 2 or k_proj.ndim != 2 or v_proj.ndim != 2:
-            raise ValueError(
-                f"Ling fused-QKV group {prefix} requires rank-2 parameters, got "
-                f"Q={tuple(q_proj.shape)}, K={tuple(k_proj.shape)}, V={tuple(v_proj.shape)}"
-            )
-        if k_proj.shape != v_proj.shape or q_proj.shape[1] != k_proj.shape[1]:
-            raise ValueError(
-                f"Ling fused-QKV group {prefix} has incompatible shapes "
-                f"Q={tuple(q_proj.shape)}, K={tuple(k_proj.shape)}, V={tuple(v_proj.shape)}"
-            )
-        if q_proj.device != k_proj.device or q_proj.device != v_proj.device:
-            raise ValueError(f"Ling fused-QKV group {prefix} spans multiple devices")
-        if q_proj.dtype != k_proj.dtype or q_proj.dtype != v_proj.dtype:
-            raise ValueError(f"Ling fused-QKV group {prefix} spans multiple dtypes")
-
-        checkpoint_name = f"{prefix}.attention.query_key_value.weight"
-        fused_destination = torch.empty(
-            (q_proj.shape[0] + k_proj.shape[0] + v_proj.shape[0], q_proj.shape[1]),
-            dtype=q_proj.dtype,
-            device=q_proj.device,
-        )
-        return _LingFusedQKVLoadGroup(
-            destinations={checkpoint_name: fused_destination},
-            native_keys=frozenset({q_name, k_name, v_name}),
-            q_proj=q_proj,
-            k_proj=k_proj,
-            v_proj=v_proj,
-        )
-
-    @staticmethod
-    def _build_expert_load_group(
-        expert_root: str,
-        builder: _LingExpertGroupBuilder,
-        n_experts: int,
-    ) -> CheckpointLoadGroup:
-        """Validate and freeze one complete layer of final-storage expert destinations.
-
-        Args:
-            expert_root: Checkpoint prefix ending in ``layers.{layer}.mlp.experts``.
-            builder: Per-expert gate, up, and down checkpoint views. Gate/up views have shape [expert_hidden, hidden],
-                down views have shape [hidden, expert_hidden], and all views use the grouped model weights as memory.
-            n_experts: Number of checkpoint experts required for the layer.
-
-        Returns:
-            No-op installation group containing exactly three destinations per expert.
-
-        Raises:
-            ValueError: If either grouped projection or any split expert destination is missing.
-        """
-        missing_projections = {"gate_and_up_projs", "down_projs"} - builder.projections
-        if missing_projections:
-            raise ValueError(f"Incomplete Ling expert group {expert_root}: missing {sorted(missing_projections)}")
-        expected_destinations = 3 * n_experts
-        if len(builder.destinations) != expected_destinations:
-            raise ValueError(
-                f"Ling expert group {expert_root} has {len(builder.destinations)} checkpoint destinations, "
-                f"expected {expected_destinations}"
-            )
-        return CheckpointLoadGroup(
-            destinations=builder.destinations,
-            native_keys=frozenset(builder.native_keys),
-        )
-
-    def iter_checkpoint_load_groups(
-        self,
-        checkpoint_state: dict[str, torch.Tensor],
-        device_mesh: DeviceMesh | None = None,
-    ) -> Iterator[CheckpointLoadGroup]:
-        """Yield Ling checkpoint tensors in groups that are installed before the next group is read.
-
-        Args:
-            checkpoint_state: Canonical native names and final Ling model tensors. Attention Q/K/V parameters have
-                shapes [q_hidden, hidden], [kv_hidden, hidden], and [kv_hidden, hidden]. Grouped gate/up experts have
-                shape [experts, hidden, 2 * expert_hidden], and grouped down experts have shape
-                [experts, expert_hidden, hidden].
-            device_mesh: Must be ``None``. Distributed Ling loads continue to use the rank-sharded DCP path.
-
-        Returns:
-            Iterator whose ordinary and expert destinations use model weight memory directly. Each attention group
-            uses one temporary fused-QKV tensor and splits it into Q, K, and V before the iterator advances.
-
-        Raises:
-            RuntimeError: If the configured expert backend does not support checkpoint load groups.
-            ValueError: If a distributed mesh, unsupported tensor layout, incomplete dependency group, or allocating
-                ordinary destination is encountered.
-        """
-        if not self.supports_checkpoint_load_groups:
-            raise RuntimeError(
-                f"{type(self).__name__} does not support checkpoint load groups for experts={self.backend.experts}, "
-                f"dispatcher={self.backend.dispatcher}"
-            )
-        if device_mesh is not None:
-            raise ValueError("Ling checkpoint load groups currently require a single-device model")
-
-        ordinary_destinations: dict[str, torch.Tensor] = {}
-        ordinary_native_keys: set[str] = set()
-        qkv_groups: dict[str, _LingQKVGroupBuilder] = {}
-        expert_groups: dict[str, _LingExpertGroupBuilder] = {}
-        n_experts = self.moe_config.n_routed_experts
-
-        for native_name, state_tensor in checkpoint_state.items():
-            qkv_match = _NATIVE_LAYER_QKV_RE.match(native_name)
-            if qkv_match is not None:
-                prefix = qkv_match.group("prefix")
-                projection = qkv_match.group("projection")
-                builder = qkv_groups.setdefault(prefix, _LingQKVGroupBuilder(projections={}))
-                if projection in builder.projections:
-                    raise ValueError(f"Duplicate Ling {projection.upper()} projection for {prefix}")
-                builder.projections[projection] = (native_name, state_tensor)
-                continue
-
-            expert_match = _STREAMABLE_EXPERT_WEIGHT.match(native_name)
-            if expert_match is not None:
-                expert_root = expert_match.group("root")
-                projection = expert_match.group("projection")
-                builder = expert_groups.setdefault(
-                    expert_root,
-                    _LingExpertGroupBuilder(destinations={}, native_keys=set(), projections=set()),
-                )
-                self._add_expert_checkpoint_views(
-                    builder,
-                    expert_root=expert_root,
-                    native_name=native_name,
-                    projection=projection,
-                    grouped_weight=state_tensor,
-                )
-                continue
-
-            checkpoint_name, checkpoint_destination = self._single_checkpoint_load_view(native_name, state_tensor)
-            if checkpoint_name in ordinary_destinations:
-                raise ValueError(f"Duplicate ordinary Ling checkpoint destination {checkpoint_name}")
-            ordinary_destinations[checkpoint_name] = checkpoint_destination
-            ordinary_native_keys.add(native_name)
-
-        if not qkv_groups:
-            raise ValueError("Ling checkpoint load plan found no fused-QKV layers")
-        if not expert_groups:
-            raise ValueError("Ling checkpoint load plan found no grouped expert layers")
-
-        if ordinary_destinations:
-            yield CheckpointLoadGroup(
-                destinations=ordinary_destinations,
-                native_keys=frozenset(ordinary_native_keys),
-            )
-
-        for prefix, builder in qkv_groups.items():
-            yield self._build_qkv_load_group(prefix, builder)
-
-        for expert_root, builder in expert_groups.items():
-            yield self._build_expert_load_group(expert_root, builder, n_experts)
+    def supports_checkpoint_load_without_full_copy(self) -> bool:
+        """Whether Ling can use DCP with only its small fused-QKV temporary tensors."""
+        return self._supports_write_through_expert_checkpoint_load and not self.moe_config.expert_bias
 
     # ---- HF -> native ----------------------------------------------------
 
diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
index 8d81633e7a..1dedec9962 100644
--- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
@@ -113,23 +113,6 @@ def __init__(
             "model.layers.{}.mixer.experts.{}.down_proj.weight": "model.layers.{}.mixer.experts.down_projs",
         }
 
-    @property
-    def supports_checkpoint_load_groups(self) -> bool:
-        """Whether the shared direct-view group can load this single-device runtime layout.
-
-        TE/DeepEP is configured for distributed runs, but the MoE layer deliberately falls back to ``GroupedExperts``
-        when world size is one. The adapter configuration still says ``experts="te"``, so the legacy converter assumes
-        the grouped tensors are TE stack copies and rebuilds them. The shared checkpoint-group implementation instead
-        verifies that every converted tensor is a view of the actual runtime model storage. Expert bias and other
-        activations keep the full-checkpoint fallback until their layouts have concrete coverage.
-        """
-        return (
-            self.moe_config is not None
-            and self.backend.experts == "te"
-            and self.moe_config.expert_activation == "relu2"
-            and not self.moe_config.expert_bias
-        )
-
     @property
     def _hf_prefix(self) -> str:
         """Return the source checkpoint's public Nemotron-H model prefix."""
diff --git a/nemo_automodel/components/moe/state_dict_mixin.py b/nemo_automodel/components/moe/state_dict_mixin.py
index 5a28467436..fe22bc12cf 100644
--- a/nemo_automodel/components/moe/state_dict_mixin.py
+++ b/nemo_automodel/components/moe/state_dict_mixin.py
@@ -19,6 +19,7 @@
 import torch
 from torch.distributed.device_mesh import DeviceMesh
 
+from nemo_automodel.components.distributed.init_utils import get_world_size_safe
 from nemo_automodel.components.moe.state_dict_utils import (
     create_dtensor_from_local,
     get_expert_range_for_rank_from_mesh,
@@ -62,7 +63,23 @@ def _supports_write_through_expert_checkpoint_load(self) -> bool:
         This covers only the shared expert conversion. A concrete adapter must also verify that all non-expert
         checkpoint tensors load directly before it enables the full-checkpoint fast path.
         """
-        return self.backend.experts != "te" and self.backend.dispatcher != "mok"
+        return self._grouped_expert_storage_is_model_weight and self.backend.dispatcher != "mok"
+
+    @property
+    def _grouped_expert_storage_is_model_weight(self) -> bool:
+        """Whether the runtime grouped expert tensors use the model's parameter storage.
+
+        TE normally exposes virtual grouped tensors backed by ``torch.stack`` copies. The MoE layer instead constructs
+        ordinary ``GroupedExperts`` when an EP dispatcher runs with world size one, so the same TE configuration uses
+        real model parameters in that runtime. Non-EP dispatchers also construct ordinary grouped experts.
+        """
+        if self.backend.experts != "te":
+            return True
+        if self.backend.dispatcher == "mok":
+            return False
+        if self.backend.dispatcher not in {"deepep", "hybridep", "uccl_ep"}:
+            return True
+        return get_world_size_safe() == 1
 
     @property
     def _is_gated_moe(self) -> bool:
@@ -829,7 +846,6 @@ def _convert_single_merged_expert_to_hf_split_experts(
         *,
         prefix_override: str | None = None,
         for_checkpoint_load: bool = False,
-        track_inplace_load: bool = True,
         **kwargs,
     ) -> list[tuple[str, torch.Tensor]]:
         """Convert one grouped expert tensor to Hugging Face's per-expert layout.
@@ -846,8 +862,6 @@ def _convert_single_merged_expert_to_hf_split_experts(
                 outside the main backbone, e.g. ``"mtp."`` for the MTP head.
             for_checkpoint_load: Return views that DCP will completely overwrite. Save/export callers leave this
                 disabled so converted tensors preserve their current values in contiguous storage.
-            track_inplace_load: Record direct-write views for a later ``from_hf`` call. Checkpoint load groups disable
-                this because they write into the model and finish the load without calling ``from_hf``.
             **kwargs: Absorbed for forward-compatibility with base callers
                 that forward arbitrary state-dict kwargs (e.g. ``exclude_key_regex``).
 
@@ -870,7 +884,7 @@ def _convert_single_merged_expert_to_hf_split_experts(
         # down_projs transpose still uses the model's weight memory. Treat the two native tensors independently so
         # MoK rebuilds gate/up after the read while DCP loads down directly into the model.
         backend = getattr(self, "backend", None)
-        grouped_storage_is_model_weight = getattr(backend, "experts", None) != "te"
+        grouped_storage_is_model_weight = self._grouped_expert_storage_is_model_weight
         gate_up_storage_is_model_weight = (
             grouped_storage_is_model_weight and getattr(backend, "dispatcher", None) != "mok"
         )
@@ -906,7 +920,7 @@ def checkpoint_load_destination(view: torch.Tensor, source: torch.Tensor) -> tor
                 and not quantization
                 and gate_up_storage_is_model_weight
             )
-            if inplace_ok and track_inplace_load:
+            if inplace_ok:
                 self._register_inplace_loaded_key(fqn, prefix_override)
 
             result = []
@@ -959,7 +973,7 @@ def checkpoint_load_destination(view: torch.Tensor, source: torch.Tensor) -> tor
                 and not quantization
                 and down_storage_is_model_weight
             )
-            if inplace_ok and track_inplace_load:
+            if inplace_ok:
                 self._register_inplace_loaded_key(fqn, prefix_override)
 
             result = []
diff --git a/tests/unit_tests/checkpoint/test_checkpointing.py b/tests/unit_tests/checkpoint/test_checkpointing.py
index 4e87ae8d9d..510fb2aed4 100644
--- a/tests/unit_tests/checkpoint/test_checkpointing.py
+++ b/tests/unit_tests/checkpoint/test_checkpointing.py
@@ -34,7 +34,6 @@
 from nemo_automodel.components.checkpoint._backports.hf_storage import (
     _DIFFUSERS_INDEX_FN,
     _extract_file_index_with_status,
-    _HuggingFaceStorageReader,
     get_fqn_to_dtype_mapping,
     get_fqn_to_file_index_mapping,
 )
@@ -64,7 +63,7 @@
     is_cloud_path,
     save_config,
 )
-from nemo_automodel.components.checkpoint.state_dict_adapter import CheckpointLoadGroup, StateDictAdapter
+from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
 from nemo_automodel.components.checkpoint.stateful_wrappers import (
     ModelState,
     OptimizerState,
@@ -1697,7 +1696,7 @@ def test_single_device_custom_model_uses_fast_path(self, mock_load_full, mock_lo
     @patch("nemo_automodel.components.checkpoint.checkpointing._is_safetensors_checkpoint", return_value=True)
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype")
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model")
-    @pytest.mark.parametrize("load_capability", ["write_through", "without_full_copy", "groups"])
+    @pytest.mark.parametrize("load_capability", ["write_through", "without_full_copy"])
     @pytest.mark.parametrize("dequantize_base_checkpoint", [False, True])
     def test_single_device_adapter_without_full_copy_routes_by_quantization(
         self,
@@ -1716,7 +1715,6 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
         model.state_dict_adapter = MagicMock(spec=StateDictAdapter)
         model.state_dict_adapter.supports_write_through_checkpoint_load = load_capability == "write_through"
         model.state_dict_adapter.supports_checkpoint_load_without_full_copy = load_capability == "without_full_copy"
-        model.state_dict_adapter.supports_checkpoint_load_groups = load_capability == "groups"
         mock_state_dict = {"layer.weight": torch.randn(4, 4), "layer.bias": torch.randn(4)}
         mock_load_hf.return_value = mock_state_dict
 
@@ -1737,7 +1735,6 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
                 side_effect=lambda model_part, state_dict, **kwargs: state_dict,
             ),
             patch.object(checkpointer, "_get_storage_reader", return_value=MagicMock()),
-            patch.object(checkpointer, "_load_model_in_checkpoint_groups") as mock_grouped_load,
             patch.object(checkpointer, "_do_load", return_value=mock_state_dict) as mock_dcp_load,
         ):
             mock_model_state = mock_model_state_cls.return_value
@@ -1749,178 +1746,14 @@ def test_single_device_adapter_without_full_copy_routes_by_quantization(
         if dequantize_base_checkpoint:
             mock_load_full.assert_called_once()
             mock_load_hf.assert_called_once()
-            mock_grouped_load.assert_not_called()
-            mock_dcp_load.assert_not_called()
-        elif load_capability == "groups":
-            mock_load_full.assert_not_called()
-            mock_load_hf.assert_not_called()
-            mock_grouped_load.assert_called_once()
             mock_dcp_load.assert_not_called()
         else:
             mock_load_full.assert_not_called()
             mock_load_hf.assert_not_called()
-            mock_grouped_load.assert_not_called()
             mock_dcp_load.assert_called_once()
             assert "load_model:" in caplog.text
 
 
-def test_checkpoint_group_executor_installs_before_requesting_next_group(caplog):
-    checkpointer = TestLoadModelCustomModelGuard()._make_checkpointer()
-    model_state = SimpleNamespace(
-        model=[torch.nn.Linear(2, 2)],
-        uses_tied_lm_head=False,
-        is_peft=False,
-    )
-    adapter = MagicMock(spec=StateDictAdapter)
-    reader = MagicMock()
-    reader.read_metadata.return_value = SimpleNamespace(
-        state_dict_metadata={"checkpoint.0": object(), "checkpoint.1": object()}
-    )
-
-    active_groups = 0
-    peak_active_groups = 0
-    installed: dict[int, torch.Tensor] = {}
-
-    checkpoint_state = {
-        "native.0": torch.empty(3),
-        "native.1": torch.empty(3),
-    }
-    adapter.get_checkpoint_load_state.return_value = checkpoint_state
-
-    def iter_groups(load_state, device_mesh=None):
-        """Yield groups after verifying the prior group's destination was installed.
-
-        Args:
-            load_state: Native model tensors of shape [elements] that the groups must cover.
-            device_mesh: Optional distributed mesh, expected to be ``None`` for this single-process test.
-
-        Returns:
-            Iterator of groups containing destination tensors of shape [elements].
-        """
-        nonlocal active_groups, peak_active_groups
-        assert load_state is checkpoint_state
-        assert device_mesh is None
-        for group_id in range(2):
-            assert active_groups == 0
-            active_groups += 1
-            peak_active_groups = max(peak_active_groups, active_groups)
-            destination = torch.empty(3)
-            group = CheckpointLoadGroup(
-                destinations={f"checkpoint.{group_id}": destination},
-                native_keys=frozenset({f"native.{group_id}"}),
-            )
-
-            def install(group_id=group_id, destination=destination):
-                """Record one loaded destination tensor.
-
-                Args:
-                    group_id: Sequential group identifier.
-                    destination: Loaded tensor of shape [elements].
-                """
-                nonlocal active_groups
-                installed[group_id] = destination.clone()
-                active_groups -= 1
-
-            group.install = install
-            yield group
-            assert group_id in installed
-
-    def load_group(destinations, **kwargs):
-        """Fill one checkpoint destination mapping as DCP would.
-
-        Args:
-            destinations: Mapping containing one tensor of shape [elements].
-            **kwargs: DCP reader arguments.
-        """
-        assert kwargs["storage_reader"] is reader
-        key = next(iter(destinations))
-        destinations[key].fill_(int(key.rsplit(".", 1)[1]) + 1)
-
-    adapter.iter_checkpoint_load_groups.side_effect = iter_groups
-    caplog.set_level(logging.INFO)
-    with patch("nemo_automodel.components.checkpoint.checkpointing.dcp.load", side_effect=load_group) as mock_load:
-        checkpointer._load_model_in_checkpoint_groups(model_state, adapter, "/checkpoint", reader)
-
-    assert peak_active_groups == 1
-    assert active_groups == 0
-    assert torch.equal(installed[0], torch.ones(3))
-    assert torch.equal(installed[1], torch.full((3,), 2.0))
-    assert mock_load.call_count == 2
-    reader.read_metadata.assert_called_once()
-    assert "loaded" in caplog.text
-    assert "2 checkpoint groups" in caplog.text
-
-
-def test_checkpoint_group_executor_reuses_hf_metadata_and_loads_exact_values(tmp_path):
-    checkpoint = {
-        "checkpoint.0": torch.arange(6, dtype=torch.float32).reshape(2, 3),
-        "checkpoint.1": torch.arange(4, dtype=torch.bfloat16).reshape(2, 2),
-    }
-    save_file(checkpoint, tmp_path / "model.safetensors")
-
-    backing_tensors = {
-        key: torch.empty(value.shape[1], value.shape[0], dtype=value.dtype) for key, value in checkpoint.items()
-    }
-    destinations = {key: backing.transpose(0, 1) for key, backing in backing_tensors.items()}
-    assert all(not destination.is_contiguous() for destination in destinations.values())
-    groups = [
-        CheckpointLoadGroup(
-            destinations={key: destinations[key]},
-            native_keys=frozenset({f"native.{index}"}),
-        )
-        for index, key in enumerate(checkpoint)
-    ]
-    adapter = MagicMock(spec=StateDictAdapter)
-    adapter.get_checkpoint_load_state.return_value = {
-        f"native.{index}": backing_tensors[key] for index, key in enumerate(checkpoint)
-    }
-    adapter.iter_checkpoint_load_groups.return_value = iter(groups)
-    model_state = SimpleNamespace(
-        model=[torch.nn.Linear(2, 2)],
-        uses_tied_lm_head=False,
-        is_peft=False,
-    )
-    reader = _HuggingFaceStorageReader(str(tmp_path))
-    checkpointer = TestLoadModelCustomModelGuard()._make_checkpointer()
-
-    with patch.object(reader.fs, "ls", wraps=reader.fs.ls) as mock_list_files:
-        checkpointer._load_model_in_checkpoint_groups(model_state, adapter, str(tmp_path), reader)
-
-    assert mock_list_files.call_count == 1
-    for key, expected in checkpoint.items():
-        torch.testing.assert_close(destinations[key], expected)
-
-
-def test_checkpoint_group_executor_rejects_incomplete_model_coverage():
-    checkpointer = TestLoadModelCustomModelGuard()._make_checkpointer()
-    model_state = SimpleNamespace(
-        model=[torch.nn.Linear(2, 2)],
-        uses_tied_lm_head=False,
-        is_peft=False,
-    )
-    adapter = MagicMock(spec=StateDictAdapter)
-    adapter.get_checkpoint_load_state.return_value = {
-        "native.loaded": torch.empty(2),
-        "native.omitted": torch.empty(2),
-    }
-    adapter.iter_checkpoint_load_groups.return_value = iter(
-        [
-            CheckpointLoadGroup(
-                destinations={"checkpoint.loaded": torch.empty(2)},
-                native_keys=frozenset({"native.loaded"}),
-            )
-        ]
-    )
-    reader = MagicMock()
-    reader.read_metadata.return_value = SimpleNamespace(state_dict_metadata={"checkpoint.loaded": object()})
-
-    with (
-        patch("nemo_automodel.components.checkpoint.checkpointing.dcp.load"),
-        pytest.raises(RuntimeError, match="omitted 1 model tensors.*native.omitted"),
-    ):
-        checkpointer._load_model_in_checkpoint_groups(model_state, adapter, "/checkpoint", reader)
-
-
 class TestLoadModelCheckpointKeySubset:
     """Test allow_checkpoint_key_subset support for torch_save exports."""
 
diff --git a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
index e642609199..cc725dbb06 100644
--- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
+++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -67,23 +68,6 @@ def convert_single_tensor_to_hf(self, fqn, tensor, **kwargs):
         return [(fqn, tensor)]
 
 
-class _TransposeLoadAdapter(_AllocatingKeysAdapter):
-    """Test adapter whose checkpoint layout transposes each model tensor."""
-
-    _supports_checkpoint_load_groups = True
-
-    def convert_single_tensor_to_hf(self, fqn, tensor, **kwargs):
-        self.checkpoint_conversion_kwargs = kwargs
-        return [(f"checkpoint.{fqn}", tensor.transpose(0, 1))]
-
-
-class _CopyingLoadAdapter(_TransposeLoadAdapter):
-    """Invalid direct-load adapter that allocates a new checkpoint tensor."""
-
-    def convert_single_tensor_to_hf(self, fqn, tensor, **kwargs):
-        return [(f"checkpoint.{fqn}", tensor.transpose(0, 1).contiguous())]
-
-
 def test_get_hf_state_dict_keys_uses_shape_only_tensors() -> None:
     adapter = _AllocatingKeysAdapter()
     state_dict = {
@@ -98,47 +82,6 @@ def test_get_hf_state_dict_keys_uses_shape_only_tensors() -> None:
     assert list(state_dict) == ["q.weight", "k.weight"]
 
 
-def test_shared_checkpoint_group_uses_model_views_and_excludes_non_base_state() -> None:
-    class _ModelWithRuntimeState(torch.nn.Module):
-        def __init__(self) -> None:
-            super().__init__()
-            self.weight = torch.nn.Parameter(torch.zeros(2, 3))
-            self.lora_A = torch.nn.Parameter(torch.ones(1, 3))
-            self.register_buffer("running_scale", torch.zeros(2, 1))
-
-        def get_extra_state(self) -> dict[str, bool]:
-            return {"runtime_only": True}
-
-        def set_extra_state(self, state: object) -> None:
-            del state
-
-    model = _ModelWithRuntimeState()
-    adapter = _TransposeLoadAdapter()
-
-    checkpoint_state = adapter.get_checkpoint_load_state(model)
-    groups = list(adapter.iter_checkpoint_load_groups(checkpoint_state))
-
-    assert set(checkpoint_state) == {"weight", "running_scale"}
-    assert len(groups) == 1
-    assert groups[0].native_keys == frozenset(checkpoint_state)
-    assert adapter.checkpoint_conversion_kwargs["for_checkpoint_load"] is True
-    assert adapter.checkpoint_conversion_kwargs["track_inplace_load"] is False
-    weight_destination = groups[0].destinations["checkpoint.weight"]
-    assert tuple(weight_destination.shape) == (3, 2)
-    assert not weight_destination.is_contiguous()
-    assert weight_destination.untyped_storage().data_ptr() == model.weight.untyped_storage().data_ptr()
-    weight_destination.fill_(4)
-    torch.testing.assert_close(model.weight, torch.full_like(model.weight, 4))
-
-
-def test_shared_checkpoint_group_rejects_allocating_conversion() -> None:
-    source = torch.zeros(2, 3)
-    adapter = _CopyingLoadAdapter()
-
-    with pytest.raises(ValueError, match="does not use model tensor weight storage"):
-        list(adapter.iter_checkpoint_load_groups({"weight": source}))
-
-
 def _assert_destinations_write_through(
     adapter: StateDictAdapter,
     state_dict: dict[str, torch.Tensor],
@@ -234,7 +177,13 @@ def test_write_through_grouped_adapters_preserve_non_expert_storage_and_require_
     )
 
     adapter.backend = SimpleNamespace(experts="te", dispatcher="torch")
-    assert adapter.supports_write_through_checkpoint_load is False
+    assert adapter.supports_write_through_checkpoint_load is True
+
+    adapter.backend = SimpleNamespace(experts="te", dispatcher="deepep")
+    with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+        assert adapter.supports_write_through_checkpoint_load is False
+    with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
+        assert adapter.supports_write_through_checkpoint_load is True
 
     adapter.backend = SimpleNamespace(experts="torch", dispatcher="mok")
     assert adapter.supports_write_through_checkpoint_load is False
@@ -269,6 +218,15 @@ def test_gemma4_moe_adapter_loads_without_a_full_checkpoint_copy():
     assert adapter.supports_checkpoint_load_without_full_copy is True
 
 
+def test_ling_adapter_loads_without_a_full_checkpoint_copy():
+    adapter = object.__new__(BailingMoeV2StateDictAdapter)
+    adapter.moe_config = SimpleNamespace(expert_bias=False)
+    adapter.backend = SimpleNamespace(experts="torch_mm", dispatcher="torch")
+
+    assert adapter.supports_write_through_checkpoint_load is False
+    assert adapter.supports_checkpoint_load_without_full_copy is True
+
+
 def test_dense_grouped_adapter_does_not_require_an_expert_backend():
     adapter = object.__new__(NemotronV3StateDictAdapter)
     adapter.moe_config = None
diff --git a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
index 5e81cdc0f5..f4447bbba6 100644
--- a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
+++ b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
@@ -12,14 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import weakref
-
 import pytest
 import torch
 
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.ling_v2.config import BailingMoeV2Config
-from nemo_automodel.components.models.ling_v2.model import BailingMoeV2ForCausalLM
 from nemo_automodel.components.models.ling_v2.state_dict_adapter import BailingMoeV2StateDictAdapter
 from nemo_automodel.components.moe.config import MoEConfig
 
@@ -227,101 +224,65 @@ def test_hf_to_native_to_hf_is_lossless_for_qkv_path(self, adapter, config):
         torch.testing.assert_close(roundtripped["model.word_embeddings.weight"], hf_sd["model.word_embeddings.weight"])
 
 
-class TestCheckpointLoadGroups:
-    def test_loads_fused_qkv_through_one_temporary_and_experts_through_final_views(
-        self,
-        adapter,
-        config,
-        moe_config,
-        backend_config,
-    ):
-        model = BailingMoeV2ForCausalLM(config=config, moe_config=moe_config, backend=backend_config)
-        model_state = adapter.get_checkpoint_load_state(model)
-        groups = adapter.iter_checkpoint_load_groups(model_state)
-        loaded_native_keys = set()
-
-        ordinary_group = next(groups)
-        loaded_native_keys.update(ordinary_group.native_keys)
-        assert "model.word_embeddings.weight" in ordinary_group.destinations
-        assert ordinary_group.destinations["model.word_embeddings.weight"].untyped_storage().data_ptr() == (
-            model_state["model.embed_tokens.weight"].untyped_storage().data_ptr()
+class TestCheckpointLoadWithoutFullCopy:
+    def test_uses_only_fused_qkv_temporaries(self, adapter, config, moe_config):
+        q_size = config.num_attention_heads * config.head_dim
+        kv_size = config.num_key_value_heads * config.head_dim
+        q_proj = torch.randn(q_size, config.hidden_size)
+        k_proj = torch.randn(kv_size, config.hidden_size)
+        v_proj = torch.randn(kv_size, config.hidden_size)
+        grouped_gate_up = torch.randn(
+            moe_config.n_routed_experts,
+            moe_config.dim,
+            2 * moe_config.moe_inter_dim,
         )
-        correction_bias = ordinary_group.destinations["model.layers.1.mlp.gate.expert_bias"]
-        correction_bias.fill_(3.0)
-        torch.testing.assert_close(
-            model_state["model.layers.1.mlp.gate.e_score_correction_bias"],
-            torch.full_like(model_state["model.layers.1.mlp.gate.e_score_correction_bias"], 3.0),
+        grouped_down = torch.randn(
+            moe_config.n_routed_experts,
+            moe_config.moe_inter_dim,
+            moe_config.dim,
         )
+        native = {
+            "model.layers.1.self_attn.q_proj.weight": q_proj,
+            "model.layers.1.self_attn.k_proj.weight": k_proj,
+            "model.layers.1.self_attn.v_proj.weight": v_proj,
+            "model.layers.1.mlp.experts.gate_and_up_projs": grouped_gate_up,
+            "model.layers.1.mlp.experts.down_projs": grouped_down,
+        }
 
-        pending_group = None
-        for layer_id in range(config.num_hidden_layers):
-            qkv_group = next(groups) if pending_group is None else pending_group
-            pending_group = None
-            checkpoint_name = f"model.layers.{layer_id}.attention.query_key_value.weight"
-            assert set(qkv_group.destinations) == {checkpoint_name}
-            loaded_native_keys.update(qkv_group.native_keys)
-
-            fused_qkv = qkv_group.destinations[checkpoint_name]
-            fused_values = torch.arange(fused_qkv.numel(), dtype=torch.float32).reshape(fused_qkv.shape)
-            fused_values = fused_values.to(fused_qkv.dtype)
-            fused_qkv.copy_(fused_values)
-            temporary_ref = weakref.ref(fused_qkv)
-            qkv_group.install()
+        destinations = adapter.to_hf(native, for_checkpoint_load=True)
 
-            q_size = config.num_attention_heads * config.head_dim
-            kv_size = config.num_key_value_heads * config.head_dim
-            torch.testing.assert_close(
-                model_state[f"model.layers.{layer_id}.self_attn.q_proj.weight"], fused_values[:q_size]
-            )
-            torch.testing.assert_close(
-                model_state[f"model.layers.{layer_id}.self_attn.k_proj.weight"],
-                fused_values[q_size : q_size + kv_size],
-            )
-            torch.testing.assert_close(
-                model_state[f"model.layers.{layer_id}.self_attn.v_proj.weight"], fused_values[q_size + kv_size :]
-            )
+        fused_qkv = destinations["model.layers.1.attention.query_key_value.weight"]
+        assert fused_qkv.shape == (q_size + 2 * kv_size, config.hidden_size)
+        assert fused_qkv.numel() == q_proj.numel() + k_proj.numel() + v_proj.numel()
+        qkv_storage = {
+            q_proj.untyped_storage().data_ptr(),
+            k_proj.untyped_storage().data_ptr(),
+            v_proj.untyped_storage().data_ptr(),
+        }
+        assert fused_qkv.untyped_storage().data_ptr() not in qkv_storage
 
-            del fused_qkv, qkv_group
-            pending_group = next(groups)
-            assert temporary_ref() is None
-
-        expert_group = pending_group
-        assert expert_group is not None
-        loaded_native_keys.update(expert_group.native_keys)
-        gate_key = "model.layers.1.mlp.experts.0.gate_proj.weight"
-        gate_destination = expert_group.destinations[gate_key]
-        gate_destination.fill_(7.0)
-        expert_group.install()
-        grouped_gate_up = model_state["model.layers.1.mlp.experts.gate_and_up_projs"]
-        torch.testing.assert_close(
-            grouped_gate_up[0, :, : config.moe_intermediate_size],
-            torch.full_like(grouped_gate_up[0, :, : config.moe_intermediate_size], 7.0),
-        )
+        gate_destination = destinations["model.layers.1.mlp.experts.0.gate_proj.weight"]
+        up_destination = destinations["model.layers.1.mlp.experts.0.up_proj.weight"]
+        down_destination = destinations["model.layers.1.mlp.experts.0.down_proj.weight"]
+        assert gate_destination.untyped_storage().data_ptr() == grouped_gate_up.untyped_storage().data_ptr()
+        assert up_destination.untyped_storage().data_ptr() == grouped_gate_up.untyped_storage().data_ptr()
+        assert down_destination.untyped_storage().data_ptr() == grouped_down.untyped_storage().data_ptr()
 
-        with pytest.raises(StopIteration):
-            next(groups)
+    def test_capability_matches_runtime_expert_storage(self, adapter, monkeypatch):
+        assert adapter.supports_checkpoint_load_without_full_copy is True
 
-        assert loaded_native_keys == set(model_state)
+        adapter.backend.experts = "te"
+        adapter.backend.dispatcher = "deepep"
+        monkeypatch.setattr("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", lambda: 1)
+        assert adapter.supports_checkpoint_load_without_full_copy is True
 
-    def test_checkpoint_load_groups_guard_backend_and_distributed_mesh(self, adapter, config, moe_config):
-        assert adapter.supports_checkpoint_load_groups is True
-        with pytest.raises(ValueError, match="single-device"):
-            next(adapter.iter_checkpoint_load_groups({}, device_mesh=object()))
+        monkeypatch.setattr("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", lambda: 8)
+        assert adapter.supports_checkpoint_load_without_full_copy is False
 
-        unsupported_backend = BackendConfig(
-            attn="sdpa",
-            linear="torch",
-            rms_norm="torch",
-            experts="te",
-            dispatcher="deepep",
-            enable_hf_state_dict_adapter=False,
-        )
-        unsupported_adapter = BailingMoeV2StateDictAdapter(
-            config=config,
-            moe_config=moe_config,
-            backend=unsupported_backend,
-            dtype=torch.float32,
-        )
-        assert unsupported_adapter.supports_checkpoint_load_groups is False
-        with pytest.raises(RuntimeError, match="does not support checkpoint load groups"):
-            next(unsupported_adapter.iter_checkpoint_load_groups({}))
+        adapter.backend.experts = "torch_mm"
+        adapter.backend.dispatcher = "mok"
+        assert adapter.supports_checkpoint_load_without_full_copy is False
+
+        adapter.backend.dispatcher = "torch"
+        adapter.moe_config.expert_bias = True
+        assert adapter.supports_checkpoint_load_without_full_copy is False
diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
index 5d68a35a8f..6f25e5a542 100644
--- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
+++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
@@ -120,91 +120,55 @@ def test_write_through_load_capability_requires_aliasing_expert_backend(self, co
         assert adapter.supports_write_through_checkpoint_load is True
 
         adapter.backend.experts = "te"
-        assert adapter.supports_write_through_checkpoint_load is False
+        adapter.backend.dispatcher = "deepep"
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+            assert adapter.supports_write_through_checkpoint_load is False
+
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
+            assert adapter.supports_write_through_checkpoint_load is True
 
         adapter.backend.experts = "gmm"
         adapter.backend.dispatcher = "mok"
         assert adapter.supports_write_through_checkpoint_load is False
 
-    def test_te_single_device_fallback_streams_through_grouped_parameter_views(self, config, moe_config, backend):
+    def test_te_single_device_fallback_loads_through_grouped_parameter_views(self, config, moe_config, backend):
         backend.experts = "te"
+        backend.dispatcher = "deepep"
         adapter = NemotronV3StateDictAdapter(config, moe_config, backend)
-
-        class _FakeSingleDeviceFallbackExperts(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.gate_and_up_projs = torch.nn.Parameter(
-                    torch.zeros(moe_config.n_routed_experts, moe_config.expert_dim, moe_config.moe_inter_dim)
-                )
-                self.down_projs = torch.nn.Parameter(
-                    torch.zeros(moe_config.n_routed_experts, moe_config.moe_inter_dim, moe_config.expert_dim)
-                )
-
-            def get_extra_state(self) -> dict[str, bool]:
-                return {"runtime_only": True}
-
-            def set_extra_state(self, state: object) -> None:
-                del state
-
-        model = torch.nn.Module()
-        model.model = torch.nn.Module()
-        model.model.embed_tokens = torch.nn.Embedding(8, moe_config.expert_dim)
-        model.model.embed_tokens.lora_A = torch.nn.Linear(moe_config.expert_dim, 2, bias=False)
-        lora_initial = model.model.embed_tokens.lora_A.weight.detach().clone()
-        layer = torch.nn.Module()
-        layer.mixer = torch.nn.Module()
-        layer.mixer.gate = torch.nn.Module()
-        layer.mixer.gate.register_buffer(
-            "e_score_correction_bias",
-            torch.zeros(moe_config.n_routed_experts, dtype=torch.float32),
+        gate_and_up = torch.zeros(
+            moe_config.n_routed_experts,
+            moe_config.expert_dim,
+            moe_config.moe_inter_dim,
+        )
+        down = torch.zeros(
+            moe_config.n_routed_experts,
+            moe_config.moe_inter_dim,
+            moe_config.expert_dim,
         )
-        layer.mixer.experts = _FakeSingleDeviceFallbackExperts()
-        experts = layer.mixer.experts
-        model.model.layers = torch.nn.ModuleList([layer])
-
-        checkpoint_state = adapter.get_checkpoint_load_state(model)
-        groups = list(adapter.iter_checkpoint_load_groups(checkpoint_state))
-        assert adapter.supports_checkpoint_load_groups is True
-        assert len(groups) == 1
-        assert groups[0].native_keys == frozenset(checkpoint_state)
-
-        all_destinations = {key: value for group in groups for key, value in group.destinations.items()}
-        assert set(all_destinations) == {
-            "backbone.embeddings.weight",
-            "backbone.layers.0.mixer.gate.e_score_correction_bias",
-            *{
-                f"backbone.layers.0.mixer.experts.{expert_id}.{projection}.weight"
-                for expert_id in range(moe_config.n_routed_experts)
-                for projection in ("up_proj", "down_proj")
-            },
+
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
+            destinations = adapter.to_hf(
+                {
+                    "model.layers.0.mixer.experts.gate_and_up_projs": gate_and_up,
+                    "model.layers.0.mixer.experts.down_projs": down,
+                },
+                for_checkpoint_load=True,
+            )
+            assert adapter.supports_write_through_checkpoint_load is True
+
+        assert set(destinations) == {
+            f"backbone.layers.0.mixer.experts.{expert_id}.{projection}.weight"
+            for expert_id in range(moe_config.n_routed_experts)
+            for projection in ("up_proj", "down_proj")
         }
         for expert_id in range(moe_config.n_routed_experts):
-            up_destination = all_destinations[f"backbone.layers.0.mixer.experts.{expert_id}.up_proj.weight"]
-            down_destination = all_destinations[f"backbone.layers.0.mixer.experts.{expert_id}.down_proj.weight"]
-            assert up_destination.untyped_storage().data_ptr() == experts.gate_and_up_projs.untyped_storage().data_ptr()
-            assert down_destination.untyped_storage().data_ptr() == experts.down_projs.untyped_storage().data_ptr()
+            up_destination = destinations[f"backbone.layers.0.mixer.experts.{expert_id}.up_proj.weight"]
+            down_destination = destinations[f"backbone.layers.0.mixer.experts.{expert_id}.down_proj.weight"]
+            assert up_destination.untyped_storage().data_ptr() == gate_and_up.untyped_storage().data_ptr()
+            assert down_destination.untyped_storage().data_ptr() == down.untyped_storage().data_ptr()
             assert not up_destination.is_contiguous()
             assert not down_destination.is_contiguous()
 
-        assert not any("lora" in key for key in all_destinations)
-        assert not any("_extra_state" in key for key in all_destinations)
-        correction_bias = all_destinations["backbone.layers.0.mixer.gate.e_score_correction_bias"]
-        correction_bias.fill_(2)
-        torch.testing.assert_close(
-            layer.mixer.gate.e_score_correction_bias,
-            torch.full_like(layer.mixer.gate.e_score_correction_bias, 2),
-        )
-        for checkpoint_name, destination in all_destinations.items():
-            if ".experts." in checkpoint_name:
-                destination.fill_(3)
-        groups[0].install()
-        assert torch.count_nonzero(experts.gate_and_up_projs) == experts.gate_and_up_projs.numel()
-        assert torch.count_nonzero(experts.down_projs) == experts.down_projs.numel()
-        torch.testing.assert_close(model.model.embed_tokens.lora_A.weight, lora_initial)
-
-        moe_config.expert_bias = True
-        assert adapter.supports_checkpoint_load_groups is False
-
     def test_from_hf_map_structure(self, config, moe_config, backend):
         """Test from_hf_map structure."""
         adapter = NemotronV3StateDictAdapter(config, moe_config, backend)
diff --git a/tests/unit_tests/moe/test_state_dict_mixin.py b/tests/unit_tests/moe/test_state_dict_mixin.py
index 3e67c7a254..307610a015 100644
--- a/tests/unit_tests/moe/test_state_dict_mixin.py
+++ b/tests/unit_tests/moe/test_state_dict_mixin.py
@@ -1053,7 +1053,14 @@ def test_expert_write_through_capability_matches_grouped_storage_aliasing(self):
         assert mixin._supports_write_through_expert_checkpoint_load is True
 
         mixin.backend.experts = "te"
-        assert mixin._supports_write_through_expert_checkpoint_load is False
+        mixin.backend.dispatcher = "deepep"
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+            assert mixin._supports_write_through_expert_checkpoint_load is False
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
+            assert mixin._supports_write_through_expert_checkpoint_load is True
+
+        mixin.backend.dispatcher = "torch"
+        assert mixin._supports_write_through_expert_checkpoint_load is True
 
         mixin.backend.experts = "gmm"
         mixin.backend.dispatcher = "mok"
@@ -1178,13 +1185,15 @@ def test_inplace_load_skips_when_backend_experts_is_te_gate_and_up(self):
         # the copy_ would write the throwaway and the experts would never be loaded.
         mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=512)
         mixin.backend.experts = "te"
+        mixin.backend.dispatcher = "deepep"
         local_storage = torch.randn(2, 1024, 1024)
         splits = [local_storage[i] for i in range(2)]
         mock_dtensor = Mock()
 
-        result = self._run_inplace_conversion(
-            mixin, "model.layers.0.mlp.experts.gate_and_up_projs", mock_dtensor, splits
-        )
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+            result = self._run_inplace_conversion(
+                mixin, "model.layers.0.mlp.experts.gate_and_up_projs", mock_dtensor, splits
+            )
 
         assert result is not None
         converted = dict(result)
@@ -1225,6 +1234,7 @@ def test_mok_noncontiguous_checkpoint_views_are_reused(self):
     def test_te_checkpoint_layout_views_are_reused(self):
         mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=3)
         mixin.backend.experts = "te"
+        mixin.backend.dispatcher = "deepep"
         gate_up_storage = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4)
         down_storage = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3)
 
@@ -1234,6 +1244,7 @@ def test_te_checkpoint_layout_views_are_reused(self):
         virtual_down = down_storage.transpose(-1, -2)
 
         with (
+            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
             patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
             patch("torch.empty_like") as empty_like,
         ):
@@ -1272,12 +1283,14 @@ def test_te_checkpoint_layout_cuda_peak_stays_within_one_buffer(self):
         """Catch allocating a second model-sized destination for TE checkpoint views."""
         mixin = MockMoEStateDictMixin(n_experts=8, inter_dim=1024, dtype=torch.bfloat16)
         mixin.backend.experts = "te"
+        mixin.backend.dispatcher = "deepep"
         # The TE stack is 64 MiB. Reusing its checkpoint-layout views fits this 70 MiB budget; allocating blank
         # destinations of the same total size would double live allocation to 128 MiB.
         gate_up_storage = torch.empty((8, 2048, 2048), dtype=torch.bfloat16, device="cuda")
         virtual_gate_up = gate_up_storage.transpose(-1, -2)
 
         with (
+            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
             patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
             patch("nemo_automodel.components.moe.state_dict_mixin.gc.collect") as collect,
             patch("torch.cuda.empty_cache") as empty_cache,
@@ -1297,10 +1310,14 @@ def test_te_checkpoint_layout_cuda_peak_stays_within_one_buffer(self):
     def test_temporary_checkpoint_views_round_trip_loaded_values(self):
         mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=3)
         mixin.backend.experts = "te"
+        mixin.backend.dispatcher = "deepep"
         initialized_gate_up = torch.full((2, 4, 6), 99.0)
         initialized_down = torch.full((2, 3, 4), 99.0)
 
-        with patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False):
+        with (
+            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
+            patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
+        ):
             destinations = dict(
                 mixin._convert_single_merged_expert_to_hf_split_experts(
                     "model.layers.0.mlp.experts.gate_and_up_projs",
@@ -1343,6 +1360,7 @@ def test_inplace_load_skips_when_backend_experts_is_te_down_projs(self):
         # Same non-aliasing reason as the gate_and_up case, for the down_projs branch.
         mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=512)
         mixin.backend.experts = "te"
+        mixin.backend.dispatcher = "deepep"
         local_storage = torch.randn(2, 512, 1024)
         splits = [local_storage[i] for i in range(2)]
         mock_dtensor = Mock(spec=["ndim", "shape", "is_meta"])
@@ -1350,7 +1368,8 @@ def test_inplace_load_skips_when_backend_experts_is_te_down_projs(self):
         mock_dtensor.shape = (2, 512, 1024)
         mock_dtensor.is_meta = False
 
-        result = self._run_inplace_conversion(mixin, "model.layers.3.mlp.experts.down_projs", mock_dtensor, splits)
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+            result = self._run_inplace_conversion(mixin, "model.layers.3.mlp.experts.down_projs", mock_dtensor, splits)
 
         assert result is not None and len(result) == 2
         for _, v in result:

From 14766e7c0f8c7ad909466acb3d3fae12250c7797 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Fri, 21 Aug 2026 10:10:05 -0700
Subject: [PATCH 10/20] refactor(checkpoint): clarify low-memory DCP capability

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../components/checkpoint/checkpointing.py    | 14 +++----
 .../checkpoint/state_dict_adapter.py          | 24 ++++-------
 .../models/bagel/state_dict_adapter.py        |  2 +-
 .../models/ernie4_5/state_dict_adapter.py     |  4 +-
 .../models/gemma4_moe/state_dict_adapter.py   |  2 +-
 .../gemma4_unified/state_dict_adapter.py      |  2 +-
 .../models/glm4_moe/state_dict_adapter.py     |  2 +-
 .../models/glm_moe_dsa/state_dict_adapter.py  |  2 +-
 .../models/hy_mt2/state_dict_adapter.py       |  2 +-
 .../models/hy_v3/state_dict_adapter.py        |  2 +-
 .../models/laguna/state_dict_adapter.py       |  2 +-
 .../models/ling_v2/state_dict_adapter.py      |  6 +--
 .../llava_onevision/state_dict_adapter.py     |  2 +-
 .../models/muse_glimmer/state_dict_adapter.py |  2 +-
 .../nemotron_omni/state_dict_adapter.py       |  6 +--
 .../models/nemotron_v3/state_dict_adapter.py  |  2 +-
 .../models/qwen2_5_omni/state_dict_adapter.py |  2 +-
 .../models/qwen3_5/state_dict_adapter.py      |  2 +-
 .../models/qwen3_moe/state_dict_adapter.py    |  2 +-
 .../models/qwen3_next/state_dict_adapter.py   |  2 +-
 .../qwen3_omni_moe/state_dict_adapter.py      |  2 +-
 .../models/qwen3_vl_moe/state_dict_adapter.py |  2 +-
 .../components/moe/state_dict_mixin.py        | 16 ++++----
 .../checkpoint/test_checkpointing.py          | 15 +++----
 .../test_state_dict_adapter_capabilities.py   | 41 +++++++++----------
 .../test_ling_v2_state_dict_adapter.py        | 12 +++---
 .../test_nemotron_v3_state_dict_adapter.py    | 14 +++----
 tests/unit_tests/moe/test_state_dict_mixin.py | 12 +++---
 28 files changed, 89 insertions(+), 109 deletions(-)

diff --git a/nemo_automodel/components/checkpoint/checkpointing.py b/nemo_automodel/components/checkpoint/checkpointing.py
index 7b9920dfe7..79d892a51a 100644
--- a/nemo_automodel/components/checkpoint/checkpointing.py
+++ b/nemo_automodel/components/checkpoint/checkpointing.py
@@ -831,9 +831,8 @@ def load_model(
         is_safetensors = _is_safetensors_checkpoint(model_path)
         is_custom_model = _is_custom_model(model_state.model[0])
         # Custom models traditionally loaded the complete checkpoint on the host because model-specific conversion
-        # could otherwise create a second full copy on the GPU. Use DCP when the adapter can place large checkpoint
-        # tensors directly in model weight memory. An adapter such as Gemma4 may also use a small temporary tensor that
-        # it applies after the read. Other custom adapters and quantized initialization keep the host fallback.
+        # could otherwise create a second full copy on the GPU. Use DCP when most tensors load into model weight memory
+        # and any temporary tensors are small. Other custom adapters and quantized initialization keep the host fallback.
         # World size inline (not via components.distributed) so the checkpoint component stays
         # independent per the import-linter contract.
         if torch.distributed.is_initialized():
@@ -841,16 +840,13 @@ def load_model(
         else:
             world_size = int(os.environ.get("WORLD_SIZE", "1"))
         state_dict_adapter = getattr(_unwrap_ddp_model(model_state.model[0]), "state_dict_adapter", None)
-        can_load_without_full_copy = (
+        can_use_low_memory_dcp = (
             isinstance(state_dict_adapter, StateDictAdapter)
-            and (
-                state_dict_adapter.supports_write_through_checkpoint_load
-                or state_dict_adapter.supports_checkpoint_load_without_full_copy
-            )
+            and state_dict_adapter.supports_low_memory_dcp_load
             and not self.config.dequantize_base_checkpoint
         )
         single_device_custom_safetensors = (
-            is_safetensors and is_custom_model and world_size == 1 and not can_load_without_full_copy
+            is_safetensors and is_custom_model and world_size == 1 and not can_use_low_memory_dcp
         )
         if (
             is_init_step
diff --git a/nemo_automodel/components/checkpoint/state_dict_adapter.py b/nemo_automodel/components/checkpoint/state_dict_adapter.py
index a1bbb4e90f..e5ee4341d0 100644
--- a/nemo_automodel/components/checkpoint/state_dict_adapter.py
+++ b/nemo_automodel/components/checkpoint/state_dict_adapter.py
@@ -28,27 +28,17 @@ class StateDictAdapter(ABC):
     state dict format and other model state dict formats.
     """
 
-    _supports_write_through_checkpoint_load: bool = False
-    _supports_checkpoint_load_without_full_copy: bool = False
+    _supports_low_memory_dcp_load: bool = False
 
     @property
-    def supports_write_through_checkpoint_load(self) -> bool:
-        """Whether every checkpoint tensor is loaded directly into the model's existing weight memory.
+    def supports_low_memory_dcp_load(self) -> bool:
+        """Whether DCP can load the checkpoint with zero or small temporary tensors.
 
-        Enable this only when writing every tensor returned by ``to_hf`` for base-checkpoint loading updates the
-        model itself. This lets the loader skip a complete CPU copy of the checkpoint.
+        Most checkpoint tensors must load directly into the model's existing weight memory. Small temporary tensors
+        are allowed when they are converted and released after the read. Enable this only when the extra memory is
+        safely below a full model copy, including when loading on one GPU.
         """
-        return self._supports_write_through_checkpoint_load
-
-    @property
-    def supports_checkpoint_load_without_full_copy(self) -> bool:
-        """Whether DCP can load this adapter without another full set of model weights.
-
-        Large checkpoint tensors must be loaded into the model's existing weight memory. Small temporary tensors are
-        allowed when they can be applied and discarded without making a model-sized copy. For example, Gemma4 loads
-        a scale tensor and applies it to already-loaded expert weights.
-        """
-        return self._supports_checkpoint_load_without_full_copy
+        return self._supports_low_memory_dcp_load
 
     @abstractmethod
     def to_hf(self, state_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
diff --git a/nemo_automodel/components/models/bagel/state_dict_adapter.py b/nemo_automodel/components/models/bagel/state_dict_adapter.py
index e1db9f8fc0..8d554eea10 100644
--- a/nemo_automodel/components/models/bagel/state_dict_adapter.py
+++ b/nemo_automodel/components/models/bagel/state_dict_adapter.py
@@ -181,7 +181,7 @@ class BagelStateDictAdapter(StateDictAdapter):
             ``1`` / ``2``.
     """
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(self, config: Any = None, *, stage: Any = "stage1") -> None:
         self.config = config
diff --git a/nemo_automodel/components/models/ernie4_5/state_dict_adapter.py b/nemo_automodel/components/models/ernie4_5/state_dict_adapter.py
index daf28117cd..47ef887b5f 100644
--- a/nemo_automodel/components/models/ernie4_5/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ernie4_5/state_dict_adapter.py
@@ -32,7 +32,7 @@
 class Ernie4_5StateDictAdapter(StateDictAdapter):
     """Passthrough adapter for dense ERNIE 4.5 checkpoints."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(self, config: Any):
         self.config = config
@@ -60,7 +60,7 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
 class Ernie4_5_MoeStateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter):
     """Convert ERNIE 4.5 MoE HF checkpoints to AutoModel grouped-expert format."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py b/nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py
index 121f16eecb..e17fcbb8b9 100644
--- a/nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py
+++ b/nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py
@@ -61,7 +61,7 @@ class Gemma4MoEStateDictAdapter(StateDictAdapter):
       4. Expert-parallel sharding when a device mesh is provided
     """
 
-    _supports_checkpoint_load_without_full_copy = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/gemma4_unified/state_dict_adapter.py b/nemo_automodel/components/models/gemma4_unified/state_dict_adapter.py
index f3351208ef..bbca820449 100644
--- a/nemo_automodel/components/models/gemma4_unified/state_dict_adapter.py
+++ b/nemo_automodel/components/models/gemma4_unified/state_dict_adapter.py
@@ -70,7 +70,7 @@ def _rename_state_dict(state_dict: dict[str, Any], renames: dict[str, str]) -> d
 class Gemma4UnifiedStateDictAdapter(StateDictAdapter):
     """Translate Gemma4 Unified keys between model and published HF names."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def to_hf(
         self,
diff --git a/nemo_automodel/components/models/glm4_moe/state_dict_adapter.py b/nemo_automodel/components/models/glm4_moe/state_dict_adapter.py
index 0602d48f7d..13cc0eb7be 100644
--- a/nemo_automodel/components/models/glm4_moe/state_dict_adapter.py
+++ b/nemo_automodel/components/models/glm4_moe/state_dict_adapter.py
@@ -45,7 +45,7 @@ class Glm4MoeStateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter):
       model.layers.{L}.mlp.shared_expert.down_proj.weight
     """
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/glm_moe_dsa/state_dict_adapter.py b/nemo_automodel/components/models/glm_moe_dsa/state_dict_adapter.py
index bfcd701401..b72cdbcb8b 100644
--- a/nemo_automodel/components/models/glm_moe_dsa/state_dict_adapter.py
+++ b/nemo_automodel/components/models/glm_moe_dsa/state_dict_adapter.py
@@ -27,7 +27,7 @@ class GlmMoeDsaStateDictAdapter(Glm4MoeStateDictAdapter):
     that should not be quantized (k_norm, weights_proj).
     """
 
-    _supports_write_through_checkpoint_load = False
+    _supports_low_memory_dcp_load = False
 
     _indexer_non_quantized_keys = [
         "indexer.k_norm.weight",
diff --git a/nemo_automodel/components/models/hy_mt2/state_dict_adapter.py b/nemo_automodel/components/models/hy_mt2/state_dict_adapter.py
index ea1c54e994..d48f0942ab 100644
--- a/nemo_automodel/components/models/hy_mt2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/hy_mt2/state_dict_adapter.py
@@ -70,7 +70,7 @@
 class HyMT2StateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter):
     """Bridges Automodel native (grouped experts) and on-disk Hy-MT2 HF format."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/hy_v3/state_dict_adapter.py b/nemo_automodel/components/models/hy_v3/state_dict_adapter.py
index f92be67d74..75ea99034a 100644
--- a/nemo_automodel/components/models/hy_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/hy_v3/state_dict_adapter.py
@@ -78,7 +78,7 @@ class HYV3StateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter):
     only the three HYV3-specific name renames + MTP-layer filtering live here.
     """
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/laguna/state_dict_adapter.py b/nemo_automodel/components/models/laguna/state_dict_adapter.py
index f991ebf82d..e52301d5bb 100644
--- a/nemo_automodel/components/models/laguna/state_dict_adapter.py
+++ b/nemo_automodel/components/models/laguna/state_dict_adapter.py
@@ -46,7 +46,7 @@ def _apply_renames(key: str, renames: tuple[tuple[re.Pattern[str], str], ...]) -
 class LagunaStateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter):
     """Convert Laguna HF checkpoints to Automodel's grouped-MoE layout."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
index cfd647cd2f..f5d9b24f0e 100644
--- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
@@ -101,9 +101,9 @@ def __init__(
         self._uses_model_prefix = True
 
     @property
-    def supports_checkpoint_load_without_full_copy(self) -> bool:
-        """Whether Ling can use DCP with only its small fused-QKV temporary tensors."""
-        return self._supports_write_through_expert_checkpoint_load and not self.moe_config.expert_bias
+    def supports_low_memory_dcp_load(self) -> bool:
+        """Whether Ling's DCP load needs only small fused-QKV temporary tensors."""
+        return self._expert_checkpoint_tensors_use_model_storage and not self.moe_config.expert_bias
 
     # ---- HF -> native ----------------------------------------------------
 
diff --git a/nemo_automodel/components/models/llava_onevision/state_dict_adapter.py b/nemo_automodel/components/models/llava_onevision/state_dict_adapter.py
index beea12bd37..b817bc50b6 100644
--- a/nemo_automodel/components/models/llava_onevision/state_dict_adapter.py
+++ b/nemo_automodel/components/models/llava_onevision/state_dict_adapter.py
@@ -45,7 +45,7 @@
 
 
 class LlavaOneVisionStateDictAdapter(StateDictAdapter):
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(self, config: Any = None, **kwargs):
         self.config = config
diff --git a/nemo_automodel/components/models/muse_glimmer/state_dict_adapter.py b/nemo_automodel/components/models/muse_glimmer/state_dict_adapter.py
index b8019d79ce..db80f8c610 100644
--- a/nemo_automodel/components/models/muse_glimmer/state_dict_adapter.py
+++ b/nemo_automodel/components/models/muse_glimmer/state_dict_adapter.py
@@ -28,7 +28,7 @@
 class MuseGlimmerStateDictAdapter(StateDictAdapter):
     """Map canonical nested MuseGlimmer checkpoint keys to the native legacy module tree."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(self, config: MuseGlimmerConfig) -> None:
         self.config = config
diff --git a/nemo_automodel/components/models/nemotron_omni/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_omni/state_dict_adapter.py
index 92beeb6f7d..62ce6be713 100644
--- a/nemo_automodel/components/models/nemotron_omni/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_omni/state_dict_adapter.py
@@ -106,9 +106,9 @@ class NemotronOmniStateDictAdapter(StateDictAdapter):
     """
 
     @property
-    def supports_write_through_checkpoint_load(self) -> bool:
-        """Whether the embedded language adapter and all wrapper renames preserve storage."""
-        return self._llm_adapter.supports_write_through_checkpoint_load
+    def supports_low_memory_dcp_load(self) -> bool:
+        """Whether the embedded language adapter supports a low-memory DCP load."""
+        return self._llm_adapter.supports_low_memory_dcp_load
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
index 1dedec9962..a430d15484 100644
--- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
@@ -89,7 +89,7 @@ class NemotronV3StateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter
     Note: NemotronV3 uses 'mixer' instead of 'mlp' in layer paths.
     """
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/qwen2_5_omni/state_dict_adapter.py b/nemo_automodel/components/models/qwen2_5_omni/state_dict_adapter.py
index c26626ef0f..1543126dbb 100644
--- a/nemo_automodel/components/models/qwen2_5_omni/state_dict_adapter.py
+++ b/nemo_automodel/components/models/qwen2_5_omni/state_dict_adapter.py
@@ -50,7 +50,7 @@ class Qwen2_5OmniStateDictAdapter(StateDictAdapter):
     needed — this is a thin key-renaming adapter.
     """
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py
index 7613a96bb0..0f8e634ac1 100644
--- a/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py
+++ b/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py
@@ -75,7 +75,7 @@ def map_qwen3_5_mtp_to_hf_key(key: str) -> str:
 class Qwen3_5DenseStateDictAdapter(StateDictAdapter):
     """Adapter that hides the ``_fp32_params`` wrapping in saved checkpoints."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(self, *, route_linear_attn_fp32_params: bool = True) -> None:
         self.route_linear_attn_fp32_params = route_linear_attn_fp32_params
diff --git a/nemo_automodel/components/models/qwen3_moe/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_moe/state_dict_adapter.py
index 3e37c7a400..dad36b7ecb 100644
--- a/nemo_automodel/components/models/qwen3_moe/state_dict_adapter.py
+++ b/nemo_automodel/components/models/qwen3_moe/state_dict_adapter.py
@@ -43,7 +43,7 @@ class Qwen3MoeStateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter):
       model.layers.{L}.mlp.experts.down_projs         # [n_experts, moe_inter_dim, dim]
     """
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/qwen3_next/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_next/state_dict_adapter.py
index 085c98ae1e..4aa9994e09 100644
--- a/nemo_automodel/components/models/qwen3_next/state_dict_adapter.py
+++ b/nemo_automodel/components/models/qwen3_next/state_dict_adapter.py
@@ -55,7 +55,7 @@ class Qwen3NextStateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter)
       model.layers.{L}.mlp.shared_experts.down_proj.weight
     """
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py
index 20fcf60bc8..c22958e535 100644
--- a/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py
+++ b/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py
@@ -29,7 +29,7 @@
 class Qwen3OmniMoeStateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter):
     """Converts between HF Qwen3OmniMoe checkpoints and grouped-experts native format."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/models/qwen3_vl_moe/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_vl_moe/state_dict_adapter.py
index 19e911b097..e4aa0f291a 100644
--- a/nemo_automodel/components/models/qwen3_vl_moe/state_dict_adapter.py
+++ b/nemo_automodel/components/models/qwen3_vl_moe/state_dict_adapter.py
@@ -42,7 +42,7 @@ class Qwen3VLMoeStateDictAdapter(StateDictAdapter):
                  shard, and wraps in DTensor via create_dtensor_from_local.
     """
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/nemo_automodel/components/moe/state_dict_mixin.py b/nemo_automodel/components/moe/state_dict_mixin.py
index fe22bc12cf..71984bdf4d 100644
--- a/nemo_automodel/components/moe/state_dict_mixin.py
+++ b/nemo_automodel/components/moe/state_dict_mixin.py
@@ -51,17 +51,17 @@ class MoESplitExpertsStateDictMixin:
     # - self.backend: Backend configuration object
 
     @property
-    def supports_write_through_checkpoint_load(self) -> bool:
-        """Whether every checkpoint tensor, including expert tensors, loads directly into model weights."""
-        experts_load_directly = self.moe_config is None or self._supports_write_through_expert_checkpoint_load
-        return self._supports_write_through_checkpoint_load and experts_load_directly
+    def supports_low_memory_dcp_load(self) -> bool:
+        """Whether DCP needs at most small temporary tensors for this MoE checkpoint."""
+        experts_use_model_storage = self.moe_config is None or self._expert_checkpoint_tensors_use_model_storage
+        return self._supports_low_memory_dcp_load and experts_use_model_storage
 
     @property
-    def _supports_write_through_expert_checkpoint_load(self) -> bool:
-        """Whether grouped expert checkpoint tensors load directly into model weight memory.
+    def _expert_checkpoint_tensors_use_model_storage(self) -> bool:
+        """Whether grouped expert checkpoint tensors use the model's weight memory.
 
-        This covers only the shared expert conversion. A concrete adapter must also verify that all non-expert
-        checkpoint tensors load directly before it enables the full-checkpoint fast path.
+        This covers only the shared expert conversion. A concrete adapter must also verify that its non-expert
+        checkpoint tensors require at most small temporary tensors before enabling low-memory DCP loading.
         """
         return self._grouped_expert_storage_is_model_weight and self.backend.dispatcher != "mok"
 
diff --git a/tests/unit_tests/checkpoint/test_checkpointing.py b/tests/unit_tests/checkpoint/test_checkpointing.py
index 510fb2aed4..b4717c8743 100644
--- a/tests/unit_tests/checkpoint/test_checkpointing.py
+++ b/tests/unit_tests/checkpoint/test_checkpointing.py
@@ -1548,8 +1548,8 @@ class TestLoadModelCustomModelGuard:
 
     Under multi-rank (sharded) loading, custom models use the standard DCP path so each
     rank slices its local DTensor shard. On a single device (world_size == 1) there is no
-    sharding, so adapters that cannot expose write-through destinations take the frugal
-    full-state path instead. Adapters with an explicit write-through guarantee use DCP.
+    sharding, so adapters that need large temporary tensors take the frugal full-state path
+    instead. Adapters that need zero or small temporary tensors use DCP.
     """
 
     def _make_checkpointer(self):
@@ -1659,7 +1659,7 @@ def test_custom_model_skips_fast_path_uses_dcp(self, mock_load_full, mock_load_h
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype")
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model")
     def test_single_device_custom_model_uses_fast_path(self, mock_load_full, mock_load_hf, mock_is_st, caplog):
-        """A custom model without a write-through guarantee uses the full-state path.
+        """A custom model without low-memory DCP support uses the full-state path.
 
         The fast path applies the state_dict_adapter from_hf conversion on CPU (via
         _maybe_adapt_state_dict_from_hf) and copies into the model, keeping device memory at
@@ -1696,25 +1696,22 @@ def test_single_device_custom_model_uses_fast_path(self, mock_load_full, mock_lo
     @patch("nemo_automodel.components.checkpoint.checkpointing._is_safetensors_checkpoint", return_value=True)
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_hf_checkpoint_preserving_dtype")
     @patch("nemo_automodel.components.checkpoint.checkpointing._load_full_state_dict_into_model")
-    @pytest.mark.parametrize("load_capability", ["write_through", "without_full_copy"])
     @pytest.mark.parametrize("dequantize_base_checkpoint", [False, True])
-    def test_single_device_adapter_without_full_copy_routes_by_quantization(
+    def test_single_device_low_memory_dcp_routes_by_quantization(
         self,
         mock_load_full,
         mock_load_hf,
         mock_is_st,
         caplog,
         dequantize_base_checkpoint,
-        load_capability,
     ):
-        """Quantized conversion keeps the full CPU fallback for both direct-load capabilities."""
+        """Quantized conversion keeps the full CPU fallback despite low-memory DCP support."""
         CustomModel = type("CustomModel", (torch.nn.Module,), {})
         CustomModel.__module__ = "nemo_automodel.components.models.nemotron_v3.model"
         model = CustomModel()
         model.layer = torch.nn.Linear(4, 4)
         model.state_dict_adapter = MagicMock(spec=StateDictAdapter)
-        model.state_dict_adapter.supports_write_through_checkpoint_load = load_capability == "write_through"
-        model.state_dict_adapter.supports_checkpoint_load_without_full_copy = load_capability == "without_full_copy"
+        model.state_dict_adapter.supports_low_memory_dcp_load = True
         mock_state_dict = {"layer.weight": torch.randn(4, 4), "layer.bias": torch.randn(4)}
         mock_load_hf.return_value = mock_state_dict
 
diff --git a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
index cc725dbb06..a80e9ddfb0 100644
--- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
+++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
@@ -130,7 +130,7 @@ def test_write_through_adapters_expose_aliasing_destinations(
     adapter = object.__new__(adapter_type)
     adapter.__dict__.update(adapter_attrs)
 
-    assert adapter.supports_write_through_checkpoint_load is True
+    assert adapter.supports_low_memory_dcp_load is True
 
     _assert_destinations_write_through(
         adapter,
@@ -156,7 +156,7 @@ def test_write_through_adapters_expose_aliasing_destinations(
         Qwen3OmniMoeStateDictAdapter,
     ],
 )
-def test_write_through_grouped_adapters_preserve_non_expert_storage_and_require_aliasing_backend(
+def test_low_memory_dcp_grouped_adapters_preserve_non_expert_storage_and_require_model_backed_experts(
     adapter_type: type[StateDictAdapter],
 ) -> None:
     moe_config = SimpleNamespace(n_routed_experts=2, moe_inter_dim=2, expert_activation="silu")
@@ -165,7 +165,7 @@ def test_write_through_grouped_adapters_preserve_non_expert_storage_and_require_
         moe_config,
         SimpleNamespace(experts="torch", dispatcher="torch"),
     )
-    assert adapter.supports_write_through_checkpoint_load is True
+    assert adapter.supports_low_memory_dcp_load is True
 
     _assert_destinations_write_through(
         adapter,
@@ -177,22 +177,21 @@ def test_write_through_grouped_adapters_preserve_non_expert_storage_and_require_
     )
 
     adapter.backend = SimpleNamespace(experts="te", dispatcher="torch")
-    assert adapter.supports_write_through_checkpoint_load is True
+    assert adapter.supports_low_memory_dcp_load is True
 
     adapter.backend = SimpleNamespace(experts="te", dispatcher="deepep")
     with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
-        assert adapter.supports_write_through_checkpoint_load is False
+        assert adapter.supports_low_memory_dcp_load is False
     with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
-        assert adapter.supports_write_through_checkpoint_load is True
+        assert adapter.supports_low_memory_dcp_load is True
 
     adapter.backend = SimpleNamespace(experts="torch", dispatcher="mok")
-    assert adapter.supports_write_through_checkpoint_load is False
+    assert adapter.supports_low_memory_dcp_load is False
 
 
 @pytest.mark.parametrize(
     "adapter_type",
     [
-        BailingMoeV2StateDictAdapter,
         DeepSeekV3StateDictAdapter,
         GlmMoeDsaStateDictAdapter,
         KimiK25VLStateDictAdapter,
@@ -203,28 +202,26 @@ def test_write_through_grouped_adapters_preserve_non_expert_storage_and_require_
         MiniMaxM3StateDictAdapter,
     ],
 )
-def test_materializing_grouped_expert_adapters_keep_frugal_load(adapter_type):
+def test_materializing_grouped_expert_adapters_do_not_enable_low_memory_dcp(adapter_type):
     adapter = object.__new__(adapter_type)
     adapter.moe_config = object()
     adapter.backend = SimpleNamespace(experts="torch", dispatcher="torch")
 
-    assert adapter.supports_write_through_checkpoint_load is False
+    assert adapter.supports_low_memory_dcp_load is False
 
 
-def test_gemma4_moe_adapter_loads_without_a_full_checkpoint_copy():
+def test_gemma4_moe_adapter_supports_low_memory_dcp():
     adapter = object.__new__(Gemma4MoEStateDictAdapter)
 
-    assert adapter.supports_write_through_checkpoint_load is False
-    assert adapter.supports_checkpoint_load_without_full_copy is True
+    assert adapter.supports_low_memory_dcp_load is True
 
 
-def test_ling_adapter_loads_without_a_full_checkpoint_copy():
+def test_ling_adapter_supports_low_memory_dcp():
     adapter = object.__new__(BailingMoeV2StateDictAdapter)
     adapter.moe_config = SimpleNamespace(expert_bias=False)
     adapter.backend = SimpleNamespace(experts="torch_mm", dispatcher="torch")
 
-    assert adapter.supports_write_through_checkpoint_load is False
-    assert adapter.supports_checkpoint_load_without_full_copy is True
+    assert adapter.supports_low_memory_dcp_load is True
 
 
 def test_dense_grouped_adapter_does_not_require_an_expert_backend():
@@ -232,13 +229,13 @@ def test_dense_grouped_adapter_does_not_require_an_expert_backend():
     adapter.moe_config = None
     adapter.backend = SimpleNamespace(experts="te", dispatcher="mok")
 
-    assert adapter.supports_write_through_checkpoint_load is True
+    assert adapter.supports_low_memory_dcp_load is True
 
 
-def test_nemotron_omni_delegates_direct_load_support_to_language_adapter():
+def test_nemotron_omni_delegates_low_memory_dcp_support_to_language_adapter():
     adapter = object.__new__(NemotronOmniStateDictAdapter)
-    adapter._llm_adapter = SimpleNamespace(supports_write_through_checkpoint_load=True)
-    assert adapter.supports_write_through_checkpoint_load is True
+    adapter._llm_adapter = SimpleNamespace(supports_low_memory_dcp_load=True)
+    assert adapter.supports_low_memory_dcp_load is True
 
-    adapter._llm_adapter.supports_write_through_checkpoint_load = False
-    assert adapter.supports_write_through_checkpoint_load is False
+    adapter._llm_adapter.supports_low_memory_dcp_load = False
+    assert adapter.supports_low_memory_dcp_load is False
diff --git a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
index f4447bbba6..74df7caef6 100644
--- a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
+++ b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
@@ -224,7 +224,7 @@ def test_hf_to_native_to_hf_is_lossless_for_qkv_path(self, adapter, config):
         torch.testing.assert_close(roundtripped["model.word_embeddings.weight"], hf_sd["model.word_embeddings.weight"])
 
 
-class TestCheckpointLoadWithoutFullCopy:
+class TestLowMemoryDcpLoad:
     def test_uses_only_fused_qkv_temporaries(self, adapter, config, moe_config):
         q_size = config.num_attention_heads * config.head_dim
         kv_size = config.num_key_value_heads * config.head_dim
@@ -269,20 +269,20 @@ def test_uses_only_fused_qkv_temporaries(self, adapter, config, moe_config):
         assert down_destination.untyped_storage().data_ptr() == grouped_down.untyped_storage().data_ptr()
 
     def test_capability_matches_runtime_expert_storage(self, adapter, monkeypatch):
-        assert adapter.supports_checkpoint_load_without_full_copy is True
+        assert adapter.supports_low_memory_dcp_load is True
 
         adapter.backend.experts = "te"
         adapter.backend.dispatcher = "deepep"
         monkeypatch.setattr("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", lambda: 1)
-        assert adapter.supports_checkpoint_load_without_full_copy is True
+        assert adapter.supports_low_memory_dcp_load is True
 
         monkeypatch.setattr("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", lambda: 8)
-        assert adapter.supports_checkpoint_load_without_full_copy is False
+        assert adapter.supports_low_memory_dcp_load is False
 
         adapter.backend.experts = "torch_mm"
         adapter.backend.dispatcher = "mok"
-        assert adapter.supports_checkpoint_load_without_full_copy is False
+        assert adapter.supports_low_memory_dcp_load is False
 
         adapter.backend.dispatcher = "torch"
         adapter.moe_config.expert_bias = True
-        assert adapter.supports_checkpoint_load_without_full_copy is False
+        assert adapter.supports_low_memory_dcp_load is False
diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
index 6f25e5a542..92a1f4d8b3 100644
--- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
+++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
@@ -115,21 +115,21 @@ def test_expert_path_segment_property(self, config, moe_config, backend):
 
         assert adapter._expert_path_segment == "mixer.experts"
 
-    def test_write_through_load_capability_requires_aliasing_expert_backend(self, config, moe_config, backend):
+    def test_low_memory_dcp_capability_requires_model_backed_experts(self, config, moe_config, backend):
         adapter = NemotronV3StateDictAdapter(config, moe_config, backend)
-        assert adapter.supports_write_through_checkpoint_load is True
+        assert adapter.supports_low_memory_dcp_load is True
 
         adapter.backend.experts = "te"
         adapter.backend.dispatcher = "deepep"
         with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
-            assert adapter.supports_write_through_checkpoint_load is False
+            assert adapter.supports_low_memory_dcp_load is False
 
         with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
-            assert adapter.supports_write_through_checkpoint_load is True
+            assert adapter.supports_low_memory_dcp_load is True
 
         adapter.backend.experts = "gmm"
         adapter.backend.dispatcher = "mok"
-        assert adapter.supports_write_through_checkpoint_load is False
+        assert adapter.supports_low_memory_dcp_load is False
 
     def test_te_single_device_fallback_loads_through_grouped_parameter_views(self, config, moe_config, backend):
         backend.experts = "te"
@@ -154,7 +154,7 @@ def test_te_single_device_fallback_loads_through_grouped_parameter_views(self, c
                 },
                 for_checkpoint_load=True,
             )
-            assert adapter.supports_write_through_checkpoint_load is True
+            assert adapter.supports_low_memory_dcp_load is True
 
         assert set(destinations) == {
             f"backbone.layers.0.mixer.experts.{expert_id}.{projection}.weight"
@@ -196,7 +196,7 @@ def adapter(self, backend):
 
     def test_init_accepts_none_moe_config(self, adapter):
         assert adapter.moe_config is None
-        assert adapter.supports_write_through_checkpoint_load is True
+        assert adapter.supports_low_memory_dcp_load is True
 
     def test_from_hf_renames_without_experts(self, adapter):
         hf_sd = {
diff --git a/tests/unit_tests/moe/test_state_dict_mixin.py b/tests/unit_tests/moe/test_state_dict_mixin.py
index 307610a015..c9176c980a 100644
--- a/tests/unit_tests/moe/test_state_dict_mixin.py
+++ b/tests/unit_tests/moe/test_state_dict_mixin.py
@@ -1048,23 +1048,23 @@ class TestInplaceLoadViews:
     conversion function, so patches must target that module path.
     """
 
-    def test_expert_write_through_capability_matches_grouped_storage_aliasing(self):
+    def test_expert_checkpoint_storage_capability_matches_grouped_storage_aliasing(self):
         mixin = MockMoEStateDictMixin()
-        assert mixin._supports_write_through_expert_checkpoint_load is True
+        assert mixin._expert_checkpoint_tensors_use_model_storage is True
 
         mixin.backend.experts = "te"
         mixin.backend.dispatcher = "deepep"
         with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
-            assert mixin._supports_write_through_expert_checkpoint_load is False
+            assert mixin._expert_checkpoint_tensors_use_model_storage is False
         with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
-            assert mixin._supports_write_through_expert_checkpoint_load is True
+            assert mixin._expert_checkpoint_tensors_use_model_storage is True
 
         mixin.backend.dispatcher = "torch"
-        assert mixin._supports_write_through_expert_checkpoint_load is True
+        assert mixin._expert_checkpoint_tensors_use_model_storage is True
 
         mixin.backend.experts = "gmm"
         mixin.backend.dispatcher = "mok"
-        assert mixin._supports_write_through_expert_checkpoint_load is False
+        assert mixin._expert_checkpoint_tensors_use_model_storage is False
 
     def _run_inplace_conversion(self, mixin, fqn, mock_dtensor, splits):
         mixin._split_experts_weights = Mock(return_value=splits)

From d62d63761fb64a650337086e9de3e4aa4bf3ac0b Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Tue, 25 Aug 2026 11:50:06 -0700
Subject: [PATCH 11/20] fix(checkpoint): avoid distributed component import

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../components/moe/state_dict_mixin.py        |  6 ++-
 .../test_state_dict_adapter_capabilities.py   | 10 ++++-
 tests/unit_tests/moe/test_state_dict_mixin.py | 43 ++++++++++++++++---
 3 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/nemo_automodel/components/moe/state_dict_mixin.py b/nemo_automodel/components/moe/state_dict_mixin.py
index 71984bdf4d..1f1efe9afb 100644
--- a/nemo_automodel/components/moe/state_dict_mixin.py
+++ b/nemo_automodel/components/moe/state_dict_mixin.py
@@ -13,13 +13,13 @@
 # limitations under the License.
 
 import gc
+import os
 import re
 from typing import Any, Optional
 
 import torch
 from torch.distributed.device_mesh import DeviceMesh
 
-from nemo_automodel.components.distributed.init_utils import get_world_size_safe
 from nemo_automodel.components.moe.state_dict_utils import (
     create_dtensor_from_local,
     get_expert_range_for_rank_from_mesh,
@@ -79,7 +79,9 @@ def _grouped_expert_storage_is_model_weight(self) -> bool:
             return False
         if self.backend.dispatcher not in {"deepep", "hybridep", "uccl_ep"}:
             return True
-        return get_world_size_safe() == 1
+        if torch.distributed.is_initialized():
+            return torch.distributed.get_world_size() == 1
+        return int(os.environ.get("WORLD_SIZE", "1")) == 1
 
     @property
     def _is_gated_moe(self) -> bool:
diff --git a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
index a80e9ddfb0..e52bf62405 100644
--- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
+++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
@@ -180,9 +180,15 @@ def test_low_memory_dcp_grouped_adapters_preserve_non_expert_storage_and_require
     assert adapter.supports_low_memory_dcp_load is True
 
     adapter.backend = SimpleNamespace(experts="te", dispatcher="deepep")
-    with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+    with (
+        patch("nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False),
+        patch.dict("os.environ", {"WORLD_SIZE": "8"}),
+    ):
         assert adapter.supports_low_memory_dcp_load is False
-    with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
+    with (
+        patch("nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False),
+        patch.dict("os.environ", {"WORLD_SIZE": "1"}),
+    ):
         assert adapter.supports_low_memory_dcp_load is True
 
     adapter.backend = SimpleNamespace(experts="torch", dispatcher="mok")
diff --git a/tests/unit_tests/moe/test_state_dict_mixin.py b/tests/unit_tests/moe/test_state_dict_mixin.py
index c9176c980a..518fb995f3 100644
--- a/tests/unit_tests/moe/test_state_dict_mixin.py
+++ b/tests/unit_tests/moe/test_state_dict_mixin.py
@@ -1054,9 +1054,19 @@ def test_expert_checkpoint_storage_capability_matches_grouped_storage_aliasing(s
 
         mixin.backend.experts = "te"
         mixin.backend.dispatcher = "deepep"
-        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+        with (
+            patch(
+                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
+            ),
+            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
+        ):
             assert mixin._expert_checkpoint_tensors_use_model_storage is False
-        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
+        with (
+            patch(
+                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
+            ),
+            patch.dict("os.environ", {"WORLD_SIZE": "1"}),
+        ):
             assert mixin._expert_checkpoint_tensors_use_model_storage is True
 
         mixin.backend.dispatcher = "torch"
@@ -1190,7 +1200,12 @@ def test_inplace_load_skips_when_backend_experts_is_te_gate_and_up(self):
         splits = [local_storage[i] for i in range(2)]
         mock_dtensor = Mock()
 
-        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+        with (
+            patch(
+                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
+            ),
+            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
+        ):
             result = self._run_inplace_conversion(
                 mixin, "model.layers.0.mlp.experts.gate_and_up_projs", mock_dtensor, splits
             )
@@ -1244,7 +1259,10 @@ def test_te_checkpoint_layout_views_are_reused(self):
         virtual_down = down_storage.transpose(-1, -2)
 
         with (
-            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
+            patch(
+                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
+            ),
+            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
             patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
             patch("torch.empty_like") as empty_like,
         ):
@@ -1290,7 +1308,10 @@ def test_te_checkpoint_layout_cuda_peak_stays_within_one_buffer(self):
         virtual_gate_up = gate_up_storage.transpose(-1, -2)
 
         with (
-            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
+            patch(
+                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
+            ),
+            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
             patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
             patch("nemo_automodel.components.moe.state_dict_mixin.gc.collect") as collect,
             patch("torch.cuda.empty_cache") as empty_cache,
@@ -1315,7 +1336,10 @@ def test_temporary_checkpoint_views_round_trip_loaded_values(self):
         initialized_down = torch.full((2, 3, 4), 99.0)
 
         with (
-            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
+            patch(
+                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
+            ),
+            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
             patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
         ):
             destinations = dict(
@@ -1368,7 +1392,12 @@ def test_inplace_load_skips_when_backend_experts_is_te_down_projs(self):
         mock_dtensor.shape = (2, 512, 1024)
         mock_dtensor.is_meta = False
 
-        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
+        with (
+            patch(
+                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
+            ),
+            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
+        ):
             result = self._run_inplace_conversion(mixin, "model.layers.3.mlp.experts.down_projs", mock_dtensor, splits)
 
         assert result is not None and len(result) == 2

From 4c6c4988888b10135ef527dbe47ad101ffceee99 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Tue, 25 Aug 2026 14:16:48 -0700
Subject: [PATCH 12/20] fix(checkpoint): preserve backend-independent LoRA
 export

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../components/moe/state_dict_mixin.py        | 38 +++++++-----
 .../test_state_dict_adapter_capabilities.py   | 10 +--
 tests/unit_tests/moe/test_state_dict_mixin.py | 61 ++++++++-----------
 3 files changed, 49 insertions(+), 60 deletions(-)

diff --git a/nemo_automodel/components/moe/state_dict_mixin.py b/nemo_automodel/components/moe/state_dict_mixin.py
index 1f1efe9afb..c21c9082b9 100644
--- a/nemo_automodel/components/moe/state_dict_mixin.py
+++ b/nemo_automodel/components/moe/state_dict_mixin.py
@@ -33,6 +33,17 @@
 _LORA_EXPERT_SUFFIXES = ("lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B")
 
 
+def get_world_size_safe() -> int:
+    """Return the distributed world size before or after process-group initialization.
+
+    Returns:
+        The initialized process-group size, or ``WORLD_SIZE`` before initialization.
+    """
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_world_size()
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
 class MoESplitExpertsStateDictMixin:
     """Mixin class providing MoE state dict conversion utilities.
 
@@ -79,9 +90,7 @@ def _grouped_expert_storage_is_model_weight(self) -> bool:
             return False
         if self.backend.dispatcher not in {"deepep", "hybridep", "uccl_ep"}:
             return True
-        if torch.distributed.is_initialized():
-            return torch.distributed.get_world_size() == 1
-        return int(os.environ.get("WORLD_SIZE", "1")) == 1
+        return get_world_size_safe() == 1
 
     @property
     def _is_gated_moe(self) -> bool:
@@ -874,6 +883,17 @@ def _convert_single_merged_expert_to_hf_split_experts(
         inter_dim = self.moe_config.moe_inter_dim
         prefix = prefix_override if prefix_override is not None else self._hf_prefix
         expert_segment = self._expert_path_segment
+
+        # MoE expert LoRA keys do not depend on the runtime backend or its checkpoint storage aliases.
+        # When v4_compatible=True, emit per-expert split keys. Otherwise, adapters that explicitly opt in emit the
+        # fused ParamWrapper format; unvalidated adapters retain the legacy per-expert format.
+        v4_compatible = kwargs.get("v4_compatible", False)
+        for suffix in _LORA_EXPERT_SUFFIXES:
+            if f".{expert_segment}.{suffix}" in fqn and fqn.endswith(f".{suffix}"):
+                if not v4_compatible and self._v5_peft_target_parameters:
+                    return self._convert_lora_to_paramwrapper(fqn, tensor)
+                return self._convert_lora_expert_to_hf(fqn, tensor, n_experts, inter_dim, expert_segment)
+
         # Quantization casts each split with ``value.to(float8_e4m3fn)``, creating storage separate from the model's
         # grouped weights. DCP must not treat that cast as model weight memory, or the model would keep its random
         # expert values. Quantized loads therefore rebuild the grouped expert tensor after the read.
@@ -999,16 +1019,4 @@ def checkpoint_load_destination(view: torch.Tensor, source: torch.Tensor) -> tor
                     torch.cuda.empty_cache()
             return result
 
-        # MoE expert LoRA keys: convert to HF-PEFT-compatible format.
-        # When v4_compatible=True: per-expert split keys (v4 format).
-        # When v4_compatible=False and the adapter explicitly opts in: fused
-        # ParamWrapper format (v5). Unvalidated adapters retain the legacy
-        # per-expert format even when v4_compatible is not requested.
-        v4_compatible = kwargs.get("v4_compatible", False)
-        for suffix in _LORA_EXPERT_SUFFIXES:
-            if f".{expert_segment}.{suffix}" in fqn and fqn.endswith(f".{suffix}"):
-                if not v4_compatible and self._v5_peft_target_parameters:
-                    return self._convert_lora_to_paramwrapper(fqn, tensor)
-                return self._convert_lora_expert_to_hf(fqn, tensor, n_experts, inter_dim, expert_segment)
-
         return None
diff --git a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
index e52bf62405..a80e9ddfb0 100644
--- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
+++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
@@ -180,15 +180,9 @@ def test_low_memory_dcp_grouped_adapters_preserve_non_expert_storage_and_require
     assert adapter.supports_low_memory_dcp_load is True
 
     adapter.backend = SimpleNamespace(experts="te", dispatcher="deepep")
-    with (
-        patch("nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False),
-        patch.dict("os.environ", {"WORLD_SIZE": "8"}),
-    ):
+    with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
         assert adapter.supports_low_memory_dcp_load is False
-    with (
-        patch("nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False),
-        patch.dict("os.environ", {"WORLD_SIZE": "1"}),
-    ):
+    with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
         assert adapter.supports_low_memory_dcp_load is True
 
     adapter.backend = SimpleNamespace(experts="torch", dispatcher="mok")
diff --git a/tests/unit_tests/moe/test_state_dict_mixin.py b/tests/unit_tests/moe/test_state_dict_mixin.py
index 518fb995f3..bb46c296c5 100644
--- a/tests/unit_tests/moe/test_state_dict_mixin.py
+++ b/tests/unit_tests/moe/test_state_dict_mixin.py
@@ -22,7 +22,23 @@
 
 skip_if_no_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for GPU operations")
 
-from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin
+from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin, get_world_size_safe
+
+
+def test_get_world_size_safe_uses_initialized_process_group():
+    with (
+        patch("nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=True),
+        patch("nemo_automodel.components.moe.state_dict_mixin.torch.distributed.get_world_size", return_value=8),
+    ):
+        assert get_world_size_safe() == 8
+
+
+def test_get_world_size_safe_uses_preinit_environment():
+    with (
+        patch("nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False),
+        patch.dict("os.environ", {"WORLD_SIZE": "8"}),
+    ):
+        assert get_world_size_safe() == 8
 
 
 class MockMoEConfig:
@@ -1054,19 +1070,9 @@ def test_expert_checkpoint_storage_capability_matches_grouped_storage_aliasing(s
 
         mixin.backend.experts = "te"
         mixin.backend.dispatcher = "deepep"
-        with (
-            patch(
-                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
-            ),
-            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
-        ):
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
             assert mixin._expert_checkpoint_tensors_use_model_storage is False
-        with (
-            patch(
-                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
-            ),
-            patch.dict("os.environ", {"WORLD_SIZE": "1"}),
-        ):
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
             assert mixin._expert_checkpoint_tensors_use_model_storage is True
 
         mixin.backend.dispatcher = "torch"
@@ -1200,12 +1206,7 @@ def test_inplace_load_skips_when_backend_experts_is_te_gate_and_up(self):
         splits = [local_storage[i] for i in range(2)]
         mock_dtensor = Mock()
 
-        with (
-            patch(
-                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
-            ),
-            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
-        ):
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
             result = self._run_inplace_conversion(
                 mixin, "model.layers.0.mlp.experts.gate_and_up_projs", mock_dtensor, splits
             )
@@ -1259,10 +1260,7 @@ def test_te_checkpoint_layout_views_are_reused(self):
         virtual_down = down_storage.transpose(-1, -2)
 
         with (
-            patch(
-                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
-            ),
-            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
+            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
             patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
             patch("torch.empty_like") as empty_like,
         ):
@@ -1308,10 +1306,7 @@ def test_te_checkpoint_layout_cuda_peak_stays_within_one_buffer(self):
         virtual_gate_up = gate_up_storage.transpose(-1, -2)
 
         with (
-            patch(
-                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
-            ),
-            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
+            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
             patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
             patch("nemo_automodel.components.moe.state_dict_mixin.gc.collect") as collect,
             patch("torch.cuda.empty_cache") as empty_cache,
@@ -1336,10 +1331,7 @@ def test_temporary_checkpoint_views_round_trip_loaded_values(self):
         initialized_down = torch.full((2, 3, 4), 99.0)
 
         with (
-            patch(
-                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
-            ),
-            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
+            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8),
             patch("nemo_automodel.components.moe.state_dict_utils.is_dtensor", return_value=False),
         ):
             destinations = dict(
@@ -1392,12 +1384,7 @@ def test_inplace_load_skips_when_backend_experts_is_te_down_projs(self):
         mock_dtensor.shape = (2, 512, 1024)
         mock_dtensor.is_meta = False
 
-        with (
-            patch(
-                "nemo_automodel.components.moe.state_dict_mixin.torch.distributed.is_initialized", return_value=False
-            ),
-            patch.dict("os.environ", {"WORLD_SIZE": "8"}),
-        ):
+        with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
             result = self._run_inplace_conversion(mixin, "model.layers.3.mlp.experts.down_projs", mock_dtensor, splits)
 
         assert result is not None and len(result) == 2

From 47f25b482ed3b055dfca19769cb5b9a89b68e2ba Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Tue, 25 Aug 2026 17:03:05 -0700
Subject: [PATCH 13/20] test(checkpoint): exercise allocating TE expert storage

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 tests/unit_tests/moe/test_state_dict_mixin.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/unit_tests/moe/test_state_dict_mixin.py b/tests/unit_tests/moe/test_state_dict_mixin.py
index bb46c296c5..3c18155eb6 100644
--- a/tests/unit_tests/moe/test_state_dict_mixin.py
+++ b/tests/unit_tests/moe/test_state_dict_mixin.py
@@ -913,6 +913,7 @@ class TestConvertSingleMergedExpertToHfSplitExperts:
     def test_allocating_cuda_conversions_use_generation_zero_collection(self):
         mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=3)
         mixin.backend.experts = "te"
+        mixin.backend.dispatcher = "deepep"
         gate_up_tensor = Mock(spec=torch.Tensor, is_meta=False, is_cuda=True)
         down_tensor = Mock(spec=torch.Tensor, ndim=3, shape=(2, 3, 4), is_meta=False, is_cuda=True)
         gate_up_splits = [
@@ -921,6 +922,7 @@ def test_allocating_cuda_conversions_use_generation_zero_collection(self):
         down_splits = [torch.arange(12, dtype=torch.float32).reshape(3, 4) + 12 * expert_id for expert_id in range(2)]
 
         with (
+            patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=2),
             patch.object(mixin, "_split_experts_weights", side_effect=[gate_up_splits, down_splits]),
             patch("torch.cuda.is_available", return_value=True),
             patch("nemo_automodel.components.moe.state_dict_mixin.gc.collect") as collect,

From d047b3a2df4d6340e2476533e5fa09105442eb5c Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Wed, 26 Aug 2026 07:54:23 -0700
Subject: [PATCH 14/20] fix(checkpoint): tighten low-memory MoE load safety

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../models/ling_v2/state_dict_adapter.py      |  4 +-
 .../nemotron_omni/state_dict_adapter.py       |  5 +
 .../models/nemotron_v3/state_dict_adapter.py  | 22 +++--
 nemo_automodel/components/moe/layers.py       |  2 +-
 .../components/moe/state_dict_mixin.py        | 41 +++++---
 .../test_state_dict_adapter_capabilities.py   | 22 ++++-
 .../test_ling_v2_state_dict_adapter.py        | 69 +++++++++++++
 .../test_nemotron_omni_state_dict_adapter.py  | 12 +++
 .../test_nemotron_v3_state_dict_adapter.py    | 51 ++++++++++
 .../test_qwen3_next_state_dict_adapter.py     |  3 +
 tests/unit_tests/moe/test_state_dict_mixin.py | 96 ++++++++++++++++---
 11 files changed, 290 insertions(+), 37 deletions(-)

diff --git a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
index f5d9b24f0e..94b5412a13 100644
--- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py
@@ -87,6 +87,8 @@ def _rename_native_to_hf(key: str) -> str:
 class BailingMoeV2StateDictAdapter(MoESplitExpertsStateDictMixin, StateDictAdapter):
     """State-dict adapter for BailingMoeV2 / Ling 2.0 checkpoints."""
 
+    _supports_low_memory_dcp_load = True
+
     def __init__(
         self,
         config: BailingMoeV2Config,
@@ -103,7 +105,7 @@ def __init__(
     @property
     def supports_low_memory_dcp_load(self) -> bool:
         """Whether Ling's DCP load needs only small fused-QKV temporary tensors."""
-        return self._expert_checkpoint_tensors_use_model_storage and not self.moe_config.expert_bias
+        return self.moe_config is not None and super().supports_low_memory_dcp_load
 
     # ---- HF -> native ----------------------------------------------------
 
diff --git a/nemo_automodel/components/models/nemotron_omni/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_omni/state_dict_adapter.py
index 62ce6be713..74c44bd379 100644
--- a/nemo_automodel/components/models/nemotron_omni/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_omni/state_dict_adapter.py
@@ -110,6 +110,11 @@ def supports_low_memory_dcp_load(self) -> bool:
         """Whether the embedded language adapter supports a low-memory DCP load."""
         return self._llm_adapter.supports_low_memory_dcp_load
 
+    @property
+    def view_loaded_native_keys(self) -> set[str]:
+        """Language-model keys loaded through model-backed checkpoint views."""
+        return {f"language_model.{key}" for key in self._llm_adapter.view_loaded_native_keys}
+
     def __init__(
         self,
         config,
diff --git a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
index a430d15484..d275468a68 100644
--- a/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/nemotron_v3/state_dict_adapter.py
@@ -272,14 +272,18 @@ def from_hf(
             # reset_view_loaded_keys=False: this is the second merge of a single from_hf (after the
             # backbone merge above), so accumulate MTP view-loaded keys onto the backbone's record.
             prior_view_keys = set(self.view_loaded_native_keys)
-            merged_mtp = self._from_hf_w_merged_experts(stripped, device_mesh, reset_view_loaded_keys=False)
+            merged_mtp = (
+                stripped
+                if self.moe_config is None
+                else self._from_hf_w_merged_experts(stripped, device_mesh, reset_view_loaded_keys=False)
+            )
             for key, value in merged_mtp.items():
                 merged[f"mtp.{key}"] = value
             # The merge loop records view-loaded keys in mtp.-stripped form (it only ever sees
             # stripped keys); re-prefix them so the checkpoint loader's key-diff matches them
             # against the model's real mtp.* parameter names instead of flagging them as
             # missing/unexpected.
-            new_view_keys = self._view_loaded_native_keys - prior_view_keys
+            new_view_keys = self.view_loaded_native_keys - prior_view_keys
             self._view_loaded_native_keys = prior_view_keys | {f"mtp.{key}" for key in new_view_keys}
 
         return merged
@@ -302,11 +306,15 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
         # emitted HF keys stay under ``mtp.`` instead of ``backbone.``.
         if fqn.startswith("mtp."):
             fqn = _strip_mamba_fp32_holder_key(fqn)
-            expert_split = self._convert_single_merged_expert_to_hf_split_experts(
-                fqn,
-                tensor,
-                prefix_override="mtp.",
-                **kwargs,
+            expert_split = (
+                None
+                if self.moe_config is None
+                else self._convert_single_merged_expert_to_hf_split_experts(
+                    fqn,
+                    tensor,
+                    prefix_override="mtp.",
+                    **kwargs,
+                )
             )
             result = expert_split if expert_split is not None else [(fqn, tensor)]
             result = [(key, _upcast_mamba_fp32_state_tensor(key, value)) for key, value in result]
diff --git a/nemo_automodel/components/moe/layers.py b/nemo_automodel/components/moe/layers.py
index 30e3a39773..dfb832bada 100644
--- a/nemo_automodel/components/moe/layers.py
+++ b/nemo_automodel/components/moe/layers.py
@@ -786,7 +786,7 @@ def __init__(self, config: MoEConfig, backend: BackendConfig):
                     dispatcher_async_dispatch=backend.dispatcher_async_dispatch,
                 )
             else:
-                # experts == "te"
+                # All other expert backends use the TE grouped implementation.
                 self.experts = GroupedExpertsTE(
                     config,
                     backend=backend,
diff --git a/nemo_automodel/components/moe/state_dict_mixin.py b/nemo_automodel/components/moe/state_dict_mixin.py
index c21c9082b9..1726ff0844 100644
--- a/nemo_automodel/components/moe/state_dict_mixin.py
+++ b/nemo_automodel/components/moe/state_dict_mixin.py
@@ -64,8 +64,9 @@ class MoESplitExpertsStateDictMixin:
     @property
     def supports_low_memory_dcp_load(self) -> bool:
         """Whether DCP needs at most small temporary tensors for this MoE checkpoint."""
-        experts_use_model_storage = self.moe_config is None or self._expert_checkpoint_tensors_use_model_storage
-        return self._supports_low_memory_dcp_load and experts_use_model_storage
+        if not self._supports_low_memory_dcp_load or self.moe_config is None:
+            return self._supports_low_memory_dcp_load
+        return self._expert_checkpoint_tensors_use_model_storage and not self.moe_config.expert_bias
 
     @property
     def _expert_checkpoint_tensors_use_model_storage(self) -> bool:
@@ -80,17 +81,17 @@ def _expert_checkpoint_tensors_use_model_storage(self) -> bool:
     def _grouped_expert_storage_is_model_weight(self) -> bool:
         """Whether the runtime grouped expert tensors use the model's parameter storage.
 
-        TE normally exposes virtual grouped tensors backed by ``torch.stack`` copies. The MoE layer instead constructs
-        ordinary ``GroupedExperts`` when an EP dispatcher runs with world size one, so the same TE configuration uses
-        real model parameters in that runtime. Non-EP dispatchers also construct ordinary grouped experts.
+        This mirrors the expert implementation selected by ``MoE.__init__``. EP dispatchers use ordinary
+        ``GroupedExperts`` at world size one, ``GroupedExpertsDeepEP`` for the grouped backends at larger world sizes,
+        and ``GroupedExpertsTE`` otherwise. Non-EP dispatchers construct ordinary grouped experts.
         """
-        if self.backend.experts != "te":
-            return True
         if self.backend.dispatcher == "mok":
-            return False
+            return self.backend.experts != "te"
         if self.backend.dispatcher not in {"deepep", "hybridep", "uccl_ep"}:
             return True
-        return get_world_size_safe() == 1
+        if get_world_size_safe() == 1:
+            return True
+        return self.backend.experts in {"gmm", "torch_mm", "torch_mm_mxfp8"}
 
     @property
     def _is_gated_moe(self) -> bool:
@@ -597,6 +598,12 @@ def _to_hf_w_split_experts(self, state_dict: dict[str, Any], **kwargs: Any) -> d
         ``_convert_single_merged_expert_to_hf_split_experts`` for adapter
         compatibility (e.g. ``exclude_key_regex``).
         """
+        if self.moe_config is not None and self.moe_config.expert_bias:
+            raise NotImplementedError(
+                "Checkpoint conversion for grouped experts with expert_bias=True is not implemented; "
+                "refusing to pass expert bias tensors through under incorrect keys."
+            )
+
         hf_state_dict: dict[str, Any] = {}
 
         for fqn, tensor in state_dict.items():
@@ -643,6 +650,12 @@ def _from_hf_w_merged_experts(
             [local_experts, hidden, expert_hidden], and down projections have shape
             [local_experts, expert_hidden, hidden].
         """
+        if self.moe_config is not None and self.moe_config.expert_bias:
+            raise NotImplementedError(
+                "Checkpoint conversion for grouped experts with expert_bias=True is not implemented; "
+                "refusing to pass expert bias tensors through under incorrect keys."
+            )
+
         if reset_view_loaded_keys:
             self._view_loaded_native_keys = set()
 
@@ -884,6 +897,12 @@ def _convert_single_merged_expert_to_hf_split_experts(
         prefix = prefix_override if prefix_override is not None else self._hf_prefix
         expert_segment = self._expert_path_segment
 
+        if self.moe_config.expert_bias and f".{expert_segment}." in fqn:
+            raise NotImplementedError(
+                "Checkpoint conversion for grouped experts with expert_bias=True is not implemented; "
+                "refusing to pass expert bias tensors through under incorrect keys."
+            )
+
         # MoE expert LoRA keys do not depend on the runtime backend or its checkpoint storage aliases.
         # When v4_compatible=True, emit per-expert split keys. Otherwise, adapters that explicitly opt in emit the
         # fused ParamWrapper format; unvalidated adapters retain the legacy per-expert format.
@@ -942,7 +961,7 @@ def checkpoint_load_destination(view: torch.Tensor, source: torch.Tensor) -> tor
                 and not quantization
                 and gate_up_storage_is_model_weight
             )
-            if inplace_ok:
+            if inplace_ok and for_checkpoint_load:
                 self._register_inplace_loaded_key(fqn, prefix_override)
 
             result = []
@@ -995,7 +1014,7 @@ def checkpoint_load_destination(view: torch.Tensor, source: torch.Tensor) -> tor
                 and not quantization
                 and down_storage_is_model_weight
             )
-            if inplace_ok:
+            if inplace_ok and for_checkpoint_load:
                 self._register_inplace_loaded_key(fqn, prefix_override)
 
             result = []
diff --git a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
index a80e9ddfb0..7c4f70a67a 100644
--- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
+++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import re
 from types import SimpleNamespace
 from unittest.mock import patch
 
@@ -59,7 +60,16 @@ def __init__(self) -> None:
 
     def to_hf(self, state_dict, **kwargs):
         self.converted_devices = [value.device for value in state_dict.values() if isinstance(value, torch.Tensor)]
-        return {"fused.weight": torch.cat((state_dict.pop("q.weight"), state_dict.pop("k.weight")), dim=0)}
+        converted = {"fused.weight": torch.cat((state_dict.pop("q.weight"), state_dict.pop("k.weight")), dim=0)}
+        exclude_key_regex = kwargs.get("exclude_key_regex")
+        converted.update(
+            {
+                key: value
+                for key, value in state_dict.items()
+                if exclude_key_regex is None or not re.match(exclude_key_regex, key)
+            }
+        )
+        return converted
 
     def from_hf(self, hf_state_dict, device_mesh=None, **kwargs):
         return hf_state_dict
@@ -73,13 +83,15 @@ def test_get_hf_state_dict_keys_uses_shape_only_tensors() -> None:
     state_dict = {
         "q.weight": torch.ones(4, 8),
         "k.weight": torch.ones(2, 8),
+        "buffer": torch.ones(1),
+        "layer._extra_state": {"recipe": "metadata"},
     }
 
     keys = adapter.get_hf_state_dict_keys(state_dict)
 
-    assert keys == ["fused.weight"]
-    assert adapter.converted_devices == [torch.device("meta"), torch.device("meta")]
-    assert list(state_dict) == ["q.weight", "k.weight"]
+    assert keys == ["fused.weight", "buffer"]
+    assert adapter.converted_devices == [torch.device("meta"), torch.device("meta"), torch.device("meta")]
+    assert list(state_dict) == ["q.weight", "k.weight", "buffer", "layer._extra_state"]
 
 
 def _assert_destinations_write_through(
@@ -159,7 +171,7 @@ def test_write_through_adapters_expose_aliasing_destinations(
 def test_low_memory_dcp_grouped_adapters_preserve_non_expert_storage_and_require_model_backed_experts(
     adapter_type: type[StateDictAdapter],
 ) -> None:
-    moe_config = SimpleNamespace(n_routed_experts=2, moe_inter_dim=2, expert_activation="silu")
+    moe_config = SimpleNamespace(n_routed_experts=2, moe_inter_dim=2, expert_activation="silu", expert_bias=False)
     adapter = adapter_type(
         SimpleNamespace(num_hidden_layers=1),
         moe_config,
diff --git a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
index 74df7caef6..962f045c83 100644
--- a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
+++ b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
@@ -14,7 +14,10 @@
 
 import pytest
 import torch
+import torch.distributed.checkpoint as dcp
+from safetensors.torch import save_file
 
+from nemo_automodel.components.checkpoint._backports.hf_storage import _HuggingFaceStorageReader
 from nemo_automodel.components.models.common import BackendConfig
 from nemo_automodel.components.models.ling_v2.config import BailingMoeV2Config
 from nemo_automodel.components.models.ling_v2.state_dict_adapter import BailingMoeV2StateDictAdapter
@@ -225,6 +228,72 @@ def test_hf_to_native_to_hf_is_lossless_for_qkv_path(self, adapter, config):
 
 
 class TestLowMemoryDcpLoad:
+    def test_hf_reader_loads_qkv_and_experts_into_native_storage(self, adapter, config, moe_config, tmp_path):
+        q_size = config.num_attention_heads * config.head_dim
+        kv_size = config.num_key_value_heads * config.head_dim
+        native = {
+            "model.layers.1.self_attn.q_proj.weight": torch.zeros(q_size, config.hidden_size),
+            "model.layers.1.self_attn.k_proj.weight": torch.zeros(kv_size, config.hidden_size),
+            "model.layers.1.self_attn.v_proj.weight": torch.zeros(kv_size, config.hidden_size),
+            "model.layers.1.mlp.experts.gate_and_up_projs": torch.zeros(
+                moe_config.n_routed_experts, moe_config.dim, 2 * moe_config.moe_inter_dim
+            ),
+            "model.layers.1.mlp.experts.down_projs": torch.zeros(
+                moe_config.n_routed_experts, moe_config.moe_inter_dim, moe_config.dim
+            ),
+        }
+        generator = torch.Generator().manual_seed(7)
+        hf_checkpoint = {
+            "model.layers.1.attention.query_key_value.weight": torch.randn(
+                q_size + 2 * kv_size, config.hidden_size, generator=generator
+            )
+        }
+        for expert_id in range(moe_config.n_routed_experts):
+            hf_checkpoint[f"model.layers.1.mlp.experts.{expert_id}.gate_proj.weight"] = torch.randn(
+                moe_config.moe_inter_dim, moe_config.dim, generator=generator
+            )
+            hf_checkpoint[f"model.layers.1.mlp.experts.{expert_id}.up_proj.weight"] = torch.randn(
+                moe_config.moe_inter_dim, moe_config.dim, generator=generator
+            )
+            hf_checkpoint[f"model.layers.1.mlp.experts.{expert_id}.down_proj.weight"] = torch.randn(
+                moe_config.dim, moe_config.moe_inter_dim, generator=generator
+            )
+        save_file(hf_checkpoint, tmp_path / "model.safetensors")
+
+        destinations = adapter.to_hf(dict(native), for_checkpoint_load=True)
+        assert set(destinations) == set(hf_checkpoint)
+
+        # CPU tensors use the same checkpoint views but cannot trigger the CUDA/DTensor registration condition.
+        # Register the two grouped destinations explicitly to simulate the production DCP load lifecycle.
+        adapter._register_inplace_loaded_key("model.layers.1.mlp.experts.gate_and_up_projs", None)
+        adapter._register_inplace_loaded_key("model.layers.1.mlp.experts.down_projs", None)
+        dcp.load(destinations, storage_reader=_HuggingFaceStorageReader(path=tmp_path))
+
+        converted = adapter.from_hf(destinations)
+        for key, value in converted.items():
+            native[key].copy_(value)
+
+        fused_qkv = hf_checkpoint["model.layers.1.attention.query_key_value.weight"]
+        torch.testing.assert_close(native["model.layers.1.self_attn.q_proj.weight"], fused_qkv[:q_size])
+        torch.testing.assert_close(
+            native["model.layers.1.self_attn.k_proj.weight"], fused_qkv[q_size : q_size + kv_size]
+        )
+        torch.testing.assert_close(native["model.layers.1.self_attn.v_proj.weight"], fused_qkv[q_size + kv_size :])
+        for expert_id in range(moe_config.n_routed_experts):
+            torch.testing.assert_close(
+                native["model.layers.1.mlp.experts.gate_and_up_projs"][expert_id, :, : moe_config.moe_inter_dim],
+                hf_checkpoint[f"model.layers.1.mlp.experts.{expert_id}.gate_proj.weight"].T,
+            )
+            torch.testing.assert_close(
+                native["model.layers.1.mlp.experts.gate_and_up_projs"][expert_id, :, moe_config.moe_inter_dim :],
+                hf_checkpoint[f"model.layers.1.mlp.experts.{expert_id}.up_proj.weight"].T,
+            )
+            torch.testing.assert_close(
+                native["model.layers.1.mlp.experts.down_projs"][expert_id],
+                hf_checkpoint[f"model.layers.1.mlp.experts.{expert_id}.down_proj.weight"].T,
+            )
+        assert set(native) == set(converted) | adapter.view_loaded_native_keys
+
     def test_uses_only_fused_qkv_temporaries(self, adapter, config, moe_config):
         q_size = config.num_attention_heads * config.head_dim
         kv_size = config.num_key_value_heads * config.head_dim
diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_state_dict_adapter.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_state_dict_adapter.py
index 1efbaa837f..3e0e05eb76 100644
--- a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_state_dict_adapter.py
+++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_state_dict_adapter.py
@@ -122,6 +122,18 @@ def test_from_hf_routes_llm_through_v3_adapter_with_prefix(adapter):
     assert not any(k.startswith("language_model.") for k in delegated)
 
 
+def test_view_loaded_llm_keys_restore_outer_prefix(adapter):
+    adapter._llm_adapter.view_loaded_native_keys = {
+        "model.layers.0.mixer.experts.gate_and_up_projs",
+        "model.layers.0.mixer.experts.down_projs",
+    }
+
+    assert adapter.view_loaded_native_keys == {
+        "language_model.model.layers.0.mixer.experts.gate_and_up_projs",
+        "language_model.model.layers.0.mixer.experts.down_projs",
+    }
+
+
 # ---------------------------------------------------------------------------
 # to_hf — Automodel layout -> HF layout (round-trip)
 # ---------------------------------------------------------------------------
diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
index 92a1f4d8b3..36c4d732ba 100644
--- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
+++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_state_dict_adapter.py
@@ -228,6 +228,18 @@ def test_round_trip_dense(self, adapter):
 
         assert set(back.keys()) == set(hf_sd.keys())
 
+    def test_dense_mtp_keys_do_not_enter_moe_conversion(self, adapter):
+        hf_tensor = torch.randn(4, dtype=torch.bfloat16)
+        native = adapter.from_hf({"mtp.layers.0.mixer.A_log": hf_tensor})
+
+        assert set(native) == {"mtp.layers.0.mixer._fp32_params.A_log"}
+        assert native["mtp.layers.0.mixer._fp32_params.A_log"].dtype == torch.float32
+
+        exported = adapter.convert_single_tensor_to_hf(
+            "mtp.layers.0.mixer._fp32_params.A_log", native["mtp.layers.0.mixer._fp32_params.A_log"]
+        )
+        assert exported[0][0] == "mtp.layers.0.mixer.A_log"
+
     def test_peft_outer_prefix_round_trip(self, adapter):
         hf_key = "base_model.model.backbone.layers.0.mixer.in_proj.lora_A.weight"
         native_key = "base_model.model.model.layers.0.mixer.in_proj.lora_A.weight"
@@ -287,6 +299,45 @@ def test_view_loaded_mtp_experts_keep_native_namespace(self):
             "mtp.layers.1.mixer.experts.down_projs",
         }
 
+    def test_mtp_expert_conversion_forwards_checkpoint_kwargs(self):
+        moe_config = MoEConfig(
+            n_routed_experts=2,
+            n_shared_experts=1,
+            n_activated_experts=1,
+            n_expert_groups=1,
+            n_limited_groups=1,
+            train_gate=True,
+            gate_bias_update_factor=0.0,
+            aux_loss_coeff=0.0,
+            score_func="sigmoid",
+            route_scale=1.0,
+            dim=256,
+            inter_dim=512,
+            moe_inter_dim=128,
+            norm_topk_prob=False,
+            expert_bias=False,
+            expert_activation="relu2",
+            dtype=torch.bfloat16,
+        )
+        adapter = NemotronV3StateDictAdapter(MockNemotronV3Config(), moe_config, BackendConfig())
+        tensor = torch.randn(moe_config.n_routed_experts, 256, 128)
+
+        with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None) as convert:
+            adapter.convert_single_tensor_to_hf(
+                "mtp.layers.1.mixer.experts.gate_and_up_projs",
+                tensor,
+                for_checkpoint_load=True,
+                v4_compatible=True,
+            )
+
+        convert.assert_called_once_with(
+            "mtp.layers.1.mixer.experts.gate_and_up_projs",
+            tensor,
+            prefix_override="mtp.",
+            for_checkpoint_load=True,
+            v4_compatible=True,
+        )
+
 
 class TestNemotronV3AdapterToHf:
     """Test to_hf conversion."""
diff --git a/tests/unit_tests/models/qwen3_next/test_qwen3_next_state_dict_adapter.py b/tests/unit_tests/models/qwen3_next/test_qwen3_next_state_dict_adapter.py
index 6ee44c9b24..e61cc27067 100644
--- a/tests/unit_tests/models/qwen3_next/test_qwen3_next_state_dict_adapter.py
+++ b/tests/unit_tests/models/qwen3_next/test_qwen3_next_state_dict_adapter.py
@@ -160,6 +160,7 @@ def create_mock_moe_config(self, **overrides):
         moe_config = Mock(spec=MoEConfig)
         moe_config.n_routed_experts = overrides.get("n_routed_experts", 8)
         moe_config.moe_inter_dim = overrides.get("moe_inter_dim", 512)
+        moe_config.expert_bias = overrides.get("expert_bias", False)
         for key, value in overrides.items():
             setattr(moe_config, key, value)
         return moe_config
@@ -457,6 +458,7 @@ def create_mock_moe_config(self):
         moe_config = Mock()
         moe_config.n_routed_experts = 8
         moe_config.moe_inter_dim = 512
+        moe_config.expert_bias = False
         return moe_config
 
     def create_mock_backend_config(self):
@@ -565,6 +567,7 @@ def _make_adapter(self):
         moe_config = Mock()
         moe_config.n_routed_experts = 8
         moe_config.moe_inter_dim = 512
+        moe_config.expert_bias = False
         backend = Mock()
         backend.dispatcher = "torch"
         backend.experts = "torch"
diff --git a/tests/unit_tests/moe/test_state_dict_mixin.py b/tests/unit_tests/moe/test_state_dict_mixin.py
index 3c18155eb6..7f8abf5ebb 100644
--- a/tests/unit_tests/moe/test_state_dict_mixin.py
+++ b/tests/unit_tests/moe/test_state_dict_mixin.py
@@ -42,10 +42,11 @@ def test_get_world_size_safe_uses_preinit_environment():
 
 
 class MockMoEConfig:
-    def __init__(self, n_routed_experts=8, moe_inter_dim=512, expert_activation="swiglu"):
+    def __init__(self, n_routed_experts=8, moe_inter_dim=512, expert_activation="swiglu", expert_bias=False):
         self.n_routed_experts = n_routed_experts
         self.moe_inter_dim = moe_inter_dim
         self.expert_activation = expert_activation
+        self.expert_bias = expert_bias
 
 
 class MockConfig:
@@ -60,6 +61,8 @@ def __init__(self):
 
 
 class MockMoEStateDictMixin(MoESplitExpertsStateDictMixin):
+    _supports_low_memory_dcp_load = True
+
     def __init__(
         self,
         n_experts=8,
@@ -107,6 +110,7 @@ def _run_ep_free_dtensor_split(rank: int, world_size: int, init_file: str) -> No
             mixin._convert_single_merged_expert_to_hf_split_experts(
                 "model.layers.0.mlp.experts.gate_and_up_projs",
                 expert_sharded_weight,
+                for_checkpoint_load=True,
             )
         )
         assert len(converted) == 4
@@ -122,6 +126,7 @@ def _run_ep_free_dtensor_split(rank: int, world_size: int, init_file: str) -> No
             mixin._convert_single_merged_expert_to_hf_split_experts(
                 "model.layers.0.mlp.experts.down_projs",
                 down_weight,
+                for_checkpoint_load=True,
             )
         )
 
@@ -1066,14 +1071,17 @@ class TestInplaceLoadViews:
     conversion function, so patches must target that module path.
     """
 
-    def test_expert_checkpoint_storage_capability_matches_grouped_storage_aliasing(self):
+    @pytest.mark.parametrize("dispatcher", ["deepep", "hybridep", "uccl_ep"])
+    @pytest.mark.parametrize(
+        ("experts", "expected"),
+        [("gmm", True), ("torch_mm", True), ("torch_mm_mxfp8", True), ("te", False), ("torch", False)],
+    )
+    def test_expert_checkpoint_storage_capability_matches_grouped_storage_aliasing(self, dispatcher, experts, expected):
         mixin = MockMoEStateDictMixin()
-        assert mixin._expert_checkpoint_tensors_use_model_storage is True
-
-        mixin.backend.experts = "te"
-        mixin.backend.dispatcher = "deepep"
+        mixin.backend.experts = experts
+        mixin.backend.dispatcher = dispatcher
         with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=8):
-            assert mixin._expert_checkpoint_tensors_use_model_storage is False
+            assert mixin._expert_checkpoint_tensors_use_model_storage is expected
         with patch("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", return_value=1):
             assert mixin._expert_checkpoint_tensors_use_model_storage is True
 
@@ -1084,7 +1092,24 @@ def test_expert_checkpoint_storage_capability_matches_grouped_storage_aliasing(s
         mixin.backend.dispatcher = "mok"
         assert mixin._expert_checkpoint_tensors_use_model_storage is False
 
-    def _run_inplace_conversion(self, mixin, fqn, mock_dtensor, splits):
+        mixin.backend.experts = "te"
+        assert mixin._grouped_expert_storage_is_model_weight is False
+        assert mixin._expert_checkpoint_tensors_use_model_storage is False
+
+    def test_grouped_expert_bias_is_rejected_instead_of_passed_through(self):
+        mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=3)
+        mixin.moe_config.expert_bias = True
+
+        assert mixin.supports_low_memory_dcp_load is False
+        with pytest.raises(NotImplementedError, match="expert_bias=True"):
+            mixin._from_hf_w_merged_experts({"model.layers.0.mlp.experts.0.gate_proj.weight": torch.randn(3, 4)})
+        with pytest.raises(NotImplementedError, match="expert_bias=True"):
+            mixin._convert_single_merged_expert_to_hf_split_experts(
+                "model.layers.0.mlp.experts.gate_and_up_projs",
+                torch.randn(2, 4, 6),
+            )
+
+    def _run_inplace_conversion(self, mixin, fqn, mock_dtensor, splits, *, for_checkpoint_load=False):
         mixin._split_experts_weights = Mock(return_value=splits)
         mixin._last_expert_ids = list(range(len(splits)))
 
@@ -1095,7 +1120,9 @@ def _run_inplace_conversion(self, mixin, fqn, mock_dtensor, splits):
             ),
             patch("nemo_automodel.components.moe.state_dict_utils.validate_dtensor_expert_sharding"),
         ):
-            return mixin._convert_single_merged_expert_to_hf_split_experts(fqn, mock_dtensor)
+            return mixin._convert_single_merged_expert_to_hf_split_experts(
+                fqn, mock_dtensor, for_checkpoint_load=for_checkpoint_load
+            )
 
     def test_inplace_load_gate_and_up_returns_views(self):
         mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=512)
@@ -1105,7 +1132,11 @@ def test_inplace_load_gate_and_up_returns_views(self):
         mock_dtensor = Mock()
 
         result = self._run_inplace_conversion(
-            mixin, "model.layers.0.mlp.experts.gate_and_up_projs", mock_dtensor, splits
+            mixin,
+            "model.layers.0.mlp.experts.gate_and_up_projs",
+            mock_dtensor,
+            splits,
+            for_checkpoint_load=True,
         )
 
         assert result is not None
@@ -1128,7 +1159,13 @@ def test_inplace_load_down_projs_returns_views(self):
         mock_dtensor.shape = (2, 512, 1024)
         mock_dtensor.is_meta = False
 
-        result = self._run_inplace_conversion(mixin, "model.layers.3.mlp.experts.down_projs", mock_dtensor, splits)
+        result = self._run_inplace_conversion(
+            mixin,
+            "model.layers.3.mlp.experts.down_projs",
+            mock_dtensor,
+            splits,
+            for_checkpoint_load=True,
+        )
 
         assert result is not None and len(result) == 2
         src_ptr = local_storage.untyped_storage().data_ptr()
@@ -1166,7 +1203,11 @@ def test_inplace_load_writes_through_to_model_storage(self):
         mock_dtensor = Mock()
 
         result = self._run_inplace_conversion(
-            mixin, "model.layers.0.mlp.experts.gate_and_up_projs", mock_dtensor, splits
+            mixin,
+            "model.layers.0.mlp.experts.gate_and_up_projs",
+            mock_dtensor,
+            splits,
+            for_checkpoint_load=True,
         )
 
         gate0 = next(v for k, v in result if k.endswith("0.gate_proj.weight"))
@@ -1196,6 +1237,36 @@ def test_from_hf_skips_rebuild_for_inplace_loaded_keys(self):
         assert "model.layers.0.mlp.experts.down_projs" not in out
         assert mixin._inplace_loaded_native_keys == set()
 
+    def test_save_conversion_does_not_poison_later_from_hf(self):
+        mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=3)
+        grouped_gate_up = torch.arange(2 * 4 * 6, dtype=torch.float32).reshape(2, 4, 6)
+        grouped_down = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4)
+        gate_up_dtensor = Mock()
+        down_dtensor = Mock(spec=["ndim", "shape", "is_meta"])
+        down_dtensor.ndim = 3
+        down_dtensor.shape = (2, 3, 4)
+        down_dtensor.is_meta = False
+
+        exported = self._run_inplace_conversion(
+            mixin,
+            "model.layers.0.mlp.experts.gate_and_up_projs",
+            gate_up_dtensor,
+            [grouped_gate_up[expert_id] for expert_id in range(2)],
+            for_checkpoint_load=False,
+        )
+        exported += self._run_inplace_conversion(
+            mixin,
+            "model.layers.0.mlp.experts.down_projs",
+            down_dtensor,
+            [grouped_down[expert_id] for expert_id in range(2)],
+            for_checkpoint_load=False,
+        )
+
+        assert not hasattr(mixin, "_inplace_loaded_native_keys")
+        restored = mixin._from_hf_w_merged_experts(dict(exported))
+        torch.testing.assert_close(restored["model.layers.0.mlp.experts.gate_and_up_projs"], grouped_gate_up)
+        torch.testing.assert_close(restored["model.layers.0.mlp.experts.down_projs"], grouped_down)
+
     def test_inplace_load_skips_when_backend_experts_is_te_gate_and_up(self):
         # GroupedExpertsTE (backend.experts == "te") exposes gate_and_up_projs as a
         # torch.stack copy of per-expert weights that does not alias the model's grouped
@@ -1429,6 +1500,7 @@ def test_inplace_load_skips_mok_gate_up_copy_but_keeps_down_view(self):
             "model.layers.3.mlp.experts.down_projs",
             down_dtensor,
             [down_storage[i] for i in range(2)],
+            for_checkpoint_load=True,
         )
 
         assert down_result is not None

From 6ee26af29192d6ebd0b6714e9c6dfd8947c5e927 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Wed, 26 Aug 2026 08:23:31 -0700
Subject: [PATCH 15/20] docs(skill): update checkpoint load capability guidance

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 skills/nemo-automodel-model-onboarding/SKILL.md | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/skills/nemo-automodel-model-onboarding/SKILL.md b/skills/nemo-automodel-model-onboarding/SKILL.md
index c2f7e9e33c..cd917f86d6 100644
--- a/skills/nemo-automodel-model-onboarding/SKILL.md
+++ b/skills/nemo-automodel-model-onboarding/SKILL.md
@@ -172,11 +172,13 @@ Implement files in dependency order:
 3. **layers.py** (if needed) -- Attention, MLP, decoder block classes
 4. **model.py** -- The main `ForCausalLM` (or `ForConditionalGeneration`) class
 5. **state_dict_adapter.py** -- HF weight conversion. Leave
-   `supports_write_through_checkpoint_load` disabled unless every load
-   destination returned by `to_hf` writes through to final model storage for
-   every supported backend and configuration. Any opt-in needs a focused
-   write-through storage test; allocating conversions can otherwise create a
-   model-sized device temporary and cause an out-of-memory failure.
+   `_supports_low_memory_dcp_load` disabled unless most checkpoint tensors load
+   directly into final model storage and every remaining allocating conversion
+   has a small, bounded temporary footprint for every supported backend and
+   configuration. Any opt-in needs focused tests that prove direct destinations
+   alias model storage, bound allocating conversions, and report
+   `supports_low_memory_dcp_load=False` for unsafe variants; a model-sized device
+   temporary can otherwise cause an out-of-memory failure.
 6. **__init__.py** -- Re-export the main model class
 
 See the pattern files for detailed implementation guidance:
@@ -419,8 +421,8 @@ that only surface in a full parity comparison.
 - [ ] Implemented layers.py (if custom layers needed)
 - [ ] Implemented rope_utils.py (if custom RoPE needed)
 - [ ] Implemented model.py with `HFCheckpointingMixin`
-- [ ] Implemented state_dict_adapter.py; any write-through opt-in proves writes reach model storage and reports
-  `False` for allocating variants
+- [ ] Implemented state_dict_adapter.py; any low-memory DCP opt-in proves direct destinations reach model storage,
+  bounds allocating conversions, and reports `False` for unsafe variants
 - [ ] Implemented __init__.py with re-export
 - [ ] Registered in `MODEL_ARCH_MAPPING` in `_transformers/registry.py`
 - [ ] Registered custom config in `_CUSTOM_CONFIG_REGISTRATIONS` (if applicable)

From d3680bc58ce154ff4bc8f6d7b822ddc56c7b96e6 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Wed, 26 Aug 2026 10:56:06 -0700
Subject: [PATCH 16/20] fix(checkpoint): preserve low-memory adapter lifecycle

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../models/hy_mt2/state_dict_adapter.py       |  2 +-
 .../models/hy_v3/state_dict_adapter.py        |  2 +-
 .../minimax_m3_vl/state_dict_adapter.py       |  5 +++
 .../qwen3_omni_moe/state_dict_adapter.py      |  2 +-
 .../components/moe/state_dict_mixin.py        | 25 ++++-------
 .../nemo-automodel-model-onboarding/SKILL.md  |  8 ++++
 .../test_state_dict_adapter_capabilities.py   | 42 ++++++++++++++-----
 .../test_ernie4_5_state_dict_adapter.py       | 22 ++++++++++
 .../test_ling_v2_state_dict_adapter.py        |  6 +++
 .../minimax_m3_vl/test_minimax_m3_vlm.py      | 13 ++++++
 tests/unit_tests/moe/test_state_dict_mixin.py | 34 +++++++++++----
 11 files changed, 124 insertions(+), 37 deletions(-)

diff --git a/nemo_automodel/components/models/hy_mt2/state_dict_adapter.py b/nemo_automodel/components/models/hy_mt2/state_dict_adapter.py
index d48f0942ab..6082146655 100644
--- a/nemo_automodel/components/models/hy_mt2/state_dict_adapter.py
+++ b/nemo_automodel/components/models/hy_mt2/state_dict_adapter.py
@@ -92,7 +92,7 @@ def to_hf(
         **kwargs,
     ) -> dict[str, Any]:
         """Native -> on-disk Hy-MT2 HF: per-expert split + name renames."""
-        hf_split: dict[str, Any] = self._to_hf_w_split_experts(state_dict)
+        hf_split: dict[str, Any] = self._to_hf_w_split_experts(state_dict, **kwargs)
 
         out: dict[str, Any] = {}
         for k, v in hf_split.items():
diff --git a/nemo_automodel/components/models/hy_v3/state_dict_adapter.py b/nemo_automodel/components/models/hy_v3/state_dict_adapter.py
index 75ea99034a..a723d9fa54 100644
--- a/nemo_automodel/components/models/hy_v3/state_dict_adapter.py
+++ b/nemo_automodel/components/models/hy_v3/state_dict_adapter.py
@@ -111,7 +111,7 @@ def to_hf(
              gate.weight -> router.gate.weight, shared_experts. -> shared_mlp.).
         """
         # Step 1: per-expert split via the mixin. Pass-through for non-expert keys.
-        hf_split: dict[str, Any] = self._to_hf_w_split_experts(state_dict)
+        hf_split: dict[str, Any] = self._to_hf_w_split_experts(state_dict, **kwargs)
 
         # Step 2: rename native -> on-disk Tencent.
         out: dict[str, Any] = {}
diff --git a/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py b/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
index b8b3204102..979c2ee664 100644
--- a/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
+++ b/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
@@ -332,6 +332,11 @@ def __init__(self, config: Any, moe_config: MoEConfig, backend: BackendConfig, d
         self.config = config
         self.text_adapter = MiniMaxM3StateDictAdapter(config.text_config, moe_config, backend, dtype=dtype)
 
+    @property
+    def view_loaded_native_keys(self) -> set[str]:
+        """Native text keys loaded through the inner adapter's checkpoint views."""
+        return set(self.text_adapter.view_loaded_native_keys)
+
     @staticmethod
     def _map_non_text_from_hf(key: str) -> str | None:
         if ".mtp." in key:
diff --git a/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py
index c22958e535..501bd18126 100644
--- a/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py
+++ b/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py
@@ -48,7 +48,7 @@ def __init__(
     def to_hf(
         self, state_dict: dict[str, Any], exclude_key_regex: str | None = None, quantization: bool = False, **kwargs
     ) -> dict[str, Any]:
-        hf_state_dict = self._to_hf_w_split_experts(state_dict)
+        hf_state_dict = self._to_hf_w_split_experts(state_dict, quantization=quantization, **kwargs)
 
         if self._uses_thinker_prefix:
             hf_state_dict_with_prefix = {}
diff --git a/nemo_automodel/components/moe/state_dict_mixin.py b/nemo_automodel/components/moe/state_dict_mixin.py
index 1726ff0844..e7077fbf4d 100644
--- a/nemo_automodel/components/moe/state_dict_mixin.py
+++ b/nemo_automodel/components/moe/state_dict_mixin.py
@@ -598,12 +598,6 @@ def _to_hf_w_split_experts(self, state_dict: dict[str, Any], **kwargs: Any) -> d
         ``_convert_single_merged_expert_to_hf_split_experts`` for adapter
         compatibility (e.g. ``exclude_key_regex``).
         """
-        if self.moe_config is not None and self.moe_config.expert_bias:
-            raise NotImplementedError(
-                "Checkpoint conversion for grouped experts with expert_bias=True is not implemented; "
-                "refusing to pass expert bias tensors through under incorrect keys."
-            )
-
         hf_state_dict: dict[str, Any] = {}
 
         for fqn, tensor in state_dict.items():
@@ -650,18 +644,23 @@ def _from_hf_w_merged_experts(
             [local_experts, hidden, expert_hidden], and down projections have shape
             [local_experts, expert_hidden, hidden].
         """
+        expert_segment = self._expert_path_segment
         if self.moe_config is not None and self.moe_config.expert_bias:
-            raise NotImplementedError(
-                "Checkpoint conversion for grouped experts with expert_bias=True is not implemented; "
-                "refusing to pass expert bias tensors through under incorrect keys."
+            split_bias_pattern = re.compile(
+                rf"(?:^|\.){re.escape(expert_segment)}\.\d+\.(?:gate_proj|up_proj|down_proj)\.bias$"
             )
+            unsupported_bias_key = next((key for key in hf_state_dict if split_bias_pattern.search(key)), None)
+            if unsupported_bias_key is not None:
+                raise NotImplementedError(
+                    "Loading Hugging Face per-expert bias tensors with expert_bias=True is not implemented; "
+                    f"refusing key {unsupported_bias_key!r}."
+                )
 
         if reset_view_loaded_keys:
             self._view_loaded_native_keys = set()
 
         n_experts = self.moe_config.n_routed_experts
         is_gated = self._is_gated_moe
-        expert_segment = self._expert_path_segment
 
         self._validate_expert_availability(hf_state_dict, n_experts, device_mesh)
 
@@ -897,12 +896,6 @@ def _convert_single_merged_expert_to_hf_split_experts(
         prefix = prefix_override if prefix_override is not None else self._hf_prefix
         expert_segment = self._expert_path_segment
 
-        if self.moe_config.expert_bias and f".{expert_segment}." in fqn:
-            raise NotImplementedError(
-                "Checkpoint conversion for grouped experts with expert_bias=True is not implemented; "
-                "refusing to pass expert bias tensors through under incorrect keys."
-            )
-
         # MoE expert LoRA keys do not depend on the runtime backend or its checkpoint storage aliases.
         # When v4_compatible=True, emit per-expert split keys. Otherwise, adapters that explicitly opt in emit the
         # fused ParamWrapper format; unvalidated adapters retain the legacy per-expert format.
diff --git a/skills/nemo-automodel-model-onboarding/SKILL.md b/skills/nemo-automodel-model-onboarding/SKILL.md
index cd917f86d6..88536627a0 100644
--- a/skills/nemo-automodel-model-onboarding/SKILL.md
+++ b/skills/nemo-automodel-model-onboarding/SKILL.md
@@ -179,6 +179,14 @@ Implement files in dependency order:
    alias model storage, bound allocating conversions, and report
    `supports_low_memory_dcp_load=False` for unsafe variants; a model-sized device
    temporary can otherwise cause an out-of-memory failure.
+
+   For adapters using `MoESplitExpertsStateDictMixin`, treat the capability as
+   runtime-dependent. With the maintained EP dispatchers (`deepep`, `hybridep`,
+   and `uccl_ep`), TE expert storage is model-backed only at world size one;
+   `gmm`, `torch_mm`, and `torch_mm_mxfp8` expert storage is model-backed at any
+   world size. Non-EP dispatchers use ordinary grouped storage. The shared mixin
+   disables low-memory DCP for `expert_bias=True`, MoK, and any variant whose
+   checkpoint expert tensors require rebuilding.
 6. **__init__.py** -- Re-export the main model class
 
 See the pattern files for detailed implementation guidance:
diff --git a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
index 7c4f70a67a..68f8b89fd8 100644
--- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
+++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py
@@ -52,6 +52,19 @@
 from nemo_automodel.components.models.qwen3_vl_moe.state_dict_adapter import Qwen3VLMoeStateDictAdapter
 
 
+_LOW_MEMORY_GROUPED_ADAPTERS = (
+    Ernie4_5_MoeStateDictAdapter,
+    Glm4MoeStateDictAdapter,
+    HYV3StateDictAdapter,
+    HyMT2StateDictAdapter,
+    LagunaStateDictAdapter,
+    NemotronV3StateDictAdapter,
+    Qwen3MoeStateDictAdapter,
+    Qwen3NextStateDictAdapter,
+    Qwen3OmniMoeStateDictAdapter,
+)
+
+
 class _AllocatingKeysAdapter(StateDictAdapter):
     """Test adapter whose normal conversion joins two tensors."""
 
@@ -156,17 +169,7 @@ def test_write_through_adapters_expose_aliasing_destinations(
 
 @pytest.mark.parametrize(
     "adapter_type",
-    [
-        Ernie4_5_MoeStateDictAdapter,
-        Glm4MoeStateDictAdapter,
-        HYV3StateDictAdapter,
-        HyMT2StateDictAdapter,
-        LagunaStateDictAdapter,
-        NemotronV3StateDictAdapter,
-        Qwen3MoeStateDictAdapter,
-        Qwen3NextStateDictAdapter,
-        Qwen3OmniMoeStateDictAdapter,
-    ],
+    _LOW_MEMORY_GROUPED_ADAPTERS,
 )
 def test_low_memory_dcp_grouped_adapters_preserve_non_expert_storage_and_require_model_backed_experts(
     adapter_type: type[StateDictAdapter],
@@ -201,6 +204,23 @@ def test_low_memory_dcp_grouped_adapters_preserve_non_expert_storage_and_require
     assert adapter.supports_low_memory_dcp_load is False
 
 
+@pytest.mark.parametrize("adapter_type", _LOW_MEMORY_GROUPED_ADAPTERS)
+def test_low_memory_grouped_adapters_forward_checkpoint_load_flag(adapter_type: type[StateDictAdapter]) -> None:
+    adapter = adapter_type(
+        SimpleNamespace(num_hidden_layers=1),
+        SimpleNamespace(n_routed_experts=2, moe_inter_dim=2, expert_activation="silu", expert_bias=False),
+        SimpleNamespace(experts="torch", dispatcher="torch"),
+    )
+    with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None) as convert:
+        adapter.to_hf(
+            {"model.layers.0.input_layernorm.weight": torch.zeros(2, dtype=torch.bfloat16)},
+            for_checkpoint_load=True,
+        )
+
+    convert.assert_called_once()
+    assert convert.call_args.kwargs.get("for_checkpoint_load") is True
+
+
 @pytest.mark.parametrize(
     "adapter_type",
     [
diff --git a/tests/unit_tests/models/ernie4_5/test_ernie4_5_state_dict_adapter.py b/tests/unit_tests/models/ernie4_5/test_ernie4_5_state_dict_adapter.py
index ce8a3b3f7a..6f8e25eb44 100644
--- a/tests/unit_tests/models/ernie4_5/test_ernie4_5_state_dict_adapter.py
+++ b/tests/unit_tests/models/ernie4_5/test_ernie4_5_state_dict_adapter.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import replace
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
@@ -290,3 +291,24 @@ def test_to_hf_iterates_all_keys(self, moe_hf_config, moe_config, backend_config
         # gate bias should be renamed back to moe_statics in HF
         assert "model.layers.1.mlp.moe_statics.e_score_correction_bias" in out
         assert "model.embed_tokens.weight" in out
+
+    def test_grouped_expert_bias_does_not_break_hf_key_discovery(
+        self, moe_hf_config, moe_config, backend_config
+    ):
+        biased_hf_config = SimpleNamespace(**{**vars(moe_hf_config), "use_bias": True})
+        biased_moe_config = replace(moe_config, expert_bias=True)
+        adapter = Ernie4_5_MoeStateDictAdapter(biased_hf_config, biased_moe_config, backend_config)
+        state = {
+            "model.layers.1.mlp.experts.gate_and_up_projs": torch.zeros(4, 64, 64),
+            "model.layers.1.mlp.experts.down_projs": torch.zeros(4, 32, 64),
+            "model.layers.1.mlp.experts.gate_up_proj_bias": torch.zeros(4, 64),
+            "model.layers.1.mlp.experts.down_proj_bias": torch.zeros(4, 64),
+        }
+
+        keys = set(adapter.get_hf_state_dict_keys(state))
+
+        assert "model.layers.1.mlp.experts.gate_up_proj_bias" in keys
+        assert "model.layers.1.mlp.experts.down_proj_bias" in keys
+        assert "model.layers.1.mlp.experts.0.gate_proj.weight" in keys
+        assert "model.layers.1.mlp.experts.0.up_proj.weight" in keys
+        assert "model.layers.1.mlp.experts.0.down_proj.weight" in keys
diff --git a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
index 962f045c83..19ba67082a 100644
--- a/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
+++ b/tests/unit_tests/models/ling_v2/test_ling_v2_state_dict_adapter.py
@@ -270,6 +270,12 @@ def test_hf_reader_loads_qkv_and_experts_into_native_storage(self, adapter, conf
         dcp.load(destinations, storage_reader=_HuggingFaceStorageReader(path=tmp_path))
 
         converted = adapter.from_hf(destinations)
+        grouped_keys = {
+            "model.layers.1.mlp.experts.gate_and_up_projs",
+            "model.layers.1.mlp.experts.down_projs",
+        }
+        assert grouped_keys.isdisjoint(converted)
+        assert grouped_keys <= adapter.view_loaded_native_keys
         for key, value in converted.items():
             native[key].copy_(value)
 
diff --git a/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py b/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py
index fd5d925134..88651be00d 100644
--- a/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py
+++ b/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_vlm.py
@@ -242,6 +242,19 @@ def test_vlm_adapter_roundtrip_and_naming(vlm_model):
         assert torch.allclose(native[key].float(), back[key].float(), atol=1e-6), key
 
 
+def test_vlm_adapter_delegates_view_loaded_native_keys_without_prefix(vlm_model):
+    adapter = vlm_model.state_dict_adapter
+    adapter.text_adapter._view_loaded_native_keys = {
+        "model.layers.0.mlp.experts.gate_and_up_projs",
+        "model.layers.0.mlp.experts.down_projs",
+    }
+
+    assert adapter.view_loaded_native_keys == {
+        "model.layers.0.mlp.experts.gate_and_up_projs",
+        "model.layers.0.mlp.experts.down_projs",
+    }
+
+
 def test_video_grid_rejected(vlm_model):
     # Image-only support: grid_t beyond vision_segment_max_frames must fail loudly
     # (video temporal segmentation is not implemented).
diff --git a/tests/unit_tests/moe/test_state_dict_mixin.py b/tests/unit_tests/moe/test_state_dict_mixin.py
index 7f8abf5ebb..e4aefc4b33 100644
--- a/tests/unit_tests/moe/test_state_dict_mixin.py
+++ b/tests/unit_tests/moe/test_state_dict_mixin.py
@@ -1096,17 +1096,37 @@ def test_expert_checkpoint_storage_capability_matches_grouped_storage_aliasing(s
         assert mixin._grouped_expert_storage_is_model_weight is False
         assert mixin._expert_checkpoint_tensors_use_model_storage is False
 
-    def test_grouped_expert_bias_is_rejected_instead_of_passed_through(self):
+    def test_grouped_expert_bias_native_roundtrip_preserves_grouped_keys(self):
         mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=3)
         mixin.moe_config.expert_bias = True
+        native = {
+            "model.layers.0.mlp.experts.gate_and_up_projs": torch.randn(2, 4, 6),
+            "model.layers.0.mlp.experts.down_projs": torch.randn(2, 3, 4),
+            "model.layers.0.mlp.experts.gate_up_proj_bias": torch.randn(2, 6),
+            "model.layers.0.mlp.experts.down_proj_bias": torch.randn(2, 4),
+        }
 
         assert mixin.supports_low_memory_dcp_load is False
-        with pytest.raises(NotImplementedError, match="expert_bias=True"):
-            mixin._from_hf_w_merged_experts({"model.layers.0.mlp.experts.0.gate_proj.weight": torch.randn(3, 4)})
-        with pytest.raises(NotImplementedError, match="expert_bias=True"):
-            mixin._convert_single_merged_expert_to_hf_split_experts(
-                "model.layers.0.mlp.experts.gate_and_up_projs",
-                torch.randn(2, 4, 6),
+        hf_state = mixin._to_hf_w_split_experts(dict(native))
+        assert hf_state["model.layers.0.mlp.experts.gate_up_proj_bias"] is native[
+            "model.layers.0.mlp.experts.gate_up_proj_bias"
+        ]
+        assert hf_state["model.layers.0.mlp.experts.down_proj_bias"] is native[
+            "model.layers.0.mlp.experts.down_proj_bias"
+        ]
+
+        restored = mixin._from_hf_w_merged_experts(dict(hf_state))
+        assert set(restored) == set(native)
+        for key in native:
+            torch.testing.assert_close(restored[key], native[key])
+
+    def test_hf_per_expert_bias_load_is_rejected(self):
+        mixin = MockMoEStateDictMixin(n_experts=2, inter_dim=3)
+        mixin.moe_config.expert_bias = True
+
+        with pytest.raises(NotImplementedError, match="per-expert bias"):
+            mixin._from_hf_w_merged_experts(
+                {"model.layers.0.mlp.experts.0.gate_proj.bias": torch.randn(3)}
             )
 
     def _run_inplace_conversion(self, mixin, fqn, mock_dtensor, splits, *, for_checkpoint_load=False):

From b61ca76a064d8a3b5c9f492c2f6ba14e34662c58 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Wed, 26 Aug 2026 12:43:17 -0700
Subject: [PATCH 17/20] docs(skill): clarify low-memory MoE capability

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../nemo-automodel-model-onboarding/SKILL.md  | 35 +++++++++++++++----
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/skills/nemo-automodel-model-onboarding/SKILL.md b/skills/nemo-automodel-model-onboarding/SKILL.md
index 88536627a0..ed87d60ebd 100644
--- a/skills/nemo-automodel-model-onboarding/SKILL.md
+++ b/skills/nemo-automodel-model-onboarding/SKILL.md
@@ -180,13 +180,34 @@ Implement files in dependency order:
    `supports_low_memory_dcp_load=False` for unsafe variants; a model-sized device
    temporary can otherwise cause an out-of-memory failure.
 
-   For adapters using `MoESplitExpertsStateDictMixin`, treat the capability as
-   runtime-dependent. With the maintained EP dispatchers (`deepep`, `hybridep`,
-   and `uccl_ep`), TE expert storage is model-backed only at world size one;
-   `gmm`, `torch_mm`, and `torch_mm_mxfp8` expert storage is model-backed at any
-   world size. Non-EP dispatchers use ordinary grouped storage. The shared mixin
-   disables low-memory DCP for `expert_bias=True`, MoK, and any variant whose
-   checkpoint expert tensors require rebuilding.
+   The checkpoint router currently uses this capability to let one-process
+   custom-model safetensors initialization avoid a full host checkpoint load.
+   Multi-process loads already use DCP, but they share the same adapter conversion
+   and expert-storage rules.
+
+   For adapters using `MoESplitExpertsStateDictMixin`, decide this capability
+   from the runtime expert storage, not only from the configured backend name:
+
+   - Non-EP dispatchers use ordinary grouped experts whose checkpoint tensors
+     alias the final model parameters.
+   - At world size one, the maintained EP dispatchers (`deepep`, `hybridep`, and
+     `uccl_ep`) fall back to ordinary grouped experts and can use the low-memory
+     path.
+   - At larger world sizes, the `gmm`, `torch_mm`, and `torch_mm_mxfp8` expert
+     backends use model-backed grouped storage. TE and other expert backends use
+     virtual grouped tensors assembled from per-expert parameters, so they must
+     not advertise low-memory DCP.
+   - The mixin reports the capability as false for `expert_bias=True` and MoK.
+     The loader must also bypass the low-memory route for quantized expert
+     conversion and any other variant that must rebuild model-sized checkpoint
+     tensors after reading them.
+
+   A false capability value does not mean checkpoint loading is unsupported; it
+   selects an allocating or rebuilding load path. A false positive is dangerous:
+   if an adapter marks a temporary tensor as loaded in-place, `from_hf` can skip
+   rebuilding the real parameter and silently leave its previous value. Test each
+   optimized destination by writing a sentinel value through it and asserting
+   that the final model parameter storage changes.
 6. **__init__.py** -- Re-export the main model class
 
 See the pattern files for detailed implementation guidance:

From c768d3136183f175a37d2ceb6b77ac6dea5d7b53 Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Wed, 26 Aug 2026 16:08:59 -0700
Subject: [PATCH 18/20] fix(checkpoint): preserve MiniMax MTP load bookkeeping

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../minimax_m3_vl/state_dict_adapter.py       | 26 ++++++++++++++++---
 .../minimax_m3_vl/test_minimax_m3_mtp.py      | 26 +++++++++++++++++++
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py b/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
index 979c2ee664..201ba51ff2 100644
--- a/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
+++ b/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
@@ -233,9 +233,7 @@ def from_hf(
         return native
 
     def _mtp_from_hf(self, mtp_keys: dict[str, Any], device_mesh: Optional["DeviceMesh"] = None) -> dict[str, Any]:
-        """Convert MTP tensors: the transformer_layer reuses the full text from_hf
-        (as a fake 1-layer model, so expert-merge / index / dequant all apply); the
-        enorm/hnorm/eh_proj/final_layernorm fusion tensors pass through (eh_proj is FP8)."""
+        """Convert MTP tensors through a temporary one-layer text namespace."""
         pattern = re.compile(r"(?P<pfx>.*?)mtp\.layers\.(?P<d>\d+)\.(?P<rest>.+)")
         tl_hf: dict[str, Any] = {}
         passthrough: dict[str, Any] = {}
@@ -249,9 +247,29 @@ def _mtp_from_hf(self, mtp_keys: dict[str, Any], device_mesh: Optional["DeviceMe
 
         self._dequantize(passthrough)  # eh_proj is MXFP8
         native = dict(passthrough)
-        for key, value in self.from_hf(tl_hf, device_mesh).items():
+        self._dequantize(tl_hf)
+        for key in list(tl_hf):
+            new_key = self._hf_key_to_native(key)
+            if new_key != key:
+                tl_hf[new_key] = tl_hf.pop(key)
+
+        # The backbone merge already reset the per-load record. Preserve it while
+        # processing MTP, then translate the temporary model.layers.* names back to
+        # the native model.mtp.layers.* namespace used by the checkpoint key diff.
+        prior_view_keys = set(self.view_loaded_native_keys)
+        tl_native = self._from_hf_w_merged_experts(tl_hf, device_mesh, reset_view_loaded_keys=False)
+        for key, value in tl_native.items():
             m = re.match(r"model\.layers\.(\d+)\.(.+)", key)
             native[f"model.mtp.layers.{m.group(1)}.transformer_layer.{m.group(2)}"] = value
+        new_view_keys = self.view_loaded_native_keys - prior_view_keys
+        self._view_loaded_native_keys = prior_view_keys | {
+            re.sub(
+                r"^model\.layers\.(\d+)\.",
+                r"model.mtp.layers.\1.transformer_layer.",
+                key,
+            )
+            for key in new_view_keys
+        }
         return native
 
     def to_hf(
diff --git a/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_mtp.py b/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_mtp.py
index 041117dbc9..b351b1f5b8 100644
--- a/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_mtp.py
+++ b/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_mtp.py
@@ -99,6 +99,32 @@ def test_mtp_adapter_roundtrip_and_naming(mtp_model):
         assert torch.allclose(native[key].float(), back[key].float(), atol=1e-6), key
 
 
+def test_mtp_view_loaded_keys_keep_backbone_and_native_mtp_names(mtp_model):
+    adapter = mtp_model.state_dict_adapter
+    hf = adapter.to_hf(mtp_model.state_dict())
+
+    native_backbone = {
+        "model.layers.1.mlp.experts.gate_and_up_projs",
+        "model.layers.1.mlp.experts.down_projs",
+    }
+    temporary_mtp = {
+        "model.layers.0.mlp.experts.gate_and_up_projs",
+        "model.layers.0.mlp.experts.down_projs",
+    }
+    native_mtp = {
+        "model.mtp.layers.0.transformer_layer.mlp.experts.gate_and_up_projs",
+        "model.mtp.layers.0.transformer_layer.mlp.experts.down_projs",
+    }
+    adapter._inplace_loaded_native_keys = native_backbone | temporary_mtp
+
+    converted = adapter.from_hf(hf)
+
+    assert not native_backbone & converted.keys()
+    assert not native_mtp & converted.keys()
+    assert adapter.view_loaded_native_keys == native_backbone | native_mtp
+    assert adapter._inplace_loaded_native_keys == set()
+
+
 def test_from_hf_drops_mtp_when_disabled(model):
     """A model without MTP (num_mtp_modules=0) drops any MTP tensors on load."""
     adapter = model.state_dict_adapter

From ca42a11653e323afb8114c308d301ceb9ed3de3b Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 27 Aug 2026 07:52:02 -0700
Subject: [PATCH 19/20] fix(checkpoint): separate MiniMax MTP load bookkeeping

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../minimax_m3_vl/state_dict_adapter.py       | 48 ++++++++++-----
 .../models/minimax_m3_vl/conftest.py          |  9 ++-
 .../minimax_m3_vl/test_minimax_m3_mtp.py      | 61 ++++++++++++++++---
 3 files changed, 92 insertions(+), 26 deletions(-)

diff --git a/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py b/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
index 201ba51ff2..3e67dda2b6 100644
--- a/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
+++ b/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py
@@ -175,6 +175,20 @@ def _dequantize(self, state_dict: dict[str, Any]) -> dict[str, Any]:
             state_dict.pop(key, None)
         return state_dict
 
+    def _prepare_hf_for_merge(self, state_dict: dict[str, Any]) -> None:
+        """Prepare Hugging Face projection tensors for grouped-expert merging.
+
+        Args:
+            state_dict: Mapping mutated in place. Expert projection tensors have shape
+                [expert_hidden, hidden] for gate/up and [hidden, expert_hidden] for down;
+                MXFP8 projection tensors are dequantized before their keys are renamed.
+        """
+        self._dequantize(state_dict)
+        for key in list(state_dict):
+            new_key = self._hf_key_to_native(key)
+            if new_key != key:
+                state_dict[new_key] = state_dict.pop(key)
+
     @property
     def _mtp_enabled(self) -> bool:
         return int(getattr(self.config, "num_mtp_modules", 0) or 0) > 0
@@ -221,11 +235,7 @@ def from_hf(
         # decoder block); dropped entirely when the model has no MTP module.
         mtp_keys = {k: hf_state_dict.pop(k) for k in list(hf_state_dict) if ".mtp." in k}
 
-        self._dequantize(hf_state_dict)
-        for key in list(hf_state_dict.keys()):
-            new_key = self._hf_key_to_native(key)
-            if new_key != key:
-                hf_state_dict[new_key] = hf_state_dict.pop(key)
+        self._prepare_hf_for_merge(hf_state_dict)
         native = self._from_hf_w_merged_experts(hf_state_dict, device_mesh)
 
         if mtp_keys and self._mtp_enabled:
@@ -239,32 +249,31 @@ def _mtp_from_hf(self, mtp_keys: dict[str, Any], device_mesh: Optional["DeviceMe
         passthrough: dict[str, Any] = {}
         for key, value in mtp_keys.items():
             m = pattern.match(key)
+            if m is None:
+                passthrough[key] = value
+                continue
             depth, rest = m.group("d"), m.group("rest")
             if rest.startswith("transformer_layer."):
-                tl_hf[f"model.layers.{depth}.{rest[len('transformer_layer.') :]}"] = value
+                tl_hf[f"layers.{depth}.{rest[len('transformer_layer.') :]}"] = value
             else:
                 passthrough[key] = value
 
         self._dequantize(passthrough)  # eh_proj is MXFP8
         native = dict(passthrough)
-        self._dequantize(tl_hf)
-        for key in list(tl_hf):
-            new_key = self._hf_key_to_native(key)
-            if new_key != key:
-                tl_hf[new_key] = tl_hf.pop(key)
+        self._prepare_hf_for_merge(tl_hf)
 
         # The backbone merge already reset the per-load record. Preserve it while
-        # processing MTP, then translate the temporary model.layers.* names back to
+        # processing MTP, then translate the temporary layers.* names back to
         # the native model.mtp.layers.* namespace used by the checkpoint key diff.
         prior_view_keys = set(self.view_loaded_native_keys)
         tl_native = self._from_hf_w_merged_experts(tl_hf, device_mesh, reset_view_loaded_keys=False)
         for key, value in tl_native.items():
-            m = re.match(r"model\.layers\.(\d+)\.(.+)", key)
+            m = re.match(r"layers\.(\d+)\.(.+)", key)
             native[f"model.mtp.layers.{m.group(1)}.transformer_layer.{m.group(2)}"] = value
         new_view_keys = self.view_loaded_native_keys - prior_view_keys
         self._view_loaded_native_keys = prior_view_keys | {
             re.sub(
-                r"^model\.layers\.(\d+)\.",
+                r"^layers\.(\d+)\.",
                 r"model.mtp.layers.\1.transformer_layer.",
                 key,
             )
@@ -321,12 +330,19 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
 
     def _mtp_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[tuple[str, Any]]:
         m = re.match(r"(?P<head>.*?)mtp\.layers\.(?P<d>\d+)\.(?P<rest>.+)", fqn)
+        if m is None:
+            exclude_key_regex = kwargs.get("exclude_key_regex", None)
+            if exclude_key_regex and re.match(exclude_key_regex, fqn):
+                return []
+            return [(fqn, tensor)]
         head, depth, rest = m.group("head"), m.group("d"), m.group("rest")
         if rest.startswith("transformer_layer."):
             suffix = rest[len("transformer_layer.") :]
-            converted = self.convert_single_tensor_to_hf(f"model.layers.{depth}.{suffix}", tensor, **kwargs)
+            converted = self.convert_single_tensor_to_hf(
+                f"mtp.layers.{depth}.{suffix}", tensor, prefix_override="mtp.", **kwargs
+            )
             tl_prefix = f"{head}mtp.layers.{depth}.transformer_layer."
-            return [(k.replace(f"model.layers.{depth}.", tl_prefix, 1), v) for k, v in converted]
+            return [(k.replace(f"mtp.layers.{depth}.", tl_prefix, 1), v) for k, v in converted]
         exclude_key_regex = kwargs.get("exclude_key_regex", None)
         if exclude_key_regex and re.match(exclude_key_regex, fqn):
             return []
diff --git a/tests/unit_tests/models/minimax_m3_vl/conftest.py b/tests/unit_tests/models/minimax_m3_vl/conftest.py
index c2b159bdff..6e21fa5889 100644
--- a/tests/unit_tests/models/minimax_m3_vl/conftest.py
+++ b/tests/unit_tests/models/minimax_m3_vl/conftest.py
@@ -128,12 +128,17 @@ def sparse_model(sparse_text_config, backend):
 
 @pytest.fixture
 def mtp_model(backend):
-    """Sparse text backbone with one MTP module (DeepSeek-V3 style)."""
+    """Default all-MoE text backbone with one MTP module (DeepSeek-V3 style)."""
     from nemo_automodel.components.models.minimax_m3_vl.model import MiniMaxM3SparseForCausalLM
 
     cfg = MiniMaxM3VLTextConfig(
         torch_dtype="float32",
-        **{**TINY_CFG, "num_mtp_modules": 1, "sparse_attention_config": dict(SPARSE_ATTENTION_CONFIG)},
+        **{
+            **TINY_CFG,
+            "moe_layer_freq": None,
+            "num_mtp_modules": 1,
+            "sparse_attention_config": dict(SPARSE_ATTENTION_CONFIG),
+        },
     )
     m = MiniMaxM3SparseForCausalLM(cfg, backend=backend).eval()
     m.initialize_weights(dtype=torch.float32)
diff --git a/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_mtp.py b/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_mtp.py
index b351b1f5b8..329f7f1666 100644
--- a/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_mtp.py
+++ b/tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_mtp.py
@@ -99,23 +99,58 @@ def test_mtp_adapter_roundtrip_and_naming(mtp_model):
         assert torch.allclose(native[key].float(), back[key].float(), atol=1e-6), key
 
 
-def test_mtp_view_loaded_keys_keep_backbone_and_native_mtp_names(mtp_model):
+def test_mtp_view_loaded_keys_keep_backbone_and_native_mtp_names(mtp_model, monkeypatch):
     adapter = mtp_model.state_dict_adapter
-    hf = adapter.to_hf(mtp_model.state_dict())
-
     native_backbone = {
-        "model.layers.1.mlp.experts.gate_and_up_projs",
-        "model.layers.1.mlp.experts.down_projs",
-    }
-    temporary_mtp = {
         "model.layers.0.mlp.experts.gate_and_up_projs",
         "model.layers.0.mlp.experts.down_projs",
     }
+    temporary_mtp = {
+        "layers.0.mlp.experts.gate_and_up_projs",
+        "layers.0.mlp.experts.down_projs",
+    }
     native_mtp = {
         "model.mtp.layers.0.transformer_layer.mlp.experts.gate_and_up_projs",
         "model.mtp.layers.0.transformer_layer.mlp.experts.down_projs",
     }
-    adapter._inplace_loaded_native_keys = native_backbone | temporary_mtp
+    model_state = mtp_model.state_dict()
+    native = {key: model_state[key] for key in native_backbone | native_mtp}
+
+    def split_local_experts(weight: torch.Tensor, n_experts: int) -> list[torch.Tensor]:
+        """Expose per-expert views of grouped checkpoint storage.
+
+        Args:
+            weight: Tensor of shape [experts, ...], with arbitrary trailing dimensions.
+            n_experts: Number of experts on the leading axis.
+
+        Returns:
+            List of tensors of shape [...], one view per expert.
+        """
+        adapter._last_expert_ids = list(range(n_experts))
+        return [weight[expert_id] for expert_id in range(n_experts)]
+
+    monkeypatch.setattr(adapter, "_split_experts_weights", split_local_experts)
+    monkeypatch.setattr(
+        "nemo_automodel.components.moe.state_dict_utils.is_dtensor",
+        lambda tensor: isinstance(tensor, torch.Tensor) and tensor.ndim == 3,
+    )
+    monkeypatch.setattr(
+        "nemo_automodel.components.moe.state_dict_utils.validate_dtensor_expert_sharding",
+        lambda *args, **kwargs: None,
+    )
+
+    hf = adapter.to_hf(native, for_checkpoint_load=True)
+
+    assert adapter._inplace_loaded_native_keys == native_backbone | temporary_mtp
+    assert any(key.startswith("model.layers.0.block_sparse_moe.experts.") for key in hf)
+    assert any(key.startswith("model.mtp.layers.0.transformer_layer.block_sparse_moe.experts.") for key in hf)
+    with torch.no_grad():
+        for key, destination in hf.items():
+            destination.fill_(3.0 if key.startswith("model.layers.0.") else 7.0)
+    assert torch.all(native["model.layers.0.mlp.experts.gate_and_up_projs"] == 3.0)
+    assert torch.all(native["model.layers.0.mlp.experts.down_projs"] == 3.0)
+    assert torch.all(native["model.mtp.layers.0.transformer_layer.mlp.experts.gate_and_up_projs"] == 7.0)
+    assert torch.all(native["model.mtp.layers.0.transformer_layer.mlp.experts.down_projs"] == 7.0)
 
     converted = adapter.from_hf(hf)
 
@@ -125,6 +160,16 @@ def test_mtp_view_loaded_keys_keep_backbone_and_native_mtp_names(mtp_model):
     assert adapter._inplace_loaded_native_keys == set()
 
 
+def test_from_hf_preserves_unrecognized_mtp_keys(mtp_model):
+    adapter = mtp_model.state_dict_adapter
+    key = "model.mtp.auxiliary.weight"
+    tensor = torch.randn(4, 4)
+
+    converted = adapter.from_hf({key: tensor})
+
+    assert converted[key] is tensor
+
+
 def test_from_hf_drops_mtp_when_disabled(model):
     """A model without MTP (num_mtp_modules=0) drops any MTP tensors on load."""
     adapter = model.state_dict_adapter

From a7f06cab6d263fee89f464af862a2b2410e2709d Mon Sep 17 00:00:00 2001
From: Yuhe Zhang <yuhez@nvidia.com>
Date: Thu, 27 Aug 2026 11:05:25 -0700
Subject: [PATCH 20/20] fix(checkpoint): migrate Qwen3.8 load capability

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
---
 .../models/qwen3_8_flash_next/state_dict_adapter.py           | 2 +-
 .../test_qwen3_8_flash_next_state_dict_adapter.py             | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/nemo_automodel/components/models/qwen3_8_flash_next/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_8_flash_next/state_dict_adapter.py
index 875a70fd25..0e544635e9 100644
--- a/nemo_automodel/components/models/qwen3_8_flash_next/state_dict_adapter.py
+++ b/nemo_automodel/components/models/qwen3_8_flash_next/state_dict_adapter.py
@@ -43,7 +43,7 @@
 class Qwen3_8_FlashNextStateDictAdapter(Qwen3_5MoeStateDictAdapter):
     """Convert Qwen3.8-Flash-Next checkpoints without gathering the global PLE table."""
 
-    _supports_write_through_checkpoint_load = True
+    _supports_low_memory_dcp_load = True
 
     def __init__(
         self,
diff --git a/tests/unit_tests/models/qwen3_8_flash_next/test_qwen3_8_flash_next_state_dict_adapter.py b/tests/unit_tests/models/qwen3_8_flash_next/test_qwen3_8_flash_next_state_dict_adapter.py
index da0cf83cdc..0a96d4c04c 100644
--- a/tests/unit_tests/models/qwen3_8_flash_next/test_qwen3_8_flash_next_state_dict_adapter.py
+++ b/tests/unit_tests/models/qwen3_8_flash_next/test_qwen3_8_flash_next_state_dict_adapter.py
@@ -545,11 +545,11 @@ def test_grouped_experts_and_fp32_gdn_use_inherited_qwen35_conversion(
     torch.testing.assert_close(restored_a_log, a_log)
 
 
-def test_adapter_advertises_nonquantized_write_through_loading(
+def test_adapter_advertises_nonquantized_low_memory_loading(
     adapter: Qwen3_8_FlashNextStateDictAdapter,
 ) -> None:
     """The loader may bypass host staging for valid BF16/fp32 model state."""
-    assert adapter.supports_write_through_checkpoint_load
+    assert adapter.supports_low_memory_dcp_load
 
 
 def test_table_export_honors_exclude_regex(