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 11371af940..e5ee4341d0 100644 --- a/nemo_automodel/components/checkpoint/state_dict_adapter.py +++ b/nemo_automodel/components/checkpoint/state_dict_adapter.py @@ -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. 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 1d60dabaea..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, @@ -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/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..6082146655 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, @@ -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 f92be67d74..a723d9fa54 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, @@ -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/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 9e1391c5b4..94b5412a13 100644 --- a/nemo_automodel/components/models/ling_v2/state_dict_adapter.py +++ b/nemo_automodel/components/models/ling_v2/state_dict_adapter.py @@ -64,6 +64,7 @@ ) _LAYER_QKV_RE = re.compile(r"^(?P(?:.*\.)?layers\.\d+)\.attention\.query_key_value\.weight$") +_NATIVE_LAYER_QKV_RE = re.compile(r"^(?P(?:.*\.)?layers\.\d+)\.self_attn\.(?P[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(?:.*\.)?layers\.\d+)\.self_attn\.(?P[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(?:.*\.)?layers\.\d+)\.self_attn\.(?P[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/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/minimax_m3_vl/state_dict_adapter.py b/nemo_automodel/components/models/minimax_m3_vl/state_dict_adapter.py index b8b3204102..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: @@ -233,25 +243,42 @@ 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.*?)mtp\.layers\.(?P\d+)\.(?P.+)") tl_hf: dict[str, Any] = {} 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) - for key, value in self.from_hf(tl_hf, device_mesh).items(): - m = re.match(r"model\.layers\.(\d+)\.(.+)", 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 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"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"^layers\.(\d+)\.", + r"model.mtp.layers.\1.transformer_layer.", + key, + ) + for key in new_view_keys + } return native def to_hf( @@ -303,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.*?)mtp\.layers\.(?P\d+)\.(?P.+)", 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 [] @@ -332,6 +366,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/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..74c44bd379 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,14 @@ 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 + + @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, 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..d275468a68 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, @@ -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,7 +306,16 @@ 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 = ( + 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] if exclude_key_regex: 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_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/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..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 @@ -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, @@ -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/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/layers.py b/nemo_automodel/components/moe/layers.py index 2afe823ec9..285b20fed8 100644 --- a/nemo_automodel/components/moe/layers.py +++ b/nemo_automodel/components/moe/layers.py @@ -799,7 +799,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 766522b9ab..e7077fbf4d 100644 --- a/nemo_automodel/components/moe/state_dict_mixin.py +++ b/nemo_automodel/components/moe/state_dict_mixin.py @@ -13,6 +13,7 @@ # limitations under the License. import gc +import os import re from typing import Any, Optional @@ -32,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. @@ -50,19 +62,36 @@ 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.""" + 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: + """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 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" @property - def _supports_write_through_expert_checkpoint_load(self) -> bool: - """Whether grouped expert checkpoint tensors load directly into model weight memory. + def _grouped_expert_storage_is_model_weight(self) -> bool: + """Whether the runtime grouped expert tensors use the model's parameter storage. - 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 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. """ - return self.backend.experts != "te" and self.backend.dispatcher != "mok" + if self.backend.dispatcher == "mok": + return self.backend.experts != "te" + if self.backend.dispatcher not in {"deepep", "hybridep", "uccl_ep"}: + return True + 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: @@ -615,12 +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: + 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) @@ -855,6 +895,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. @@ -867,7 +918,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" ) @@ -903,7 +954,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 = [] @@ -956,7 +1007,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 = [] @@ -980,16 +1031,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/skills/nemo-automodel-model-onboarding/SKILL.md b/skills/nemo-automodel-model-onboarding/SKILL.md index c2f7e9e33c..ed87d60ebd 100644 --- a/skills/nemo-automodel-model-onboarding/SKILL.md +++ b/skills/nemo-automodel-model-onboarding/SKILL.md @@ -172,11 +172,42 @@ 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. + + 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: @@ -419,8 +450,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) diff --git a/tests/unit_tests/checkpoint/test_checkpointing.py b/tests/unit_tests/checkpoint/test_checkpointing.py index 36af0c99e4..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 @@ -1734,7 +1731,7 @@ 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, "_do_load", return_value=mock_state_dict) as mock_dcp_load, ): mock_model_state = mock_model_state_cls.return_value 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..68f8b89fd8 100644 --- a/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py +++ b/tests/unit_tests/checkpoint/test_state_dict_adapter_capabilities.py @@ -12,7 +12,9 @@ # 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 import pytest import torch @@ -50,6 +52,61 @@ 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.""" + + 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)] + 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 + + 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), + "buffer": torch.ones(1), + "layer._extra_state": {"recipe": "metadata"}, + } + + keys = adapter.get_hf_state_dict_keys(state_dict) + + 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( adapter: StateDictAdapter, state_dict: dict[str, torch.Tensor], @@ -98,7 +155,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, @@ -112,28 +169,18 @@ 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_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") + 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, 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, @@ -145,16 +192,38 @@ 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_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_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_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", _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", [ - BailingMoeV2StateDictAdapter, DeepSeekV3StateDictAdapter, GlmMoeDsaStateDictAdapter, KimiK25VLStateDictAdapter, @@ -165,19 +234,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_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_low_memory_dcp_load is True def test_dense_grouped_adapter_does_not_require_an_expert_backend(): @@ -185,13 +261,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/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 a25ac29d56..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 @@ -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 @@ -78,7 +81,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 +225,139 @@ 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 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) + 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) + + 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 + 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, + ) + 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, + } + + destinations = adapter.to_hf(native, for_checkpoint_load=True) + + 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 + + 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() + + def test_capability_matches_runtime_expert_storage(self, adapter, monkeypatch): + 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_low_memory_dcp_load is True + + monkeypatch.setattr("nemo_automodel.components.moe.state_dict_mixin.get_world_size_safe", lambda: 8) + assert adapter.supports_low_memory_dcp_load is False + + adapter.backend.experts = "torch_mm" + adapter.backend.dispatcher = "mok" + assert adapter.supports_low_memory_dcp_load is False + + adapter.backend.dispatcher = "torch" + adapter.moe_config.expert_bias = True + assert adapter.supports_low_memory_dcp_load is False 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 041117dbc9..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,6 +99,77 @@ 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, monkeypatch): + adapter = mtp_model.state_dict_adapter + native_backbone = { + "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", + } + 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) + + 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_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 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/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 e5b133a821..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 @@ -115,16 +115,59 @@ 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" - 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_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_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" + backend.dispatcher = "deepep" + adapter = NemotronV3StateDictAdapter(config, moe_config, backend) + 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, + ) + + 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_low_memory_dcp_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 = 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() def test_from_hf_map_structure(self, config, moe_config, backend): """Test from_hf_map structure.""" @@ -153,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 = { @@ -185,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" @@ -244,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_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( 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 3e67c7a254..e4aefc4b33 100644 --- a/tests/unit_tests/moe/test_state_dict_mixin.py +++ b/tests/unit_tests/moe/test_state_dict_mixin.py @@ -22,14 +22,31 @@ 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: - 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: @@ -44,6 +61,8 @@ def __init__(self): class MockMoEStateDictMixin(MoESplitExpertsStateDictMixin): + _supports_low_memory_dcp_load = True + def __init__( self, n_experts=8, @@ -91,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 @@ -106,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, ) ) @@ -897,6 +918,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 = [ @@ -905,6 +927,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, @@ -1048,18 +1071,65 @@ class TestInplaceLoadViews: conversion function, so patches must target that module path. """ - def test_expert_write_through_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._supports_write_through_expert_checkpoint_load is True + 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 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 - mixin.backend.experts = "te" - assert mixin._supports_write_through_expert_checkpoint_load is False + mixin.backend.dispatcher = "torch" + 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 + + 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_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 + 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 - def _run_inplace_conversion(self, mixin, fqn, mock_dtensor, splits): + 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): mixin._split_experts_weights = Mock(return_value=splits) mixin._last_expert_ids = list(range(len(splits))) @@ -1070,7 +1140,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) @@ -1080,7 +1152,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 @@ -1103,7 +1179,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() @@ -1141,7 +1223,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")) @@ -1171,6 +1257,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 @@ -1178,13 +1294,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 +1343,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 +1353,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 +1392,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 +1419,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 +1469,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 +1477,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: @@ -1392,6 +1520,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