diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index ab047dc32f..7bf9b3699d 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -9,6 +9,7 @@ import subprocess import sys import pathlib +from unittest.mock import patch sys.path.append(str(pathlib.Path(__file__).resolve().parent.parent)) from utils import run_distributed @@ -958,6 +959,86 @@ def _test_mxfp8_empty_master_shard(dp_group): ) +def _test_rowwise_master_cast(group, 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) + + for step in range(2): + # Also exercise empty ranks in the existing two-rank CI launch. + if step == 1 and world > 1: + lo, hi = (0, full.numel()) if rank == 0 else (0, 0) + 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 + with patch("torch.distributed.all_reduce", side_effect=record_reduce): + 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 + + def run_parallel_tests() -> None: """Run parallel tests""" @@ -1000,6 +1081,9 @@ def run_parallel_tests() -> None: if is_mxfp8_available(): print("starting mxfp8 empty master shard test") _test_mxfp8_empty_master_shard(dp_group) + for mixed in (False, True): + for fragments in (False, True): + _test_rowwise_master_cast(dp_group, mixed, fragments) nvfp4_available, _ = is_nvfp4_available(return_reason=True) if nvfp4_available: print("starting cast master weights to nvfp4 test") diff --git a/tests/pytorch/test_partial_cast.py b/tests/pytorch/test_partial_cast.py index bbb18503b1..b7f683b222 100644 --- a/tests/pytorch/test_partial_cast.py +++ b/tests/pytorch/test_partial_cast.py @@ -128,6 +128,7 @@ def _pad(x, padding): def test_mxfp8_scaling_partial_cast(): torch.cuda.manual_seed(1234) + run_one_case(0, 64, 64, 0) run_one_case(3, 32, 64, 31) run_one_case(64 * 64 - 2, 64, 64, 1) run_one_case(16384 * 6144, 16384, 6144, 0) @@ -135,3 +136,39 @@ def test_mxfp8_scaling_partial_cast(): run_one_case(131072, 768, 256, 0) run_one_case(65536, 768, 256, 131072) run_one_case(98304, 128, 768, 0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("offset,length", [(0, 0), (31, 3), (1, 4094)]) +@pytest.mark.parametrize("with_columnwise_scales", [False, True]) +def test_rowwise_partial_kernels(offset, length, with_columnwise_scales): + 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) + # The PyTorch extension must normalize empty views with backing storage. + 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") + columnwise_scales = ( + torch.full((4, 128), 126, dtype=torch.uint8, device="cuda") + if with_columnwise_scales + else scales.flatten()[:0].view(0, 0) + ) + out = torch.empty(length, dtype=torch.uint8, device="cuda") + tex.mxfp8_scaling_partial_cast(inp, out, out[:0], scales, columnwise_scales, 64, 64, offset) + torch.testing.assert_close(out, inp.to(torch.float8_e4m3fn).view(torch.uint8), rtol=0, atol=0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_mxfp8_partial_cast_requires_columnwise_scales(): + inp = torch.randn(3, device="cuda", dtype=torch.bfloat16) + row = torch.empty(3, device="cuda", dtype=torch.uint8) + col = torch.empty_like(row) + scales = torch.full((128, 4), 127, device="cuda", dtype=torch.uint8) + with pytest.raises(RuntimeError, match="Columnwise output requires columnwise scales"): + tex.mxfp8_scaling_partial_cast( + inp, row, col, scales, scales.flatten()[:0].view(0, 0), 64, 64, 31 + ) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index b3feeba6b7..98c5eed098 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. + * An empty NVTETensor 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,9 @@ 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 an empty NVTETensor for + * output_colwise; scale_inv_colwise is ignored in this case. Column-wise output + * requires column-wise scales. 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..b25e3d4b00 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -23,7 +23,7 @@ constexpr int kColsPerTile = 128; // Columns each block processes constexpr int kThreadsPerBlock = 128; -template +template __global__ void __launch_bounds__(kThreadsPerBlock) mxfp8_scaling_compute_partial_amax_kernel(const IType *input, IType *amax_rowwise, IType *amax_colwise, int amax_rowwise_stride, @@ -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 constexpr (kWithColumnwise) 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 constexpr (kWithColumnwise) { + amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); + } __syncthreads(); @@ -74,7 +76,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) amax_rowwise[r * amax_rowwise_stride + c] = static_cast(smem_amax_rowwise[r_][c_]); } -template +template __global__ void __launch_bounds__(kThreadsPerBlock) mxfp8_scaling_partial_cast_kernel(const IType *input, OType *output_rowwise, OType *output_colwise, const e8m0_t *scale_inv_rowwise, @@ -95,7 +97,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) } // Load scales_colwise - { + if constexpr (kWithColumnwise) { int c_ = threadIdx.x; int r = blockIdx.y * kRowsPerTile / 32; int c = blockIdx.x * kColsPerTile + c_; @@ -108,7 +110,7 @@ __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 = kWithColumnwise ? 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 +123,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 constexpr (kWithColumnwise) { + OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); + output_colwise_minus_offset[idx] = out_colwise; + } } r++; @@ -149,14 +153,17 @@ 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"); + const bool with_columnwise = amax_colwise.has_data(); + 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; @@ -164,11 +171,15 @@ void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, - 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]);) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_columnwise, kWithColumnwise, + mxfp8_scaling_compute_partial_amax_kernel + <<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(amax_rowwise.data.dptr), + 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 +193,23 @@ 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 = output_colwise.has_data(); + if (with_columnwise) { + NVTE_CHECK(scale_inv_colwise.has_data(), "Columnwise output requires columnwise scales"); + 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 output must have a data pointer"); + 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 +220,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; @@ -217,14 +237,19 @@ void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tenso TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( input.dtype(), IType, - mxfp8_scaling_partial_cast_kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(output_rowwise.data.dptr), - reinterpret_cast(output_colwise.data.dptr), - 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]);) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_columnwise, kWithColumnwise, + mxfp8_scaling_partial_cast_kernel + <<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output_rowwise.data.dptr), + with_columnwise ? reinterpret_cast(output_colwise.data.dptr) : nullptr, + reinterpret_cast(scale_inv_rowwise.data.dptr), + 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/csrc/extensions/fp8_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp index d6693a485e..dc02e67e5f 100644 --- a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp @@ -57,7 +57,9 @@ void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax const TensorWrapper input_cu = makeTransformerEngineTensor(input); TensorWrapper amax_rowwise_cu = makeTransformerEngineTensor(amax_rowwise); - TensorWrapper amax_colwise_cu = makeTransformerEngineTensor(amax_colwise); + // Empty PyTorch views may have a backing pointer. Normalize omission at this boundary. + TensorWrapper amax_colwise_cu = + amax_colwise.numel() == 0 ? TensorWrapper() : makeTransformerEngineTensor(amax_colwise); nvte_mxfp8_scaling_compute_partial_amax(input_cu.data(), amax_rowwise_cu.data(), amax_colwise_cu.data(), rows, cols, start_offset, @@ -72,13 +74,19 @@ void mxfp8_scaling_partial_cast(const at::Tensor &input, at::Tensor output_rowwi TORCH_CHECK(output_rowwise.is_contiguous(), "output_rowwise must be contiguous"); TORCH_CHECK(output_colwise.is_contiguous(), "output_colwise must be contiguous"); TORCH_CHECK(scale_inv_rowwise.is_contiguous(), "scale_inv_rowwise must be contiguous"); - TORCH_CHECK(scale_inv_colwise.is_contiguous(), "scale_inv_colwise must be contiguous"); const TensorWrapper input_cu = makeTransformerEngineTensor(input); TensorWrapper output_rowwise_cu = makeTransformerEngineTensor(output_rowwise); - TensorWrapper output_colwise_cu = makeTransformerEngineTensor(output_colwise); + const bool with_columnwise = output_colwise.numel() != 0; + if (with_columnwise) { + TORCH_CHECK(scale_inv_colwise.numel() != 0, "Columnwise output requires columnwise scales"); + TORCH_CHECK(scale_inv_colwise.is_contiguous(), "scale_inv_colwise must be contiguous"); + } + TensorWrapper output_colwise_cu = + with_columnwise ? makeTransformerEngineTensor(output_colwise) : TensorWrapper(); const TensorWrapper scale_inv_rowwise_cu = makeTransformerEngineTensor(scale_inv_rowwise); - const TensorWrapper scale_inv_colwise_cu = makeTransformerEngineTensor(scale_inv_colwise); + const TensorWrapper scale_inv_colwise_cu = + with_columnwise ? makeTransformerEngineTensor(scale_inv_colwise) : TensorWrapper(); nvte_mxfp8_scaling_partial_cast(input_cu.data(), output_rowwise_cu.data(), output_colwise_cu.data(), scale_inv_rowwise_cu.data(), diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index ccad043a65..caaf9cf5aa 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1029,6 +1029,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 @@ -1058,7 +1060,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}" @@ -1087,7 +1095,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] @@ -1097,7 +1109,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: @@ -1123,8 +1140,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()], ], ) @@ -1145,10 +1162,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, "