Route non-Muon params through DistributedOptimizer

megatron/core/distributed/distributed_data_parallel.py

@@ -473,6 +473,57 @@ def no_sync(self):
             for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
                 bucket_group.is_last_microbatch = True
 
+    def _start_bucket_group_param_sync(
+        self, bucket_group: '_ParamAndGradBucketGroup', force_sync: bool
+    ) -> None:
+        """Dispatch one bucket group's param all-gather + run the FP8 / MXFP8
+        post-all-gather work the synchronous path needs.
+
+        Factored out of :meth:`start_param_sync` so callers that own a subset
+        of bucket groups (e.g. a chained ``LayerWiseDistributedOptimizer`` +
+        ``DistributedOptimizer`` pair) can sync only their own buckets without
+        losing the FP8 post-processing that follows the collective.
+        """
+        bucket_group.start_param_sync(force_sync=force_sync)
+
+        if self.ddp_config.overlap_param_gather:
+            return
+
+        # For MXFP8 params, we need to copy the all-gathered param data from the buffer to
+        # the param.data, since param buffer is not mapped to model params for MXFP8 case.
+        # The paramaters are cast from bf16 to MXFP8 during copy.
+        # In the case of "overlap_param_gather=True", the param copy is done
+        # in "finish_param_sync" stage after zeroing the shared gardient buffers.
+        if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
+            for bucket in bucket_group.buckets:
+                is_bf16_weight_bucket = False
+                for param in bucket.params:
+                    # Skip copying since bf16 weights in the mxfp8 model
+                    # are already mapped to param.data.
+                    if not is_float8tensor(param):
+                        is_bf16_weight_bucket = True
+                        break
+                    param_start, param_end = bucket.param_to_index[param]
+                    param_slice = bucket.param_data.view(-1)[param_start:param_end]
+                    param.data.copy_(param_slice.view(param.data.shape))
+                if is_bf16_weight_bucket:
+                    continue
+                # All-gathered params are not needed after being copied to param.data.
+                # Zero out the param buffer (shared with grad buffer) for gradient
+                # accumulation. We cannot zero out the entire grad buffer because one grad
+                # buffer may correspond to multiple param buffers. If we zero out the entire
+                # grad buffer, it would clear the data of those param buffers that have not
+                # yet completed AG.
+                bucket.param_data.zero_()
+        else:
+            fp8_params = []
+            for bucket in bucket_group.buckets:
+                for param in bucket.params:
+                    if is_float8tensor(param):
+                        fp8_params.append(param)
+            if len(fp8_params) > 0:
+                post_all_gather_processing(fp8_params)
+
     def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bool = False):
         """
         Initiates param sync (all-gather) communication operations for all model parameters.
@@ -493,43 +544,7 @@ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bo
                 return
 
         for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
-            bucket_group.start_param_sync(force_sync=force_sync)
-
-            if not self.ddp_config.overlap_param_gather:
-                # For MXFP8 params, we need to copy the all-gathered param data from the buffer to
-                # the param.data, since param buffer is not mapped to model params for MXFP8 case.
-                # The paramaters are cast from bf16 to MXFP8 during copy.
-                # In the case of "overlap_param_gather=True", the param copy is done
-                # in "finish_param_sync" stage after zeroing the shared gardient buffers.
-                if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
-                    for bucket in bucket_group.buckets:
-                        is_bf16_weight_bucket = False
-                        for param in bucket.params:
-                            # Skip copying since bf16 weights in the mxfp8 model
-                            # are already mapped to param.data.
-                            if not is_float8tensor(param):
-                                is_bf16_weight_bucket = True
-                                break
-                            param_start, param_end = bucket.param_to_index[param]
-                            param_slice = bucket.param_data.view(-1)[param_start:param_end]
-                            param.data.copy_(param_slice.view(param.data.shape))
-                        if is_bf16_weight_bucket:
-                            continue
-                        # All-gathered params are not needed after being copied to param.data.
-                        # Zero out the param buffer (shared with grad buffer) for gradient
-                        # accumulation. We cannot zero out the entire grad buffer because one grad
-                        # buffer may correspond to multiple param buffers. If we zero out the entire
-                        # grad buffer, it would clear the data of those param buffers that have not
-                        # yet completed AG.
-                        bucket.param_data.zero_()
-                else:
-                    fp8_params = []
-                    for bucket in bucket_group.buckets:
-                        for param in bucket.params:
-                            if is_float8tensor(param):
-                                fp8_params.append(param)
-                    if len(fp8_params) > 0:
-                        post_all_gather_processing(fp8_params)
+            self._start_bucket_group_param_sync(bucket_group, force_sync=force_sync)
 
     def start_grad_sync(self, *unused):
         """

