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

@@ -830,25 +830,27 @@ def load_model(
         # fall behind other ranks' async allocations.
         is_safetensors = _is_safetensors_checkpoint(model_path)
         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.
+        # 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.
         # 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 = (
+        can_load_without_full_copy = (
             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_checkpoint_load_without_full_copy
+            )
             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 can_load_without_full_copy
         )
         if (
             is_init_step
@@ -868,6 +870,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 +899,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 +934,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,
+            for_checkpoint_load=True,
         )
         destinations_ready = time.monotonic()
         requested_bytes = sum(
@@ -1346,7 +1351,7 @@ def _consolidate() -> None:
                     _maybe_rename_index_for_diffusers(consolidated_dir)
                 if is_rank_0():
                     logger.info("Successfully exported consolidated HF safetensors to %s.", consolidated_dir)
-            except BaseException as e:  # noqa: B036 - re-raised on the main thread in async_wait
+            except BaseException as e:  # Re-raised on the main thread in async_wait.
                 self._consolidation_error = e
 
         self._consolidation_thread = threading.Thread(

nemo_automodel/components/checkpoint/state_dict_adapter.py

@@ -27,18 +27,27 @@ class StateDictAdapter(ABC):
     """
 
     _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 all checkpoint-load destinations write through to final model storage.
+        """Whether every checkpoint tensor is loaded directly into the model's existing weight memory.
 
-        Adapters should set ``_supports_write_through_checkpoint_load`` only when every tensor produced by
-        ``to_hf`` for base-checkpoint loading is either the original model tensor or a view that writes through
-        to it. The checkpoint loader uses this guarantee to avoid materializing a converted full state dict on
-        the host.
+        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
 
+    @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
+
     @abstractmethod
     def to_hf(self, state_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
         """Convert from native model state dict to HuggingFace format.