[main] Muon fp8 param gather for decoupled layout

megatron/core/distributed/distributed_data_parallel.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import logging
 import weakref
@@ -182,8 +182,16 @@ def __init__(
             param_to_name[param] = name
             all_params.append(param)
 
-        # Group parameters by (param_dtype, grad_dtype, is_expert_parallel).
-        buffer_groups = group_params_for_buffers(all_params, self.ddp_config.grad_reduce_in_fp32)
+        # Group parameters by (param_dtype, grad_dtype, is_expert_parallel). fp8 params key to
+        # uint8 (own buffer); partition_buckets later merges the small non-fp8 bucket groups into
+        # the fp8 group to aggregate their communication.
+        buffer_groups = group_params_for_buffers(
+            all_params,
+            self.ddp_config.grad_reduce_in_fp32,
+            merge_layerwise_fp8_grads=not getattr(
+                self.ddp_config, 'use_layer_wise_param_layout', True
+            ),
+        )
 
         # Auto-compute layouts when using distributed optimizer but no layout was provided.
         # This maintains backward compatibility for callers that create DDP directly

megatron/core/distributed/distributed_data_parallel_config.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 from dataclasses import dataclass
 from typing import Optional, Tuple
@@ -36,6 +36,12 @@ class DistributedDataParallelConfig:
        enabled. Defaults to 1, which means DistOpt is across entire DP domain.
     """
 
+    use_layer_wise_param_layout: bool = False
+    """Layer-wise (Muon) optimizer only. When True, LayerWise-managed buffers use
+       the shard-aligned padded LayerWise param layout. When False (default), the compact
+       decoupled layout is selected instead.
+    """
+
     check_for_nan_in_grad: bool = False
     """
     If true, check for NaNs and Infs in gradients _before_ communication collective.

megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. 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.
@@ -66,6 +66,7 @@
     from megatron.core.distributed.distributed_data_parallel_config import (
         DistributedDataParallelConfig,
     )
+    from megatron.core.fp8_utils import pop_high_precision_init_val
     from megatron.core.tensor_parallel import get_cuda_rng_tracker
     from megatron.core.utils import is_submodule
 
@@ -77,6 +78,18 @@
     from .distributed_data_parallel_config import DistributedDataParallelConfig
     from .utils import get_cuda_rng_tracker, is_submodule
 
+    # Standalone compatibility: MCore's shared helper is intentionally unavailable when the
+    # independently installable megatron_fsdp package is used without Megatron Core.
+    def pop_high_precision_init_val(param: torch.Tensor) -> Optional[torch.Tensor]:
+        """Return and clear a TE preserved high-precision initial value, if present."""
+        getter = getattr(param, "get_high_precision_init_val", None)
+        if getter is None:
+            return None
+
+        high_precision_init_val = getter()
+        param.clear_high_precision_init_val()
+        return high_precision_init_val
+
     HAVE_MCORE = False
     logger.info("Megatron Core is not installed, Megatron-FSDP will run without Megatron Core.")
 
@@ -3002,8 +3015,9 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params):
                         )
                         # Needed to instantiate FP8 parameters. Requires installing
                         # TransformerEngine.
-                        mbuf.set_item(item_id, p.get_high_precision_init_val())
-                        p.clear_high_precision_init_val()
+                        high_precision_init_val = pop_high_precision_init_val(p)
+                        assert high_precision_init_val is not None
+                        mbuf.set_item(item_id, high_precision_init_val)
                     else:
                         # Insert a copy of the model weight parameter tensor into
                         # the (high-precision) main weight buffer.

megatron/core/fp8_utils.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 """Utility functions related to FP8 that are used throughout Megatron core"""
 
@@ -60,8 +60,17 @@
     HAVE_TE_MXFP8TENSOR = True
 except (ImportError, ModuleNotFoundError):
     # MXFP8Tensor not found
