state_dict_mixin.py

# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gc
import re
from typing import Any, Optional

import torch
from torch.distributed.device_mesh import DeviceMesh

from nemo_automodel.components.moe.state_dict_utils import (
    create_dtensor_from_local,
    get_expert_range_for_rank_from_mesh,
    get_submesh,
    is_dtensor,
    should_load_expert_for_rank,
    split_experts_weights_dtensor_aware,
)


class MoESplitExpertsStateDictMixin:
    """Mixin class providing MoE state dict conversion utilities.

    This mixin provides methods for:
    - Expert parallelism calculations (ranges, assignment)
    - Format conversion between HuggingFace and native formats
    - Both GroupedExperts and DeepEP format support
    - DTensor-aware expert loading and conversion

    Can be used by any MoE model that needs expert parallelism and format conversion.
    """

    # These attributes must be set by subclasses in their __init__ method:
    # - self.moe_config: MoE configuration object with expert settings
    # - self.config: Model configuration object
    # - self.backend: Backend configuration object

    @property
    def _is_gated_moe(self) -> bool:
        """Check if the MoE uses gated activation (e.g., SwiGLU) or non-gated (e.g., ReLU²)."""
        from nemo_automodel.components.moe.experts import is_gated_activation

        return is_gated_activation(self.moe_config.expert_activation)

    def _register_inplace_loaded_key(self, fqn: str, prefix_override: str | None) -> None:
        """Mark ``fqn`` as loaded via in-place views so ``_from_hf_w_merged_experts`` skips its rebuild.

        The tracked key must match the native_key that the from_hf merge loop
        reconstructs from the HF per-expert keys. For backbone tensors the
        native_key equals ``fqn``; for MTP tensors (``prefix_override="mtp."``)
        the HF keys live under the ``mtp.`` namespace and from_hf processes
        them with that prefix stripped, so the tracked key is also the
        ``mtp.``-less form. The user of this set (``_from_hf_w_merged_experts``)
        receives the matching stripped key when called via the adapter's
        per-namespace dispatch.
        """
        if prefix_override is not None and prefix_override.endswith("."):
            tracked = fqn[len(prefix_override) :] if fqn.startswith(prefix_override) else fqn
        else:
            tracked = fqn
        if not hasattr(self, "_inplace_loaded_native_keys") or self._inplace_loaded_native_keys is None:
            self._inplace_loaded_native_keys = set()
        self._inplace_loaded_native_keys.add(tracked)

    @property
    def view_loaded_native_keys(self) -> set[str]:
        """Native keys loaded in-place via strided views during the most recent ``from_hf``.

        MoE experts with a plain local split are loaded by DCP writing the checkpoint tensors
        straight through non-contiguous strided views into the model's grouped expert storage.
        Such keys are intentionally absent from the dict ``from_hf`` returns (the data is already
        in the model) but are NOT missing. ``_from_hf_w_merged_experts`` records them here so the
        checkpoint loader can exclude them from false "missing" key-diff warnings. The record is
        reset at the start of each load by ``_from_hf_w_merged_experts(reset_view_loaded_keys=True)``.
        """
        return getattr(self, "_view_loaded_native_keys", None) or set()

    @property
    def _hf_prefix(self) -> str:
        """Prefix for HuggingFace format keys. Override in subclass."""
        return "model." if self._uses_model_prefix else ""

    @property
    def _expert_path_segment(self) -> str:
        """Path segment for experts (e.g., 'mlp.experts' or 'mixer.experts'). Override in subclass."""
        return "mlp.experts"

    def _validate_expert_availability(
        self,
        hf_state_dict: dict[str, Any],
        n_experts: int,
        device_mesh: Optional["DeviceMesh"] = None,
    ) -> None:
        """Validate that all required experts are available in the HF state dict before loading.
        Only validates experts needed for the current rank and layers present in the state dict.

        Args:
            hf_state_dict: HuggingFace format state dict
            n_experts: Total number of experts
            device_mesh: Optional device mesh for expert parallelism

        Raises:
            RuntimeError: If required expert weights are missing from the checkpoint
        """
        if device_mesh is not None:
            start_expert, end_expert = get_expert_range_for_rank_from_mesh(device_mesh, n_experts)
            required_experts = list(range(start_expert, end_expert))
            rank = (
                get_submesh(device_mesh, ("ep",)).get_rank()
                if "ep" in device_mesh.mesh_dim_names
                else device_mesh.get_rank()
            )
            rank_info = f" (rank {rank})"
        else:
            required_experts = list(range(n_experts))
            rank_info = ""

        expert_segment = self._expert_path_segment

        # Detect actual prefix from keys (handles both HF format and pre-renamed internal format)
        key_prefix = ""
        for key in hf_state_dict.keys():
            if f".{expert_segment}." in key and "layers." in key:
                key_prefix = key[: key.index("layers.")]
                break

        # Build list of all possible prefixes
        prefixes = ["model.language_model.", "model.", "language_model.", ""]
        if key_prefix and key_prefix not in prefixes:
            prefixes.insert(0, key_prefix)

        layers_with_experts: dict[int, set[str]] = {}
        # Create pattern with all prefixes
        escaped_prefixes = [re.escape(p) for p in prefixes]
        prefix_pattern = "(?P<prefix>" + "|".join(escaped_prefixes) + ")"
        pattern = (
            rf"{prefix_pattern}layers\.(\d+)\.{re.escape(expert_segment)}\.\d+\.(gate_proj|up_proj|down_proj)\.weight"
        )
        for key in hf_state_dict.keys():
            match = re.match(pattern, key)
            if match:
                prefix = match.group("prefix") or ""
                layer_num = int(match.group(2))
                layers_with_experts.setdefault(layer_num, set()).add(prefix)

        if not layers_with_experts:
            return

        missing_weights = []
        projection_types = ["gate_proj", "up_proj", "down_proj"] if self._is_gated_moe else ["up_proj", "down_proj"]

        for layer_num, prefixes in layers_with_experts.items():
            for prefix in prefixes:
                for expert_id in required_experts:
                    for proj_type in projection_types:
                        expected_key = f"{prefix}layers.{layer_num}.{expert_segment}.{expert_id}.{proj_type}.weight"
                        if expected_key not in hf_state_dict:
                            missing_weights.append(expected_key)

        if missing_weights:
            missing_count = len(missing_weights)
            total_required = len(required_experts) * len(layers_with_experts) * len(projection_types)
            raise RuntimeError(
                f"Expert weights missing from checkpoint{rank_info}: {missing_count}/{total_required} required weights not found. "
                f"Cannot load experts - checkpoint may be incomplete or corrupted. "
                f"Layers with experts: {sorted(layers_with_experts)}, Required experts: {required_experts}. "
                f"First few missing keys: {missing_weights[:5]}"
                + (f" (and {missing_count - 5} more)" if missing_count > 5 else "")
            )

    def _split_experts_weights(self, weight: torch.Tensor, n_experts: int) -> list[torch.Tensor]:
        """Split grouped expert weights into individual expert weights.
        For grouped expert weights with shape [n_experts, ...], split into n_experts tensors each with shape [...].
        Supports both regular tensors and DTensors.
        """
        if is_dtensor(weight):
            split_weights, expert_ids = split_experts_weights_dtensor_aware(weight, n_experts)
            self._last_expert_ids = expert_ids
            return split_weights
        else:
            if weight.shape[0] != n_experts:
                raise ValueError(f"Expected first dimension to be {n_experts}, got {weight.shape[0]}")

            split_weights = []
            expert_ids = []
            for i in range(n_experts):
                expert_weight = weight[i]  # Shape: [...] (expert dimension removed)
                split_weights.append(expert_weight)
                expert_ids.append(i)

            self._last_expert_ids = expert_ids
            return split_weights

    def _concatenate_expert_weights(
        self, expert_weights_by_layer: dict[str, Any], n_experts: int
    ) -> Optional[torch.Tensor]:
        """Concatenate the weights of separate experts into GroupedExpert weights.

        Args:
            expert_weights_by_layer: Nested dict structure containing expert weights
            n_experts: Total number of experts expected

        Returns:
            Stacked tensor if all experts are available for a layer, None otherwise
        """
        for layer, abstract_keys in list(expert_weights_by_layer.items()):
            for abstract_key, experts in list(abstract_keys.items()):
                if len(experts) == n_experts:
                    sorted_expert_ids = sorted(experts.keys())
                    sorted_experts = [experts[i] for i in sorted_expert_ids]
                    stacked_tensor = torch.stack(sorted_experts, dim=0)

                    del expert_weights_by_layer[layer][abstract_key]
                    if not expert_weights_by_layer[layer]:
                        del expert_weights_by_layer[layer]

                    return stacked_tensor

        return None

    def _convert_lora_expert_to_hf(
        self,
        fqn: str,
        tensor: torch.Tensor,
        n_experts: int,
        inter_dim: int,
        expert_segment: str,
    ) -> list[tuple[str, torch.Tensor]]:
        """Convert a grouped MoE expert LoRA tensor to per-expert HF PEFT format.

        Handles the four LoRA parameter types produced by GroupedExpertsLoRA /
        GroupedExpertsDeepEPLoRA and converts them to per-expert ``lora_A.weight``
        / ``lora_B.weight`` keys that HF PEFT understands.

        The prefix (e.g. ``base_model.model.model.``) is preserved from the
        incoming *fqn* so that both PEFT and FFT save paths work correctly.
        """
        match = re.search(r"(.*)layers\.(\d+)\.", fqn)
        if not match:
            return None
        fqn_prefix = match.group(1)
        layer_num = match.group(2)
        suffix = fqn.rsplit(".", 1)[-1]

        splits = self._split_experts_weights(tensor, n_experts)
        result: list[tuple[str, torch.Tensor]] = []

        for i, w in enumerate(splits):
            expert_id = self._last_expert_ids[i]
            base = f"{fqn_prefix}layers.{layer_num}.{expert_segment}.{expert_id}"

            if suffix == "lora_gate_and_up_A":
                # [dim, lora_dim] -> [lora_dim, dim] (nn.Linear convention)
                w_t = w.transpose(0, 1).contiguous()
                if self._is_gated_moe:
                    result.append((f"{base}.gate_proj.lora_A.weight", w_t))
                    result.append((f"{base}.up_proj.lora_A.weight", w_t.clone()))
                else:
                    result.append((f"{base}.up_proj.lora_A.weight", w_t))

            elif suffix == "lora_gate_and_up_B":
                # [lora_dim, 2*inter] (gated) or [lora_dim, inter] (non-gated)
                if self._is_gated_moe:
                    w_gate = w[:, :inter_dim].transpose(0, 1).contiguous()
                    w_up = w[:, inter_dim:].transpose(0, 1).contiguous()
                    result.append((f"{base}.gate_proj.lora_B.weight", w_gate))
                    result.append((f"{base}.up_proj.lora_B.weight", w_up))
                else:
                    result.append((f"{base}.up_proj.lora_B.weight", w.transpose(0, 1).contiguous()))

            elif suffix == "lora_down_A":
                # [inter_dim, lora_dim] -> [lora_dim, inter_dim]
                result.append((f"{base}.down_proj.lora_A.weight", w.transpose(0, 1).contiguous()))

            elif suffix == "lora_down_B":
                # [lora_dim, dim] -> [dim, lora_dim]
                result.append((f"{base}.down_proj.lora_B.weight", w.transpose(0, 1).contiguous()))

        return result

    def _recombine_lora_expert_keys(self, state_dict: dict[str, Any]) -> dict[str, Any]:
        """Recombine per-expert HF LoRA keys back to grouped MoE LoRA format.

        This is the reverse of ``_convert_lora_expert_to_hf``.  It detects
        per-expert LoRA keys (e.g.
        ``layers.0.mlp.experts.0.gate_proj.lora_A.weight``) and recombines
        them into the grouped tensors expected by GroupedExpertsLoRA /
        GroupedExpertsDeepEPLoRA (e.g. ``layers.0.mlp.experts.lora_gate_and_up_A``).
        """
        expert_segment = re.escape(self._expert_path_segment)
        n_experts = self.moe_config.n_routed_experts

        lora_pattern = re.compile(
            rf"(?P<prefix>.*)layers\.(?P<layer>\d+)\.{expert_segment}\."
            rf"(?P<expert>\d+)\.(?P<proj>gate_proj|up_proj|down_proj)\.(?P<lora>lora_[AB])\.weight"
        )

        # Group: (prefix, layer, proj, lora) -> {expert_id: tensor}
        lora_groups: dict[tuple, dict[int, torch.Tensor]] = {}
        consumed_keys: set[str] = set()

        for key, value in state_dict.items():
            m = lora_pattern.match(key)
            if m:
                group_key = (m.group("prefix"), m.group("layer"), m.group("proj"), m.group("lora"))
                lora_groups.setdefault(group_key, {})[int(m.group("expert"))] = value
                consumed_keys.add(key)

        if not consumed_keys:
            return state_dict

        result = {k: v for k, v in state_dict.items() if k not in consumed_keys}
        processed: set[tuple] = set()

        for (prefix, layer, proj, lora), experts in lora_groups.items():
            group_id = (prefix, layer, proj, lora)
            if group_id in processed:
                continue
            processed.add(group_id)

            if len(experts) != n_experts:
                for eid, t in experts.items():
                    orig_seg = self._expert_path_segment
                    result[f"{prefix}layers.{layer}.{orig_seg}.{eid}.{proj}.{lora}.weight"] = t
                continue

            sorted_ids = sorted(experts.keys())
            base_key = f"{prefix}layers.{layer}.{self._expert_path_segment}"

            if proj in ("gate_proj", "up_proj") and lora == "lora_A":
                if self._is_gated_moe and proj == "up_proj":
                    continue  # gate_proj.lora_A already produces this (deduplicate)
                tensors = [experts[eid].transpose(0, 1).contiguous() for eid in sorted_ids]
                result[f"{base_key}.lora_gate_and_up_A"] = torch.stack(tensors, dim=0)
                if self._is_gated_moe:
                    processed.add((prefix, layer, "up_proj", "lora_A"))

            elif proj in ("gate_proj", "up_proj") and lora == "lora_B":
                if self._is_gated_moe and proj == "up_proj":
                    continue  # handled by gate_proj.lora_B below
                if self._is_gated_moe:
                    up_key = (prefix, layer, "up_proj", "lora_B")
                    up_experts = lora_groups.get(up_key, {})
                    if len(up_experts) != n_experts:
                        for eid, t in experts.items():
                            orig_seg = self._expert_path_segment
                            result[f"{prefix}layers.{layer}.{orig_seg}.{eid}.{proj}.{lora}.weight"] = t
                        continue
                    gate_ts = [experts[eid].transpose(0, 1).contiguous() for eid in sorted_ids]
                    up_ts = [up_experts[eid].transpose(0, 1).contiguous() for eid in sorted_ids]
                    combined = torch.cat([torch.stack(gate_ts, dim=0), torch.stack(up_ts, dim=0)], dim=-1)
                    result[f"{base_key}.lora_gate_and_up_B"] = combined
                    processed.add(up_key)
                else:
                    tensors = [experts[eid].transpose(0, 1).contiguous() for eid in sorted_ids]
                    result[f"{base_key}.lora_gate_and_up_B"] = torch.stack(tensors, dim=0)

            elif proj == "down_proj":
                native_suffix = "lora_down_A" if lora == "lora_A" else "lora_down_B"
                tensors = [experts[eid].transpose(0, 1).contiguous() for eid in sorted_ids]
                result[f"{base_key}.{native_suffix}"] = torch.stack(tensors, dim=0)

        return result

    def _to_hf_w_split_experts(self, state_dict: dict[str, Any], **kwargs: Any) -> dict[str, Any]:
        """Convert DeepEP format to HuggingFace format.

        Handles ``gate_and_up_projs`` / ``down_projs`` -> individual expert
        weights. Forwards ``**kwargs`` to
        ``_convert_single_merged_expert_to_hf_split_experts`` for adapter
        compatibility (e.g. ``exclude_key_regex``).
        """
        hf_state_dict: dict[str, Any] = {}

        for fqn, tensor in state_dict.items():
            converted = self._convert_single_merged_expert_to_hf_split_experts(fqn, tensor, **kwargs)
            if converted is not None:
                for key, value in converted:
                    hf_state_dict[key] = value
            else:
                hf_state_dict[fqn] = tensor

        return hf_state_dict

    def _from_hf_w_merged_experts(
        self,
        hf_state_dict: dict[str, Any],
        device_mesh: Optional["DeviceMesh"] = None,
        reset_view_loaded_keys: bool = True,
    ) -> dict[str, Any]:
        """Convert HF checkpoint to native format.

        For gated activations (SwiGLU, Quick-GEGLU):
            Creates combined gate_and_up_projs [n_experts, dim, 2*inter_dim] and
            transposed down_projs tensors.

        For non-gated activations (ReLU²):
            Creates gate_and_up_projs [n_experts, dim, inter_dim] and transposed down_projs tensors.

        Args:
            reset_view_loaded_keys: Clear the in-place (strided-view) loaded-key record at the
                start of this call. A single ``from_hf`` may invoke this method more than once
                (e.g. backbone then MTP merge); the later call(s) pass ``False`` so the view-loaded
                keys accumulate across one logical load. Resetting here (rather than in the loader)
                keeps the whole view-key lifecycle inside the adapter and ensures each load starts
                clean (no leak from a prior load such as an init-time partial load).
        """
        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)

        if device_mesh is not None:
            start_expert, end_expert = get_expert_range_for_rank_from_mesh(device_mesh, n_experts)
            expected_experts_per_rank = end_expert - start_expert
            rank = (
                get_submesh(device_mesh, ("ep",)).get_rank()
                if "ep" in device_mesh.mesh_dim_names
                else device_mesh.get_rank()
            )
        else:
            start_expert, end_expert = 0, n_experts
            expected_experts_per_rank = n_experts
            rank = None

        state_dict: dict[str, Any] = {}
        expert_weights_by_layer: dict[str, dict[str, dict[int, torch.Tensor]]] = {}

        # Handle both formats:
        # - model.layers.{L}.{expert_segment}.{E}.gate_proj.weight (with model prefix)
        # - language_model.layers.{L}.{expert_segment}.{E}.gate_proj.weight (with language_model prefix)
        # - layers.{L}.{expert_segment}.{E}.gate_proj.weight (without model prefix)
        expert_pattern = re.compile(
            rf"(?P<prefix>(?:model\.)?(?:language_model\.)?)layers\.(\d+)\.{re.escape(expert_segment)}\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight"
        )

        inplace_loaded_keys: set = getattr(self, "_inplace_loaded_native_keys", None) or set()
        consumed_inplace_keys: set = set()

        for key in list(hf_state_dict.keys()):
            value = hf_state_dict.pop(key)
            if f".{expert_segment}." in key and key.endswith(".weight"):
                m = expert_pattern.match(key)
                if m is None:
                    state_dict[key] = value
                    continue

                prefix = m.group("prefix") or ""
                layer_num, expert_num, which = m.group(2), m.group(3), m.group(4)
                expert_num = int(expert_num)

                if which in ["gate_proj", "up_proj"]:
                    native_key = f"{prefix}layers.{layer_num}.{expert_segment}.gate_and_up_projs"
                else:  # down_proj
                    native_key = f"{prefix}layers.{layer_num}.{expert_segment}.down_projs"

                # Skip rebuild: DCP wrote through the view; model already holds the data.
                if native_key in inplace_loaded_keys:
                    consumed_inplace_keys.add(native_key)
                    del value
                    continue

                if not should_load_expert_for_rank(expert_num, device_mesh, n_experts):
                    continue

                if layer_num not in expert_weights_by_layer:
                    expert_weights_by_layer[layer_num] = {}

                if native_key not in expert_weights_by_layer[layer_num]:
                    expert_weights_by_layer[layer_num][native_key] = {}

                if which in ["gate_proj", "up_proj"]:
                    # Non-gated models only use up_proj, skip gate_proj
                    if not is_gated and which == "gate_proj":
                        continue

                    # Store weight: gated uses dict for gate+up, non-gated stores tensor directly
                    if is_gated:
                        if expert_num not in expert_weights_by_layer[layer_num][native_key]:
                            expert_weights_by_layer[layer_num][native_key][expert_num] = {}
                        expert_weights_by_layer[layer_num][native_key][expert_num][which] = value
                    else:
                        expert_weights_by_layer[layer_num][native_key][expert_num] = value

                    # Check if all experts are complete
                    all_complete = len(expert_weights_by_layer[layer_num][native_key]) == expected_experts_per_rank
                    if is_gated:
                        all_complete = all_complete and all(
                            isinstance(d, dict) and "gate_proj" in d and "up_proj" in d
                            for d in expert_weights_by_layer[layer_num][native_key].values()
                        )

                    if all_complete:
                        expert_ids = sorted(expert_weights_by_layer[layer_num][native_key].keys())
                        tensors = []
                        for expert_id in expert_ids:
                            expert_data = expert_weights_by_layer[layer_num][native_key][expert_id]

                            if is_gated:
                                gate_weight = expert_data["gate_proj"]
                                up_weight = expert_data["up_proj"]
                                if is_dtensor(gate_weight):
                                    gate_weight = gate_weight.to_local()
                                if is_dtensor(up_weight):
                                    up_weight = up_weight.to_local()
                                gate_t = gate_weight.transpose(0, 1)
                                up_t = up_weight.transpose(0, 1)
                                tensors.append(torch.cat([gate_t, up_t], dim=-1))
                            else:
                                up_weight = expert_data
                                if is_dtensor(up_weight):
                                    up_weight = up_weight.to_local()
                                tensors.append(up_weight.transpose(0, 1))

                        stacked = torch.stack(tensors, dim=0).to(self.dtype)
                        state_dict[native_key] = create_dtensor_from_local(stacked, device_mesh, rank)

                        # Aggressively release intermediates so the per-layer
                        # transient does not pile on top of the model's
                        # already-materialized GPU DTensors. Without this,
                        # ``tensors``/``stacked`` and the per-expert dict
                        # entries hang around until Python's refcount GC
                        # eventually runs — too late under tight GPU budgets
                        # (e.g. a large MoE on 2 nodes / 8 GPUs).
                        del tensors, stacked
                        del expert_weights_by_layer[layer_num][native_key]
                        if not expert_weights_by_layer[layer_num]:
                            del expert_weights_by_layer[layer_num]
                        gc.collect()
                        if torch.cuda.is_available():
                            torch.cuda.empty_cache()

                else:  # down_proj
                    expert_weights_by_layer[layer_num][native_key][expert_num] = value

                    if len(expert_weights_by_layer[layer_num][native_key]) == expected_experts_per_rank:
                        expert_ids = sorted(expert_weights_by_layer[layer_num][native_key].keys())

                        ordered = []
                        for expert_id in expert_ids:
                            down_weight = expert_weights_by_layer[layer_num][native_key][expert_id]  # [dim, inter_dim]

                            # Extract local tensor if input is already a DTensor
                            if is_dtensor(down_weight):
                                down_weight = down_weight.to_local()

                            down_t = down_weight.transpose(0, 1)  # [inter_dim, dim]
                            ordered.append(down_t)

                        stacked = torch.stack(ordered, dim=0)
                        stacked = stacked.to(self.dtype)

                        state_dict[native_key] = create_dtensor_from_local(stacked, device_mesh, rank)

                        # See gate/up branch above for the cleanup rationale.
                        del ordered, stacked
                        del expert_weights_by_layer[layer_num][native_key]
                        if not expert_weights_by_layer[layer_num]:
                            del expert_weights_by_layer[layer_num]
                        gc.collect()
                        if torch.cuda.is_available():
                            torch.cuda.empty_cache()

            else:
                if not key.endswith("_scale_inv"):
                    state_dict[key] = value

        # Drop consumed entries so a subsequent from_hf (e.g. MTP merge after backbone) starts clean.
        if consumed_inplace_keys:
            self._inplace_loaded_native_keys -= consumed_inplace_keys
            # These native keys were loaded in-place via strided views into model storage (DCP
            # writes the checkpoint tensors straight through them), so they are intentionally
            # absent from the returned state_dict but are NOT missing. Record them so the
            # checkpoint loader's key-diff can exclude them and only flag genuinely unloaded params.
            self._view_loaded_native_keys = (
                getattr(self, "_view_loaded_native_keys", None) or set()
            ) | consumed_inplace_keys

        # Recombine any per-expert HF LoRA keys back to grouped format
        state_dict = self._recombine_lora_expert_keys(state_dict)

        return state_dict

    def _convert_single_merged_expert_to_hf_split_experts(
        self,
        fqn: str,
        tensor: torch.Tensor,
        *,
        prefix_override: str | None = None,
        **kwargs,
    ) -> list[tuple[str, torch.Tensor]]:
        """Convert a single merged expert tensor from native format to split HuggingFace format.

        When ``tensor`` is a model DTensor with a plain (non-DTensor) local
        split — i.e. ``ep_shard == 1`` — the per-expert outputs are returned
        as **non-contiguous strided views** into the local storage of the
        model's grouped DTensor instead of newly-allocated contiguous copies.
        DCP's ``target.copy_(source)`` then writes safetensors data directly
        through the views into the model's storage, and
        ``_from_hf_w_merged_experts`` skips the rebuild for the corresponding
        native key (tracked in ``_inplace_loaded_native_keys``). For loads of
        large MoE checkpoints this avoids tens of GB of per-expert
        scratch on top of the already-materialized model.

        Save callers must materialize the views before serializing —
        ``safetensors.torch.save`` rejects non-contiguous tensors. See
        ``_materialize_to_hf_views_for_save`` in ``checkpointing.py``.

        Args:
            fqn: Fully qualified name of the tensor in native format.
            tensor: The tensor to convert.
            prefix_override: When provided, replaces ``self._hf_prefix`` in
                emitted HF keys. Used to route conversions through namespaces
                outside the main backbone, e.g. ``"mtp."`` for the MTP head.
            **kwargs: Absorbed for forward-compatibility with base callers
                that forward arbitrary state-dict kwargs (e.g. ``exclude_key_regex``).

        Returns:
            List of (fqn, tensor) tuples in HuggingFace format, or None if not an expert tensor.
        """
        n_experts = self.moe_config.n_routed_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
        # When quantizing, the adapter casts each split with ``value.to(float8_e4m3fn)``,
        # which ALLOCATES a new tensor that no longer aliases the model's grouped storage.
        # An in-place DCP ``copy_`` into that throwaway buffer would silently never reach the
        # model (experts stay at random init -> garbage loss). So fp8 loads must take the
        # rebuild path: never mark these keys in-place-loaded when ``quantization`` is set.
        quantization = kwargs.get("quantization", False)

        # GroupedExpertsTE (backend.experts == "te") exposes gate_and_up_projs/down_projs as a
        # torch.stack COPY of TE's per-expert weight{i} params; it does not alias the model's
        # grouped storage, so an in-place copy_ would write that throwaway buffer and leave the TE
        # experts at their initial values. Force the rebuild path for it. Other expert backends
        # keep a real aliasing stacked Parameter and are unaffected.
        experts_alias_grouped_storage = getattr(getattr(self, "backend", None), "experts", None) != "te"

        from nemo_automodel.components.moe.state_dict_utils import (
            is_dtensor,
            validate_dtensor_expert_sharding,
        )

        if f".{expert_segment}.gate_and_up_projs" in fqn and fqn.endswith(".gate_and_up_projs"):
            layer_num = re.search(r"layers\.(\d+)", fqn).group(1)

            if is_dtensor(tensor):
                validate_dtensor_expert_sharding(tensor, n_experts, f"gate_and_up_projs layer {layer_num}")

            splits = self._split_experts_weights(tensor, n_experts)

            inplace_ok = (
                (is_dtensor(tensor) or (isinstance(tensor, torch.Tensor) and tensor.is_cuda and not tensor.is_meta))
                and len(splits) > 0
                and not is_dtensor(splits[0])
                and not quantization
                and experts_alias_grouped_storage
            )
            if inplace_ok:
                self._register_inplace_loaded_key(fqn, prefix_override)

            result = []
            for i, w in enumerate(splits):
                expert_id = self._last_expert_ids[i]
                if self._is_gated_moe:
                    # Gated: split into gate_proj and up_proj
                    if inplace_ok:
                        w_gate = w[:, :inter_dim].transpose(0, 1)
                        w_up = w[:, inter_dim:].transpose(0, 1)
                    else:
                        w_gate = w[:, :inter_dim].transpose(0, 1).contiguous()
                        w_up = w[:, inter_dim:].transpose(0, 1).contiguous()
                    result.append((f"{prefix}layers.{layer_num}.{expert_segment}.{expert_id}.gate_proj.weight", w_gate))
                    result.append((f"{prefix}layers.{layer_num}.{expert_segment}.{expert_id}.up_proj.weight", w_up))
                else:
                    # Non-gated: only up_proj (tensor is [dim, inter_dim], not [dim, 2*inter_dim])
                    if inplace_ok:
                        w_up = w.transpose(0, 1)
                    else:
                        w_up = w.transpose(0, 1).contiguous()
                    result.append((f"{prefix}layers.{layer_num}.{expert_segment}.{expert_id}.up_proj.weight", w_up))
            del splits
            if not inplace_ok and isinstance(tensor, torch.Tensor) and not tensor.is_meta and torch.cuda.is_available():
                gc.collect()
                torch.cuda.empty_cache()
            return result

        elif (
            f".{expert_segment}.down_projs" in fqn
            and fqn.endswith(".down_projs")
            and tensor.ndim == 3
            and tensor.shape[1] == inter_dim
        ):
            layer_num = re.search(r"layers\.(\d+)", fqn).group(1)

            if is_dtensor(tensor):
                validate_dtensor_expert_sharding(tensor, n_experts, f"down_projs (DeepEP) layer {layer_num}")

            splits = self._split_experts_weights(tensor, n_experts)
            inplace_ok = (
                (is_dtensor(tensor) or (isinstance(tensor, torch.Tensor) and tensor.is_cuda and not tensor.is_meta))
                and len(splits) > 0
                and not is_dtensor(splits[0])
                and not quantization
                and experts_alias_grouped_storage
            )
            if inplace_ok:
                self._register_inplace_loaded_key(fqn, prefix_override)

            result = []
            for i, w in enumerate(splits):
                expert_id = self._last_expert_ids[i]
                if inplace_ok:
                    w_down = w.transpose(0, 1)
                else:
                    w_down = w.transpose(0, 1).contiguous()
                result.append(
                    (
                        f"{prefix}layers.{layer_num}.{expert_segment}.{expert_id}.down_proj.weight",
                        w_down,
                    )
                )
            # See gate_and_up branch above for the cleanup rationale.
            del splits
            if not inplace_ok and isinstance(tensor, torch.Tensor) and not tensor.is_meta and torch.cuda.is_available():
                gc.collect()
                torch.cuda.empty_cache()
            return result

        # MoE expert LoRA keys: split grouped 3-D adapter tensors into per-expert
        # HF-PEFT-compatible keys so that AutoPeftModelForCausalLM can load & merge.
        _LORA_EXPERT_SUFFIXES = ("lora_gate_and_up_A", "lora_gate_and_up_B", "lora_down_A", "lora_down_B")
        for suffix in _LORA_EXPERT_SUFFIXES:
            if f".{expert_segment}.{suffix}" in fqn and fqn.endswith(f".{suffix}"):
                return self._convert_lora_expert_to_hf(fqn, tensor, n_experts, inter_dim, expert_segment)

        return None