diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 15247c2be53..406449e256c 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/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. diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 5411b676d83..895d46e9b3d 100644 --- a/megatron/core/fp8_utils.py +++ b/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) diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index 5be17f03c9f..dc65d541455 100644 --- a/tests/unit_tests/test_fp8_utils.py +++ b/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([], [])