megatron/core/distributed/param_and_grad_buffer.py

@@ -811,15 +811,22 @@ def group_params_for_buffers(
             param_dtype = torch.uint8
         grad_dtype = torch.float if grad_reduce_in_fp32 else param.dtype
         is_expert_parallel = not getattr(param, 'allreduce', True)
+        is_managed_by_layer_wise_optimizer = getattr(
+            param, 'is_managed_by_layer_wise_optimizer', False
+        )
 
-        key = BufferKey(param_dtype, grad_dtype, is_expert_parallel)
+        key = BufferKey(
+            param_dtype, grad_dtype, is_expert_parallel, is_managed_by_layer_wise_optimizer
+        )
         param_list = key_to_params.get(key, [])
         param_list.append(param)
         key_to_params[key] = param_list
 
         # Use param.dtype (not param_dtype) so FP8/NVFP4 params share offsets with their
         # logical high-precision dtype, needed for checkpoint compatibility.
-        offset_key = BufferKey(param.dtype, grad_dtype, is_expert_parallel)
+        offset_key = BufferKey(
+            param.dtype, grad_dtype, is_expert_parallel, is_managed_by_layer_wise_optimizer
+        )
         offset = dtype_to_offsets.get(offset_key, 0)
         dtype_to_offsets[offset_key] = offset + 1
         indices = key_to_indices.get(key, [])
@@ -1498,12 +1505,16 @@ def partition_buckets(
     if len(buffers) == 0:
         return []
 
-    dtype_to_buffer_map = {}
+    # At most one fp8 (uint8) buffer is allowed; Cases 2 and 3 below branch on
+    # whether one is present. Non-uint8 dtypes can legitimately appear in
+    # multiple buffers (e.g. LayerWise-managed bf16 weights + Adam-managed bf16
+    # biases share the bf16 ``param_dtype`` but live in separate buffers), so
+    # the uniqueness check is restricted to uint8.
+    fp8_buffer = None
     for buffer in buffers:
-        dtype = buffer.param_dtype
-        # Make sure that the param_dtype of any two buffers is different.
-        assert dtype not in dtype_to_buffer_map
-        dtype_to_buffer_map[dtype] = buffer
+        if buffer.param_dtype == torch.uint8:
+            assert fp8_buffer is None
+            fp8_buffer = buffer
 
     # Case 1: Put all buckets into a single bucket group if force_single_bucket_group is True.
     if force_single_bucket_group:
@@ -1522,7 +1533,7 @@ def partition_buckets(
         )
         return [bucket_group]
 
-    if torch.uint8 not in dtype_to_buffer_map:
+    if fp8_buffer is None:
         # Case 2: When there is no fp8 buffer in the input buffers, let each bucket group have
         #         only one bucket.
         bucket_groups = []
@@ -1546,7 +1557,6 @@ def partition_buckets(
                     non_fp8_buckets.append(bucket)
 
         bucket_groups = []
-        fp8_buffer = dtype_to_buffer_map[torch.uint8]
         for bucket in fp8_buffer.buckets:
             if len(bucket_groups) == len(fp8_buffer.buckets) - 1:
                 # reduce_scatter_with_fp32_accumulation requires exactly one bucket