diff --git a/tests/pytorch/mxfp8/test_rowwise_master_cast.py b/tests/pytorch/mxfp8/test_rowwise_master_cast.py new file mode 100644 index 0000000000..a876a8fb23 --- /dev/null +++ b/tests/pytorch/mxfp8/test_rowwise_master_cast.py @@ -0,0 +1,132 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# See LICENSE for license information. + +"""Run with pytest on one GPU or torchrun -m pytest on multiple GPUs.""" + +import os + +import pytest +import torch +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch.tensor.utils import ( + cast_master_weights_to_fp8, + quantize_master_weights, +) + +available, reason = te.is_mxfp8_available(return_reason=True) +pytestmark = pytest.mark.skipif(not available, reason=reason) + + +@pytest.fixture(scope="module") +def group(): + owned = not torch.distributed.is_initialized() + if owned: + torch.cuda.set_device(int(os.getenv("LOCAL_RANK", "0"))) + if "RANK" in os.environ: + torch.distributed.init_process_group("nccl") + else: + torch.distributed.init_process_group( + "nccl", store=torch.distributed.HashStore(), rank=0, world_size=1 + ) + yield torch.distributed.group.WORLD + if owned: + torch.distributed.destroy_process_group() + + +@pytest.mark.parametrize("mixed", [False, True]) +@pytest.mark.parametrize("fragments", [False, True]) +def test_rowwise_master_cast(group, monkeypatch, mixed, fragments): + rank = torch.distributed.get_rank(group) + world = torch.distributed.get_world_size(group) + shape = (160, 96) # Both scale padding and non-tile-aligned logical shapes. + full = torch.linspace(-4, 7, 160 * 96, device="cuda").reshape(shape) + quantizers = [te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=False)] + if mixed: + quantizers.append(te.MXFP8Quantizer(te.DType.kFloat8E4M3, rowwise=True, columnwise=True)) + weights = [q(full.to(torch.bfloat16)) for q in quantizers] + pointers = [(w._rowwise_data.data_ptr(), w._rowwise_scale_inv.data_ptr()) for w in weights] + if world == 1: + lo, hi = 0, full.numel() + elif rank == 0: + lo, hi = 0, 17 # Split inside a 32-value block. + elif rank == 1: + lo, hi = 17, full.numel() + else: + lo, hi = 0, 0 # No master shard on tail ranks. + real_reduce = torch.distributed.all_reduce + reduced_sizes = [] + + def record_reduce(tensor, *args, **kwargs): + reduced_sizes.append(tensor.numel()) + return real_reduce(tensor, *args, **kwargs) + + monkeypatch.setattr(torch.distributed, "all_reduce", record_reduce) + for step in range(2): + master = full + step * 0.25 + shard = master.flatten()[lo:hi] if lo < hi else None + outputs = [ + ( + torch.empty(hi - lo, dtype=torch.uint8, device="cuda"), + torch.empty(hi - lo, dtype=torch.uint8, device="cuda") if i else None, + ) + for i in range(len(weights)) + ] + reduced_sizes.clear() + cast = quantize_master_weights if step == 0 else cast_master_weights_to_fp8 + cast( + weights, + [shard] * len(weights), + [lo if shard is not None else None] * len(weights), + group, + fsdp_shard_model_weights=outputs if fragments else None, + ) + expected_amax = sum(w._rowwise_scale_inv.numel() for w in weights) + if mixed: + expected_amax += weights[1]._columnwise_scale_inv.numel() + assert reduced_sizes == [expected_amax] + for idx, (weight, quantizer) in enumerate(zip(weights, quantizers)): + expected = quantizer(master.to(torch.bfloat16)) + for direction in ["rowwise", "columnwise"] if idx else ["rowwise"]: + data = getattr(weight, f"_{direction}_data") + actual = torch.zeros_like(data).flatten() + if lo < hi: + src = ( + outputs[idx][direction == "columnwise"] + if fragments + else data.flatten()[lo:hi] + ) + actual[lo:hi].copy_(src) + real_reduce(actual, op=torch.distributed.ReduceOp.MAX, group=group) + torch.testing.assert_close( + actual.view(shape), getattr(expected, f"_{direction}_data"), rtol=0, atol=0 + ) + scale = getattr(weight, f"_{direction}_scale_inv") + ref = getattr(expected, f"_{direction}_scale_inv") + rows, cols = (160, 3) if direction == "rowwise" else (5, 96) + torch.testing.assert_close(scale[:rows, :cols], ref[:rows, :cols], rtol=0, atol=0) + assert pointers[idx] == ( + weight._rowwise_data.data_ptr(), + weight._rowwise_scale_inv.data_ptr(), + ) + assert weights[0]._columnwise_data is None + assert weights[0]._columnwise_scale_inv is None + + +@pytest.mark.parametrize("offset,length", [(0, 0), (31, 3), (1, 4094)]) +def test_rowwise_partial_kernels(offset, length): + inp = torch.randn(length, device="cuda", dtype=torch.bfloat16) + row = torch.zeros((128, 4), device="cuda", dtype=inp.dtype) + ref = torch.zeros_like(row) + col = torch.zeros((4, 128), device="cuda", dtype=inp.dtype) + # Empty view may still have backing storage. Omission is shape-based. + omitted = col.flatten()[:0].view(0, 0) + tex.mxfp8_scaling_compute_partial_amax(inp, row, omitted, 64, 64, offset) + tex.mxfp8_scaling_compute_partial_amax(inp, ref, col, 64, 64, offset) + torch.testing.assert_close(row, ref, rtol=0, atol=0) + scales = torch.full((128, 4), 127, dtype=torch.uint8, device="cuda") + out = torch.empty(length, dtype=torch.uint8, device="cuda") + tex.mxfp8_scaling_partial_cast( + inp, out, out[:0], scales, scales.flatten()[:0].view(0, 0), 64, 64, offset + ) + torch.testing.assert_close(out, inp.to(torch.float8_e4m3fn).view(torch.uint8), rtol=0, atol=0) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 47539a89a1..1ce8923a96 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -290,6 +290,7 @@ void nvte_fp8_block_scaling_partial_cast(const NVTETensor inp, NVTETensor out, * \param[in] input Input tensor (continuous segment of flattened original tensor). * \param[in,out] amax_rowwise Output tensor for row-wise maximum absolute values. * \param[in,out] amax_colwise Output tensor for column-wise maximum absolute values. + * Shape [0, 0] skips column-wise computation. * \param[in] rows Number of rows in the logical tensor. * \param[in] cols Number of columns in the logical tensor. * \param[in] start_offset Starting offset in the flattened tensor. @@ -303,7 +304,8 @@ void nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor * * This function casts the input tensor to MXFP8 format, producing both row-wise and * column-wise scaled outputs. input contains a continuous segment from the flattened - * original tensor. + * original tensor. To skip column-wise computation, pass output_colwise with shape [0] + * and scale_inv_colwise with shape [0, 0]. Row-wise outputs remain required. * * \param[in] input Input (continuous segment of flattened original tensor). * \param[out] output_rowwise Output tensor with row-wise scaling (MXFP8 format). diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu index be692d4563..5d40d5475e 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -47,7 +47,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { float abs_input = fabs(static_cast(input_minus_offset[idx])); row_amax = fmaxf(row_amax, abs_input); - col_amax = fmaxf(col_amax, abs_input); + if (amax_colwise != nullptr) col_amax = fmaxf(col_amax, abs_input); } #pragma unroll @@ -63,7 +63,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock) r++; } - amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); + if (amax_colwise != nullptr) { + amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); + } __syncthreads(); @@ -95,7 +97,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) } // Load scales_colwise - { + if (output_colwise != nullptr) { int c_ = threadIdx.x; int r = blockIdx.y * kRowsPerTile / 32; int c = blockIdx.x * kColsPerTile + c_; @@ -108,7 +110,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) size_t end_offset = start_offset + len; const IType *input_minus_offset = input - start_offset; OType *output_rowwise_minus_offset = output_rowwise - start_offset; - OType *output_colwise_minus_offset = output_colwise - start_offset; + OType *output_colwise_minus_offset = + output_colwise != nullptr ? output_colwise - start_offset : nullptr; int warp_idx = threadIdx.x / 32; // int lane_idx = threadIdx.x % 32; int c = blockIdx.x * kColsPerTile + threadIdx.x; @@ -121,9 +124,11 @@ __global__ void __launch_bounds__(kThreadsPerBlock) if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { float inp = static_cast(input_minus_offset[idx]); OType out_rowwise = static_cast(inp * smem_scales_rowwise[i][warp_idx]); - OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); output_rowwise_minus_offset[idx] = out_rowwise; - output_colwise_minus_offset[idx] = out_colwise; + if (output_colwise != nullptr) { + OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); + output_colwise_minus_offset[idx] = out_colwise; + } } r++; @@ -149,14 +154,20 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, NVTE_CHECK(amax_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); NVTE_CHECK(amax_rowwise.dtype() == input.dtype(), "Wrong dtype of amax_rowwise"); - NVTE_CHECK(amax_colwise.data.shape.size() == 2, "amax_colwise must be a 2D tensor"); - NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0, - "Wrong padding of amax_colwise's rows"); - NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, "Invalid rows"); - NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0, - "Wrong padding of amax_colwise's cols"); - NVTE_CHECK(amax_colwise.data.shape[1] >= cols, "Invalid cols"); - NVTE_CHECK(amax_colwise.dtype() == input.dtype(), "Wrong dtype of amax_colwise"); + // Empty views can have a non-null data pointer: use shape to detect omission. + const bool with_columnwise = + !(amax_colwise.data.shape.size() == 2 && amax_colwise.data.shape[0] == 0 && + amax_colwise.data.shape[1] == 0); + if (with_columnwise) { + NVTE_CHECK(amax_colwise.data.shape.size() == 2, "amax_colwise must be a 2D tensor"); + NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0, + "Wrong padding of amax_colwise's rows"); + NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, "Invalid rows"); + NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0, + "Wrong padding of amax_colwise's cols"); + NVTE_CHECK(amax_colwise.data.shape[1] >= cols, "Invalid cols"); + NVTE_CHECK(amax_colwise.dtype() == input.dtype(), "Wrong dtype of amax_colwise"); + } int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile; int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile; @@ -167,8 +178,9 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, mxfp8_scaling_compute_partial_amax_kernel<<>>( reinterpret_cast(input.data.dptr), reinterpret_cast(amax_rowwise.data.dptr), - reinterpret_cast(amax_colwise.data.dptr), amax_rowwise.data.shape[1], - amax_colwise.data.shape[1], rows, cols, start_offset, input.data.shape[0]);) + with_columnwise ? reinterpret_cast(amax_colwise.data.dptr) : nullptr, + amax_rowwise.data.shape[1], with_columnwise ? amax_colwise.data.shape[1] : 0, rows, cols, + start_offset, input.data.shape[0]);) } void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tensor output_colwise, @@ -182,16 +194,27 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso "Invalid start_offset"); NVTE_CHECK(output_rowwise.data.shape.size() == 1, "output_rowwise must be a 1D tensor"); - NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); NVTE_CHECK(output_rowwise.data.shape[0] == input.data.shape[0], "Size of input and output_rowwise mismatch"); - NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0], - "Size of input and output_colwise mismatch"); NVTE_CHECK(output_rowwise.dtype() == DType::kFloat8E4M3 || output_rowwise.dtype() == DType::kByte, "output_rowwise should be e4m3 or uint8"); - NVTE_CHECK(output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte, - "output_colwise should be e4m3 or uint8"); + const bool with_columnwise = + !(scale_inv_colwise.data.shape.size() == 2 && scale_inv_colwise.data.shape[0] == 0 && + scale_inv_colwise.data.shape[1] == 0); + NVTE_CHECK(with_columnwise || + (output_colwise.data.shape.size() == 1 && output_colwise.data.shape[0] == 0), + "Omitted columnwise scales require an empty columnwise output"); + if (with_columnwise) { + NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); + NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0], + "Size of input and output_colwise mismatch"); + NVTE_CHECK(input.data.shape[0] == 0 || output_colwise.data.dptr != nullptr, + "Columnwise scales require columnwise output"); + NVTE_CHECK( + output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte, + "output_colwise should be e4m3 or uint8"); + } NVTE_CHECK(scale_inv_rowwise.data.shape.size() == 2, "scale_inv_rowwise must be a 2D tensor"); NVTE_CHECK(scale_inv_rowwise.data.shape[0] % rowwise_row_padding == 0, @@ -202,14 +225,16 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso NVTE_CHECK(scale_inv_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); NVTE_CHECK(scale_inv_rowwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_rowwise"); - NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, "scale_inv_colwise must be a 2D tensor"); - NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0, - "Wrong padding of scale_inv_colwise's rows"); - NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, "Invalid rows"); - NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0, - "Wrong padding of scale_inv_colwise's cols"); - NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, "Invalid cols"); - NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_colwise"); + if (with_columnwise) { + NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, "scale_inv_colwise must be a 2D tensor"); + NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0, + "Wrong padding of scale_inv_colwise's rows"); + NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, "Invalid rows"); + NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0, + "Wrong padding of scale_inv_colwise's cols"); + NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, "Invalid cols"); + NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_colwise"); + } int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile; int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile; @@ -220,11 +245,11 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso mxfp8_scaling_partial_cast_kernel<<>>( reinterpret_cast(input.data.dptr), reinterpret_cast(output_rowwise.data.dptr), - reinterpret_cast(output_colwise.data.dptr), + with_columnwise ? reinterpret_cast(output_colwise.data.dptr) : nullptr, reinterpret_cast(scale_inv_rowwise.data.dptr), - reinterpret_cast(scale_inv_colwise.data.dptr), - scale_inv_rowwise.data.shape[1], scale_inv_colwise.data.shape[1], rows, cols, - start_offset, input.data.shape[0]);) + with_columnwise ? reinterpret_cast(scale_inv_colwise.data.dptr) : nullptr, + scale_inv_rowwise.data.shape[1], with_columnwise ? scale_inv_colwise.data.shape[1] : 0, + rows, cols, start_offset, input.data.shape[0]);) } } // namespace mxfp8_scaling_recipe diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index cef45c0223..6da03679a1 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1017,6 +1017,8 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( ): # pylint: disable=unused-argument r"""Helper function to cast master weights to FP8 primary weights for mxfp8 scaling. + Row-only primaries omit columnwise amax, scale updates, and output writes. + Parameters ---------- params : List of tuple, each tuple contains a model weight, a master weight, and an offset @@ -1046,7 +1048,13 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( raise ValueError( f"rowwise_scale_inv must be 2D, got {len(rowwise_shape)}D shape {rowwise_shape}" ) - colwise_shape = model_weight._columnwise_scale_inv.shape + if (model_weight._columnwise_data is None) != (model_weight._columnwise_scale_inv is None): + raise ValueError("MXFP8 columnwise data and scales must both be present or absent") + colwise_shape = ( + model_weight._columnwise_scale_inv.shape + if model_weight._columnwise_scale_inv is not None + else (0, 0) + ) if len(colwise_shape) != 2: raise ValueError( f"columnwise_scale_inv must be 2D, got {len(colwise_shape)}D shape {colwise_shape}" @@ -1075,7 +1083,11 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( amaxes_colwise, scale_invs_colwise = [], [] for i, (model_weight, master_weight, start_offset, _) in enumerate(params): rowwise_shape = model_weight._rowwise_scale_inv.shape - colwise_shape = model_weight._columnwise_scale_inv.shape + colwise_shape = ( + model_weight._columnwise_scale_inv.shape + if model_weight._columnwise_scale_inv is not None + else (0, 0) + ) rowwise_start = cu_rowwise_amax_sizes[i] rowwise_end = cu_rowwise_amax_sizes[i + 1] colwise_start = cu_rowwise_amax_sizes[-1] + cu_colwise_amax_sizes[i] @@ -1085,7 +1097,12 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( amaxes_rowwise.append(amax_rowwise) amaxes_colwise.append(amax_colwise) scale_invs_rowwise.append(model_weight._rowwise_scale_inv) - scale_invs_colwise.append(model_weight._columnwise_scale_inv) + # Empty tensors mark an omitted direction without allocating a GPU payload. + scale_invs_colwise.append( + model_weight._columnwise_scale_inv + if model_weight._columnwise_scale_inv is not None + else torch.empty((0, 0), dtype=torch.uint8, device=device) + ) # Compute amax of the master weight and store it in packed_amaxes. if master_weight is not None: @@ -1111,8 +1128,8 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( multi_tensor_compute_scale_inv_e8m0, None, # dummy_overflow_buf [ - amaxes_rowwise + amaxes_colwise, - scale_invs_rowwise + scale_invs_colwise, + amaxes_rowwise + [amax for amax in amaxes_colwise if amax.numel()], + scale_invs_rowwise + [scale for scale in scale_invs_colwise if scale.numel()], ], ) @@ -1133,10 +1150,18 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( end_offset = start_offset + master_weight.numel() if use_fsdp_shard_model_weights: rowwise_fragment = model_weight_fragment[0] - colwise_fragment = model_weight_fragment[1] + colwise_fragment = ( + model_weight_fragment[1] + if scale_inv_colwise.numel() + else torch.empty(0, dtype=torch.uint8, device=device) + ) else: rowwise_fragment = model_weight._rowwise_data.reshape(-1)[start_offset:end_offset] - colwise_fragment = model_weight._columnwise_data.reshape(-1)[start_offset:end_offset] + colwise_fragment = ( + model_weight._columnwise_data.reshape(-1)[start_offset:end_offset] + if scale_inv_colwise.numel() + else torch.empty(0, dtype=torch.uint8, device=device) + ) if len(model_weight.shape) != 2: raise ValueError( "model_weight must be 2D for MXFP8 scaling partial cast, "