Perf: skip per-param copy_ dispatch in the MXFP8 param copy-back

megatron/core/distributed/param_and_grad_buffer.py

@@ -28,7 +28,7 @@
     modify_nvfp4_rowwise_storage,
 )
 from ..fp8_utils import (
-    copy_tensor_to_quantized_param,
+    copy_tensors_to_quantized_params,
     is_float8tensor,
     is_grouped_mxfp8tensor,
     is_grouped_tensor,
@@ -308,6 +308,9 @@ def _post_param_sync(self):
                     # buffer to copy back from.
                     continue
                 has_non_quantized_weight = False
+                quantized_params = []
+                param_slices = []
+                flat_param_data = bucket.param_data.view(-1)
                 for param in bucket.params:
                     # Non-quantized weights are already mapped to param.data. Skip
                     # mixed buckets because zeroing bucket.param_data would also
@@ -316,8 +319,11 @@ def _post_param_sync(self):
                         has_non_quantized_weight = True
                         break
                     param_start, param_end = bucket.param_to_index[param]
-                    param_slice = bucket.param_data.view(-1)[param_start:param_end]
-                    copy_tensor_to_quantized_param(param, param_slice)
+                    quantized_params.append(param)
+                    param_slices.append(flat_param_data[param_start:param_end])
+                # Cast the bucket in one call: these casts are small, so the per-param cost of
+                # issuing them is worth avoiding.
+                copy_tensors_to_quantized_params(quantized_params, param_slices)
                 if has_non_quantized_weight:
                     continue
                 # All-gathered params are not needed after being copied to param.data.

megatron/core/fp8_utils.py

@@ -226,6 +226,46 @@ def copy_tensor_to_quantized_param(param: torch.Tensor, src: torch.Tensor) -> No
     dst.copy_(src.view(dst.shape))
 
 
+def copy_tensors_to_quantized_params(params: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
+    """List form of :func:`copy_tensor_to_quantized_param`, for a whole bucket of params.
+
+    Same values, minus the per-param ``copy_`` and tensor-subclass dispatch: the quantizer is
+    resolved up front and called directly. Cast kernels are unchanged, one per param. Worth it
+    because those casts are small and issuing them is expensive, and under
+    --reuse-grad-buf-for-mxfp8-param-ag they run inside the forward pass.
+
+    Args:
+        params: quantized model params to write into.
+        srcs: high-precision source values, one per param, in the same order.
+    """
+    if len(params) == 0:
+        return
+
+    srcs_to_cast = []
+    dsts_to_cast = []
+    quantizers = []
+    for param, src in zip(params, srcs):
+        dst = _unwrap_parameter_data(param)
+        quantizer = (
+            None
+            if is_grouped_tensor_with_quantized_storage(dst)
+            else getattr(dst, "_quantizer", None)
+        )
+        if quantizer is None:
+            # Grouped storage quantizes per member; a missing quantizer has to be built. Both
+            # cases are handled by the single-param path.
+            copy_tensor_to_quantized_param(param, src)
+            continue
+        srcs_to_cast.append(src.view(dst.shape))
+        dsts_to_cast.append(dst)
+        quantizers.append(quantizer)
+
+    # Equivalent to dst.copy_(src), but entered directly instead of via the aten::copy_ op,
+    # QuantizedTensor.__torch_dispatch__ (type and usage checks) and dst.quantize_(src).
+    for src, quantizer, dst in zip(srcs_to_cast, quantizers, dsts_to_cast):
+        quantizer.update_quantized(src, dst)
+
+
 def modify_grouped_tensor_rowwise_storage(tensor: torch.Tensor, new_storage: torch.Tensor) -> None:
     """Replace a high-precision Transformer Engine GroupedTensor's rowwise storage."""
     tensor = _unwrap_parameter_data(tensor)

tests/unit_tests/test_fp8_utils.py

@@ -7,8 +7,21 @@
 import torch.nn as nn
 
 from megatron.core import fp8_utils
+from megatron.training.utils import get_device_arch_version
 from tests.unit_tests.test_utilities import Utils
 
+try:
+    import transformer_engine_torch as tex
+    from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer
+
+    HAVE_MXFP8_TENSOR = True
+except ImportError:
+    HAVE_MXFP8_TENSOR = False
+
+# MXFP8 needs Blackwell or newer.
+mxfp8_available = HAVE_MXFP8_TENSOR and get_device_arch_version() >= 10
+reason_for_no_mxfp8 = "MXFP8 requires Transformer Engine and device arch >= 10"
+
 
 class MockTELinear(nn.Module):
     """Mock TE Linear module for testing."""
@@ -130,3 +143,66 @@ def track_forward(x):
 
             # Verify output has original shape
             assert output.shape == (6, 2, 4096)  # Back to original seq_len
+
+
+@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+class TestCopyTensorsToQuantizedParams:
+    """Cover the batched MXFP8 param copy-back used by _post_param_sync.
+
+    ``copy_tensors_to_quantized_params`` bypasses ``copy_`` and calls the destination quantizer
+    directly, so the contract to protect is that it still writes exactly what the per-param
+    ``copy_tensor_to_quantized_param`` would have written.
+    """
+
+    SHAPES = [(1024, 512), (2048, 256), (512, 1024)]
+
+    def _make_param(self, shape):
+        quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True)
+        tensor = quantizer.make_empty(shape, dtype=torch.bfloat16, device="cuda")
+        return torch.nn.Parameter(tensor, requires_grad=False)
+
+    def _raw_buffers(self, param):
+        """The four buffers MXFP8 storage is made of, i.e. everything a cast writes."""
+        data = param.data
+        return (
+            data._rowwise_data,
+            data._rowwise_scale_inv,
+            data._columnwise_data,
+            data._columnwise_scale_inv,
+        )
+
+    def test_matches_per_param_copy(self):
+        """Batched copy-back is bitwise identical to copying one param at a time."""
+        torch.manual_seed(0)
+        reference_params = [self._make_param(shape) for shape in self.SHAPES]
+        batched_params = [self._make_param(shape) for shape in self.SHAPES]
+        # Sources are flat slices, matching how _post_param_sync views the param buffer.
+        srcs = [
+            torch.randn(shape, dtype=torch.bfloat16, device="cuda").view(-1)
+            for shape in self.SHAPES
+        ]
+
+        for param, src in zip(reference_params, srcs):
+            fp8_utils.copy_tensor_to_quantized_param(param, src)
+        fp8_utils.copy_tensors_to_quantized_params(batched_params, srcs)
+        torch.cuda.synchronize()
+
+        for reference, batched in zip(reference_params, batched_params):
+            for expected, actual in zip(self._raw_buffers(reference), self._raw_buffers(batched)):
+                assert torch.equal(expected, actual)
+
+    def test_falls_back_without_quantizer(self):
+        """A destination with no quantizer of its own still gets written."""
+        param = self._make_param(self.SHAPES[0])
+        param.data._quantizer = None
+        src = torch.randn(self.SHAPES[0], dtype=torch.bfloat16, device="cuda").view(-1)
+
+        fp8_utils.copy_tensors_to_quantized_params([param], [src])
+        torch.cuda.synchronize()
+
+        # A quantized copy of a non-zero source cannot be all zeros.
+        assert param.data._rowwise_data.any()
+
+    def test_empty_input(self):
+        """No params is a no-op rather than an error."""
+        fp8_utils.copy_tensors_to_quantized_params([], [])