perf(checkpoint): reduce allocating grouped MoE load overhead

examples/vlm_finetune/gemma4/gemma4_26b_a4b_moe.yaml

@@ -119,8 +119,8 @@ freeze_config:
 ci:
   recipe_owner: athitten
   nodes: 1
-  # Full robustness takes about 50 minutes.
-  time: "01:15:00"
+  # Full robustness measured 32 minutes; keep 13 minutes for filesystem variance.
+  time: "00:45:00"
   # Keep failures blocking; residual MoE drift is tracked by AM-707/AM-719.
   checkpoint_robustness:
     # The shared robustness batch of 2 exceeds this 26B MoE recipe's memory

nemo_automodel/components/checkpoint/checkpointing.py

@@ -832,23 +832,27 @@ def load_model(
         is_custom_model = _is_custom_model(model_state.model[0])
         # Custom adapters traditionally took the frugal full-state path here because converting
         # grouped experts into load destinations could materialize a second on-device copy and OOM.
-        # Adapters whose destinations all alias final model storage can now opt into the standard
-        # DCP path below, which writes checkpoint tensors directly through those views. Keep the CPU
-        # path for non-aliasing backends and quantized initialization, whose conversion allocates.
+        # Adapters whose model-sized destinations alias final model storage can now opt into the standard DCP path
+        # below, which writes checkpoint tensors directly through those views. A bounded adapter may also retain
+        # small auxiliary destinations that it consumes in-place after the read. Keep the CPU path for other
+        # non-aliasing adapters and quantized initialization, whose conversion allocates model-sized tensors.
         # 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)
-        supports_write_through_checkpoint_load = (
+        supports_memory_bounded_checkpoint_load = (
             isinstance(state_dict_adapter, StateDictAdapter)
-            and state_dict_adapter.supports_write_through_checkpoint_load
+            and (
+                state_dict_adapter.supports_write_through_checkpoint_load
+                or state_dict_adapter.supports_bounded_checkpoint_load
+            )
             and not self.config.dequantize_base_checkpoint
         )
         single_device_custom_safetensors = (
-            is_safetensors and is_custom_model and world_size == 1 and not supports_write_through_checkpoint_load
+            is_safetensors and is_custom_model and world_size == 1 and not supports_memory_bounded_checkpoint_load
         )
         if (
             is_init_step
@@ -868,6 +872,7 @@ def load_model(
                 )
             else:
                 state_dict_from_disk = {}
+            t_adapt = time.monotonic()
 
             # Apply key_mapping (e.g. _checkpoint_conversion_mapping) so that
             # HF checkpoint keys are renamed to match the model's parameter FQNs.
@@ -896,13 +901,14 @@ def load_model(
             t_end = time.monotonic()
 
             disk_s = t_disk - t0
-            dist_s = t_end - t_disk
+            adapt_s = t_adapt - t_disk
+            install_s = t_end - t_adapt
             total_s = t_end - t0
             gb = total_bytes / (1 << 30)
             logging.info(
                 f"load_model: {gb:.2f} GB loaded in {total_s:.2f}s "
                 f"({gb / total_s:.2f} GB/s overall | "
-                f"disk read {disk_s:.2f}s, distribute {dist_s:.2f}s)"
+                f"disk read {disk_s:.2f}s, adapt {adapt_s:.2f}s, install {install_s:.2f}s)"
             )
             del state_dict_from_disk
             gc.collect()
@@ -930,6 +936,7 @@ def load_model(
             # Only base-checkpoint initialization needs FP8 scale destinations.
             quantization=bool(is_init_step and self.config.dequantize_base_checkpoint),
             device_mesh=self.moe_mesh,
+            load_into_empty_destinations=True,
         )
         destinations_ready = time.monotonic()
         requested_bytes = sum(

nemo_automodel/components/checkpoint/state_dict_adapter.py

@@ -27,6 +27,7 @@ class StateDictAdapter(ABC):
     """
 
     _supports_write_through_checkpoint_load: bool = False
+    _supports_bounded_checkpoint_load: bool = False
 
     @property
     def supports_write_through_checkpoint_load(self) -> bool:
@@ -39,6 +40,18 @@ def supports_write_through_checkpoint_load(self) -> bool:
         """
         return self._supports_write_through_checkpoint_load
 
+    @property
+    def supports_bounded_checkpoint_load(self) -> bool:
+        """Whether single-device checkpoint loading avoids a full converted state dict.
+
+        Adapters should set ``_supports_bounded_checkpoint_load`` only when every model-sized destination produced by
+        ``to_hf(..., load_into_empty_destinations=True)`` writes through to final model storage. Any allocating
+        destinations must be bounded auxiliary tensors, and ``from_hf`` must complete their conversion without
+        allocating another model-sized tensor. This is weaker than ``supports_write_through_checkpoint_load`` because
+        bounded auxiliary destinations do not need to alias model storage.
+        """
+        return self._supports_bounded_checkpoint_load
+
     @abstractmethod
     def to_hf(self, state_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
         """Convert from native model state dict to HuggingFace format.

nemo_automodel/components/models/gemma4_moe/state_dict_adapter.py

@@ -42,6 +42,8 @@
 import torch
 import torch.distributed as dist
 from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor.placement_types import Replicate, Shard
 
 from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
 from nemo_automodel.components.models.common import BackendConfig
@@ -59,6 +61,8 @@ class Gemma4MoEStateDictAdapter(StateDictAdapter):
       4. Expert-parallel sharding when a device mesh is provided
     """
 
+    _supports_bounded_checkpoint_load = True
+
     def __init__(
         self,
         config: Any,
@@ -71,6 +75,7 @@ def __init__(
         self.backend = backend
         self.dtype = dtype
         self._uses_model_prefix = True
+        self._load_destinations_alias_model = False
 
     # ------------------------------------------------------------------
     # HF -> NeMo
@@ -81,6 +86,25 @@ def from_hf(
         device_mesh: Optional[DeviceMesh] = None,
         **kwargs,
     ) -> dict[str, Any]:
+        """Convert Hugging Face Gemma4 weights into native model layout.
+
+        Args:
+            hf_state_dict: Hugging Face state mapping. Expert gate/up tensors have shape
+                ``[experts, 2 * expert_hidden, hidden]`` and down tensors have shape
+                ``[experts, hidden, expert_hidden]``. When prepared by ``to_hf`` for a bounded checkpoint load,
+                those tensors are transposed views of final model storage and are mutated in place.
+            device_mesh: Optional expert-parallel mesh. Distributed conversion slices the global expert axis and may
+                shard the native feature axis according to the mesh.
+            **kwargs: Additional adapter-interface arguments.
+
+        Returns:
+            Native state mapping. Expert gate/up tensors have shape ``[local_experts, hidden, 2 * expert_hidden]``
+            and down tensors have shape ``[local_experts, expert_hidden, hidden]``. Sharded bounded-load outputs are
+            DTensors with global expert shape and ``Shard(0)`` on the ``ep`` mesh; their local tensors alias final
+            model storage. Materialized conversion outputs own their storage.
+        """
+        load_destinations_alias_model = self._load_destinations_alias_model
+        self._load_destinations_alias_model = False
         self._uses_model_prefix = any(key.startswith("model.") for key in hf_state_dict)
         model_prefix = "model." if self._uses_model_prefix else ""
 
@@ -125,7 +149,7 @@ def from_hf(
             state_dict[key] = value
 
         # Process collected expert weights per layer
-        _REQUIRED_EXPERT_KEYS = {"gate_up_proj", "down_proj"}
+        _REQUIRED_EXPERT_KEYS = {"gate_up_proj", "down_proj", "per_expert_scale"}
         for layer_path, tensors in expert_buffers.items():
             missing = _REQUIRED_EXPERT_KEYS - tensors.keys()
             if missing:
@@ -139,15 +163,52 @@ def from_hf(
             per_expert_scale = tensors["per_expert_scale"]  # [E]
 
             # Transpose gate_up_proj from HF [E, 2*inter, hidden] to NeMo [E, hidden, 2*inter]
-            gate_and_up = gate_up_proj.transpose(-2, -1)  # [E, hidden, 2*inter]
-
-            # Transpose down_proj from HF [E, hidden, inter] to NeMo [E, inter, hidden]
-            # and absorb per_expert_scale
-            down = down_proj.transpose(-2, -1) * per_expert_scale[:, None, None]  # [E, inter, hidden]
-
-            # Slice for EP
-            gate_and_up_local = gate_and_up[start_expert:end_expert].to(self.dtype)
-            down_local = down[start_expert:end_expert].to(self.dtype)
+            if load_destinations_alias_model and state_dict_utils.is_dtensor(gate_up_proj):
+                if not state_dict_utils.is_dtensor(down_proj) or not state_dict_utils.is_dtensor(per_expert_scale):
+                    raise RuntimeError(
+                        f"Inconsistent sharded load destinations for {layer_path}: gate/up, down, and scale must "
+                        "all be DTensors."
+                    )
+
+                gate_and_up_local = gate_up_proj.transpose(-2, -1)
+                down_local_tensor = down_proj.to_local()
+                scale_local_tensor = per_expert_scale.to_local()
+                if down_local_tensor.shape[0] != scale_local_tensor.shape[0]:
+                    raise RuntimeError(
+                        f"Sharded down projection for {layer_path} has {down_local_tensor.shape[0]} local experts, "
+                        f"but its scale destination has {scale_local_tensor.shape[0]}."
+                    )
+                # Checkpoint destinations are detached state-dict views owned exclusively by this load. The raw HF
+                # down tensor already occupies final native storage, so absorb its local scales without allocating a
+                # second expert tensor or communicating across the EP mesh.
+                with torch.no_grad():
+                    down_local_tensor.mul_(scale_local_tensor[:, None, None])
+                down_local = down_proj.transpose(-2, -1)
+            elif load_destinations_alias_model:
+                # DCP loaded the checkpoint through transposed views of final model storage. Gate/up only needs its
+                # native view restored; down additionally absorbs the small per-expert scale destination in place.
+                gate_and_up_local = gate_up_proj.transpose(-2, -1)
+                down_proj.mul_(per_expert_scale[:, None, None])
+                down_local = down_proj.transpose(-2, -1)
+            else:
+                gate_and_up = gate_up_proj.transpose(-2, -1)  # [E, hidden, 2*inter]
+
+                # Transpose down_proj from HF [E, hidden, inter] to NeMo [E, inter, hidden]
+                # and absorb per_expert_scale
+                down = down_proj.transpose(-2, -1) * per_expert_scale[:, None, None]  # [E, inter, hidden]
+
+                # Slice for EP
+                gate_and_up_local = gate_and_up[start_expert:end_expert].to(self.dtype)
+                down_local = down[start_expert:end_expert].to(self.dtype)
+
+            if load_destinations_alias_model and state_dict_utils.is_dtensor(gate_and_up_local):
+                # The HF-layout load DTensors and these transposed native views both alias the original model
+                # parameters. Their Shard(0) placement already selects this rank's experts, so slicing or wrapping
+                # again would corrupt the global offsets DCP used for the read.
+                prefix = f"{model_prefix}language_model.{layer_path}"
+                state_dict[f"{prefix}.moe.experts.gate_and_up_projs"] = gate_and_up_local
+                state_dict[f"{prefix}.moe.experts.down_projs"] = down_local
+                continue
 
             # Slice for EP_SHARD across the feature dimension before wrapping as DTensor.
             if device_mesh is not None and "ep_shard" in device_mesh.mesh_dim_names:
@@ -185,10 +246,49 @@ def to_hf(
         quantization: bool = False,
         **kwargs,
     ) -> dict[str, Any]:
+        """Convert native Gemma4 weights to Hugging Face keys and layouts.
+
+        Args:
+            state_dict: Native state mapping. Expert gate/up tensors have shape
+                ``[local_experts, hidden, 2 * expert_hidden]`` and down tensors have shape
+                ``[local_experts, expert_hidden, hidden]``.
+            exclude_key_regex: Optional pattern selecting keys to omit.
+            quantization: Whether checkpoint initialization requires a precision conversion. Quantized loads do not
+                use aliasing destinations.
+            **kwargs: Adapter-interface arguments. ``device_mesh`` describes expert sharding.
+                ``load_into_empty_destinations=True`` requests memory-bounded load destinations.
+
+        Returns:
+            Hugging Face state mapping. Expert gate/up tensors have shape
+            ``[experts, 2 * expert_hidden, hidden]`` and down tensors have shape
+            ``[experts, hidden, expert_hidden]``. For a bounded load, model-sized outputs are transposed views of
+            final model storage. EP outputs retain ``Shard(0)`` on the global experts axis and each rank owns a local
+            ``per_expert_scale`` slice; unsharded scale outputs are independently owned ``[experts]`` tensors.
+        """
         self._uses_model_prefix = any(key.startswith("model.") for key in state_dict)
         prefix = "model." if self._uses_model_prefix else ""
         device_mesh: Optional[DeviceMesh] = kwargs.get("device_mesh")
         n_experts = self.moe_config.n_routed_experts
+        expert_tensors = [tensor for fqn, tensor in state_dict.items() if ".moe.experts." in fqn]
+        single_device_destinations_alias = (
+            device_mesh is None
+            and bool(expert_tensors)
+            and all(
+                isinstance(tensor, torch.Tensor) and not tensor.is_meta and not state_dict_utils.is_dtensor(tensor)
+                for tensor in expert_tensors
+            )
+        )
+        ep_destinations_alias = (
+            device_mesh is not None
+            and bool(expert_tensors)
+            and all(self._supports_ep_load_destination(tensor, n_experts) for tensor in expert_tensors)
+        )
+        load_into_model_storage = (
+            bool(kwargs.get("load_into_empty_destinations", False))
+            and not quantization
+            and (single_device_destinations_alias or ep_destinations_alias)
+        )
+        self._load_destinations_alias_model = load_into_model_storage
 
         hf_state_dict: dict[str, Any] = {}
 
@@ -205,20 +305,42 @@ def to_hf(
             # --- Expert: gate_and_up_projs -> experts.gate_up_proj ---
             if ".moe.experts.gate_and_up_projs" in fqn:
                 layer_num = re.search(r"layers\.(\d+)", fqn).group(1)
-                global_tensor = self._gather_expert_tensor(tensor, device_mesh, n_experts)
                 layer_prefix = f"{prefix}language_model.layers.{layer_num}"
-                # Transpose from NeMo [E, hidden, 2*inter] to HF [E, 2*inter, hidden]
-                hf_state_dict[f"{layer_prefix}.experts.gate_up_proj"] = global_tensor.transpose(-2, -1).contiguous()
+                if load_into_model_storage:
+                    # DCP writes the HF [E, 2*inter, hidden] checkpoint directly through this transposed view into
+                    # the native [E, hidden, 2*inter] parameter storage.
+                    hf_state_dict[f"{layer_prefix}.experts.gate_up_proj"] = tensor.transpose(-2, -1)
+                else:
+                    global_tensor = self._gather_expert_tensor(tensor, device_mesh, n_experts)
+                    hf_state_dict[f"{layer_prefix}.experts.gate_up_proj"] = global_tensor.transpose(-2, -1).contiguous()
                 continue
 
             # --- Expert: down_projs -> experts.down_proj + router.per_expert_scale ---
             if ".moe.experts.down_projs" in fqn:
                 layer_num = re.search(r"layers\.(\d+)", fqn).group(1)
-                global_tensor = self._gather_expert_tensor(tensor, device_mesh, n_experts)
                 layer_prefix = f"{prefix}language_model.layers.{layer_num}"
-                # Transpose from NeMo [E, inter, hidden] to HF [E, hidden, inter]
-                hf_state_dict[f"{layer_prefix}.experts.down_proj"] = global_tensor.transpose(-2, -1).contiguous()
-                hf_state_dict[f"{layer_prefix}.router.per_expert_scale"] = torch.ones(n_experts, dtype=self.dtype)
+                if load_into_model_storage:
+                    # The raw HF down weight first lands in final native storage through a transposed view. from_hf
+                    # then multiplies that storage in place by this bounded auxiliary scale vector.
+                    hf_state_dict[f"{layer_prefix}.experts.down_proj"] = tensor.transpose(-2, -1)
+                    if state_dict_utils.is_dtensor(tensor):
+                        local_tensor = tensor.to_local()
+                        local_scale = torch.empty(local_tensor.shape[0], dtype=self.dtype, device=local_tensor.device)
+                        hf_state_dict[f"{layer_prefix}.router.per_expert_scale"] = DTensor.from_local(
+                            local_scale,
+                            tensor.device_mesh,
+                            tensor.placements,
+                            shape=torch.Size([n_experts]),
+                            stride=(1,),
+                        )
+                    else:
+                        hf_state_dict[f"{layer_prefix}.router.per_expert_scale"] = torch.empty(
+                            n_experts, dtype=self.dtype, device=tensor.device
+                        )
+                else:
+                    global_tensor = self._gather_expert_tensor(tensor, device_mesh, n_experts)
+                    hf_state_dict[f"{layer_prefix}.experts.down_proj"] = global_tensor.transpose(-2, -1).contiguous()
+                    hf_state_dict[f"{layer_prefix}.router.per_expert_scale"] = torch.ones(n_experts, dtype=self.dtype)
                 continue
 
             # --- Pass-through ---
@@ -229,6 +351,38 @@ def to_hf(
 
         return hf_state_dict
 
+    @staticmethod
+    def _supports_ep_load_destination(tensor: Any, n_experts: int) -> bool:
+        """Return whether a grouped expert DTensor can receive its HF checkpoint slice in place.
+
+        Args:
+            tensor: Native expert DTensor with global shape ``[experts, ...]`` and local shape
+                ``[local_experts, ...]``. The ``ep`` mesh dimension must use ``Shard(0)``; every other mesh dimension
+                must replicate the tensor. Inner-axis expert sharding is intentionally unsupported by this path.
+            n_experts: Total number of routed experts in the checkpoint.
+
+        Returns:
+            ``True`` when transposing the final local tensor preserves an HF-layout ``Shard(0)`` destination that
+            DCP can fill without gathering global experts.
+        """
+        if (
+            not state_dict_utils.is_dtensor(tensor)
+            or tensor.is_meta
+            or tensor.shape[0] != n_experts
+            or "ep" not in tensor.device_mesh.mesh_dim_names
+        ):
+            return False
+
+        return all(
+            (
+                isinstance(placement, Shard)
+                and placement.dim == 0
+                and tensor.device_mesh.mesh_dim_names[mesh_dim] == "ep"
+            )
+            or (isinstance(placement, Replicate) and tensor.device_mesh.mesh_dim_names[mesh_dim] != "ep")
+            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.