From a53f850ca8945b9759264d2385c4fb83a39d9ba9 Mon Sep 17 00:00:00 2001 From: Layali Rashid Date: Tue, 13 Oct 2026 07:10:04 -0700 Subject: [PATCH 1/6] Fuse projection boundary quantization and gradient scaling Signed-off-by: Layali Rashid --- tests/pytorch/test_mxfp8_segments.py | 139 +++++++++++ .../test_nonmaterialized_norm_gradient.py | 58 +++++ tests/pytorch/test_weight_gradient_scale.py | 103 ++++++++ transformer_engine/common/CMakeLists.txt | 1 + .../common/cast/quantize_mxfp8_segments.cu | 219 ++++++++++++++++++ .../common/include/transformer_engine/cast.h | 14 ++ transformer_engine/pytorch/csrc/extensions.h | 3 + .../pytorch/csrc/extensions/cast.cpp | 39 ++++ .../pytorch/csrc/extensions/pybind.cpp | 2 + .../pytorch/distributed_weight.py | 45 ++++ .../pytorch/module/layernorm_linear.py | 33 ++- transformer_engine/pytorch/module/linear.py | 8 + .../pytorch/tensor/mxfp8_tensor.py | 54 +++++ 13 files changed, 717 insertions(+), 1 deletion(-) create mode 100644 tests/pytorch/test_mxfp8_segments.py create mode 100644 tests/pytorch/test_nonmaterialized_norm_gradient.py create mode 100644 tests/pytorch/test_weight_gradient_scale.py create mode 100644 transformer_engine/common/cast/quantize_mxfp8_segments.cu diff --git a/tests/pytorch/test_mxfp8_segments.py b/tests/pytorch/test_mxfp8_segments.py new file mode 100644 index 00000000000..e74353163a7 --- /dev/null +++ b/tests/pytorch/test_mxfp8_segments.py @@ -0,0 +1,139 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Exact MXFP8 payload parity for segmented projection gradients.""" + +import math + +import pytest +import torch + +from transformer_engine.pytorch import DType, MXFP8Quantizer + + +def assert_payload_equal(actual, expected): + rows = math.prod(expected.shape[:-1]) + columns = expected.shape[-1] + for name in ( + "_rowwise_data", + "_rowwise_scale_inv", + "_columnwise_data", + "_columnwise_scale_inv", + ): + before, after = getattr(expected, name), getattr(actual, name) + assert (before is None) == (after is None) + if before is None: + continue + assert before.shape == after.shape + before, after = before.view(torch.uint8), after.view(torch.uint8) + if "scale_inv" in name: + columnwise = name.startswith("_columnwise") + outer, inner = (columns, rows // 32) if columnwise else (rows, columns // 32) + if expected._with_gemm_swizzled_scales: + i = torch.arange(outer, device=before.device).view(-1, 1) + j = torch.arange(inner, device=before.device).view(1, -1) + index = ((i // 128) * math.ceil(inner / 4) + j // 4) * 512 + index = index + (i % 32) * 16 + ((i % 128) // 32) * 4 + j % 4 + before, after = before.flatten()[index], after.flatten()[index] + elif columnwise: + before, after = before[:inner, :outer], after[:inner, :outer] + else: + before, after = before[:outer, :inner], after[:outer, :inner] + if not torch.equal(before, after): + mismatch = (before != after).nonzero()[:8] + indexes = tuple(mismatch[:, i] for i in range(mismatch.shape[1])) + pytest.fail( + f"{name}: positions={mismatch.tolist()}, expected={before[indexes].tolist()}, actual={after[indexes].tolist()}" + ) + + +@pytest.fixture(autouse=True) +def require_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("MXFP8 requires Blackwell") + + +@pytest.mark.parametrize("usage", [(True, False), (False, True), (True, True)]) +@pytest.mark.parametrize("swizzle", [False, True]) +@pytest.mark.parametrize("layout", ["contiguous", "strided", "singleton", "fallback", "broadcast"]) +@pytest.mark.parametrize("pattern", ["random", "signed_zero", "extreme", "all_nan", "all_inf", "all_zero"]) +def test_segmented_payload(usage, swizzle, layout, pattern): + torch.manual_seed(1234) + rows = 128 if layout != "fallback" else 96 + widths = (256, 128, 128) + source = torch.randn(rows, sum(widths), device="cuda", dtype=torch.bfloat16) + if pattern == "signed_zero": + source[:, ::2] = 0.0 + source[:, 1::2] = -0.0 + elif pattern == "extreme": + values = torch.tensor( + [0.0, -0.0, 2**-100, -(2**-100), 2**100, -(2**100), float("inf"), float("nan")], + device="cuda", + dtype=torch.bfloat16, + ) + source.copy_(values.repeat(rows, sum(widths) // values.numel())) + if pattern in ("all_nan", "all_inf", "all_zero"): + source.fill_({"all_nan": float("nan"), "all_inf": float("inf"), "all_zero": 0.0}[pattern]) + parts = torch.split(source, widths, dim=-1) + if layout == "contiguous": + parts = tuple(p.contiguous() for p in parts) + elif layout == "singleton": + parts = tuple(p[:, None, None, :] for p in parts) + elif layout == "broadcast": + parts = tuple(p[:1].expand(rows, -1) for p in parts) + quantizer = MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3, rowwise=usage[0], columnwise=usage[1]) + quantizer.optimize_for_gemm = swizzle + reference = quantizer(torch.cat(parts, dim=-1)) + actual = quantizer.quantize_segments(parts) + assert_payload_equal(actual, reference) + replay = quantizer.quantize_segments(parts) + assert_payload_equal(replay, actual) + + +@pytest.mark.parametrize("segments", [1, 3]) +def test_segmented_stream_and_graph(segments): + parts = tuple( + torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) for _ in range(segments) + ) + quantizer = MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=True) + quantizer.optimize_for_gemm = True + expected = quantizer(torch.cat(parts, dim=-1)) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + actual = quantizer.quantize_segments(parts) + torch.cuda.current_stream().wait_stream(stream) + assert_payload_equal(actual, expected) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = quantizer.quantize_segments(parts) + graph.replay() + torch.cuda.synchronize() + assert_payload_equal(captured, expected) + + +def test_segmented_rejects_invalid_shapes(): + quantizer = MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=True) + part = torch.empty(128, 128, device="cuda", dtype=torch.bfloat16) + with pytest.raises(ValueError, match="one or three"): + quantizer.quantize_segments((part, part)) + with pytest.raises(ValueError, match="output shape"): + quantizer.quantize_segments((part,), (128, 256)) + with pytest.raises(ValueError, match="row counts"): + quantizer.quantize_segments((part, part[:96], part)) + + +@pytest.mark.parametrize("usage", [(True, False), (False, True)]) +def test_native_segmented_binding_rejects_unqualified_usage(usage): + if torch.cuda.get_device_capability() != (10, 3): + pytest.skip("The direct segmented binding is qualified for SM103") + import transformer_engine_torch as tex + + part = torch.ones(128, 128, device="cuda", dtype=torch.bfloat16) + quantizer = MXFP8Quantizer(DType.kFloat8E4M3, rowwise=usage[0], columnwise=usage[1]) + quantizer.optimize_for_gemm = True + with pytest.raises(RuntimeError, match="both orientations"): + tex.quantize_mxfp8_segments((part,), quantizer, tuple(part.shape)) diff --git a/tests/pytorch/test_nonmaterialized_norm_gradient.py b/tests/pytorch/test_nonmaterialized_norm_gradient.py new file mode 100644 index 00000000000..221f09f6b26 --- /dev/null +++ b/tests/pytorch/test_nonmaterialized_norm_gradient.py @@ -0,0 +1,58 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Gradient parity when the returned normalization output is unused.""" + +import pytest +import torch + +from transformer_engine.common.recipe import MXFP8BlockScaling +from transformer_engine.pytorch import LayerNormLinear, autocast + + +@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) +@pytest.mark.parametrize("consumer", ["linear", "norm", "both"]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_nonmaterialized_norm_gradient(normalization, consumer, fp8): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if fp8 and torch.cuda.get_device_capability() < (10, 0): + pytest.skip("MXFP8 requires Blackwell") + torch.manual_seed(1234) + options = dict( + in_features=128, + out_features=256, + normalization=normalization, + params_dtype=torch.bfloat16, + return_layernorm_output=True, + bias=False, + ) + reference = LayerNormLinear(**options) + fused = LayerNormLinear(**options, materialize_norm_output_grads=False) + fused.load_state_dict(reference.state_dict()) + source = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + linear_upstream = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + norm_upstream = torch.randn_like(source) + results = [] + for module in (reference, fused): + inp = source.detach().clone().requires_grad_(True) + with autocast(enabled=fp8, recipe=MXFP8BlockScaling()): + linear, norm = module(inp) + outputs, upstream = [], [] + if consumer in ("linear", "both"): + outputs.append(linear) + upstream.append(linear_upstream) + if consumer in ("norm", "both"): + outputs.append(norm) + upstream.append(norm_upstream) + gradients = torch.autograd.grad( + outputs, (inp, *module.parameters()), upstream, allow_unused=True + ) + results.append((linear.detach(), norm.detach(), gradients)) + for before, after in zip(results[0][:2], results[1][:2]): + torch.testing.assert_close(before, after, rtol=0, atol=0) + for before, after in zip(results[0][2], results[1][2]): + assert (before is None) == (after is None) + if before is not None: + torch.testing.assert_close(before, after, rtol=0, atol=0) diff --git a/tests/pytorch/test_weight_gradient_scale.py b/tests/pytorch/test_weight_gradient_scale.py new file mode 100644 index 00000000000..803a2edb27e --- /dev/null +++ b/tests/pytorch/test_weight_gradient_scale.py @@ -0,0 +1,103 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Opt-in distributed-weight gradient scaling contract.""" + +from types import SimpleNamespace + +import pytest +import torch +import transformer_engine_torch as tex + +from transformer_engine.pytorch.distributed_weight import ( + get_weight_gradient_scale, + mark_weight_gradient_scale, +) +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + +def scaling_inputs(monkeypatch, scale=1.0 / 64): + """Use a tensor implementer and the real MXFP8 quantizer type.""" + weight = torch.empty(1) + weight.is_distributed_weight = True + weight.fuse_wgrad_mean_scaling = True + weight.get_weight_gradient_scale = lambda: scale + args = SimpleNamespace( + fp8=True, + debug=False, + use_bias=False, + backward_override=None, + activation_dtype=torch.bfloat16, + grad_weight_quantizer=None, + cpu_offloading=False, + ub_bulk_wgrad=False, + wgrad_store=None, + input_quantizer=MXFP8Quantizer(tex.DType.kFloat8E4M3), + ) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return weight, args + + +def test_plain_and_existing_weights_keep_unit_scale(): + weight = torch.empty(1) + assert get_weight_gradient_scale(weight, None) == 1.0 + weight.is_distributed_weight = True + assert get_weight_gradient_scale(weight, None) == 1.0 + weight.fuse_wgrad_mean_scaling = True + assert get_weight_gradient_scale(weight, None) == 1.0 + + +def test_opt_in_scale(monkeypatch): + weight, args = scaling_inputs(monkeypatch) + assert get_weight_gradient_scale(weight, args) == 1.0 / 64 + + +@pytest.mark.parametrize( + "field,value", + [ + ("fp8", False), + ("debug", True), + ("use_bias", True), + ("backward_override", "unfused"), + ("activation_dtype", torch.float16), + ("grad_weight_quantizer", object()), + ("cpu_offloading", True), + ("ub_bulk_wgrad", True), + ("input_quantizer", None), + ("wgrad_store", SimpleNamespace(delay_wgrad_compute=lambda: True)), + ], +) +def test_unsupported_routes_keep_unit_scale(monkeypatch, field, value): + weight, args = scaling_inputs(monkeypatch) + setattr(args, field, value) + assert get_weight_gradient_scale(weight, args) == 1.0 + + +def test_graph_capture_keeps_unit_scale(monkeypatch): + weight, args = scaling_inputs(monkeypatch) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert get_weight_gradient_scale(weight, args) == 1.0 + + +@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf")]) +def test_invalid_provider_scale_rejected(monkeypatch, scale): + weight, args = scaling_inputs(monkeypatch, scale) + with pytest.raises(ValueError, match="finite and positive"): + get_weight_gradient_scale(weight, args) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_marker_preserves_values_and_rejects_duplicate_scaling(dtype): + gradient = torch.arange(8, dtype=dtype) / 64 + expected = gradient.clone() + mark_weight_gradient_scale(gradient, 1.0 / 64) + assert torch.equal(gradient, expected) + assert gradient._te_wgrad_scale == 1.0 / 64 + with pytest.raises(ValueError, match="already carries"): + mark_weight_gradient_scale(gradient, 1.0 / 64) + + +def test_unsupported_gradient_dtype_rejected(): + with pytest.raises(ValueError, match="FP32 or BF16"): + mark_weight_gradient_scale(torch.empty(1, dtype=torch.float16), 1.0 / 64) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 077a9a2b3a1..0243fdc5026 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -359,6 +359,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources activation/swiglu_grouped_bwd.cu activation/swiglu_grouped_dbias.cu cast/cast.cu + cast/quantize_mxfp8_segments.cu cast/cast_dbias.cu cast/cast_grouped.cu cast/cast_grouped_dbias.cu diff --git a/transformer_engine/common/cast/quantize_mxfp8_segments.cu b/transformer_engine/common/cast/quantize_mxfp8_segments.cu new file mode 100644 index 00000000000..37b78cd207b --- /dev/null +++ b/transformer_engine/common/cast/quantize_mxfp8_segments.cu @@ -0,0 +1,219 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include + +#include +#include + +#include "../common.h" +#include "../util/cuda_runtime.h" +#include "../util/ptx.cuh" + +namespace transformer_engine { +namespace segmented_mxfp8 { + +struct Segments { + const __nv_bfloat16 *q, *k, *v; + int dq, dk, dv, qs, ks, vs; +}; + +__device__ __forceinline__ const __nv_bfloat16 *base(Segments s, int col, int &offset, + int &stride) { + if (col < s.dq) { + offset = col; + stride = s.qs; + return s.q; + } + if (col < s.dq + s.dk) { + offset = col - s.dq; + stride = s.ks; + return s.k; + } + offset = col - s.dq - s.dk; + stride = s.vs; + return s.v; +} + +__device__ __forceinline__ unsigned char scale_exponent(float amax) { + float normalized = amax * (1.0f / 448.0f); + unsigned int packed; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + asm volatile( + "{ .reg .b16 pair,zero; mov.b16 zero,0; cvt.rp.satfinite.ue8m0x2.f32 pair,%1,%1; mov.b32 " + "%0,{pair,zero}; }" + : "=r"(packed) + : "f"(normalized)); +#else + packed = 0; +#endif + return packed & 255; +} + +__device__ __forceinline__ float inverse(unsigned char e) { + unsigned int bits = e == 254 ? 0x00400000u : (254u - e) << 23; + if (e == 255) bits = 0x7fffffffu; + return __uint_as_float(bits); +} + +__device__ __forceinline__ int64_t sf(int i, int j, int tiles) { + return (static_cast(i / 128) * tiles + j / 4) * 512 + (i % 32) * 16 + + ((i % 128) / 32) * 4 + j % 4; +} + +__device__ __forceinline__ uint16_t pair_cast(float a, float b) { + uint16_t result; + asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0,%2,%1;" : "=h"(result) : "f"(a), "f"(b)); + return result; +} + +__global__ void row_kernel(Segments s, unsigned char *out, unsigned char *scales, int n, int cols, + bool swizzle) { + int group = blockIdx.x * blockDim.x + threadIdx.x; + if (group >= n * (cols / 32)) return; + int row = group / (cols / 32), col = (group % (cols / 32)) * 32, offset, stride; + const __nv_bfloat16 *p = base(s, col, offset, stride); + p += static_cast(row) * stride + offset; + uint4 input[4]; +#pragma unroll + for (int i = 0; i < 4; i++) input[i] = reinterpret_cast(p)[i]; + const __nv_bfloat16 *values = reinterpret_cast(input); + float x[32], amax = 0.0f; +#pragma unroll + for (int i = 0; i < 32; i++) { + x[i] = __bfloat162float(values[i]); + amax = fmaxf(amax, fabsf(x[i])); + } + unsigned char e = scale_exponent(amax); + alignas(16) unsigned int packed[8]; +#pragma unroll + for (int i = 0; i < 8; i++) { + ptx::floatx4 input_group; +#pragma unroll + for (int j = 0; j < 4; ++j) reinterpret_cast(&input_group)[j] = x[i * 4 + j]; + ptx::floatx2 scale; + reinterpret_cast(&scale)[0] = inverse(e); + reinterpret_cast(&scale)[1] = inverse(e); + ptx::fp8e4m3x4 converted; + ptx::mul_cvt_4x(converted, input_group, scale); + const unsigned int word = reinterpret_cast(converted); + packed[i] = word; + } + auto dst = reinterpret_cast(out + static_cast(row) * cols + col); +#pragma unroll + for (int i = 0; i < 2; i++) + reinterpret_cast(dst)[i] = reinterpret_cast(packed)[i]; + scales[swizzle ? sf(row, col / 32, cols / 128) + : static_cast(row) * (cols / 32) + col / 32] = e; +} + +__global__ void column_kernel(Segments s, unsigned char *out, unsigned char *scales, int n, + int cols, bool swizzle, bool canonical_zero) { + int group = blockIdx.x * blockDim.x + threadIdx.x; + if (group >= (n / 32) * cols) return; + int row = (group / cols) * 32, col = group % cols, offset, stride; + const __nv_bfloat16 *p = base(s, col, offset, stride); + p += static_cast(row) * stride + offset; + float x[32], amax = 0.0f; +#pragma unroll + for (int i = 0; i < 32; i++) { + x[i] = __bfloat162float(p[static_cast(i) * stride]); + amax = i == 0 ? fabsf(x[i]) : fmaxf(amax, fabsf(x[i])); + } + unsigned char e = scale_exponent(amax); + float inv = inverse(e); +#pragma unroll + for (int i = 0; i < 16; i++) { + uint16_t packed = pair_cast(__fmaf_rn(x[i * 2], inv, 0.0f), __fmaf_rn(x[i * 2 + 1], inv, 0.0f)); + unsigned char a = packed & 255, b = (packed >> 8) & 255; + if (canonical_zero && x[i * 2] == 0.0f) a = 0; + if (canonical_zero && x[i * 2 + 1] == 0.0f) b = 0; + out[static_cast(row + i * 2) * cols + col] = a; + out[static_cast(row + i * 2 + 1) * cols + col] = b; + } + scales[swizzle ? sf(col, row / 32, n / 128) : static_cast(row / 32) * cols + col] = e; +} + +} // namespace segmented_mxfp8 +} // namespace transformer_engine + +void nvte_quantize_mxfp8_segments(const NVTETensor *inputs, size_t num_inputs, + const int64_t *row_strides, NVTETensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_mxfp8_segments); + using namespace transformer_engine; + using namespace segmented_mxfp8; + NVTE_CHECK(cuda::sm_arch() == 103, "Segmented MXFP8 is qualified for SM103."); + NVTE_CHECK(num_inputs == 1 || num_inputs == 3, "Expected one or three column segments."); + NVTE_CHECK(inputs != nullptr && row_strides != nullptr, "Missing input segments or strides."); + auto *out = convertNVTETensorCheck(output); + NVTE_CHECK(out->scaling_mode == NVTE_MXFP8_1D_SCALING && out->dtype() == DType::kFloat8E4M3, + "Expected 1D E4M3 MXFP8 output."); + Segments segments{}; + const __nv_bfloat16 **pointers[] = {&segments.q, &segments.k, &segments.v}; + int *widths[] = {&segments.dq, &segments.dk, &segments.dv}; + int *strides[] = {&segments.qs, &segments.ks, &segments.vs}; + int64_t rows = -1, columns = 0; + for (size_t i = 0; i < num_inputs; ++i) { + const auto *input = convertNVTETensorCheck(inputs[i]); + NVTE_CHECK(input->data.dtype == DType::kBFloat16 && input->data.shape.size() == 2, + "Segments must be BF16 matrices."); + const int64_t n = input->data.shape[0], width = input->data.shape[1]; + NVTE_CHECK(n > 0 && n % 128 == 0 && width > 0 && width % 128 == 0, + "Rows and each segment width must be positive multiples of 128."); + NVTE_CHECK(rows == -1 || rows == n, "Segment row counts must match."); + NVTE_CHECK(row_strides[i] >= width && row_strides[i] % 8 == 0 && + row_strides[i] <= std::numeric_limits::max(), + "Invalid row stride."); + NVTE_CHECK(reinterpret_cast(input->data.dptr) % 16 == 0, + "Segment addresses must be 16-byte aligned."); + rows = n; + columns += width; + NVTE_CHECK(width <= std::numeric_limits::max(), "Segment is too wide."); + *pointers[i] = static_cast(input->data.dptr); + *widths[i] = static_cast(width); + *strides[i] = static_cast(row_strides[i]); + } + if (num_inputs == 1) { + segments.k = segments.v = segments.q; + segments.ks = segments.vs = segments.qs; + } + NVTE_CHECK(rows <= std::numeric_limits::max() && + columns <= std::numeric_limits::max() && + rows * columns / 32 <= std::numeric_limits::max() - 128, + "Matrix is too large."); + NVTE_CHECK(out->numel() == static_cast(rows * columns), "Output shape mismatch."); + const bool rw = out->data.dptr != nullptr, cw = out->columnwise_data.dptr != nullptr; + NVTE_CHECK(rw && cw && out->with_gemm_swizzled_scales, + "Segmented MXFP8 requires both orientations and GEMM-swizzled scales."); + const size_t scales = rows * columns / 32; + if (rw) { + NVTE_CHECK(out->scale_inv.dptr != nullptr && out->scale_inv.numel() >= scales, + "Missing rowwise scale storage."); + } + if (cw) { + NVTE_CHECK( + out->columnwise_scale_inv.dptr != nullptr && out->columnwise_scale_inv.numel() >= scales, + "Missing columnwise scale storage."); + } + constexpr int threads = 128; + const int groups = rows * columns / 32; + if (rw) { + row_kernel<<<(groups + threads - 1) / threads, threads, 0, stream>>>( + segments, static_cast(out->data.dptr), + static_cast(out->scale_inv.dptr), rows, columns, + out->with_gemm_swizzled_scales); + } + if (cw) { + column_kernel<<<(groups + threads - 1) / threads, threads, 0, stream>>>( + segments, static_cast(out->columnwise_data.dptr), + static_cast(out->columnwise_scale_inv.dptr), rows, columns, + out->with_gemm_swizzled_scales, rw); + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index ca3eb13ed24..0f3c16d8c23 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -125,6 +125,20 @@ void nvte_quantize_noop(const NVTETensor input, NVTETensor output, NVTETensor no void nvte_quantize_v2(const NVTETensor input, NVTETensor output, const NVTEQuantizationConfig quant_config, cudaStream_t stream); +/*! \brief Quantize one or three BF16 column segments directly into an MXFP8 tensor. + * + * SM103 only. Rows and segment widths must be multiples of 128. Feature rows + * must be contiguous, with 16-byte aligned addresses and token strides. + * \param[in] inputs Segment tensors with logical shape [rows, columns]. + * \param[in] num_inputs One or three segments in their final column order. + * \param[in] row_strides Row strides in BF16 elements. + * \param[in,out] output E4M3 MXFP8 output with both orientations and GEMM-swizzled scales. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_quantize_mxfp8_segments(const NVTETensor *inputs, size_t num_inputs, + const int64_t *row_strides, NVTETensor output, + cudaStream_t stream); + /*! \brief Casts input tensor to MXFP8. Additionally, reduces the input along columns. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 5a4378ea328..215a18fecc0 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -382,6 +382,9 @@ std::vector bulk_allocate(const std::vector> &sh * Quantize **************************************************************************************************/ +py::object quantize_mxfp8_segments(const std::vector &segments, + py::handle quantizer, const std::vector &shape); + py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector &shape, at::ScalarType dtype, at::Device device, bool pin_memory); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 06d03cd6e1c..479b4ffc2cf 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -6,6 +6,8 @@ #include "transformer_engine/cast.h" +#include + #include #include #include @@ -135,6 +137,43 @@ py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantiz return output_py; } +py::object quantize_mxfp8_segments(const std::vector &segments, py::handle quantizer, + const std::vector &shape) { + NVTE_CHECK(segments.size() == 1 || segments.size() == 3, + "Expected one or three MXFP8 column segments."); + NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()), "Expected a native MXFP8 quantizer."); + const auto device = segments.front().device(); + NVTE_CHECK(device.is_cuda(), "Segments must be CUDA tensors."); + c10::cuda::CUDAGuard guard(device); + auto quantizer_cpp = convert_quantizer(quantizer); + auto *mxfp8 = static_cast(quantizer_cpp.get()); + NVTE_CHECK(!mxfp8->with_2d_quantization, "Segmented conversion requires 1D scaling."); + std::vector wrappers; + std::vector inputs; + std::vector strides; + wrappers.reserve(segments.size()); + inputs.reserve(segments.size()); + for (const auto &segment : segments) { + NVTE_CHECK(segment.device() == device && segment.scalar_type() == at::kBFloat16 && + segment.dim() >= 2 && segment.stride(-1) == 1, + "Invalid segmented MXFP8 input layout."); + for (int dim = 1; dim < segment.dim() - 1; ++dim) { + NVTE_CHECK(segment.size(dim) == 1, "Only singleton batch dimensions are supported."); + } + wrappers.emplace_back(makeTransformerEngineTensor( + segment.data_ptr(), + std::vector{static_cast(segment.size(0)), static_cast(segment.size(-1))}, + DType::kBFloat16)); + inputs.push_back(wrappers.back().data()); + strides.push_back(segment.stride(0)); + } + auto [output_cpp, output_py] = quantizer_cpp->create_tensor(shape, DType::kBFloat16, device); + output_cpp.set_with_gemm_swizzled_scales(quantizer_cpp->optimize_for_gemm); + nvte_quantize_mxfp8_segments(inputs.data(), inputs.size(), strides.data(), output_cpp.data(), + at::cuda::getCurrentCUDAStream(device.index())); + return output_py; +} + py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector &shape, at::ScalarType dtype, at::Device device, bool pin_memory) { auto quantizer_cpp = convert_quantizer(quantizer); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c7f70182f36..1330d3c0fb5 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -210,6 +210,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("output") = py::none(), py::arg("noop") = py::none()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), py::arg("otype")); + m.def("quantize_mxfp8_segments", &transformer_engine::pytorch::quantize_mxfp8_segments, + py::arg("segments"), py::arg("quantizer"), py::arg("shape")); m.def("create_empty_quantized_tensor", &transformer_engine::pytorch::create_empty_quantized_tensor, "Create an empty quantized tensor", py::arg("quantizer"), py::arg("shape"), diff --git a/transformer_engine/pytorch/distributed_weight.py b/transformer_engine/pytorch/distributed_weight.py index 0d491430e52..5ecf398d949 100644 --- a/transformer_engine/pytorch/distributed_weight.py +++ b/transformer_engine/pytorch/distributed_weight.py @@ -9,6 +9,7 @@ GroupedLinear -> N; leader is ``weights[0]``) and no-op on plain tensors. """ +import math from typing import Any, List, Protocol, runtime_checkable import torch @@ -19,6 +20,8 @@ "materialize_weight_for_forward", "materialize_weight_for_backward", "finalize_weight_grads", + "get_weight_gradient_scale", + "mark_weight_gradient_scale", ] @@ -68,6 +71,48 @@ def is_distributed_weight(weight: Any) -> bool: return flag +def get_weight_gradient_scale(weight: Any, args: Any) -> float: + """Request an optional distributed-weight mean scale for an MXFP8 wgrad GEMM. + + The weight opts in through ``fuse_wgrad_mean_scaling``. Its optional + ``get_weight_gradient_scale()`` callback owns the reduction + semantics. Its finalizer must consume the gradient's ``_te_wgrad_scale`` marker + exactly once. Existing implementers without the callback retain a scale of one. + """ + if not getattr(weight, "fuse_wgrad_mean_scaling", False): + return 1.0 + callback = getattr(weight, "get_weight_gradient_scale", None) + if callback is None or not is_distributed_weight(weight): + return 1.0 + if not args.fp8 or args.debug or args.use_bias: + return 1.0 + if args.backward_override is not None or args.activation_dtype != torch.bfloat16: + return 1.0 + if args.grad_weight_quantizer is not None or args.cpu_offloading or args.ub_bulk_wgrad: + return 1.0 + if args.wgrad_store is not None and args.wgrad_store.delay_wgrad_compute(): + return 1.0 + from .tensor.mxfp8_tensor import MXFP8Quantizer + + if not isinstance(args.input_quantizer, MXFP8Quantizer): + return 1.0 + if torch.cuda.is_current_stream_capturing(): + return 1.0 + scale = float(callback()) + if not math.isfinite(scale) or scale <= 0: + raise ValueError("Distributed weight gradient scale must be finite and positive") + return scale + + +def mark_weight_gradient_scale(gradient: torch.Tensor, scale: float) -> None: + """Mark a freshly produced wgrad so its finalizer can avoid duplicate scaling.""" + if gradient.dtype not in (torch.float32, torch.bfloat16): + raise ValueError("Scaled weight gradients must use FP32 or BF16") + if hasattr(gradient, "_te_wgrad_scale"): + raise ValueError("Weight gradient already carries a scaling marker") + gradient._te_wgrad_scale = scale + + def materialize_weight_for_forward(weights: Any) -> List[Any]: """Prepare the weight(s) fed to the forward GEMM, always returned as a list. diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index e93d8c9641b..4f4a7c4411f 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """LayerNormLinear API""" + import os import warnings import weakref @@ -58,6 +59,8 @@ materialize_weight_for_forward, materialize_weight_for_backward, finalize_weight_grads, + get_weight_gradient_scale, + mark_weight_gradient_scale, ) from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo @@ -179,6 +182,7 @@ class LayerNormLinearFwdArgs: # --- Misc --- cpu_offloading: bool is_grad_enabled: bool + materialize_norm_output_grads: bool = True def any_requires_grad(self) -> bool: """Whether any differentiable input requires a gradient.""" @@ -1309,6 +1313,10 @@ def _layernorm_linear_backward_impl( "bulk_overlap": args.ub_bulk_wgrad, } + wgrad_scale = get_weight_gradient_scale(saved_weight, args) + if wgrad_scale != 1.0: + wgrad_gemm_kwargs["alpha"] = wgrad_scale + def wgrad_gemm( x: torch.Tensor, dy: torch.Tensor, @@ -1323,6 +1331,8 @@ def wgrad_gemm( """ nvtx_range_push(f"{nvtx_label}.wgrad_gemm") dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) + if wgrad_scale != 1.0: + mark_weight_gradient_scale(dw, wgrad_scale) nvtx_range_pop(f"{nvtx_label}.wgrad_gemm") return dw, db @@ -1392,7 +1402,11 @@ def wgrad_gemm( dgrad_work = None # Residual gradient - if args.return_layernorm_output and not args.return_layernorm_output_gathered: + if ( + args.return_layernorm_output + and not args.return_layernorm_output_gathered + and args.grad_ln_out is not None + ): dgrad = dgrad + args.grad_ln_out # Norm gradient @@ -1500,6 +1514,14 @@ def forward( ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects ctx.backward_objects = bwd_args + if ( + not fwd_args.materialize_norm_output_grads + and fwd_args.is_grad_enabled + and fwd_args.return_layernorm_output + and not fwd_args.return_layernorm_output_gathered + ): + ctx.set_materialize_grads(False) + ctx.main_output_spec = (tuple(out.shape), out.dtype, out.device) if fwd_args.fp8 and fwd_args.any_requires_grad(): bwd_args.reduce_and_update_bwd_fp8_tensors = check_fp8_reduce_and_update() if fwd_args.backward_override is not None: @@ -1516,6 +1538,9 @@ def backward( ) -> Tuple[Union[torch.Tensor, None], ...]: """Backward pass: compute gradients and reduce FP8 scaling factors.""" bwd_args: LayerNormLinearBwdArgs = ctx.backward_objects + if grad_output is None: + shape, dtype, device = ctx.main_output_spec + grad_output = torch.zeros(shape, dtype=dtype, device=device) bwd_args.grad_output = grad_output bwd_args.grad_ln_out = grad_ln_out bwd_args.setup_saved_tensors(ctx) @@ -1626,6 +1651,9 @@ class LayerNormLinear(TransformerEngineBaseModule): instead return the bias value during the forward pass together with the output of the linear transformation :math:`y = xA^T`. This is useful when the bias addition can be fused to subsequent operations. + materialize_norm_output_grads : bool, default = True + Set to ``False`` to skip zero gradients for an unused, ungathered + normalization output. Norm-only and combined consumers retain their gradients. params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters @@ -1671,6 +1699,7 @@ def __init__( delay_wgrad_compute: bool = False, symmetric_ar_type: Optional[str] = None, name: Optional[str] = None, + materialize_norm_output_grads: bool = True, ) -> None: super().__init__(name) @@ -1684,6 +1713,7 @@ def __init__( self.return_bias = return_bias self.apply_bias = self.use_bias and not return_bias self.return_layernorm_output = return_layernorm_output + self.materialize_norm_output_grads = materialize_norm_output_grads self.return_layernorm_output_gathered = ( return_layernorm_output_gathered if return_layernorm_output else False ) @@ -2164,6 +2194,7 @@ def forward( bwd_ln_sm_margin=self.bwd_ln_sm_margin, return_layernorm_output=self.return_layernorm_output, return_layernorm_output_gathered=self.return_layernorm_output_gathered, + materialize_norm_output_grads=self.materialize_norm_output_grads, # numerical / dtype config activation_dtype=self.activation_dtype, fp8=self.fp8, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 7f00e19644a..b02a5b7461e 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -69,6 +69,8 @@ materialize_weight_for_forward, materialize_weight_for_backward, finalize_weight_grads, + get_weight_gradient_scale, + mark_weight_gradient_scale, ) from ..cpp_extensions import ( general_gemm, @@ -1603,6 +1605,10 @@ def _linear_backward_impl(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None "bulk_overlap": bwd_args.ub_bulk_wgrad, } + wgrad_scale = get_weight_gradient_scale(saved_weight, bwd_args) + if wgrad_scale != 1.0: + wgrad_gemm_kwargs["alpha"] = wgrad_scale + def wgrad_gemm( x: torch.Tensor, dy: torch.Tensor, @@ -1617,6 +1623,8 @@ def wgrad_gemm( """ nvtx_range_push(f"{nvtx_label}.wgrad_gemm") dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) + if wgrad_scale != 1.0: + mark_weight_gradient_scale(dw, wgrad_scale) nvtx_range_pop(f"{nvtx_label}.wgrad_gemm") return dw, db diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 267806e43ee..dbaa3e8a872 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Tensor class with MXFP8 data""" + from __future__ import annotations from collections.abc import Iterable import math @@ -122,6 +123,59 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) + def quantize_segments( + self, parts: Tuple[torch.Tensor, ...], shape: Optional[Tuple[int, ...]] = None + ) -> QuantizedTensor: + """Quantize one or three column segments without a BF16 join buffer. + + Aligned BF16 matrices with both GEMM orientations on SM103 use direct conversion and + scale swizzling. Other supported inputs use the native quantizer. + """ + if len(parts) not in (1, 3): + raise ValueError("Expected one or three column segments") + if any(p.ndim < 2 for p in parts): + raise ValueError("Column segments must have at least two dimensions") + rows = math.prod(parts[0].shape[:-1]) + columns = sum(p.shape[-1] for p in parts) + shape = tuple(shape or (*parts[0].shape[:-1], columns)) + if shape[-1] != columns or math.prod(shape) != rows * columns: + raise ValueError("Segmented output shape does not match its inputs") + if any(math.prod(p.shape[:-1]) != rows for p in parts): + raise ValueError("Segment row counts must match") + direct = ( + self.dtype == DType.kFloat8E4M3 + and self.rowwise_usage + and self.columnwise_usage + and self.optimize_for_gemm + and not self.with_2d_quantization + and parts[0].is_cuda + and torch.cuda.get_device_capability(parts[0].device) == (10, 3) + and rows > 0 + and rows % 128 == 0 + and rows * columns // 32 < 2**31 - 128 + and all( + all(size == 1 for size in p.shape[1:-1]) + and p.shape[0] == rows + and p.dtype == torch.bfloat16 + and p.device == parts[0].device + and p.stride(-1) == 1 + and p.shape[-1] > 0 + and p.shape[-1] % 128 == 0 + and p.data_ptr() % 16 == 0 + and p.shape[-1] <= p.stride(0) <= 2**31 - 1 + and p.stride(0) % 8 == 0 + for p in parts + ) + ) + if direct: + # The native binding accepts the exact Python quantizer type. + # pylint: disable=unidiomatic-typecheck + allocator = self if type(self) is MXFP8Quantizer else self.copy() + # pylint: enable=unidiomatic-typecheck + return tex.quantize_mxfp8_segments(parts, allocator, shape) + joined = parts[0] if len(parts) == 1 else torch.cat(parts, dim=-1) + return MXFP8Quantizer.quantize_impl(self.copy(), joined.reshape(shape)) + def is_quantizable(self, inp: torch.Tensor) -> bool: """Returns whether or not given inp can be quantized""" if inp.ndim < 2: From 647b9dddabc881fb3d1c0d1fdd59aaf621aa6ad7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:15:34 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_mxfp8_segments.py | 7 +++++-- .../common/include/transformer_engine/cast.h | 4 ++-- transformer_engine/pytorch/csrc/extensions.h | 4 ++-- transformer_engine/pytorch/csrc/extensions/cast.cpp | 9 +++++---- 4 files changed, 14 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/test_mxfp8_segments.py b/tests/pytorch/test_mxfp8_segments.py index e74353163a7..1e72f136c70 100644 --- a/tests/pytorch/test_mxfp8_segments.py +++ b/tests/pytorch/test_mxfp8_segments.py @@ -44,7 +44,8 @@ def assert_payload_equal(actual, expected): mismatch = (before != after).nonzero()[:8] indexes = tuple(mismatch[:, i] for i in range(mismatch.shape[1])) pytest.fail( - f"{name}: positions={mismatch.tolist()}, expected={before[indexes].tolist()}, actual={after[indexes].tolist()}" + f"{name}: positions={mismatch.tolist()}, expected={before[indexes].tolist()}," + f" actual={after[indexes].tolist()}" ) @@ -59,7 +60,9 @@ def require_cuda(): @pytest.mark.parametrize("usage", [(True, False), (False, True), (True, True)]) @pytest.mark.parametrize("swizzle", [False, True]) @pytest.mark.parametrize("layout", ["contiguous", "strided", "singleton", "fallback", "broadcast"]) -@pytest.mark.parametrize("pattern", ["random", "signed_zero", "extreme", "all_nan", "all_inf", "all_zero"]) +@pytest.mark.parametrize( + "pattern", ["random", "signed_zero", "extreme", "all_nan", "all_inf", "all_zero"] +) def test_segmented_payload(usage, swizzle, layout, pattern): torch.manual_seed(1234) rows = 128 if layout != "fallback" else 96 diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 0f3c16d8c23..6457f05ab20 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -136,8 +136,8 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output, * \param[in] stream CUDA stream used for the operation. */ void nvte_quantize_mxfp8_segments(const NVTETensor *inputs, size_t num_inputs, - const int64_t *row_strides, NVTETensor output, - cudaStream_t stream); + const int64_t *row_strides, NVTETensor output, + cudaStream_t stream); /*! \brief Casts input tensor to MXFP8. Additionally, reduces the input along columns. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 215a18fecc0..bb6f7c0ce1b 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -382,8 +382,8 @@ std::vector bulk_allocate(const std::vector> &sh * Quantize **************************************************************************************************/ -py::object quantize_mxfp8_segments(const std::vector &segments, - py::handle quantizer, const std::vector &shape); +py::object quantize_mxfp8_segments(const std::vector &segments, py::handle quantizer, + const std::vector &shape); py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector &shape, at::ScalarType dtype, at::Device device, bool pin_memory); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 479b4ffc2cf..4920962e30a 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -160,10 +160,11 @@ py::object quantize_mxfp8_segments(const std::vector &segments, py:: for (int dim = 1; dim < segment.dim() - 1; ++dim) { NVTE_CHECK(segment.size(dim) == 1, "Only singleton batch dimensions are supported."); } - wrappers.emplace_back(makeTransformerEngineTensor( - segment.data_ptr(), - std::vector{static_cast(segment.size(0)), static_cast(segment.size(-1))}, - DType::kBFloat16)); + wrappers.emplace_back( + makeTransformerEngineTensor(segment.data_ptr(), + std::vector{static_cast(segment.size(0)), + static_cast(segment.size(-1))}, + DType::kBFloat16)); inputs.push_back(wrappers.back().data()); strides.push_back(segment.stride(0)); } From ea45788aba70a858324c32e2342b78b4f62f43c9 Mon Sep 17 00:00:00 2001 From: Layali Rashid Date: Tue, 13 Oct 2026 10:01:08 -0700 Subject: [PATCH 3/6] Guard segmented FP8 conversion on older GPU targets Signed-off-by: Layali Rashid --- transformer_engine/common/cast/quantize_mxfp8_segments.cu | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/cast/quantize_mxfp8_segments.cu b/transformer_engine/common/cast/quantize_mxfp8_segments.cu index 37b78cd207b..2b18947bc37 100644 --- a/transformer_engine/common/cast/quantize_mxfp8_segments.cu +++ b/transformer_engine/common/cast/quantize_mxfp8_segments.cu @@ -41,9 +41,9 @@ __device__ __forceinline__ const __nv_bfloat16 *base(Segments s, int col, int &o } __device__ __forceinline__ unsigned char scale_exponent(float amax) { - float normalized = amax * (1.0f / 448.0f); unsigned int packed; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + float normalized = amax * (1.0f / 448.0f); asm volatile( "{ .reg .b16 pair,zero; mov.b16 zero,0; cvt.rp.satfinite.ue8m0x2.f32 pair,%1,%1; mov.b32 " "%0,{pair,zero}; }" @@ -67,9 +67,14 @@ __device__ __forceinline__ int64_t sf(int i, int j, int tiles) { } __device__ __forceinline__ uint16_t pair_cast(float a, float b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890 uint16_t result; asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0,%2,%1;" : "=h"(result) : "f"(a), "f"(b)); return result; +#else + NVTE_DEVICE_ERROR("FP8 pair conversion is only supported on SM 8.9+."); + return 0; +#endif } __global__ void row_kernel(Segments s, unsigned char *out, unsigned char *scales, int n, int cols, From 329399bf184397d6afcee45fcacc4cdaacbde809 Mon Sep 17 00:00:00 2001 From: Layali Rashid Date: Tue, 13 Oct 2026 10:41:13 -0700 Subject: [PATCH 4/6] Consolidate attention boundary fusion tests Signed-off-by: Layali Rashid --- qa/L0_pytorch_unittest/test.sh | 1 + ....py => test_attention_boundary_fusions.py} | 150 +++++++++++++++++- .../test_nonmaterialized_norm_gradient.py | 58 ------- tests/pytorch/test_weight_gradient_scale.py | 103 ------------ 4 files changed, 148 insertions(+), 164 deletions(-) rename tests/pytorch/{test_mxfp8_segments.py => test_attention_boundary_fusions.py} (50%) delete mode 100644 tests/pytorch/test_nonmaterialized_norm_gradient.py delete mode 100644 tests/pytorch/test_weight_gradient_scale.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 7ace265048c..9ffdc8079af 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -55,6 +55,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_selective_activation_checkpoint.xml $TE_PATH/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py || test_fail "test_selective_activation_checkpoint.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_distributed_weight.xml $TE_PATH/tests/pytorch/test_distributed_weight.py || test_fail "test_distributed_weight.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_boundary_fusions.xml $TE_PATH/tests/pytorch/test_attention_boundary_fusions.py || test_fail "test_attention_boundary_fusions.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cross_entropy.xml $TE_PATH/tests/pytorch/test_cross_entropy.py || test_fail "test_cross_entropy.py" diff --git a/tests/pytorch/test_mxfp8_segments.py b/tests/pytorch/test_attention_boundary_fusions.py similarity index 50% rename from tests/pytorch/test_mxfp8_segments.py rename to tests/pytorch/test_attention_boundary_fusions.py index 1e72f136c70..077bfdf9eee 100644 --- a/tests/pytorch/test_mxfp8_segments.py +++ b/tests/pytorch/test_attention_boundary_fusions.py @@ -2,14 +2,21 @@ # # See LICENSE for license information. -"""Exact MXFP8 payload parity for segmented projection gradients.""" +"""Attention-boundary payload, gradient scaling, and normalization contracts.""" import math +from types import SimpleNamespace import pytest import torch +import transformer_engine_torch as tex -from transformer_engine.pytorch import DType, MXFP8Quantizer +from transformer_engine.common.recipe import MXFP8BlockScaling +from transformer_engine.pytorch import DType, LayerNormLinear, MXFP8Quantizer, autocast +from transformer_engine.pytorch.distributed_weight import ( + get_weight_gradient_scale, + mark_weight_gradient_scale, +) def assert_payload_equal(actual, expected): @@ -49,7 +56,7 @@ def assert_payload_equal(actual, expected): ) -@pytest.fixture(autouse=True) +@pytest.fixture def require_cuda(): if not torch.cuda.is_available(): pytest.skip("CUDA is required") @@ -63,6 +70,7 @@ def require_cuda(): @pytest.mark.parametrize( "pattern", ["random", "signed_zero", "extreme", "all_nan", "all_inf", "all_zero"] ) +@pytest.mark.usefixtures("require_cuda") def test_segmented_payload(usage, swizzle, layout, pattern): torch.manual_seed(1234) rows = 128 if layout != "fallback" else 96 @@ -97,6 +105,7 @@ def test_segmented_payload(usage, swizzle, layout, pattern): @pytest.mark.parametrize("segments", [1, 3]) +@pytest.mark.usefixtures("require_cuda") def test_segmented_stream_and_graph(segments): parts = tuple( torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) for _ in range(segments) @@ -118,6 +127,7 @@ def test_segmented_stream_and_graph(segments): assert_payload_equal(captured, expected) +@pytest.mark.usefixtures("require_cuda") def test_segmented_rejects_invalid_shapes(): quantizer = MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=True) part = torch.empty(128, 128, device="cuda", dtype=torch.bfloat16) @@ -130,6 +140,7 @@ def test_segmented_rejects_invalid_shapes(): @pytest.mark.parametrize("usage", [(True, False), (False, True)]) +@pytest.mark.usefixtures("require_cuda") def test_native_segmented_binding_rejects_unqualified_usage(usage): if torch.cuda.get_device_capability() != (10, 3): pytest.skip("The direct segmented binding is qualified for SM103") @@ -140,3 +151,136 @@ def test_native_segmented_binding_rejects_unqualified_usage(usage): quantizer.optimize_for_gemm = True with pytest.raises(RuntimeError, match="both orientations"): tex.quantize_mxfp8_segments((part,), quantizer, tuple(part.shape)) + + +def scaling_inputs(monkeypatch, scale=1.0 / 64): + """Use a tensor implementer and the real MXFP8 quantizer type.""" + weight = torch.empty(1) + weight.is_distributed_weight = True + weight.fuse_wgrad_mean_scaling = True + weight.get_weight_gradient_scale = lambda: scale + args = SimpleNamespace( + fp8=True, + debug=False, + use_bias=False, + backward_override=None, + activation_dtype=torch.bfloat16, + grad_weight_quantizer=None, + cpu_offloading=False, + ub_bulk_wgrad=False, + wgrad_store=None, + input_quantizer=MXFP8Quantizer(tex.DType.kFloat8E4M3), + ) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return weight, args + + +def test_plain_and_existing_weights_keep_unit_scale(): + weight = torch.empty(1) + assert get_weight_gradient_scale(weight, None) == 1.0 + weight.is_distributed_weight = True + assert get_weight_gradient_scale(weight, None) == 1.0 + weight.fuse_wgrad_mean_scaling = True + assert get_weight_gradient_scale(weight, None) == 1.0 + + +def test_opt_in_scale(monkeypatch): + weight, args = scaling_inputs(monkeypatch) + assert get_weight_gradient_scale(weight, args) == 1.0 / 64 + + +@pytest.mark.parametrize( + "field,value", + [ + ("fp8", False), + ("debug", True), + ("use_bias", True), + ("backward_override", "unfused"), + ("activation_dtype", torch.float16), + ("grad_weight_quantizer", object()), + ("cpu_offloading", True), + ("ub_bulk_wgrad", True), + ("input_quantizer", None), + ("wgrad_store", SimpleNamespace(delay_wgrad_compute=lambda: True)), + ], +) +def test_unsupported_routes_keep_unit_scale(monkeypatch, field, value): + weight, args = scaling_inputs(monkeypatch) + setattr(args, field, value) + assert get_weight_gradient_scale(weight, args) == 1.0 + + +def test_graph_capture_keeps_unit_scale(monkeypatch): + weight, args = scaling_inputs(monkeypatch) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert get_weight_gradient_scale(weight, args) == 1.0 + + +@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf")]) +def test_invalid_provider_scale_rejected(monkeypatch, scale): + weight, args = scaling_inputs(monkeypatch, scale) + with pytest.raises(ValueError, match="finite and positive"): + get_weight_gradient_scale(weight, args) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_marker_preserves_values_and_rejects_duplicate_scaling(dtype): + gradient = torch.arange(8, dtype=dtype) / 64 + expected = gradient.clone() + mark_weight_gradient_scale(gradient, 1.0 / 64) + assert torch.equal(gradient, expected) + assert gradient._te_wgrad_scale == 1.0 / 64 + with pytest.raises(ValueError, match="already carries"): + mark_weight_gradient_scale(gradient, 1.0 / 64) + + +def test_unsupported_gradient_dtype_rejected(): + with pytest.raises(ValueError, match="FP32 or BF16"): + mark_weight_gradient_scale(torch.empty(1, dtype=torch.float16), 1.0 / 64) + + +@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) +@pytest.mark.parametrize("consumer", ["linear", "norm", "both"]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_nonmaterialized_norm_gradient(normalization, consumer, fp8): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + if fp8 and torch.cuda.get_device_capability() < (10, 0): + pytest.skip("MXFP8 requires Blackwell") + torch.manual_seed(1234) + options = dict( + in_features=128, + out_features=256, + normalization=normalization, + params_dtype=torch.bfloat16, + return_layernorm_output=True, + bias=False, + ) + reference = LayerNormLinear(**options) + fused = LayerNormLinear(**options, materialize_norm_output_grads=False) + fused.load_state_dict(reference.state_dict()) + source = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + linear_upstream = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + norm_upstream = torch.randn_like(source) + results = [] + for module in (reference, fused): + inp = source.detach().clone().requires_grad_(True) + with autocast(enabled=fp8, recipe=MXFP8BlockScaling()): + linear, norm = module(inp) + outputs, upstream = [], [] + if consumer in ("linear", "both"): + outputs.append(linear) + upstream.append(linear_upstream) + if consumer in ("norm", "both"): + outputs.append(norm) + upstream.append(norm_upstream) + gradients = torch.autograd.grad( + outputs, (inp, *module.parameters()), upstream, allow_unused=True + ) + results.append((linear.detach(), norm.detach(), gradients)) + for before, after in zip(results[0][:2], results[1][:2]): + torch.testing.assert_close(before, after, rtol=0, atol=0) + for before, after in zip(results[0][2], results[1][2]): + assert (before is None) == (after is None) + if before is not None: + torch.testing.assert_close(before, after, rtol=0, atol=0) diff --git a/tests/pytorch/test_nonmaterialized_norm_gradient.py b/tests/pytorch/test_nonmaterialized_norm_gradient.py deleted file mode 100644 index 221f09f6b26..00000000000 --- a/tests/pytorch/test_nonmaterialized_norm_gradient.py +++ /dev/null @@ -1,58 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Gradient parity when the returned normalization output is unused.""" - -import pytest -import torch - -from transformer_engine.common.recipe import MXFP8BlockScaling -from transformer_engine.pytorch import LayerNormLinear, autocast - - -@pytest.mark.parametrize("normalization", ["LayerNorm", "RMSNorm"]) -@pytest.mark.parametrize("consumer", ["linear", "norm", "both"]) -@pytest.mark.parametrize("fp8", [False, True]) -def test_nonmaterialized_norm_gradient(normalization, consumer, fp8): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required") - if fp8 and torch.cuda.get_device_capability() < (10, 0): - pytest.skip("MXFP8 requires Blackwell") - torch.manual_seed(1234) - options = dict( - in_features=128, - out_features=256, - normalization=normalization, - params_dtype=torch.bfloat16, - return_layernorm_output=True, - bias=False, - ) - reference = LayerNormLinear(**options) - fused = LayerNormLinear(**options, materialize_norm_output_grads=False) - fused.load_state_dict(reference.state_dict()) - source = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") - linear_upstream = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") - norm_upstream = torch.randn_like(source) - results = [] - for module in (reference, fused): - inp = source.detach().clone().requires_grad_(True) - with autocast(enabled=fp8, recipe=MXFP8BlockScaling()): - linear, norm = module(inp) - outputs, upstream = [], [] - if consumer in ("linear", "both"): - outputs.append(linear) - upstream.append(linear_upstream) - if consumer in ("norm", "both"): - outputs.append(norm) - upstream.append(norm_upstream) - gradients = torch.autograd.grad( - outputs, (inp, *module.parameters()), upstream, allow_unused=True - ) - results.append((linear.detach(), norm.detach(), gradients)) - for before, after in zip(results[0][:2], results[1][:2]): - torch.testing.assert_close(before, after, rtol=0, atol=0) - for before, after in zip(results[0][2], results[1][2]): - assert (before is None) == (after is None) - if before is not None: - torch.testing.assert_close(before, after, rtol=0, atol=0) diff --git a/tests/pytorch/test_weight_gradient_scale.py b/tests/pytorch/test_weight_gradient_scale.py deleted file mode 100644 index 803a2edb27e..00000000000 --- a/tests/pytorch/test_weight_gradient_scale.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Opt-in distributed-weight gradient scaling contract.""" - -from types import SimpleNamespace - -import pytest -import torch -import transformer_engine_torch as tex - -from transformer_engine.pytorch.distributed_weight import ( - get_weight_gradient_scale, - mark_weight_gradient_scale, -) -from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - -def scaling_inputs(monkeypatch, scale=1.0 / 64): - """Use a tensor implementer and the real MXFP8 quantizer type.""" - weight = torch.empty(1) - weight.is_distributed_weight = True - weight.fuse_wgrad_mean_scaling = True - weight.get_weight_gradient_scale = lambda: scale - args = SimpleNamespace( - fp8=True, - debug=False, - use_bias=False, - backward_override=None, - activation_dtype=torch.bfloat16, - grad_weight_quantizer=None, - cpu_offloading=False, - ub_bulk_wgrad=False, - wgrad_store=None, - input_quantizer=MXFP8Quantizer(tex.DType.kFloat8E4M3), - ) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - return weight, args - - -def test_plain_and_existing_weights_keep_unit_scale(): - weight = torch.empty(1) - assert get_weight_gradient_scale(weight, None) == 1.0 - weight.is_distributed_weight = True - assert get_weight_gradient_scale(weight, None) == 1.0 - weight.fuse_wgrad_mean_scaling = True - assert get_weight_gradient_scale(weight, None) == 1.0 - - -def test_opt_in_scale(monkeypatch): - weight, args = scaling_inputs(monkeypatch) - assert get_weight_gradient_scale(weight, args) == 1.0 / 64 - - -@pytest.mark.parametrize( - "field,value", - [ - ("fp8", False), - ("debug", True), - ("use_bias", True), - ("backward_override", "unfused"), - ("activation_dtype", torch.float16), - ("grad_weight_quantizer", object()), - ("cpu_offloading", True), - ("ub_bulk_wgrad", True), - ("input_quantizer", None), - ("wgrad_store", SimpleNamespace(delay_wgrad_compute=lambda: True)), - ], -) -def test_unsupported_routes_keep_unit_scale(monkeypatch, field, value): - weight, args = scaling_inputs(monkeypatch) - setattr(args, field, value) - assert get_weight_gradient_scale(weight, args) == 1.0 - - -def test_graph_capture_keeps_unit_scale(monkeypatch): - weight, args = scaling_inputs(monkeypatch) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) - assert get_weight_gradient_scale(weight, args) == 1.0 - - -@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf")]) -def test_invalid_provider_scale_rejected(monkeypatch, scale): - weight, args = scaling_inputs(monkeypatch, scale) - with pytest.raises(ValueError, match="finite and positive"): - get_weight_gradient_scale(weight, args) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_marker_preserves_values_and_rejects_duplicate_scaling(dtype): - gradient = torch.arange(8, dtype=dtype) / 64 - expected = gradient.clone() - mark_weight_gradient_scale(gradient, 1.0 / 64) - assert torch.equal(gradient, expected) - assert gradient._te_wgrad_scale == 1.0 / 64 - with pytest.raises(ValueError, match="already carries"): - mark_weight_gradient_scale(gradient, 1.0 / 64) - - -def test_unsupported_gradient_dtype_rejected(): - with pytest.raises(ValueError, match="FP32 or BF16"): - mark_weight_gradient_scale(torch.empty(1, dtype=torch.float16), 1.0 / 64) From e7b93973804781d7180d76a3eb760c61111da688 Mon Sep 17 00:00:00 2001 From: Layali Rashid Date: Tue, 20 Oct 2026 14:18:18 -0700 Subject: [PATCH 5/6] Initialize segmented MXFP8 quantization before type checks Signed-off-by: Layali Rashid --- .../test_attention_boundary_fusions.py | 32 +++++++++++++++++-- .../pytorch/csrc/extensions/cast.cpp | 1 + 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_attention_boundary_fusions.py b/tests/pytorch/test_attention_boundary_fusions.py index 077bfdf9eee..43c2da62353 100644 --- a/tests/pytorch/test_attention_boundary_fusions.py +++ b/tests/pytorch/test_attention_boundary_fusions.py @@ -5,11 +5,12 @@ """Attention-boundary payload, gradient scaling, and normalization contracts.""" import math +import subprocess +import sys from types import SimpleNamespace import pytest import torch -import transformer_engine_torch as tex from transformer_engine.common.recipe import MXFP8BlockScaling from transformer_engine.pytorch import DType, LayerNormLinear, MXFP8Quantizer, autocast @@ -64,6 +65,33 @@ def require_cuda(): pytest.skip("MXFP8 requires Blackwell") +@pytest.mark.usefixtures("require_cuda") +def test_segmented_first_call(): + if torch.cuda.get_device_capability() != (10, 3): + pytest.skip("The direct segmented binding is qualified for SM103") + # A fresh process must reach segmented conversion before any native TE call. + subprocess.run( + [ + sys.executable, + "-c", + """ +import torch +from transformer_engine.pytorch import DType, MXFP8Quantizer +x = torch.ones(128, 128, device="cuda", dtype=torch.bfloat16) +q = MXFP8Quantizer(DType.kFloat8E4M3, rowwise=True, columnwise=True) +q.optimize_for_gemm = True +actual = q.quantize_segments((x,)) +expected = q(x) +torch.testing.assert_close(actual.dequantize(), expected.dequantize(), rtol=0, atol=0) +for name in ("_rowwise_data", "_columnwise_data"): + assert torch.equal(getattr(actual, name), getattr(expected, name)), name +""", + ], + check=True, + timeout=120, + ) + + @pytest.mark.parametrize("usage", [(True, False), (False, True), (True, True)]) @pytest.mark.parametrize("swizzle", [False, True]) @pytest.mark.parametrize("layout", ["contiguous", "strided", "singleton", "fallback", "broadcast"]) @@ -169,7 +197,7 @@ def scaling_inputs(monkeypatch, scale=1.0 / 64): cpu_offloading=False, ub_bulk_wgrad=False, wgrad_store=None, - input_quantizer=MXFP8Quantizer(tex.DType.kFloat8E4M3), + input_quantizer=MXFP8Quantizer(DType.kFloat8E4M3), ) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) return weight, args diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 4920962e30a..a9e4c61a81e 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -139,6 +139,7 @@ py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantiz py::object quantize_mxfp8_segments(const std::vector &segments, py::handle quantizer, const std::vector &shape) { + init_extension(); NVTE_CHECK(segments.size() == 1 || segments.size() == 3, "Expected one or three MXFP8 column segments."); NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()), "Expected a native MXFP8 quantizer."); From b2d54252d281af630515ad09d5ccc8f7fff478b8 Mon Sep 17 00:00:00 2001 From: Layali Rashid Date: Tue, 20 Oct 2026 15:28:56 -0700 Subject: [PATCH 6/6] Preserve distributed weight finalizer gradients Signed-off-by: Layali Rashid --- .../test_attention_boundary_fusions.py | 51 ++++++++++++++++++- .../pytorch/module/layernorm_linear.py | 4 +- transformer_engine/pytorch/module/linear.py | 4 +- 3 files changed, 54 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_attention_boundary_fusions.py b/tests/pytorch/test_attention_boundary_fusions.py index 43c2da62353..1af2d7f4059 100644 --- a/tests/pytorch/test_attention_boundary_fusions.py +++ b/tests/pytorch/test_attention_boundary_fusions.py @@ -13,7 +13,7 @@ import torch from transformer_engine.common.recipe import MXFP8BlockScaling -from transformer_engine.pytorch import DType, LayerNormLinear, MXFP8Quantizer, autocast +from transformer_engine.pytorch import DType, LayerNormLinear, Linear, MXFP8Quantizer, autocast from transformer_engine.pytorch.distributed_weight import ( get_weight_gradient_scale, mark_weight_gradient_scale, @@ -92,6 +92,55 @@ def test_segmented_first_call(): ) +@pytest.mark.parametrize("module_type", [Linear, LayerNormLinear]) +@pytest.mark.parametrize("return_dummy", [False, True]) +def test_distributed_finalizer_owns_gradient(module_type, return_dummy): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + torch.manual_seed(1234) + options = dict(in_features=128, out_features=256, bias=False, params_dtype=torch.bfloat16) + reference = module_type(**options) + module = module_type(**options, fuse_wgrad_accumulation=True) + module.load_state_dict(reference.state_dict()) + full_weight = module.weight.detach().clone() + weight = torch.nn.Parameter(full_weight[:128].clone()) + module.weight = weight + weight.is_distributed_weight = True + weight.grad_added_to_main_grad = False + weight.main_grad = torch.zeros_like(weight, dtype=torch.float32) + workspace = torch.zeros_like(full_weight, dtype=torch.float32) + weight.grad_buffer = lambda: workspace + weight.materialize_group_for_forward = lambda: full_weight + weight.materialize_group_for_backward = lambda: full_weight + dummy = torch.full_like(weight, 3.25) if return_dummy else None + finalized = [] + + def finalize(gradient): + finalized.append(tuple(gradient.shape)) + weight.main_grad.copy_(gradient[:128]) + weight.grad_added_to_main_grad = True + return dummy + + weight.finalize_group_grads = finalize + source = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) + upstream = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + expected_input = source.clone().requires_grad_(True) + actual_input = source.clone().requires_grad_(True) + expected = reference(expected_input) + actual = module(actual_input) + expected.backward(upstream) + actual.backward(upstream) + assert finalized == [(256, 128)] + assert (weight.grad is None) == (dummy is None) + if dummy is not None: + torch.testing.assert_close(weight.grad, dummy, rtol=0, atol=0) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(actual_input.grad, expected_input.grad, rtol=0, atol=0) + torch.testing.assert_close( + weight.main_grad.to(torch.bfloat16), reference.weight.grad[:128], rtol=0.01, atol=0.01 + ) + + @pytest.mark.parametrize("usage", [(True, False), (False, True), (True, True)]) @pytest.mark.parametrize("swizzle", [False, True]) @pytest.mark.parametrize("layout", ["contiguous", "strided", "singleton", "fallback", "broadcast"]) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 4f4a7c4411f..f14714a8dcb 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1437,7 +1437,7 @@ def wgrad_gemm( clear_tensor_data(mu) clear_tensor_data(rsigma) - if args.requires_wgrad: + if args.requires_wgrad and not is_dist_weight: # Handle custom DDP from mcore. if args.fuse_wgrad_accumulation and hasattr(origin_weight, "grad_added_to_main_grad"): origin_weight.grad_added_to_main_grad = True @@ -1454,7 +1454,7 @@ def wgrad_gemm( ) elif args.fuse_wgrad_accumulation: wgrad = None - else: + elif not args.requires_wgrad: wgrad = None return ( diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index b02a5b7461e..924705cdf70 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1693,7 +1693,7 @@ def wgrad_gemm( dgrad_work.wait() dgrad_work = None - if bwd_args.requires_wgrad: + if bwd_args.requires_wgrad and not is_dist_weight: # Handle custom DDP from mcore. if bwd_args.fuse_wgrad_accumulation and hasattr( origin_weight_python_object, "grad_added_to_main_grad" @@ -1716,7 +1716,7 @@ def wgrad_gemm( ) elif bwd_args.fuse_wgrad_accumulation: wgrad = None - else: + elif not bwd_args.requires_wgrad: wgrad = None # Scatter fp8 weight buffers