+    MXFP8Tensor = None
     HAVE_TE_MXFP8TENSOR = False
 
+try:
+    from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockwiseQTensor
+
+    HAVE_TE_BLOCKWISE_FP8TENSOR = True
+except (ImportError, ModuleNotFoundError):
+    Float8BlockwiseQTensor = None
+    HAVE_TE_BLOCKWISE_FP8TENSOR = False
+
 try:
     from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor
 
@@ -136,6 +145,23 @@ def is_mxfp8tensor(tensor: torch.Tensor) -> bool:
     return HAVE_TE_MXFP8TENSOR and _is_instance_or_param_data(tensor, MXFP8Tensor)
 
 
+def is_blockwise_float8tensor(tensor: torch.Tensor) -> bool:
+    """Check if a tensor is a Transformer Engine Float8BlockwiseQTensor."""
+    return HAVE_TE_BLOCKWISE_FP8TENSOR and _is_instance_or_param_data(
+        tensor, Float8BlockwiseQTensor
+    )
+
+
+def is_layerwise_fp8_param(tensor: torch.Tensor) -> bool:
+    """Check if an FP8 parameter uses storage supported by LayerWise parameter gather.
+
+    The compact LayerWise path stages whole parameters in BF16 and requantizes them on
+    copy-back. It supports plain MXFP8 and blockwise tensors only; generic Float8Tensor,
+    NVFP4, and GroupedTensor storage require different handling.
+    """
+    return is_mxfp8tensor(tensor) or is_blockwise_float8tensor(tensor)
+
+
 def is_grouped_tensor(tensor: torch.Tensor) -> bool:
     """Check if a tensor is a Transformer Engine GroupedTensor."""
     return HAVE_TE_GROUPED_TENSOR_CLASS and _is_instance_or_param_data(tensor, GroupedTensor)
@@ -150,6 +176,21 @@ def is_grouped_tensor_with_quantized_storage(tensor: torch.Tensor) -> bool:
     return rowwise_data is not None and rowwise_data.dtype == torch.uint8
 
 
+def pop_high_precision_init_val(param: torch.Tensor) -> Optional[torch.Tensor]:
+    """Return and clear a Transformer Engine preserved high-precision initial value.
+
+    The returned tensor is left unmodified so each optimizer path can preserve its
+    existing slicing, cloning, device placement, and dtype conversion behavior.
+    """
+    getter = getattr(param, "get_high_precision_init_val", None)
+    if getter is None:
+        return None
+
+    high_precision_init_val = getter()
+    param.clear_high_precision_init_val()
+    return high_precision_init_val
+
+
 def _get_grouped_quantized_recipe(tensor: torch.Tensor):
     """Return TE recipe for grouped quantized storage, or None if unavailable."""
     tensor = _unwrap_parameter_data(tensor)
@@ -308,6 +349,48 @@ def dequantize_fp8_tensor(fp8_tensor: torch.Tensor) -> torch.Tensor:
         return fp8_tensor.from_float8()
 
 
