fix(refit): harden native grouped expert export

Signed-off-by: seonjinn <sna@nvidia.com>

Author: seonjinn

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -248,6 +248,19 @@ def _is_mtp_megatron_param(param_name: str) -> bool:
     return param_name.startswith("mtp.") or ".mtp." in param_name
 
 
+def _grouped_expert_member_views(weight: torch.Tensor) -> list[torch.Tensor]:
+    """Return cached TE grouped members without invoking unsupported indexing."""
+    storage = weight.data if isinstance(weight, torch.nn.Parameter) else weight
+    splitter = getattr(storage, "split_into_quantized_tensors", None)
+    if callable(splitter):
+        members = getattr(storage, "quantized_tensors", None)
+        if members is None:
+            members = splitter()
+            storage.quantized_tensors = members
+        return list(members)
+    return list(storage.unbind(0))
+
+
 def _collect_mtp_hf_layer_names(conversion_tasks: Optional[list]) -> set[str]:
     """Return HF layer names whose weights originate from Megatron's MTP module.
 
@@ -2474,6 +2487,44 @@ def _build_native_mxfp8_conversion_tasks(self) -> list[Any]:
         )
         self._native_grouped_mxfp8_tasks = grouped_tasks
         grouped_names = {task.global_param_name for task in grouped_tasks}
+        grouped_suffixes = (
+            ".mlp.experts.linear_fc1.weight",
+            ".mlp.experts.linear_fc2.weight",
+        )
+        num_experts = int(getattr(self.model.config, "num_moe_experts", 0) or 0)
+        ep_size = int(getattr(self.model.config, "expert_model_parallel_size", 1) or 1)
+        if num_experts and num_experts % ep_size:
+            raise ValueError(
+                f"num_moe_experts={num_experts} must be divisible by "
+                f"expert_model_parallel_size={ep_size}"
+            )
+        local_expert_count = num_experts // ep_size if num_experts else 0
+        grouped_misc_names: dict[str, list[str]] = {}
+        expanded_global_names: list[str] = []
+        for global_name in global_names:
+            if (
+                global_name.endswith(grouped_suffixes)
+                and global_name not in grouped_names
+            ):
+                if _is_mtp_megatron_param(global_name):
+                    raise ValueError(
+                        "native MXFP8 refit does not yet support co-trained MTP "
+                        "grouped experts"
+                    )
+                if local_expert_count <= 0:
+                    raise ValueError(
+                        f"Cannot expand grouped expert parameter {global_name!r} "
+                        "without num_moe_experts"
+                    )
+                expanded = [
+                    f"{global_name}{expert_id}"
+                    for expert_id in range(local_expert_count)
+                ]
+                grouped_misc_names[global_name] = expanded
+                expanded_global_names.extend(expanded)
+            else:
+                expanded_global_names.append(global_name)
+        global_names = expanded_global_names
         remaining_names = [name for name in global_names if name not in grouped_names]
         if not remaining_names:
             return grouped_tasks
@@ -2501,6 +2552,35 @@ def _build_native_mxfp8_conversion_tasks(self) -> list[Any]:
             global_name = _megatron_local_name_to_global(
                 models, self.model.config, local_name, 0
             )
+            expanded_names = grouped_misc_names.get(global_name)
+            if expanded_names is not None:
+                local_module, local_weight = get_module_and_param_from_name(
+                    models, local_name, 0
+                )
+                members = (
+                    []
+                    if local_weight is None
+                    else _grouped_expert_member_views(local_weight)
+                )
+                if len(members) != len(expanded_names):
+                    raise ValueError(
+                        f"Grouped expert parameter {global_name!r} has local shape "
+                        f"{getattr(local_weight, 'shape', None)}, expected "
+                        f"{len(expanded_names)} local experts"
+                    )
+                if local_module is not None and not hasattr(local_module, "config"):
+                    setattr(local_module, "config", self.model.config)
+                for expert_id, expanded_name in enumerate(expanded_names):
+                    local_tasks[expanded_name] = WeightConversionTask(
+                        pp_rank=pp_rank,
+                        vp_stage=0,
+                        param_name=f"{local_name}{expert_id}",
+                        global_param_name=expanded_name,
+                        megatron_module=local_module,
+                        param_weight=members[expert_id],
+                        mapping=mappings[expanded_name],
+                    )
+                continue
             if global_name not in remaining_set:
                 continue
             local_module, local_weight = get_module_and_param_from_name(
@@ -2887,8 +2967,18 @@ def _task_uses_native_mxfp8_storage(self, task: Any, *, grouped: bool) -> bool:
                     members = get_grouped_quantized_members(
                         task.param_weight, create_if_missing=False
                     )
-                except (RuntimeError, ValueError):
-                    return False
+                except RuntimeError:
+                    members = get_grouped_quantized_members(
+                        task.param_weight, create_if_missing=True
+                    )
+                except ValueError as error:
+                    logical_name = self._native_task_projections(task, grouped=True)[0][
+                        0
+                    ]
+                    raise ValueError(
+                        f"Invalid grouped MXFP8 source {logical_name!r} role "
+                        f"'weight': {error}"
+                    ) from error
                 if not members:
                     logical_name = self._native_task_projections(task, grouped=True)[0][
                         0

nemo_rl/weight_sync/nccl_reshard_utils.py

@@ -749,10 +749,25 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None:
         #   BF16 storage → MXFP8 gen  (receiver quantizes the resharded BF16 shard)
         # FP8→BF16 has no consumer (vLLM doesn't accept FP8 bytes into a BF16 param).
         fp8_cfg = megatron_cfg.get("fp8_cfg", {}) or {}
-        fp8_param = fp8_cfg.get("fp8_param", False)
+        fp8_param = bool(
+            fp8_cfg.get("enabled", False) and fp8_cfg.get("fp8_param", False)
+        )
         fp8_recipe = fp8_cfg.get("fp8_recipe", None)
         trainer_precision = policy.get("precision")
         gen_precision = vllm_cfg.get("precision", None)
+        native_mxfp8 = bool(
+            fp8_param
+            and fp8_recipe == "mxfp8"
+            and gen_precision == "fp8"
+            and vllm_cfg.get("is_mx") is True
+        )
+
+        if native_mxfp8 and (megatron_cfg.get("mtp_num_layers", 0) or 0) > 0:
+            violations.append(
+                "native MXFP8 refit does not yet support co-trained MTP layers; "
+                "set policy.megatron_cfg.mtp_num_layers=0 and load static MTP "
+                "weights from the generation checkpoint"
+            )
 
         # The refit byte-copies weights train -> gen, so gen dtype must match
         # train: BF16 (unset / "auto" / "bf16" / "bfloat16") or FP8 ("fp8").  A
@@ -770,7 +785,16 @@ def check_nccl_reshard_refit_support(master_config: Any) -> None:
 
         if gen_precision == "fp8":
             if fp8_param:
-                if vllm_cfg.get("is_mx"):
+                if native_mxfp8:
+                    pass
+                elif fp8_recipe == "mxfp8":
+                    violations.append(
+                        "native MXFP8 storage requires "
+                        "policy.generation.vllm_cfg.is_mx=True "
+                        "(native MXFP8 values and E8M0 scales cannot be "
+                        "loaded by a blockwise-FP8 target)."
+                    )
+                elif vllm_cfg.get("is_mx"):
                     violations.append(
                         "policy.generation.vllm_cfg.is_mx=True does not support "
                         "blockwise-FP8 storage from "
@@ -954,6 +978,8 @@ def _build_dst_meshes(num_gpus: int, rank_offset: int):
 
     per_layer_params: dict[str, list] = OrderedDict()
     for name, meta in state_dict_metadata.items():
+        normalized_components = normalize_refit_components(name, meta)
+        logical_weight = normalized_components[0]
         layer = _extract_layer_name(name)
         expert = is_expert_param(name)
         # Pick the gen (dst) mesh: experts go to the EP/TP-expert mesh, all other
@@ -980,8 +1006,8 @@ def _build_dst_meshes(num_gpus: int, rank_offset: int):
             src_dim_map = stage_src_dim_map
             info = {
                 "name": name,
-                "global_shape": tuple(meta["shape"]),
-                "dtype": meta["dtype"],
+                "global_shape": logical_weight.global_shape,
+                "dtype": logical_weight.dtype,
                 "pp_stage": stage,
                 "src_mesh_info": src_mesh,
                 "dst_mesh_info": dst_mesh,
@@ -996,13 +1022,12 @@ def _build_dst_meshes(num_gpus: int, rank_offset: int):
             src_dim_map = this_src_dim_map
             info = {
                 "name": name,
-                "global_shape": tuple(meta["shape"]),
-                "dtype": meta["dtype"],
+                "global_shape": logical_weight.global_shape,
+                "dtype": logical_weight.dtype,
                 "src_mesh_info": src_mesh,
                 "dst_mesh_info": dst_mesh,
             }
 
-        normalized_components = normalize_refit_components(name, meta)
         component_infos = []
         for component in normalized_components:
             src_placements = get_placements(

nemo_rl/weight_sync/refit_components.py

@@ -187,23 +187,83 @@ def component_plan_digest(refit_info: Mapping[str, Any]) -> str:
     return hashlib.sha256(payload.encode("utf-8")).hexdigest()
 
 
-def native_mxfp8_param_names(refit_info: Mapping[str, Any]) -> set[str]:
-    """Return parameters represented by a canonical MXFP8 value/scale pair."""
+def native_mxfp8_param_names(
+    refit_info: Mapping[str, Any],
+    *,
+    strict: bool = False,
+) -> set[str]:
+    """Return parameters represented by a canonical MXFP8 value/scale pair.
+
+    ``strict=False`` is suitable for feature detection: malformed entries are
+    ignored conservatively. ``strict=True`` validates every serialized
+    value/scale pair and reports malformed native metadata to the caller.
+    """
     result: set[str] = set()
     per_layer_params = refit_info.get("per_layer_params", {})
     for layer_name in refit_info.get("layer_names", []):
         for param_info in per_layer_params.get(layer_name, []):
-            components = param_info.get("components", [])
-            if [component.get("role") for component in components] != [
-                "weight",
-                "weight_scale",
-            ]:
+            if not isinstance(param_info, Mapping):
+                if strict:
+                    raise ValueError("refit parameter metadata must be a mapping")
+                continue
+            logical_name = param_info.get("name")
+            if not isinstance(logical_name, str):
+                if strict:
+                    raise ValueError("refit parameter metadata must contain a name")
+                continue
+            try:
+                serialized = param_info.get("components")
+                metadata: dict[str, Any] = {
+                    "shape": param_info.get("global_shape"),
+                    "dtype": param_info.get("dtype"),
+                }
+                if serialized is not None:
+                    if not isinstance(serialized, Sequence) or isinstance(
+                        serialized, (str, bytes)
+                    ):
+                        raise ValueError(
+                            f"{logical_name} components must be a sequence"
+                        )
+                    normalized_components: list[dict[str, Any]] = []
+                    for component in serialized:
+                        if not isinstance(component, Mapping):
+                            raise ValueError(
+                                f"{logical_name} component metadata must be mappings"
+                            )
+                        normalized_components.append(
+                            {
+                                "role": component.get("role"),
+                                "shape": component.get("global_shape"),
+                                "dtype": component.get("dtype"),
+                            }
+                        )
+                    metadata["components"] = normalized_components
+                normalized = normalize_refit_components(
+                    logical_name,
+                    metadata,
+                )
+            except ValueError:
+                if strict:
+                    raise
+                continue
+            if len(normalized) == 1:
+                continue
+            weight, scale = normalized
+            if weight.dtype != "torch.float8_e4m3fn":
+                if strict:
+                    raise ValueError(
+                        f"{logical_name!r} native weight dtype must be "
+                        "torch.float8_e4m3fn"
+                    )
+                continue
+            if scale.dtype != "torch.uint8":
+                if strict:
+                    raise ValueError(
+                        f"{logical_name!r} weight_scale dtype must be torch.uint8"
+                    )
                 continue
-            if (
-                str(components[0].get("dtype")) == "torch.float8_e4m3fn"
-                and str(components[1].get("dtype")) == "torch.uint8"
-            ):
-                result.add(param_info["name"])
+            if weight.role == "weight" and scale.role == "weight_scale":
+                result.add(logical_name)
     return result