NVIDIA-NeMo · yuhezhang-ai · Aug 20, 2026 · Aug 20, 2026 · Aug 20, 2026 · Aug 20, 2026
@@ -831,26 +831,22 @@ 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():
             world_size = torch.distributed.get_world_size()
         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

@@ -15,6 +15,8 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Optional
 
+import torch
+
 if TYPE_CHECKING:
     from torch.distributed.device_mesh import DeviceMesh
 
@@ -26,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
-
-    @property
-    def supports_write_through_checkpoint_load(self) -> bool:
-        """Whether every checkpoint tensor is loaded directly into the model's existing weight memory.
-
-        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.
-        """
-        return self._supports_write_through_checkpoint_load
+    _supports_low_memory_dcp_load: bool = False
 
     @property
-    def supports_checkpoint_load_without_full_copy(self) -> bool:
-        """Whether DCP can load this adapter without another full set of model weights.
+    def supports_low_memory_dcp_load(self) -> bool:
+        """Whether DCP can load the checkpoint with zero or small temporary tensors.
 
-        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.
+        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_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]:
@@ -95,7 +87,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
@@ -105,7 +97,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.

@@ -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

@@ -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,

@@ -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,
@@ -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,

@@ -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,

@@ -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,

@@ -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",

@@ -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,
@@ -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():

@@ -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,
@@ -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] = {}

@@ -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,

@@ -64,6 +64,7 @@
 )
 
 _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$")
 
 
 def _rename_hf_to_native(key: str) -> str:
@@ -86,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,
@@ -99,6 +102,11 @@ def __init__(
         self.dtype = dtype
         self._uses_model_prefix = True
 
+    @property
+    def supports_low_memory_dcp_load(self) -> bool:
+        """Whether Ling's DCP load needs only small fused-QKV temporary tensors."""
+        return self.moe_config is not None and super().supports_low_memory_dcp_load
+
     # ---- HF -> native ----------------------------------------------------
 
     def from_hf(
@@ -167,9 +175,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 +209,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)]
@@ -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