+def copy_back_gathered_bf16_into_fp8_params(
+    model_params: List[torch.Tensor], srcs_bf16: List[torch.Tensor]
+) -> None:
+    """Copy gathered BF16 whole-params into compact LayerWise bucket parameters.
+
+    Plain BF16 parameters are allowed because a LayerWise bucket can contain BF16 siblings of
+    supported FP8 parameters. Quantized destinations are limited to MXFP8 and blockwise tensors.
+    MXFP8 columnwise data cannot be derived from rowwise data, so force both usages before the
+    batched quantized copy. Validate the entire batch before mutating any quantizer usage.
+    """
+    if len(model_params) != len(srcs_bf16):
+        raise ValueError(
+            "LayerWise FP8 parameter gather copy-back requires one source per parameter: "
+            f"got {len(model_params)} parameters and {len(srcs_bf16)} sources."
+        )
+
+    mxfp8_quantizers = []
+    for model_p in model_params:
+        if is_grouped_tensor_with_quantized_storage(model_p):
+            raise TypeError(
+                "LayerWise FP8 parameter gather does not support Transformer Engine "
+                "GroupedTensor quantized storage. Disable --moe-single-grouped-weight."
+            )
+        if is_float8tensor(model_p) and not is_layerwise_fp8_param(model_p):
+            raise TypeError(
+                "LayerWise FP8 parameter gather supports only MXFP8Tensor and "
+                "Float8BlockwiseQTensor destinations."
+            )
+        if is_mxfp8tensor(model_p):
+            mxfp8_quantizers.append(model_p.data._get_quantizer())
+
+    for quantizer in mxfp8_quantizers:
+        quantizer.set_usage(rowwise=True, columnwise=True)
+
+    copy_tensors_to_quantized_params(model_params, srcs_bf16)
+
+
+def copy_back_gathered_bf16_into_fp8_param(model_p: torch.Tensor, src_bf16: torch.Tensor) -> None:
+    """Single-parameter compatibility wrapper for LayerWise BF16 copy-back."""
+    copy_back_gathered_bf16_into_fp8_params([model_p], [src_bf16])
+
+
 def _resolve_callable_from_python_import_path(dotted_path: str):
     """Resolve a Python import path like 'pkg.mod.func' to a callable.
 
@@ -347,11 +430,9 @@ def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, F
     try:
         custom_recipe = transformer_engine.common.recipe.CustomRecipe(qfactory=quantizer_factory)
     except AttributeError:
-        raise ValueError(
-            """CustomRecipe recipe is not available in this version of 
-            Transformer Engine. Please make sure you are using TE version 
-            >= 2.9.0.dev0."""
-        )
+        raise ValueError("""CustomRecipe recipe is not available in this version of
+            Transformer Engine. Please make sure you are using TE version
+            >= 2.9.0.dev0.""")
     return custom_recipe
 
 

megatron/core/optimizer/distrib_optimizer.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 """Megatron distributed optimizer."""
 
@@ -58,6 +58,7 @@
     get_grouped_quantized_members,
     is_float8tensor,
     is_grouped_tensor_with_quantized_storage,
+    pop_high_precision_init_val,
     quantize_param_shard,
 )
 from ..transformer.fsdp_dtensor_checkpoint import handle_experts_in_state_dict
@@ -428,15 +429,16 @@ def _build_model_and_main_param_groups(
                         if is_nvfp4tensor(model_param) or cls._is_distopt_quantized_param(
                             model_param
                         ):
-                            if hasattr(model_param, 'get_high_precision_init_val'):
+                            high_precision_init_val = pop_high_precision_init_val(model_param)
+                            if high_precision_init_val is not None:
                                 shard_main_param = (
-                                    model_param.get_high_precision_init_val()
-                                    .view(-1)[param_range.start : param_range.end]
+                                    high_precision_init_val.view(-1)[
+                                        param_range.start : param_range.end
+                                    ]
                                     .clone()
                                     .to(model_param.device)
                                     .float()
                                 )
-                                model_param.clear_high_precision_init_val()
                             else:
                                 shard_main_param = model_param.float().view(-1)[
                                     param_range.start : param_range.end
@@ -613,7 +615,11 @@ def compute_full_param_layout(
         Returns:
             FullParamLayout with a PerBufferParamLayout per buffer group.
         """
-        buffer_groups = group_params_for_buffers(params, ddp_config.grad_reduce_in_fp32)
+        buffer_groups = group_params_for_buffers(
+            params,
+            ddp_config.grad_reduce_in_fp32,
+            merge_layerwise_fp8_grads=not getattr(ddp_config, 'use_layer_wise_param_layout', True),
+        )
         layouts = {}
         for buffer_key, (group_params, param_indices) in buffer_groups.items():
             if buffer_key.is_expert_parallel: