From 0c1a98fc85e3ff064e97c57b35c5031724e3cacb Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Fri, 14 Mar 2025 21:14:21 +0000 Subject: [PATCH 1/7] Coalesce NCCL all-gathers for MXFP8 all-gather Signed-off-by: Tim Moon --- transformer_engine/pytorch/distributed.py | 161 ++++++++++-------- .../pytorch/module/layernorm_linear.py | 2 +- .../pytorch/module/layernorm_mlp.py | 2 +- transformer_engine/pytorch/module/linear.py | 2 +- 4 files changed, 95 insertions(+), 72 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 2a614f67d7..c1431c2ab8 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -863,10 +863,14 @@ def _all_gather_fp8( # we cannot directly gather the transposed fp8 tensor # so we need to disable columnwise usage for the quantizer # and then set it back to the original value after quantizing + init_rowwise_usage = quantizer.rowwise_usage init_columnwise_usage = quantizer.columnwise_usage - quantizer.set_usage(columnwise=False) + quantizer.set_usage(rowwise=True, columnwise=False) inp = quantizer(inp) - quantizer.set_usage(columnwise=init_columnwise_usage) + quantizer.set_usage( + rowwise=init_rowwise_usage, + columnwise=init_columnwise_usage, + ) # Construct output tensor out: Float8TensorBase @@ -923,9 +927,34 @@ def _all_gather_mxfp8( ) -> tuple[MXFP8TensorBase, Optional[torch.distributed.Work]]: """All-gather MXFP8 tensor along first dimension.""" - # Tensor dims + # Input tensor attributes + in_shape: Iterable[int] + device: torch.device + dtype: torch.dtype + if isinstance(inp, torch.Tensor): + in_shape = inp.size() + device = inp.device + dtype = inp.dtype + elif isinstance(inp, MXFP8TensorBase): + if inp._rowwise_data is not None: + in_shape = inp._rowwise_data.device.size() + device = inp._rowwise_data.device + dtype = inp._rowwise_data.dtype + elif inp._columnwise_data is not None: + in_shape = inp._columnwise_data.device.size() + device = inp._columnwise_data.device + dtype = inp._columnwise_data.dtype + else: + raise ValueError("Got MXFP8 input tensor without any data") + dtype = torch.bfloat16 + else: + raise ValueError( + "Invalid type for input tensor (expected torch.Tensor or MXFP8TensorBase, " + f"found {inp.__class__.__name__})" + ) + + # Output tensor shape world_size = get_distributed_world_size(process_group) - in_shape = list(inp.size()) if out_shape is None: out_shape = [in_shape[0] * world_size] + in_shape[1:] @@ -938,25 +967,20 @@ def _all_gather_mxfp8( ): out = torch.empty( out_shape, - dtype=inp.dtype, - device=inp.device, + dtype=dtype, + device=device, memory_format=torch.contiguous_format, ) torch.distributed.all_gather_into_tensor(out, inp, group=process_group) out = quantizer(out) return out, None - inp_dtype = inp.dtype - inp_device = inp.device - # Cast input tensor to MXFP8 with required data if not isinstance(inp, MXFP8TensorBase): inp = quantizer(inp) elif ( - inp.rowwise_data is None - and quantizer.rowwise_usage - or inp.columnwise_data is None - and quantizer.columnwise_usage + (quantizer.rowwise_usage and inp._rowwise_data is None) + or (quantizer.columnwise_usage and inp._columnwise_data is None) ): warnings.warn( "Input and quantizer do not have matching usages. " @@ -965,65 +989,64 @@ def _all_gather_mxfp8( inp = quantizer(inp.dequantize()) # Construct MXFP8 output tensor - out = quantizer.make_empty(out_shape, dtype=inp_dtype, device=inp_device) - - # Async op handle - handle = None - - # Gather MXFP8 data for row-wise usage - if quantizer.rowwise_usage: - - # Remove padding from MXFP8 scale-inverses - in_scale_inv = inp._rowwise_scale_inv - out_scale_inv = out._rowwise_scale_inv - flattened_in_shape0 = math.prod(in_shape[:-1]) - if in_scale_inv.size(0) != flattened_in_shape0: - in_scale_inv = in_scale_inv[:flattened_in_shape0] - out_scale_inv[flattened_in_shape0 * world_size :].zero_() - out_scale_inv = out_scale_inv[: flattened_in_shape0 * world_size] - - # Launch all-gathers - if handle is not None: - handle.wait() - torch.distributed.all_gather_into_tensor( - out_scale_inv, - in_scale_inv, - group=process_group, - ) - handle = torch.distributed.all_gather_into_tensor( - out._rowwise_data, - inp._rowwise_data, - group=process_group, - async_op=async_op, - ) - - # Gather MXFP8 data for column-wise usage - if quantizer.columnwise_usage: + out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - # Remove padding from MXFP8 scale-inverses - in_scale_inv = inp._columnwise_scale_inv - out_scale_inv = out._columnwise_scale_inv - flattened_in_shape0 = math.prod(in_shape[:-1]) // 32 - if in_scale_inv.size(0) != flattened_in_shape0: - in_scale_inv = in_scale_inv[:flattened_in_shape0] - out_scale_inv[flattened_in_shape0 * world_size :].zero_() - out_scale_inv = out_scale_inv[: flattened_in_shape0 * world_size] + # Coalesce NCCL collectives + with torch.distributed._coalescing_manager( + group=process_group, + device=device, + async_ops=async_op, + ) as coalescing_manager: + + # Gather MXFP8 data for row-wise usage + if quantizer.rowwise_usage: + + # Remove padding from MXFP8 scale-inverses + in_scale_inv = inp._rowwise_scale_inv + out_scale_inv = out._rowwise_scale_inv + flattened_in_shape0 = math.prod(in_shape[:-1]) + if in_scale_inv.size(0) != flattened_in_shape0: + in_scale_inv = in_scale_inv[:flattened_in_shape0] + out_scale_inv[flattened_in_shape0 * world_size :].zero_() + out_scale_inv = out_scale_inv[: flattened_in_shape0 * world_size] + + # Launch all-gathers + torch.distributed.all_gather_into_tensor( + out_scale_inv, + in_scale_inv, + group=process_group, + ) + torch.distributed.all_gather_into_tensor( + out._rowwise_data, + inp._rowwise_data, + group=process_group, + ) - # Launch all-gathers - if handle is not None: - handle.wait() - torch.distributed.all_gather_into_tensor( - out_scale_inv, - in_scale_inv, - group=process_group, - ) - handle = torch.distributed.all_gather_into_tensor( - out._columnwise_data, - inp._columnwise_data, - group=process_group, - async_op=async_op, - ) + # Gather MXFP8 data for column-wise usage + if quantizer.columnwise_usage: + + # Remove padding from MXFP8 scale-inverses + in_scale_inv = inp._columnwise_scale_inv + out_scale_inv = out._columnwise_scale_inv + flattened_in_shape0 = math.prod(in_shape[:-1]) // 32 + if in_scale_inv.size(0) != flattened_in_shape0: + in_scale_inv = in_scale_inv[:flattened_in_shape0] + out_scale_inv[flattened_in_shape0 * world_size :].zero_() + out_scale_inv = out_scale_inv[: flattened_in_shape0 * world_size] + + # Launch all-gathers + torch.distributed.all_gather_into_tensor( + out_scale_inv, + in_scale_inv, + group=process_group, + ) + torch.distributed.all_gather_into_tensor( + out._columnwise_data, + inp._columnwise_data, + group=process_group, + ) + handle = coalescing_manager if async_op else None return out, handle diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 1b62f8d777..f8f12348b0 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -574,7 +574,7 @@ def backward( quantizer = None if ctx.fp8: quantizer = ctx.input_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) + quantizer.set_usage(rowwise=False, columnwise=True) nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") ln_out_total, ln_out_total_work = gather_along_first_dim( ln_out, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 1f167b5a7e..b01b2bc6a7 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -687,7 +687,7 @@ def backward( quantizer = None if ctx.fp8: quantizer = ctx.fc1_input_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) + quantizer.set_usage(rowwise=False, columnwise=True) ln_out_total, ln_out_total_work = gather_along_first_dim( ln_out, ctx.tp_group, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 4c87396e3c..7dc76032d7 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -466,7 +466,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], quantizer = None if ctx.fp8: quantizer = ctx.input_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) + quantizer.set_usage(rowwise=False, columnwise=True) nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") inputmat_total, inputmat_total_work = gather_along_first_dim( inputmat, From 79e0dad0d53eb00d6e0b1bb23299b4a426746947 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Sat, 15 Mar 2025 02:35:15 +0000 Subject: [PATCH 2/7] Add missing import Signed-off-by: Tim Moon --- transformer_engine/pytorch/distributed.py | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index c1431c2ab8..832b55b598 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -5,6 +5,7 @@ """Methods needed for distributed training (DP/TP).""" from __future__ import annotations +from collections.abc import Iterable from contextlib import contextmanager, AbstractContextManager, ContextDecorator from functools import lru_cache import math From 32763b1ebcb3c7417fa057d45221e14330b59938 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 15 Mar 2025 02:36:11 +0000 Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/distributed.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 832b55b598..a9b29541fa 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -979,9 +979,8 @@ def _all_gather_mxfp8( # Cast input tensor to MXFP8 with required data if not isinstance(inp, MXFP8TensorBase): inp = quantizer(inp) - elif ( - (quantizer.rowwise_usage and inp._rowwise_data is None) - or (quantizer.columnwise_usage and inp._columnwise_data is None) + elif (quantizer.rowwise_usage and inp._rowwise_data is None) or ( + quantizer.columnwise_usage and inp._columnwise_data is None ): warnings.warn( "Input and quantizer do not have matching usages. " From de2f7de669b138a51cbd9cf48bc153947ac18b66 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Sat, 22 Mar 2025 02:44:03 +0000 Subject: [PATCH 4/7] Cache quantized input tensor after linear module forward pass Signed-off-by: Tim Moon --- .../pytorch/module/layernorm_linear.py | 111 +++++++-------- .../pytorch/module/layernorm_mlp.py | 127 ++++++++---------- transformer_engine/pytorch/module/linear.py | 14 +- 3 files changed, 113 insertions(+), 139 deletions(-) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 9771231002..1a6dd8b9cf 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -146,6 +146,7 @@ def forward( backward_needs_input = is_grad_enabled and weight_requires_grad with_input_all_gather = parallel_mode == "column" and sequence_parallel + # Check if Userbuffers is supported if fp8: if any([ub_overlap_ag_fprop, ub_overlap_rs_fprop]) and not ( FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling() @@ -155,40 +156,27 @@ def forward( " current scaling" ) + # Configure quantizer for norm output + if fp8: if input_quantizer is None: raise ValueError("Missing quantizer for input tensor") + columnwise_usage = backward_needs_input + if ( + columnwise_usage + and with_input_all_gather + and not isinstance(input_quantizer, MXFP8Quantizer) + ): + columnwise_usage = False + input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - # Configure quantizer for normalization output - with_quantized_norm = fp8 and not return_layernorm_output - if with_quantized_norm: - if with_input_all_gather: - input_quantizer.set_usage(rowwise=True, columnwise=False) - if isinstance(input_quantizer, MXFP8Quantizer): - with_quantized_norm = False - else: - input_quantizer.set_usage( - rowwise=True, - columnwise=backward_needs_input, - ) - - # Reduce duplicated transpose in `_fix_gathered_fp8_transpose` - if ( + # Construct norm output + with_quantized_norm = ( fp8 - and FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling() - and ub_bulk_dgrad - ): - input_quantizer.set_usage(rowwise=True, columnwise=False) - - ub_obj_fprop = None + and not return_layernorm_output + and not return_layernorm_output_gathered + ) ln_out = None - # For DelayScaling, output of normalization will be in fp8. - # For Float8CurrentScaling, we want the output of normalization in high precision, then quantize to fp8. - if ub_overlap_ag_fprop and not isinstance(input_quantizer, Float8CurrentScalingQuantizer): - ub_obj_fprop = get_ub(ub_name + "_fprop") - ln_out = ub_obj_fprop.get_buffer(input_quantizer, local_chunk=True) - elif with_quantized_norm: - if with_input_all_gather: - input_quantizer.set_usage(rowwise=True, columnwise=False) + if with_quantized_norm: ln_out = input_quantizer.make_empty(inputmat.shape, dtype=inputmat.dtype, device="cuda") else: ln_out = torch.empty_like( @@ -212,47 +200,42 @@ def forward( ln_out_return = ln_out if return_layernorm_output else None nvtx_range_pop(f"{nvtx_label}.norm") - # For Float8CurrentScalingQuantizer, layernorm/rmsnorm has not been fused with quantizer. - # So the output of normalization is in high precision, and we need to quantize it to FP8 and put in the buffer. - if ub_overlap_ag_fprop and isinstance(input_quantizer, Float8CurrentScalingQuantizer): - ub_obj_fprop = get_ub(ub_name + "_fprop") - ln_out_local = ln_out - ln_out = ub_obj_fprop.get_buffer(input_quantizer, local_chunk=True) - input_quantizer.quantize(ln_out_local, out=ln_out) - # Prepare GEMM input # Note: Cast to expected dtype and perform tensor-parallel communication nvtx_range_push(f"{nvtx_label}.gemm_input_cast_comm") - if with_input_all_gather and not ub_overlap_ag_fprop: - with_quantized_all_gather = fp8 - if return_layernorm_output and return_layernorm_output_gathered: - with_quantized_all_gather = False - if fp8: - input_quantizer.set_usage(rowwise=True, columnwise=False) - # ln_out in this has two possibilities: - # 1. in FP8 low precision, the cast was done by fusing quantization into layernorm kernel - # 2. in high precision, then we need to cast it and then gather in FP8 - # the output ln_out_total will be in FP8, and it's a full tensor - ln_out_total, _ = gather_along_first_dim( - ln_out, - tp_group, - quantizer=(input_quantizer if with_quantized_all_gather else None), - ) - if return_layernorm_output and return_layernorm_output_gathered: + ln_out_total = None + ub_obj_fprop = None + if with_input_all_gather: + if return_layernorm_output_gathered: + # Perform all-gather in high precision if gathered + # norm output will be returned + ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) ln_out_return = ln_out_total - if fp8 and not with_quantized_all_gather: - ln_out_total = input_quantizer(ln_out_total) - else: - if ub_overlap_ag_fprop: - ln_out_total = ub_obj_fprop.get_buffer(input_quantizer) + if fp8: + ln_out = input_quantizer(ln_out) + input_quantizer.set_usage(rowwise=True, columnwise=False) + ln_out_total = input_quantizer(ln_out_total) else: if fp8: - if not isinstance(ln_out, QuantizedTensor): - input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) + if not with_quantized_norm: ln_out = input_quantizer(ln_out) - elif backward_needs_input: - ln_out.update_usage(rowwise_usage=True, columnwise_usage=True) - ln_out_total = ln_out + input_quantizer.set_usage(rowwise=True, columnwise=False) + if ub_overlap_ag_fprop: + # Copy into Userbuffers buffer + ub_obj_fprop = get_ub(ub_name + "_fprop") + ub_obj_fprop.get_buffer(input_quantizer, local_chunk=True).copy_(ln_out) + ln_out_total = ub_obj_fprop.get_buffer(input_quantizer) + else: + # All-gather with NCCL + ln_out_total, _ = gather_along_first_dim( + ln_out, + tp_group, + quantizer=(input_quantizer if fp8 else None), + ) + else: + if fp8 and not with_quantized_norm: + ln_out = input_quantizer(ln_out) + ln_out_total = ln_out nvtx_range_pop(f"{nvtx_label}.gemm_input_cast_comm") # Cast weight to expected dtype @@ -394,7 +377,7 @@ def forward( weight, bias, ln_weight, - ln_out.clone() if ub_overlap_ag_fprop else ln_out, # avoid saving a UB buffer + ln_out, mu, rsigma, ) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 354179f081..e9b09a0f46 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -214,43 +214,38 @@ def forward( with_quantized_norm = fp8 and not return_layernorm_output tp_world_size = get_distributed_world_size(tp_group) - ub_overlap_ag = ub_overlap_ag and is_grad_enabled and not return_layernorm_output + ub_overlap_ag = ub_overlap_ag and is_grad_enabled and not return_layernorm_output_gathered ub_overlap_rs = ub_overlap_rs and is_grad_enabled with_input_all_gather_nccl = sequence_parallel and not ub_overlap_ag backwards_needs_fc1_input = is_grad_enabled and fc1_weight.requires_grad - # Configure quantizer for normalization output - if fp8 and fc1_input_quantizer is None: - raise ValueError("Missing quantizer for input tensor") - if with_quantized_norm: - if with_input_all_gather_nccl: - fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) - if isinstance(fc1_input_quantizer, MXFP8Quantizer): - with_quantized_norm = False - else: - fc1_input_quantizer.set_usage( - rowwise=True, - columnwise=backwards_needs_fc1_input, - ) - - # Reduce duplicated transpose in `_fix_gathered_fp8_transpose` - if ( - fp8 - and FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling() - and ub_bulk_dgrad - ): - fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + # Configure quantizer for norm output + if fp8: + if fc1_input_quantizer is None: + raise ValueError("Missing quantizer for FC1 input tensor") + columnwise_usage = backwards_needs_fc1_input + if ( + columnwise_usage + and sequence_parallel + and not isinstance(fc1_input_quantizer, MXFP8Quantizer) + ): + columnwise_usage = False + fc1_input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - ub_obj_lnout = None + # Construct norm output ln_out = None - # For DelayScaling, output of normalization will be in fp8. - # For Float8CurrentScaling, we want the output of normalization in high precision, then quantize to fp8. - if ub_overlap_ag and not isinstance(fc1_input_quantizer, Float8CurrentScalingQuantizer): - ub_obj_lnout = get_ub("fc1_fprop") - ln_out = ub_obj_lnout.get_buffer(fc1_input_quantizer, local_chunk=True) - elif not with_quantized_norm: + if with_quantized_norm: + ln_out = fc1_input_quantizer.make_empty( + inputmat.shape, + dtype=inputmat.dtype, + device="cuda", + ) + else: ln_out = torch.empty_like( - inputmat, dtype=inputmat.dtype, memory_format=torch.contiguous_format, device="cuda" + inputmat, + dtype=inputmat.dtype, + memory_format=torch.contiguous_format, + device="cuda", ) # Apply normalization @@ -266,53 +261,43 @@ def forward( fwd_ln_sm_margin, zero_centered_gamma, ) - ln_out_return = ln_out if return_layernorm_output else None - # For Float8CurrentScalingQuantizer, layernorm/rmsnorm has not been fused with quantizer. - # So the output of normalization is in high precision, and we need to quantize it to FP8 and put in the buffer. - if ub_overlap_ag and isinstance(fc1_input_quantizer, Float8CurrentScalingQuantizer): - ub_obj_lnout = get_ub("fc1_fprop") - ln_out_local = ln_out - ln_out = ub_obj_lnout.get_buffer(fc1_input_quantizer, local_chunk=True) - fc1_input_quantizer.quantize(ln_out_local, out=ln_out) - # Prepare GEMM input # Note: Cast to expected dtype and perform tensor-parallel communication - ln_out_gathered = False - with_quantized_all_gather = fp8 - if with_input_all_gather_nccl: - if return_layernorm_output and return_layernorm_output_gathered: - with_quantized_all_gather = False - if fp8: - fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) - # ln_out in this has two possibilities: - # 1. in FP8 low precision, the cast was done by fusing quantization into layernorm kernel - # 2. in high precision, then we need to cast it and then gather in FP8 - # the output ln_out_total will be in FP8, and it's a full tensor - ln_out_total, _ = gather_along_first_dim( - ln_out, - tp_group, - quantizer=(fc1_input_quantizer if with_quantized_all_gather else None), - ) - ln_out_gathered = True - else: - with_quantized_all_gather = False - if ub_overlap_ag: - ln_out_total = ub_obj_lnout.get_buffer(fc1_input_quantizer, False) + ln_out_total = None + ub_obj_lnout = None + if sequence_parallel: + if return_layernorm_output_gathered: + # Perform all-gather in high precision if gathered + # norm output will be returned + ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) + ln_out_return = ln_out_total + if fp8: + ln_out = fc1_input_quantizer(ln_out) + fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + ln_out_total = fc1_input_quantizer(ln_out_total) else: if fp8: - if not isinstance(ln_out, QuantizedTensor): - fc1_input_quantizer.set_usage( - rowwise=True, columnwise=backwards_needs_fc1_input - ) + if not with_quantized_norm: ln_out = fc1_input_quantizer(ln_out) - elif backwards_needs_fc1_input: - ln_out.update_usage(rowwise_usage=True, columnwise_usage=True) - # here ln_out is in FP8 low precision, the cast was either done by fc1_input_quantizer - # or fused into the layernorm kernel - # ln_out_total represents the full fp8 tensor, in this case, it's the same as ln_out - ln_out_total = ln_out + fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) + if ub_overlap_ag: + # Copy into Userbuffers buffer + ub_obj_lnout = get_ub("fc1_fprop") + ub_obj_lnout.get_buffer(fc1_input_quantizer, local_chunk=True).copy_(ln_out) + ln_out_total = ub_obj_lnout.get_buffer(fc1_input_quantizer) + else: + # All-gather with NCCL + ln_out_total, _ = gather_along_first_dim( + ln_out, + tp_group, + quantizer=(fc1_input_quantizer if fp8 else None), + ) + else: + if fp8 and not with_quantized_norm: + ln_out = fc1_input_quantizer(ln_out) + ln_out_total = ln_out # Cast weights to expected dtype fc1_weight_final = fc1_weight @@ -493,7 +478,7 @@ def forward( tensors_to_save, tensor_objects = prepare_for_saving( inputmat, ln_weight, - ln_out.clone() if ub_overlap_ag else ln_out, # avoid saving a UB buffer + ln_out, fc1_weight_final, fc1_bias, fc1_out, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 314fc38372..e2d4ff5ff5 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -56,6 +56,7 @@ prepare_for_saving, restore_from_saved, ) +from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor._internal.mxfp8_tensor_base import MXFP8TensorBase from ..cpu_offload import is_cpu_offload_enabled, set_offloading_param @@ -140,9 +141,14 @@ def forward( if input_quantizer is None: raise ValueError("Missing quantizer for input tensor") if with_input_all_gather_nccl: - assert not isinstance( - inputmat, QuantizedTensor - ), "All gather of fp8 input is not supported" + if not isinstance(inputmat, QuantizedTensor): + columnwise_usage = ( + backward_needs_input + and isinstance(input_quantizer, MXFP8Quantizer) + ) + input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + inputmat = input_quantizer(inputmat) + own_quantized_input = True input_quantizer.set_usage(rowwise=True, columnwise=False) inputmat_total, _ = gather_along_first_dim( inputmat, @@ -271,7 +277,7 @@ def forward( # to gather the input. For MXFP8, columnwise only data # can be allgathered. if isinstance(inputmat, MXFP8TensorBase) or not ctx.backward_input_needs_gather: - inputmat.update_usage(rowwise_usage=False) + inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) saved_inputmat = inputmat if cpu_offloading: From 38a4c09936064551fd47981fe7f221e5f8eb9ab4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 22 Mar 2025 02:48:54 +0000 Subject: [PATCH 5/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/module/layernorm_linear.py | 4 +--- transformer_engine/pytorch/module/linear.py | 5 ++--- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 1a6dd8b9cf..a842390328 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -171,9 +171,7 @@ def forward( # Construct norm output with_quantized_norm = ( - fp8 - and not return_layernorm_output - and not return_layernorm_output_gathered + fp8 and not return_layernorm_output and not return_layernorm_output_gathered ) ln_out = None if with_quantized_norm: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index e2d4ff5ff5..0eb95f9f4b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -142,9 +142,8 @@ def forward( raise ValueError("Missing quantizer for input tensor") if with_input_all_gather_nccl: if not isinstance(inputmat, QuantizedTensor): - columnwise_usage = ( - backward_needs_input - and isinstance(input_quantizer, MXFP8Quantizer) + columnwise_usage = backward_needs_input and isinstance( + input_quantizer, MXFP8Quantizer ) input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) inputmat = input_quantizer(inputmat) From 9acd9cf890abe05998ce052f6c3db32f104c7cd1 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Sat, 22 Mar 2025 02:58:43 +0000 Subject: [PATCH 6/7] Fix linter warnings Signed-off-by: Tim Moon --- transformer_engine/pytorch/module/layernorm_linear.py | 5 +++-- transformer_engine/pytorch/module/layernorm_mlp.py | 8 ++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a842390328..3fb94a117e 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -55,7 +55,6 @@ prepare_for_saving, restore_from_saved, ) -from ..tensor.float8_tensor import Float8CurrentScalingQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor._internal.mxfp8_tensor_base import MXFP8TensorBase from ..cpu_offload import is_cpu_offload_enabled, set_offloading_param @@ -195,7 +194,9 @@ def forward( fwd_ln_sm_margin, zero_centered_gamma, ) - ln_out_return = ln_out if return_layernorm_output else None + ln_out_return = None + if return_layernorm_output or return_layernorm_output_gathered: + ln_out_return = ln_out nvtx_range_pop(f"{nvtx_label}.norm") # Prepare GEMM input diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index e9b09a0f46..e7b1dbbaec 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -60,7 +60,6 @@ from ..tensor.mxfp8_tensor import MXFP8Quantizer from ._common import apply_normalization, _fix_gathered_fp8_transpose from ..cpu_offload import is_cpu_offload_enabled, set_offloading_param -from ..tensor.float8_tensor import Float8CurrentScalingQuantizer from ..tensor.quantized_tensor import ( QuantizedTensor, Quantizer, @@ -216,7 +215,6 @@ def forward( tp_world_size = get_distributed_world_size(tp_group) ub_overlap_ag = ub_overlap_ag and is_grad_enabled and not return_layernorm_output_gathered ub_overlap_rs = ub_overlap_rs and is_grad_enabled - with_input_all_gather_nccl = sequence_parallel and not ub_overlap_ag backwards_needs_fc1_input = is_grad_enabled and fc1_weight.requires_grad # Configure quantizer for norm output @@ -261,7 +259,9 @@ def forward( fwd_ln_sm_margin, zero_centered_gamma, ) - ln_out_return = ln_out if return_layernorm_output else None + ln_out_return = None + if return_layernorm_output or return_layernorm_output_gathered: + ln_out_return = ln_out # Prepare GEMM input # Note: Cast to expected dtype and perform tensor-parallel communication @@ -525,7 +525,7 @@ def forward( ctx.bias_gelu_fusion = bias_gelu_fusion ctx.return_layernorm_output = return_layernorm_output ctx.return_layernorm_output_gathered = ( - return_layernorm_output_gathered and ln_out_gathered + return_layernorm_output_gathered and sequence_parallel ) ctx.set_parallel_mode = set_parallel_mode ctx.bwd_ln_sm_margin = bwd_ln_sm_margin From 9e3f39213582f1d9a94ab70991ec4ab22f2dd1c4 Mon Sep 17 00:00:00 2001 From: Tim Moon Date: Tue, 25 Mar 2025 00:54:46 +0000 Subject: [PATCH 7/7] Avoid unnecessarily allocating layernorm output in LayerNormLinear/LayerNormMLP Signed-off-by: Tim Moon --- .../pytorch/module/layernorm_linear.py | 21 +++++--------- .../pytorch/module/layernorm_mlp.py | 29 ++++--------------- 2 files changed, 13 insertions(+), 37 deletions(-) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 3fb94a117e..a22745f17d 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -136,6 +136,11 @@ def forward( ln_bias = cast_if_needed(ln_bias, activation_dtype) nvtx_range_pop(f"{nvtx_label}.norm_input_cast") + # Avoid quantized norm kernel if norm output will be returned + with_quantized_norm = ( + fp8 and not return_layernorm_output and not return_layernorm_output_gathered + ) + tp_world_size = get_distributed_world_size(tp_group) ub_overlap_ag_fprop = ( ub_overlap_ag_fprop and is_grad_enabled and not return_layernorm_output @@ -168,28 +173,16 @@ def forward( columnwise_usage = False input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - # Construct norm output - with_quantized_norm = ( - fp8 and not return_layernorm_output and not return_layernorm_output_gathered - ) - ln_out = None - if with_quantized_norm: - ln_out = input_quantizer.make_empty(inputmat.shape, dtype=inputmat.dtype, device="cuda") - else: - ln_out = torch.empty_like( - inputmat, dtype=inputmat.dtype, memory_format=torch.contiguous_format, device="cuda" - ) - # Apply normalization nvtx_range_push(f"{nvtx_label}.norm") ln_out, mu, rsigma = apply_normalization( inputmat, - ln_out, + None, # ln_out ln_weight, ln_bias, eps, input_quantizer if with_quantized_norm else None, - inp.dtype, + inputmat.dtype, normalization, fwd_ln_sm_margin, zero_centered_gamma, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index e7b1dbbaec..c040b4d42a 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -206,11 +206,10 @@ def forward( if ln_bias is not None: ln_bias = cast_if_needed(ln_bias, activation_dtype) - # for fp8 DelayedScaling: layernorm output = FP8 - # only output of the linear is returned - # for return_layernorm_output: layernorm output = High precision, then cast to FP8 - # high precision layernorm output and output of the linear are returned - with_quantized_norm = fp8 and not return_layernorm_output + # Avoid quantized norm kernel if norm output will be returned + with_quantized_norm = ( + fp8 and not return_layernorm_output and not return_layernorm_output_gathered + ) tp_world_size = get_distributed_world_size(tp_group) ub_overlap_ag = ub_overlap_ag and is_grad_enabled and not return_layernorm_output_gathered @@ -230,31 +229,15 @@ def forward( columnwise_usage = False fc1_input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - # Construct norm output - ln_out = None - if with_quantized_norm: - ln_out = fc1_input_quantizer.make_empty( - inputmat.shape, - dtype=inputmat.dtype, - device="cuda", - ) - else: - ln_out = torch.empty_like( - inputmat, - dtype=inputmat.dtype, - memory_format=torch.contiguous_format, - device="cuda", - ) - # Apply normalization ln_out, mu, rsigma = apply_normalization( inputmat, - ln_out, + None, # ln_out ln_weight, ln_bias, eps, fc1_input_quantizer if with_quantized_norm else None, - inp.dtype, + inputmat.dtype, normalization, fwd_ln_sm_margin, zero_centered_gamma,