From 7553e6a1bf8a41ea605d3732bb184bcdaaea943f Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 31 Mar 2026 14:51:07 +0000 Subject: [PATCH 01/97] Stage 1&2: Python containers + quantize/gemm dispatch/unwrap Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 1561 +++++++++++++++++ .../quantize_transpose_square_blockwise.cu | 64 +- transformer_engine/pytorch/__init__.py | 3 + .../pytorch/cpp_extensions/gemm.py | 37 + transformer_engine/pytorch/module/base.py | 6 +- .../pytorch/module/grouped_linear.py | 57 +- .../pytorch/module/layernorm_linear.py | 3 + .../pytorch/module/layernorm_mlp.py | 3 + transformer_engine/pytorch/tensor/__init__.py | 7 + .../pytorch/tensor/hybrid_tensor.py | 193 ++ .../tensor/storage/hybrid_tensor_storage.py | 157 ++ 11 files changed, 2064 insertions(+), 27 deletions(-) create mode 100644 tests/pytorch/test_hybrid_quantization.py create mode 100644 transformer_engine/pytorch/tensor/hybrid_tensor.py create mode 100644 transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py new file mode 100644 index 00000000000..96a28744f33 --- /dev/null +++ b/tests/pytorch/test_hybrid_quantization.py @@ -0,0 +1,1561 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for hybrid quantization (mixed rowwise/columnwise formats).""" + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex + +from transformer_engine.common import recipe +from transformer_engine.pytorch import ( + autocast, + Linear, + LayerNormLinear, + LayerNormMLP, + TransformerLayer, + GroupedLinear, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, + HybridQuantizer, + HybridQuantizedTensor, + HybridQuantizedTensorStorage, + Float8Tensor, + Float8TensorStorage, + NVFP4Tensor, + NVFP4TensorStorage, +) +from transformer_engine.pytorch.cpp_extensions.gemm import ( + _unwrap_hybrid_A, + _unwrap_hybrid_B, +) + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) + +requires_fp8 = pytest.mark.skipif( + not fp8_available, + reason=f"FP8: {reason_for_no_fp8}", +) + +requires_fp8_and_nvfp4 = pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) + + +def _make_fp8_quantizer(*, rowwise=True, columnwise=True): + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_nvfp4_quantizer(*, rowwise=True, columnwise=True): + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_hybrid_quantizer_fp8_row_fp4_col(): + """FP8 rowwise + NVFP4 columnwise.""" + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + + +def _make_hybrid_quantizer_fp4_row_fp8_col(): + """NVFP4 rowwise + FP8 columnwise (reversed direction).""" + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + + +@requires_fp8_and_nvfp4 +class TestHybridQuantizerConstruction: + """Test construction and basic properties of hybrid quantizer.""" + + def test_creation(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + assert isinstance(hq, HybridQuantizer) + assert hq.rowwise_usage is True + assert hq.columnwise_usage is True + assert isinstance(hq.rowwise_quantizer, Float8CurrentScalingQuantizer) + assert isinstance(hq.columnwise_quantizer, NVFP4Quantizer) + + def test_compatible_recipe_is_none(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + assert hq._get_compatible_recipe() is None + + +@requires_fp8_and_nvfp4 +class TestHybridQuantize: + """Test quantization via HybridQuantizer.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_quantize_returns_hybrid_tensor(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert isinstance(result, HybridQuantizedTensor) + + def test_quantize_shape_preserved(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert result.shape == input_tensor.shape + + def test_quantize_dtype_preserved(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + assert result.dtype == input_tensor.dtype + + def test_sub_storage_types_fp8_row_fp4_col(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + row_storage = result.rowwise_sub_storage + col_storage = result.columnwise_sub_storage + assert isinstance(row_storage, (Float8TensorStorage, Float8Tensor)) + assert isinstance(col_storage, (NVFP4TensorStorage, NVFP4Tensor)) + + def test_sub_storage_types_reversed(self, input_tensor): + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + result = hq.quantize(input_tensor) + row_storage = result.rowwise_sub_storage + col_storage = result.columnwise_sub_storage + assert isinstance(row_storage, (NVFP4TensorStorage, NVFP4Tensor)) + assert isinstance(col_storage, (Float8TensorStorage, Float8Tensor)) + + def test_quantize_internal_returns_storage(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(input_tensor) + assert isinstance(result, HybridQuantizedTensorStorage) + assert not isinstance(result, HybridQuantizedTensor) + hq.internal = False + + +@requires_fp8_and_nvfp4 +class TestHybridDequantize: + """Test dequantization round-trip.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_dequantize_shape(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + assert dequantized.shape == input_tensor.shape + + def test_dequantize_dtype(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + assert dequantized.dtype == input_tensor.dtype + + def test_dequantize_explicit_dtype(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize(dtype=torch.float32) + assert dequantized.dtype == torch.float32 + assert dequantized.shape == input_tensor.shape + + def test_dequantize_close_to_original(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + torch.testing.assert_close( + dequantized.float(), input_tensor.float(), rtol=0.125, atol=0.0675 + ) + + def test_dequantize_reversed_close_to_original(self, input_tensor): + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + result = hq.quantize(input_tensor) + dequantized = result.dequantize() + torch.testing.assert_close( + dequantized.float(), input_tensor.float(), rtol=0.5, atol=1.0 + ) + + def test_storage_dequantize(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(input_tensor) + dequantized = result.dequantize(dtype=torch.bfloat16) + assert dequantized.shape == input_tensor.shape + hq.internal = False + + +@requires_fp8_and_nvfp4 +class TestHybridUpdateUsage: + """Test update_usage semantics and sub-storage cleanup.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_initial_usages(self, hybrid_tensor): + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is True + assert usages["columnwise"] is True + + def test_drop_rowwise(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + assert hybrid_tensor.rowwise_sub_storage is None + assert hybrid_tensor.columnwise_sub_storage is not None + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is False + assert usages["columnwise"] is True + + def test_drop_columnwise(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + assert hybrid_tensor.columnwise_sub_storage is None + assert hybrid_tensor.rowwise_sub_storage is not None + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is True + assert usages["columnwise"] is False + + def test_drop_both(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False, columnwise_usage=False) + usages = hybrid_tensor.get_usages() + assert usages["rowwise"] is False + assert usages["columnwise"] is False + + def test_request_true_is_noop(self, hybrid_tensor): + row_before = hybrid_tensor.rowwise_sub_storage + col_before = hybrid_tensor.columnwise_sub_storage + hybrid_tensor.update_usage(rowwise_usage=True, columnwise_usage=True) + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is col_before + + def test_repr_after_drop(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + r = repr(hybrid_tensor) + assert "HybridQuantizedTensor" in r + assert "rowwise=None" in r + + hybrid_tensor.update_usage(columnwise_usage=False) + r = repr(hybrid_tensor) + assert "rowwise=None" in r + assert "columnwise=None" in r + + +@requires_fp8_and_nvfp4 +class TestHybridSaveRestore: + """Test prepare_for_saving / restore_from_saved round-trip.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_save_restore_roundtrip(self, hybrid_tensor): + dq_before = hybrid_tensor.dequantize() + tensors, obj = hybrid_tensor.prepare_for_saving() + assert isinstance(tensors, list) + assert all(t is None or isinstance(t, torch.Tensor) for t in tensors) + + remainder = obj.restore_from_saved(tensors) + assert isinstance(remainder, list) + assert len(remainder) == 0 + + dq_after = hybrid_tensor.dequantize() + torch.testing.assert_close(dq_before, dq_after) + + def test_save_clears_data(self, hybrid_tensor): + tensors, obj = hybrid_tensor.prepare_for_saving() + row_storage = hybrid_tensor.rowwise_sub_storage + row_data_tensors = row_storage.get_data_tensors() + if isinstance(row_data_tensors, tuple): + assert all(t is None for t in row_data_tensors) + else: + assert row_data_tensors is None + # Restore to clean up + obj.restore_from_saved(tensors) + + +@requires_fp8_and_nvfp4 +class TestHybridMakeEmpty: + """Test HybridQuantizer.make_empty().""" + + def test_make_empty_shape(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert isinstance(empty, HybridQuantizedTensor) + assert empty.shape == torch.Size(shape) + + def test_make_empty_dtype(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert empty.dtype == torch.bfloat16 + + def test_make_empty_has_sub_storages(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + shape = (128, 256) + empty = hq.make_empty(shape, dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is not None + assert empty.columnwise_sub_storage is not None + + +@requires_fp8_and_nvfp4 +class TestHybridTorchDispatch: + """Test torch dispatch operations.""" + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_detach(self, hybrid_tensor): + detached = hybrid_tensor.detach() + assert isinstance(detached, HybridQuantizedTensor) + assert not detached.requires_grad + + def test_repr(self, hybrid_tensor): + r = repr(hybrid_tensor) + assert "HybridQuantizedTensor" in r + + +@requires_fp8_and_nvfp4 +class TestHybridGetDataTensors: + """Test get_data_tensors returns data from both sub-storages.""" + + def test_get_data_tensors(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(inp) + data_tensors = result.get_data_tensors() + assert isinstance(data_tensors, tuple) + assert len(data_tensors) > 0 + has_non_none = any(t is not None for t in data_tensors) + assert has_non_none + + +@requires_fp8_and_nvfp4 +class TestHybridDeviceAndSize: + """Test device and size properties.""" + + def test_device(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + result = hq.quantize(inp) + assert result.device.type == "cuda" + + def test_size_from_storage(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + result = hq.quantize(inp) + size = result.size() + assert size == torch.Size([128, 256]) + hq.internal = False + + +@requires_fp8 +class TestHybridGemmBitwiseIdentical: + """Hybrid quantizer with same FP8 format in both directions must produce + bitwise-identical results to the vanilla Float8CurrentScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_fp8(self): + torch.manual_seed(123) + + in_features = 64 + out_features = 64 + batch = 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.Float8CurrentScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + loss_ref = out_ref.float().sum() + loss_ref.backward() + + def hybrid_fp8_factory(role): + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda", + ) + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_fp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + loss_hybrid = out_hybrid.float().sum() + loss_hybrid.backward() + + # Forward outputs must be bitwise identical + assert torch.equal(out_ref, out_hybrid), ( + f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + ) + + # Input gradients must be bitwise identical + assert inp_ref.grad is not None and inp_hybrid.grad is not None + assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( + f"Input grad mismatch: max diff = " + f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + ) + + # Parameter gradients must be bitwise identical + ref_params = dict(model_ref.named_parameters()) + hybrid_params = dict(model_hybrid.named_parameters()) + for name, p_ref in ref_params.items(): + p_hyb = hybrid_params[name] + assert p_ref.grad is not None and p_hyb.grad is not None, ( + f"Missing gradient for param '{name}'" + ) + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestHybridGemmBitwiseIdenticalMXFP8: + """Hybrid quantizer with MXFP8 in both directions must produce + bitwise-identical results to the vanilla MXFP8BlockScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_mxfp8(self): + torch.manual_seed(200) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.MXFP8BlockScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_mxfp8_factory(role): + if role in ("linear_grad_output", "linear_grad_input"): + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_mxfp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal(out_ref, out_hybrid), ( + f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + ) + assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( + f"Input grad mismatch: max diff = " + f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + ) + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert p_ref.grad is not None and p_hyb.grad is not None, ( + f"Missing gradient for param '{name}'" + ) + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +class TestHybridGemmBitwiseIdenticalBlockFP8: + """Hybrid quantizer with Block FP8 in both directions must produce + bitwise-identical results to the vanilla Float8BlockScaling recipe.""" + + def test_linear_fwd_bwd_matches_vanilla_block_fp8(self): + torch.manual_seed(201) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.Float8BlockScaling() + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_block_fp8_factory(role): + dim = 2 if role == "linear_weight" else 1 + if role in ("linear_grad_output", "linear_grad_input"): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, columnwise=True, + block_scaling_dim=dim, + ) + return HybridQuantizer( + rowwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, columnwise=True, + block_scaling_dim=dim, + ), + columnwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, columnwise=True, + block_scaling_dim=dim, + ), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_block_fp8_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal(out_ref, out_hybrid), ( + f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + ) + assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( + f"Input grad mismatch: max diff = " + f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + ) + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert p_ref.grad is not None and p_hyb.grad is not None, ( + f"Missing gradient for param '{name}'" + ) + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestHybridGemmBitwiseIdenticalNVFP4: + """Hybrid quantizer with NVFP4 in both directions must produce + bitwise-identical results to the vanilla NVFP4BlockScaling recipe. + + RHT, stochastic rounding, and 2D quantization are disabled so the + test is fully deterministic and two independent quantizer instances + produce the same output. + """ + + def test_linear_fwd_bwd_matches_vanilla_nvfp4(self): + torch.manual_seed(202) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + ) + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_nvfp4_factory(role): + if role in ("linear_grad_output", "linear_grad_input"): + return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + return HybridQuantizer( + rowwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + columnwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal(out_ref, out_hybrid), ( + f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + ) + assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( + f"Input grad mismatch: max diff = " + f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + ) + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert p_ref.grad is not None and p_hyb.grad is not None, ( + f"Missing gradient for param '{name}'" + ) + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + def test_linear_fwd_bwd_all_roles_hybrid(self): + """All roles (including grad_output) use HybridQuantizer with NVFP4 both + directions. Validates per-operand unwrap produces bitwise-identical results + when grad_output is hybrid.""" + torch.manual_seed(203) + + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_hybrid = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + ) + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + def hybrid_nvfp4_all_roles_factory(role): + return HybridQuantizer( + rowwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + columnwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + ) + + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_all_roles_factory) + with autocast(enabled=True, recipe=hybrid_recipe): + out_hybrid = model_hybrid(inp_hybrid) + out_hybrid.float().sum().backward() + + assert torch.equal(out_ref, out_hybrid), ( + f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + ) + assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( + f"Input grad mismatch: max diff = " + f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" + ) + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_hyb = dict(model_hybrid.named_parameters())[name] + assert p_ref.grad is not None and p_hyb.grad is not None, ( + f"Missing gradient for param '{name}'" + ) + assert torch.equal(p_ref.grad, p_hyb.grad), ( + f"Param '{name}' grad mismatch: max diff = " + f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" + ) + + +@requires_fp8_and_nvfp4 +class TestHybridGemmMixedFormat: + """FP8 rowwise + NVFP4 columnwise through te.Linear forward+backward.""" + + def test_linear_fwd_bwd_fp8_row_nvfp4_col(self): + torch.manual_seed(42) + + in_features = 128 + out_features = 128 + batch = 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True, + ) + + def mixed_factory(role): + if role in ("linear_input", "linear_weight"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return _make_nvfp4_quantizer() + return None + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert out.dtype == torch.bfloat16 + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert inp.grad.shape == inp.shape + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + assert not torch.isinf(inp.grad).any(), "Input gradient contains Inf" + + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + assert not torch.isinf(p.grad).any(), f"Gradient for '{name}' contains Inf" + + def test_numerical_sanity_against_bf16(self): + """Mixed-format output should be within reasonable tolerance of BF16 baseline.""" + torch.manual_seed(42) + + in_features = 128 + out_features = 128 + batch = 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + # BF16 baseline (no quantization) + with torch.no_grad(): + out_bf16 = model(inp) + + def mixed_factory(role): + if role in ("linear_input", "linear_weight"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + return None + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + with torch.no_grad(): + with autocast(enabled=True, recipe=mixed_recipe): + out_mixed = model(inp) + + # FP8/FP4 quantization introduces error, but the result should be + # in the same ballpark as BF16 + torch.testing.assert_close( + out_mixed.float(), out_bf16.float(), rtol=0.25, atol=0.5, + ) + + +@requires_fp8_and_nvfp4 +class TestUnwrapHybridDirection: + """Test per-operand unwrap selects the correct sub-storage. + + Operand A: transposed (layout[0]=='T') → rowwise, else → columnwise + Operand B: not-transposed (layout[1]=='N') → rowwise, else → columnwise + """ + + @pytest.fixture + def hybrid_tensor(self): + torch.manual_seed(42) + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + return hq.quantize(inp) + + def test_A_tn_returns_rowwise(self, hybrid_tensor): + assert _unwrap_hybrid_A(hybrid_tensor, "TN") is hybrid_tensor.rowwise_sub_storage + + def test_A_nn_returns_columnwise(self, hybrid_tensor): + assert _unwrap_hybrid_A(hybrid_tensor, "NN") is hybrid_tensor.columnwise_sub_storage + + def test_A_nt_returns_columnwise(self, hybrid_tensor): + assert _unwrap_hybrid_A(hybrid_tensor, "NT") is hybrid_tensor.columnwise_sub_storage + + def test_B_tn_returns_rowwise(self, hybrid_tensor): + assert _unwrap_hybrid_B(hybrid_tensor, "TN") is hybrid_tensor.rowwise_sub_storage + + def test_B_nn_returns_rowwise(self, hybrid_tensor): + assert _unwrap_hybrid_B(hybrid_tensor, "NN") is hybrid_tensor.rowwise_sub_storage + + def test_B_nt_returns_columnwise(self, hybrid_tensor): + assert _unwrap_hybrid_B(hybrid_tensor, "NT") is hybrid_tensor.columnwise_sub_storage + + def test_tn_sub_storage_type(self, hybrid_tensor): + assert isinstance( + _unwrap_hybrid_A(hybrid_tensor, "TN"), (Float8TensorStorage, Float8Tensor), + ) + + def test_nt_sub_storage_type(self, hybrid_tensor): + assert isinstance( + _unwrap_hybrid_B(hybrid_tensor, "NT"), (NVFP4TensorStorage, NVFP4Tensor), + ) + + def test_non_hybrid_passthrough(self): + plain = torch.randn(4, 4, device="cuda") + for layout in ("TN", "NN", "NT"): + assert _unwrap_hybrid_A(plain, layout) is plain + assert _unwrap_hybrid_B(plain, layout) is plain + + def test_fp8_tensor_passthrough(self): + quantizer = _make_fp8_quantizer() + inp = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") + fp8 = quantizer.quantize(inp) + for layout in ("TN", "NN", "NT"): + assert _unwrap_hybrid_A(fp8, layout) is fp8 + assert _unwrap_hybrid_B(fp8, layout) is fp8 + + +@requires_fp8 +class TestHybridBiasGradient: + """Verify bias gradients are computed correctly with HybridQuantizer. + + tex.bgrad_quantize doesn't recognize HybridQuantizer, so the unfused + bgrad path is used instead. + """ + + def _make_uniform_hybrid_factory(self): + def factory(role): + if role in ("linear_grad_output", "linear_grad_input"): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda", + ) + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ), + ) + return factory + + def test_bias_grad_matches_vanilla_fp8(self): + torch.manual_seed(456) + in_features, out_features, batch = 64, 64, 16 + + model_ref = Linear(in_features, out_features, bias=True, + params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear(in_features, out_features, bias=True, + params_dtype=torch.bfloat16).cuda() + model_hybrid.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + # Reference + inp_ref = base_inp.clone().detach().requires_grad_(True) + with autocast(enabled=True, recipe=recipe.Float8CurrentScaling()): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + + # Hybrid + inp_hyb = base_inp.clone().detach().requires_grad_(True) + with autocast(enabled=True, + recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory())): + out_hyb = model_hybrid(inp_hyb) + out_hyb.float().sum().backward() + + ref_bias_grad = dict(model_ref.named_parameters())["bias"].grad + hyb_bias_grad = dict(model_hybrid.named_parameters())["bias"].grad + assert ref_bias_grad is not None and hyb_bias_grad is not None + assert torch.equal(ref_bias_grad, hyb_bias_grad), ( + f"Bias grad mismatch: max diff = " + f"{(ref_bias_grad - hyb_bias_grad).abs().max().item()}" + ) + + def test_no_bias_fwd_bwd(self): + """Linear with bias=False skips bgrad_quantize entirely.""" + torch.manual_seed(42) + in_features, out_features, batch = 64, 64, 16 + + model = Linear(in_features, out_features, bias=False, + params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", + dtype=torch.bfloat16, requires_grad=True) + + with autocast(enabled=True, + recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory())): + out = model(inp) + out.float().sum().backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + + +@requires_fp8_and_nvfp4 +class TestHybridScalingModeCompatibility: + """cuBLAS requires matching scaling modes within a single GEMM. + + For hybrid quantization, this means the columnwise format for + linear_input/linear_weight must match the columnwise format for + linear_grad_output — otherwise the wgrad GEMM (NT layout) fails. + """ + + def test_matching_columnwise_formats_succeed(self): + """Both operands use NVFP4 columnwise → wgrad GEMM succeeds.""" + torch.manual_seed(42) + # NVFP4 GEMM requires dimensions ≥ 128 for cuBLAS support. + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, + requires_grad=True) + + def factory(role): + if role in ("linear_input", "linear_weight"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return _make_nvfp4_quantizer() + return None + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + out.float().sum().backward() + assert inp.grad is not None + + def test_mismatched_columnwise_formats_raise(self): + """NVFP4 input × FP8 grad_output columnwise → cuBLAS rejects.""" + torch.manual_seed(42) + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, + requires_grad=True) + + def factory(role): + if role in ("linear_input", "linear_weight"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda", + ) + return None + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + with pytest.raises(RuntimeError, match="scaling_mode"): + out.float().sum().backward() + + +@requires_fp8_and_nvfp4 +class TestHybridReversedDirection: + """Reversed hybrid: NVFP4 rowwise (fprop) + FP8 columnwise (backward). + + Exercises NVFP4×NVFP4 in the fprop (TN) GEMM and FP8×FP8 in the + dgrad (NN) and wgrad (NT) GEMMs — the opposite of the primary + FP8-row/NVFP4-col configuration. + """ + + def test_nvfp4_row_fp8_col_forward_only(self): + """Forward (TN) with NVFP4×NVFP4 rowwise succeeds.""" + torch.manual_seed(99) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + + def factory(role): + if role in ("linear_input", "linear_weight"): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + return None + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with torch.no_grad(): + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + + def test_nvfp4_row_fp8_col_full_fwd_bwd(self): + """Full fwd+bwd with NVFP4 rowwise (fprop) + FP8 columnwise (backward).""" + torch.manual_seed(99) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", + dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if role in ("linear_input", "linear_weight"): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return _make_fp8_quantizer() + return None + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), "Output contains NaN" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + + +@requires_fp8 +class TestHybridMixedWithNonHybrid: + """Only one operand is hybrid; the other uses a plain TE quantizer. + + Exercises _unwrap_hybrid passthrough for the non-hybrid operand. + All roles must use compatible scaling modes for each GEMM: + fprop (TN): all rowwise formats must match + dgrad (NN): weight rowwise must match grad_output rowwise + wgrad (NT): input columnwise must match grad_output columnwise + """ + + def test_hybrid_input_plain_weight_fwd_bwd(self): + """Input is hybrid (FP8 row / FP8 col), weight + grad_output plain FP8. + + Wgrad columnwise: FP8 (input.col) × FP8 (grad_output.col) → compatible. + """ + torch.manual_seed(77) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", + dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if role == "linear_input": + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + if role == "linear_weight": + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ) + if role in ("linear_grad_output", "linear_grad_input"): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda", + ) + return None + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert not torch.isnan(out).any() + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + + def test_plain_input_hybrid_weight_fwd_bwd(self): + """Input is plain FP8, weight is hybrid (FP8 row / FP8 col).""" + torch.manual_seed(88) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", + dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if role == "linear_input": + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ) + if role == "linear_weight": + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda", + ) + return None + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert not torch.isnan(out).any() + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + + +# --------------------------------------------------------------------------- +# Parametrized cross-format tests (stateless quantizers) +# --------------------------------------------------------------------------- + +def _make_mxfp8_quantizer(*, rowwise=True, columnwise=True): + return MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_mxfp8_quantizer_e5m2(*, rowwise=True, columnwise=True): + return MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E5M2, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_block_quantizer(*, rowwise=True, columnwise=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def _make_block_quantizer_e5m2(*, rowwise=True, columnwise=True): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E5M2, + rowwise=rowwise, + columnwise=columnwise, + ) + + +# (fwd_e4m3_factory, bwd_e5m2_factory, skip_condition, skip_reason) +_QUANTIZER_CONFIGS = { + "fp8_current": ( + _make_fp8_quantizer, + lambda **kw: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda", **kw), + not fp8_available, + f"FP8: {reason_for_no_fp8}", + ), + "mxfp8": ( + _make_mxfp8_quantizer, + _make_mxfp8_quantizer_e5m2, + not mxfp8_available, + f"MXFP8: {reason_for_no_mxfp8}", + ), + "block_fp8": ( + _make_block_quantizer, + _make_block_quantizer_e5m2, + not fp8_block_scaling_available, + reason_for_no_fp8_block_scaling, + ), + "nvfp4": ( + _make_nvfp4_quantizer, + None, # NVFP4 has no E5M2 variant + not (fp8_available and nvfp4_available), + f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), +} + + +def _build_cross_format_params(): + """Build parametrize list for all stateless cross-format hybrid combos.""" + combos = [ + ("fp8_current", "mxfp8"), + ("fp8_current", "nvfp4"), + ("fp8_current", "block_fp8"), + ("mxfp8", "fp8_current"), + ("mxfp8", "mxfp8"), + ("mxfp8", "nvfp4"), + ("mxfp8", "block_fp8"), + ("block_fp8", "fp8_current"), + ("block_fp8", "mxfp8"), + ("block_fp8", "nvfp4"), + ("block_fp8", "block_fp8"), + ("nvfp4", "fp8_current"), + ("nvfp4", "mxfp8"), + ("nvfp4", "block_fp8"), + ] + params = [] + for row, col in combos: + row_cfg = _QUANTIZER_CONFIGS[row] + col_cfg = _QUANTIZER_CONFIGS[col] + hw_skip = row_cfg[2] or col_cfg[2] + hw_reason = "; ".join(filter(None, [row_cfg[3] if row_cfg[2] else "", + col_cfg[3] if col_cfg[2] else ""])) + marks = [] + if hw_skip: + marks.append(pytest.mark.skipif(True, reason=hw_reason or "N/A")) + params.append(pytest.param(row, col, id=f"{row}_row_x_{col}_col", marks=marks)) + return params + + +class TestHybridCrossFormatParametrized: + """Parametrized fwd+bwd over all stateless quantizer cross-format pairs.""" + + @pytest.mark.parametrize("row_name,col_name", _build_cross_format_params()) + def test_fwd_bwd(self, row_name, col_name): + torch.manual_seed(42) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", + dtype=torch.bfloat16, requires_grad=True) + + row_cfg = _QUANTIZER_CONFIGS[row_name] + col_cfg = _QUANTIZER_CONFIGS[col_name] + make_row_e4m3 = row_cfg[0] + make_col_e4m3 = col_cfg[0] + make_col_grad = col_cfg[1] if col_cfg[1] is not None else col_cfg[0] + + def factory(role): + if role in ("linear_input", "linear_weight"): + return HybridQuantizer( + rowwise_quantizer=make_row_e4m3(), + columnwise_quantizer=make_col_e4m3(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return make_col_grad() + return None + + mixed_recipe = recipe.CustomRecipe(qfactory=factory) + with autocast(enabled=True, recipe=mixed_recipe): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), f"Output NaN ({row_name} row × {col_name} col)" + assert not torch.isinf(out).any(), f"Output Inf ({row_name} row × {col_name} col)" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert not torch.isnan(inp.grad).any(), ( + f"Input grad NaN ({row_name} row × {col_name} col)" + ) + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), ( + f"Gradient for '{name}' NaN ({row_name} row × {col_name} col)" + ) + + +# --------------------------------------------------------------------------- +# 3-format hybrid: different quantization for fprop, dgrad, wgrad +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (fp8_available and mxfp8_available and nvfp4_available), + reason="Requires FP8 + MXFP8 + NVFP4", +) +class TestHybridThreeFormats: + """Three distinct formats: FormatA (fprop), FormatB (dgrad), FormatC (wgrad). + + Per-operand unwrap selects the correct sub-storage per GEMM: + fprop TN: weight.row(A) × input.row(A) → FormatA × FormatA + dgrad NN: weight.col(B) × grad_output.row(B) → FormatB × FormatB + wgrad NT: input.col(C) × grad_output.col(C) → FormatC × FormatC + + grad_output is itself hybrid (FormatB row + FormatC col) when B ≠ C. + """ + + def test_fp8_fprop_mxfp8_dgrad_nvfp4_wgrad(self): + """FP8 current (fprop) + MXFP8 (dgrad) + NVFP4 (wgrad).""" + torch.manual_seed(300) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", + dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if role == "linear_weight": + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_mxfp8_quantizer(), + ) + if role == "linear_input": + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return HybridQuantizer( + rowwise_quantizer=_make_mxfp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + return None + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), "Output contains NaN" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + + def test_nvfp4_fprop_fp8_dgrad_mxfp8_wgrad(self): + """NVFP4 (fprop) + FP8 current (dgrad) + MXFP8 (wgrad).""" + torch.manual_seed(301) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", + dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if role == "linear_weight": + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + if role == "linear_input": + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_mxfp8_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_mxfp8_quantizer(), + ) + return None + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + + assert out.shape == (batch, out_features) + assert not torch.isnan(out).any(), "Output contains NaN" + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient is None" + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + + def test_same_dgrad_wgrad_reduces_to_plain_grad(self): + """When dgrad format == wgrad format, grad_output can be a plain quantizer.""" + torch.manual_seed(302) + in_features, out_features, batch = 128, 128, 32 + + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(batch, in_features, device="cuda", + dtype=torch.bfloat16, requires_grad=True) + + def factory(role): + if role == "linear_weight": + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_mxfp8_quantizer(), + ) + if role == "linear_input": + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_mxfp8_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return _make_mxfp8_quantizer() + return None + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + out = model(inp) + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"Gradient for '{name}' is None" + + +# --------------------------------------------------------------------------- +# All-modules test: hybrid quantization through every TE module type +# --------------------------------------------------------------------------- + + +def _make_hybrid_fp8_factory(): + """Factory returning HybridQuantizer(FP8 row + FP8 col) for fwd roles, + plain FP8 E5M2 for bwd roles.""" + def factory(role): + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda", + ) + return Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", + ) + return factory + + +@requires_fp8 +class TestHybridAllModules: + """Hybrid quantization through all TE module types (not just Linear). + + Uses FP8 in both hybrid directions so the test validates module integration + without introducing cross-format scaling-mode concerns. + """ + + hidden_size = 128 + ffn_hidden_size = 128 + num_heads = 4 + batch = 16 + seq_len = 8 + + def _run_fwd_bwd(self, model, inp): + hybrid_recipe = recipe.CustomRecipe(qfactory=_make_hybrid_fp8_factory()) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + loss = out.float().sum() + loss.backward() + + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + assert inp.grad is not None, "Input gradient is None" + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + for name, p in model.named_parameters(): + if p.requires_grad: + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + + def test_linear(self): + torch.manual_seed(500) + model = Linear( + self.hidden_size, self.ffn_hidden_size, params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, self.hidden_size, device="cuda", + dtype=torch.bfloat16, requires_grad=True, + ) + self._run_fwd_bwd(model, inp) + + def test_layernorm_linear(self): + torch.manual_seed(501) + model = LayerNormLinear( + self.hidden_size, self.ffn_hidden_size, params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, self.hidden_size, device="cuda", + dtype=torch.bfloat16, requires_grad=True, + ) + self._run_fwd_bwd(model, inp) + + def test_layernorm_mlp(self): + torch.manual_seed(502) + model = LayerNormMLP( + hidden_size=self.hidden_size, ffn_hidden_size=self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, self.hidden_size, device="cuda", + dtype=torch.bfloat16, requires_grad=True, + ) + self._run_fwd_bwd(model, inp) + + def test_grouped_linear(self): + torch.manual_seed(504) + num_gemms = 3 + model = GroupedLinear( + num_gemms, self.hidden_size, self.ffn_hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.batch, self.hidden_size, device="cuda", + dtype=torch.bfloat16, requires_grad=True, + ) + base = self.batch // num_gemms + rem = self.batch % num_gemms + m_splits = [base + (1 if i < rem else 0) for i in range(num_gemms)] + + hybrid_recipe = recipe.CustomRecipe(qfactory=_make_hybrid_fp8_factory()) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp, m_splits) + loss = out.float().sum() + loss.backward() + + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + assert inp.grad is not None, "Input gradient is None" + assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" + for name, p in model.named_parameters(): + if p.requires_grad: + assert p.grad is not None, f"Gradient for '{name}' is None" + assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + + def test_transformer_layer(self): + torch.manual_seed(503) + model = TransformerLayer( + self.hidden_size, self.ffn_hidden_size, self.num_heads, + hidden_dropout=0.0, attention_dropout=0.0, + fuse_qkv_params=True, params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + self.seq_len, self.batch, self.hidden_size, device="cuda", + dtype=torch.bfloat16, requires_grad=True, + ) + self._run_fwd_bwd(model, inp) diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu index 3a8536587c2..9c5cdb91efe 100644 --- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu @@ -161,13 +161,15 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) static_assert(std::is_same::value); const CType scale_inv = 1.0f / block_tile_scale; - size_t row_idx = tile_id_y; - size_t col_idx = tile_id_x; - tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + if (tile_scales_inv_c != nullptr) { + size_t row_idx = tile_id_y; + size_t col_idx = tile_id_x; + tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + } if constexpr (kReturnTranspose) { - row_idx = tile_id_x; - col_idx = tile_id_y; + size_t row_idx = tile_id_x; + size_t col_idx = tile_id_y; tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv; } } @@ -189,7 +191,9 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) thrd_tile_out_trans[j].data.elt[i] = scaled_elt; } } - tmp_output_c.store_to(output_c + thread_tile_start_idx + i * row_length); + if (output_c != nullptr) { + tmp_output_c.store_to(output_c + thread_tile_start_idx + i * row_length); + } } // Step 4: store transpose into shared memory @@ -388,13 +392,15 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose static_assert(std::is_same::value); const CType scale_inv = 1.0f / block_tile_scale; - size_t row_idx = tile_id_y; - size_t col_idx = tile_id_x; - tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + if (tile_scales_inv_c != nullptr) { + size_t row_idx = tile_id_y; + size_t col_idx = tile_id_x; + tile_scales_inv_c[row_idx * scale_stride_y + col_idx * scale_stride_x] = scale_inv; + } if constexpr (kReturnTranspose) { - row_idx = tile_id_x; - col_idx = tile_id_y; + size_t row_idx = tile_id_x; + size_t col_idx = tile_id_y; tile_scales_inv_t[row_idx * scale_t_stride_y + col_idx * scale_t_stride_x] = scale_inv; } } @@ -433,8 +439,10 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) block_scaled_cast_transpose thrd_tile_out_trans[j].data.elt[i] = scaled_elt; } } - tmp_output_c.store_to_elts(output_c + thread_tile_start_idx + i * row_length, 0, - thread_tile_ncols); + if (output_c != nullptr) { + tmp_output_c.store_to_elts(output_c + thread_tile_start_idx + i * row_length, 0, + thread_tile_ncols); + } } if constexpr (kReturnTranspose) { @@ -492,19 +500,26 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor "with MXFP8, which requires using power of two scaling factors."); } - NVTE_CHECK(input.shape == output.shape, "Input and output must have the same shape."); + const bool return_identity = output.dptr != nullptr; + if (return_identity) { + NVTE_CHECK(input.shape == output.shape, "Input and output must have the same shape."); + } + NVTE_CHECK(return_identity || return_transpose, + "At least one of rowwise or columnwise output must be requested."); const size_t row_length = input.shape.size() > 0 ? input.shape.back() : 1; size_t num_rows = 1; for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) { num_rows *= input.shape.at(i); } - NVTE_CHECK(scale_inv.shape.size() == 2, "scale_inv must have 2 dimensions."); - - size_t scale_k = scale_inv.shape[1]; - - const size_t scale_stride_x = 1; - const size_t scale_stride_y = scale_k; + size_t scale_k = 0; + const size_t scale_stride_x = return_identity ? 1 : 0; + size_t scale_stride_y = 0; + if (return_identity) { + NVTE_CHECK(scale_inv.shape.size() == 2, "scale_inv must have 2 dimensions."); + scale_k = scale_inv.shape[1]; + scale_stride_y = scale_k; + } size_t scale_t_stride_x = 0; size_t scale_t_stride_y = 0; @@ -522,7 +537,10 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor ") and output_t (shape=", output_t.shape, ") have incompatible dims."); } } - NVTE_CHECK(output.dtype == output_t.dtype, "output and output_t need to have the same type."); + if (return_identity) { + NVTE_CHECK(output.dtype == output_t.dtype, + "output and output_t need to have the same type."); + } NVTE_CHECK(scale_inv_t.shape.size() == 2, "scale_inv_t must have 2 dimensions."); @@ -530,6 +548,8 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor scale_t_stride_y = scale_inv_t.shape[1]; } + const auto out_dtype = return_identity ? output.dtype : output_t.dtype; + const size_t num_blocks_x = DIVUP(row_length, BLOCK_TILE_DIM); const size_t num_blocks_y = DIVUP(num_rows, BLOCK_TILE_DIM); @@ -537,7 +557,7 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor input.dtype, InputType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - output.dtype, OutputType, + out_dtype, OutputType, TRANSFORMER_ENGINE_SWITCH_CONDITION( return_transpose, kReturnTranspose, diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index cd18ca75ade..6c31b076c1c 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -81,6 +81,9 @@ from transformer_engine.pytorch.tensor import MXFP8Tensor from transformer_engine.pytorch.tensor import Float8BlockwiseQTensor from transformer_engine.pytorch.tensor import NVFP4Tensor +from transformer_engine.pytorch.tensor import HybridQuantizer +from transformer_engine.pytorch.tensor import HybridQuantizedTensorStorage +from transformer_engine.pytorch.tensor import HybridQuantizedTensor try: torch._dynamo.config.error_on_nested_jit_trace = False diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 115569ccbaa..9991ffed191 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -16,6 +16,7 @@ from ..quantized_tensor import Quantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..tensor.utils import is_custom +from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizer @@ -69,6 +70,36 @@ def validate_gemm_scale(scale: Optional[float], required: bool) -> float: return 0.0 +def _unwrap_hybrid_A(tensor, layout): + """Extract the direction-appropriate native sub-storage for GEMM operand A. + + Operand A's data direction is determined by its transpose flag (layout[0]): + T (transposed) → rowwise sub-storage (.data consumed by C++) + N (not-transposed) → columnwise sub-storage (.columnwise_data consumed by C++) + For non-hybrid tensors this is a no-op passthrough. + """ + if not isinstance(tensor, HybridQuantizedTensorStorage): + return tensor + if layout[0] == "T": + return tensor.rowwise_sub_storage + return tensor.columnwise_sub_storage + + +def _unwrap_hybrid_B(tensor, layout): + """Extract the direction-appropriate native sub-storage for GEMM operand B. + + Operand B's data direction is determined by its transpose flag (layout[1]): + N (not-transposed) → rowwise sub-storage (.data consumed by C++) + T (transposed) → columnwise sub-storage (.columnwise_data consumed by C++) + For non-hybrid tensors this is a no-op passthrough. + """ + if not isinstance(tensor, HybridQuantizedTensorStorage): + return tensor + if layout[1] == "N": + return tensor.rowwise_sub_storage + return tensor.columnwise_sub_storage + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -95,6 +126,9 @@ def general_gemm( transa = layout[0] == "T" transb = layout[1] == "T" + A = _unwrap_hybrid_A(A, layout) + B = _unwrap_hybrid_B(B, layout) + alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) workspace = get_cublas_workspace(A.device.index, ub is not None, False) @@ -204,6 +238,9 @@ def general_grouped_gemm( """ num_gemms = len(A) + A = [_unwrap_hybrid_A(a, layout) for a in A] + B = [_unwrap_hybrid_B(b, layout) for b in B] + transa = layout[0] == "T" transb = layout[1] == "T" diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 28da4873f08..1124ce80038 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -42,6 +42,7 @@ from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer +from ..tensor.hybrid_tensor import HybridQuantizer from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage @@ -1258,8 +1259,9 @@ def grad_output_preprocess( ): grad_bias = grad_output.dequantize().view(-1, grad_output.shape[-1]).sum(dim=0) else: - if isinstance(quantizer, Float8BlockQuantizer): - # unfuse bgrad for now until cast_transpose + dgrad calculation is ready for Float8BlockQuantizer. + if isinstance(quantizer, (Float8BlockQuantizer, HybridQuantizer)): + # Float8BlockQuantizer: unfused until cast_transpose + dgrad is ready. + # HybridQuantizer: tex.bgrad_quantize doesn't recognize hybrid quantizers. grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0) else: grad_bias, grad_output = tex.bgrad_quantize(grad_output, quantizer) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 30c1dbf4089..26c8e65af03 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -51,9 +51,45 @@ prepare_for_saving, restore_from_saved, ) +from ..tensor.hybrid_tensor import HybridQuantizer from ...debug.pytorch.debug_quantization import DebugQuantizer from ...debug.pytorch.debug_state import TEDebugState + +def _has_hybrid_quantizer(quantizers): + """Check if any quantizer in the list is a HybridQuantizer.""" + return any(isinstance(q, HybridQuantizer) for q in quantizers if q is not None) + + +def _hybrid_split_quantize(tensor, m_splits, quantizers): + """Grouped split+quantize for HybridQuantizer lists. + + Runs tex.split_quantize twice (once per direction with the native + sub-quantizers), then zips the results into HybridQuantizedTensorStorage. + Non-hybrid quantizers in the list fall back to per-split Python quantize. + """ + from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage + + row_quantizers = [q.rowwise_quantizer for q in quantizers] + col_quantizers = [q.columnwise_quantizer for q in quantizers] + + row_results = tex.split_quantize(tensor, m_splits, row_quantizers) + col_results = tex.split_quantize(tensor, m_splits, col_quantizers) + + return [ + HybridStorage( + rowwise_storage=row, + columnwise_storage=col, + rowwise_quantizer=rq, + columnwise_quantizer=cq, + quantizer=q, + fake_dtype=tensor.dtype, + ) + for row, col, rq, cq, q in zip( + row_results, col_results, row_quantizers, col_quantizers, quantizers, + ) + ] + __all__ = ["GroupedLinear"] @@ -144,7 +180,8 @@ def forward( ) inp_view = inp.reshape(-1, in_features) inputmats: list - if fp8 and not debug: + hybrid = _has_hybrid_quantizer(input_quantizers) + if fp8 and not debug and not hybrid: # Disable bulk allocation when CPU offloading is active: offloading skips small # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. @@ -154,6 +191,8 @@ def forward( input_quantizers, disable_bulk_allocation=cpu_offloading, ) + elif fp8 and hybrid: + inputmats = _hybrid_split_quantize(inp_view, m_splits, input_quantizers) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, input_quantizers, m_splits, activation_dtype @@ -338,7 +377,8 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) grad_output = [None] * ctx.num_gemms grad_biases = [None] * ctx.num_gemms - if ctx.fp8 and not ctx.debug: + grad_output_hybrid = _has_hybrid_quantizer(ctx.grad_output_quantizers) + if ctx.fp8 and not ctx.debug and not grad_output_hybrid: if ctx.use_bias: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) recipe = ctx.fp8_recipe @@ -365,6 +405,14 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ctx.m_splits, ctx.grad_output_quantizers, ) + elif ctx.fp8 and grad_output_hybrid: + if ctx.use_bias: + grad_output_mats = torch.split(grad_output_view, ctx.m_splits) + for i in range(ctx.num_gemms): + grad_biases[i] = grad_output_mats[i].sum(dim=0) + grad_output = _hybrid_split_quantize( + grad_output_view, ctx.m_splits, ctx.grad_output_quantizers, + ) elif ctx.debug: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) for i in range(ctx.num_gemms): @@ -451,8 +499,11 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], else: input_quantizer.set_usage(rowwise=False, columnwise=True) inputmats: list - if ctx.fp8 and not ctx.debug: + input_hybrid = _has_hybrid_quantizer(ctx.input_quantizers) + if ctx.fp8 and not ctx.debug and not input_hybrid: inputmats = tex.split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) + elif ctx.fp8 and input_hybrid: + inputmats = _hybrid_split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) elif ctx.debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index d775dc3e8ed..b1d78afc848 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -64,6 +64,7 @@ ) from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.hybrid_tensor import HybridQuantizer from ..cpu_offload import ( is_cpu_offload_enabled, start_offload, @@ -206,12 +207,14 @@ def forward( # Avoid quantized norm kernel if norm output will be returned # or if a gather of ln_out must be in high precision. custom = is_custom(input_quantizer) + hybrid = isinstance(input_quantizer, HybridQuantizer) with_quantized_norm = ( fp8 and not debug and not return_layernorm_output and not return_layernorm_output_gathered and not custom # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() + and not hybrid ) # Apply normalization diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 037fb6c858c..219c61ddd1b 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -69,6 +69,7 @@ from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer +from ..tensor.hybrid_tensor import HybridQuantizer from ._common import apply_normalization, WeightGradStore from ..cpu_offload import ( is_cpu_offload_enabled, @@ -390,12 +391,14 @@ def _forward( # for debug: : layernorm output = High precision to enable processing of this norm custom = is_custom(fc1_input_quantizer) + hybrid = isinstance(fc1_input_quantizer, HybridQuantizer) with_quantized_norm = ( fp8 and not debug and not return_layernorm_output and not return_layernorm_output_gathered and not custom + and not hybrid ) # Apply normalization diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index 56680567002..fbf725047d8 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -18,11 +18,13 @@ from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from .storage.nvfp4_tensor_storage import NVFP4TensorStorage from .storage.grouped_tensor_storage import GroupedTensorStorage +from .storage.hybrid_tensor_storage import HybridQuantizedTensorStorage from .float8_tensor import Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer from .grouped_tensor import GroupedTensor +from .hybrid_tensor import HybridQuantizedTensor, HybridQuantizer from .utils import cast_master_weights_to_fp8, replace_raw_data __all__ = [ @@ -32,18 +34,21 @@ "MXFP8Quantizer", "Float8BlockQuantizer", "NVFP4Quantizer", + "HybridQuantizer", "QuantizedTensorStorage", "Float8TensorStorage", "MXFP8TensorStorage", "Float8BlockwiseQTensorStorage", "NVFP4TensorStorage", "GroupedTensorStorage", + "HybridQuantizedTensorStorage", "QuantizedTensor", "Float8Tensor", "MXFP8Tensor", "Float8BlockwiseQTensor", "NVFP4Tensor", "GroupedTensor", + "HybridQuantizedTensor", "prepare_for_saving", "restore_from_saved", ] @@ -95,5 +100,7 @@ def get_all_tensor_types(): NVFP4TensorStorage, GroupedTensor, GroupedTensorStorage, + HybridQuantizedTensor, + HybridQuantizedTensorStorage, ] return all_tensor_types diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py new file mode 100644 index 00000000000..c47cd925750 --- /dev/null +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -0,0 +1,193 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tensor class with hybrid quantized data (different formats for rowwise vs columnwise)""" + +from __future__ import annotations +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch + +from .storage.hybrid_tensor_storage import HybridQuantizedTensorStorage +from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer + +aten = torch.ops.aten + + +class HybridQuantizer(Quantizer): + """Quantizer that composes two existing quantizers for different directions. + + Performs two-pass quantization: the rowwise_quantizer produces rowwise + quantized data and the columnwise_quantizer produces columnwise quantized + data. The results are wrapped in a HybridQuantizedTensor. + + Parameters + ---------- + rowwise_quantizer : Quantizer + Quantizer for the rowwise direction (e.g. MXFP8Quantizer). + columnwise_quantizer : Quantizer + Quantizer for the columnwise direction (e.g. NVFP4Quantizer). + + """ + + rowwise_quantizer: Quantizer + columnwise_quantizer: Quantizer + + def __init__( + self, + *, + rowwise_quantizer: Quantizer, + columnwise_quantizer: Quantizer, + ) -> None: + super().__init__(rowwise=True, columnwise=True) + self.rowwise_quantizer = rowwise_quantizer + self.columnwise_quantizer = columnwise_quantizer + + # Pin each sub-quantizer to its designated direction + self.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) + self.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: + rowwise_result = self.rowwise_quantizer.quantize(tensor) + columnwise_result = self.columnwise_quantizer.quantize(tensor) + + if self.internal: + return HybridQuantizedTensorStorage( + rowwise_storage=rowwise_result, + columnwise_storage=columnwise_result, + rowwise_quantizer=self.rowwise_quantizer, + columnwise_quantizer=self.columnwise_quantizer, + quantizer=self, + fake_dtype=tensor.dtype, + ) + + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=rowwise_result, + columnwise_storage=columnwise_result, + rowwise_quantizer=self.rowwise_quantizer, + columnwise_quantizer=self.columnwise_quantizer, + quantizer=self, + ) + + def make_empty( + self, + shape: Iterable[int], + *, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + requires_grad: bool = False, + pin_memory: bool = False, + ) -> HybridQuantizedTensor: + self.rowwise_quantizer.internal = True + rowwise_empty = self.rowwise_quantizer.make_empty( + shape, dtype=dtype, device=device, pin_memory=pin_memory, + ) + self.rowwise_quantizer.internal = False + + self.columnwise_quantizer.internal = True + columnwise_empty = self.columnwise_quantizer.make_empty( + shape, dtype=dtype, device=device, pin_memory=pin_memory, + ) + self.columnwise_quantizer.internal = False + + return HybridQuantizedTensor( + shape=shape, + dtype=dtype, + requires_grad=requires_grad, + device=device, + rowwise_storage=rowwise_empty, + columnwise_storage=columnwise_empty, + rowwise_quantizer=self.rowwise_quantizer, + columnwise_quantizer=self.columnwise_quantizer, + quantizer=self, + ) + + def set_usage( + self, *, rowwise: Optional[bool] = None, columnwise: Optional[bool] = None + ) -> None: + super().set_usage(rowwise=rowwise, columnwise=columnwise) + + def _get_compatible_recipe(self): + return None + + +class HybridQuantizedTensor(HybridQuantizedTensorStorage, QuantizedTensor): + """Quantized tensor holding data in two different formats per direction. + + The tensor presents as having a standard, higher-precision dtype, but + internally stores rowwise data in one quantized format and columnwise + data in another. + + Parameters + ---------- + shape : iterable of int + Tensor dimensions. + dtype : torch.dtype + Nominal tensor datatype. + rowwise_storage : QuantizedTensorStorage + Sub-storage for rowwise quantized data. + columnwise_storage : QuantizedTensorStorage + Sub-storage for columnwise quantized data. + rowwise_quantizer : Quantizer, optional + Quantizer used for the rowwise sub-storage. + columnwise_quantizer : Quantizer, optional + Quantizer used for the columnwise sub-storage. + quantizer : HybridQuantizer, optional + Parent hybrid quantizer. + requires_grad : bool, default = False + Whether to compute gradients for this tensor. + + """ + + def __new__( + cls, + *args, + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + rowwise_quantizer: Optional[Quantizer] = None, + columnwise_quantizer: Optional[Quantizer] = None, + quantizer: Optional[Quantizer] = None, + **kwargs, + ): + instance = super().__new__( + cls, + *args, + rowwise_storage=rowwise_storage, + columnwise_storage=columnwise_storage, + rowwise_quantizer=rowwise_quantizer, + columnwise_quantizer=columnwise_quantizer, + quantizer=quantizer, + **kwargs, + ) + return instance + + def __repr__(self, *, tensor_contents=None): + row_type = type(self._rowwise_storage).__name__ if self._rowwise_storage is not None else "None" + col_type = type(self._columnwise_storage).__name__ if self._columnwise_storage is not None else "None" + return ( + f"HybridQuantizedTensor(" + f"rowwise={row_type}, " + f"columnwise={col_type}, " + f"dtype={self.dtype})" + ) + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if dtype is None: + dtype = self.dtype + return HybridQuantizedTensorStorage.dequantize(self, dtype=dtype) + + def detach(self) -> HybridQuantizedTensor: + return HybridQuantizedTensor.make_like(self) + + def get_metadata(self) -> Dict[str, Any]: + return HybridQuantizedTensorStorage.get_metadata(self) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if func == aten.detach.default: + return args[0].detach() + + return super().__torch_dispatch__(func, types, args, kwargs) diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py new file mode 100644 index 00000000000..3668809e73e --- /dev/null +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -0,0 +1,157 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Mixin class holding data specific for HybridQuantizedTensor""" + +from __future__ import annotations +from typing import Any, Dict, Optional, Tuple + +import torch + +from ...quantized_tensor import QuantizedTensorStorage, Quantizer + + +class HybridQuantizedTensorStorage(QuantizedTensorStorage): + """Storage that composes two QuantizedTensorStorage instances. + + One sub-storage provides rowwise quantized data and the other provides + columnwise quantized data. This enables mixed-precision quantization + where, for example, rowwise data is FP8 and columnwise data is FP4. + + """ + + _rowwise_storage: Optional[QuantizedTensorStorage] + _columnwise_storage: Optional[QuantizedTensorStorage] + _rowwise_quantizer: Optional[Quantizer] + _columnwise_quantizer: Optional[Quantizer] + _quantizer: Optional[Quantizer] + + def __new__( + cls, + *args, + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + rowwise_quantizer: Optional[Quantizer] = None, + columnwise_quantizer: Optional[Quantizer] = None, + quantizer: Optional[Quantizer] = None, + fake_dtype: Optional[torch.dtype] = None, + **kwargs, + ): + if cls is HybridQuantizedTensorStorage: + instance = object.__new__(cls) + instance._dtype = fake_dtype if fake_dtype is not None else torch.float32 + else: + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) + + instance._rowwise_storage = rowwise_storage + instance._columnwise_storage = columnwise_storage + instance._rowwise_quantizer = rowwise_quantizer + instance._columnwise_quantizer = columnwise_quantizer + instance._quantizer = quantizer + return instance + + @property + def rowwise_sub_storage(self) -> Optional[QuantizedTensorStorage]: + """The sub-storage providing rowwise quantized data.""" + return self._rowwise_storage + + @property + def columnwise_sub_storage(self) -> Optional[QuantizedTensorStorage]: + """The sub-storage providing columnwise quantized data.""" + return self._columnwise_storage + + def update_usage( + self, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ): + if rowwise_usage is not None and not rowwise_usage: + self._rowwise_storage = None + if columnwise_usage is not None and not columnwise_usage: + self._columnwise_storage = None + + def get_usages(self) -> Dict[str, bool]: + return { + "rowwise": self._rowwise_storage is not None, + "columnwise": self._columnwise_storage is not None, + } + + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], HybridQuantizedTensorStorage]: + tensors = [] + if self._rowwise_storage is not None: + row_tensors, _ = self._rowwise_storage.prepare_for_saving() + tensors.extend(row_tensors) + if self._columnwise_storage is not None: + col_tensors, _ = self._columnwise_storage.prepare_for_saving() + tensors.extend(col_tensors) + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + if self._rowwise_storage is not None: + tensors = self._rowwise_storage.restore_from_saved(tensors) + if self._columnwise_storage is not None: + tensors = self._columnwise_storage.restore_from_saved(tensors) + return tensors + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if dtype is None: + dtype = self._dtype + if self._rowwise_storage is not None: + return self._rowwise_storage.dequantize(dtype=dtype) + if self._columnwise_storage is not None: + return self._columnwise_storage.dequantize(dtype=dtype) + raise RuntimeError("HybridQuantizedTensorStorage has no data to dequantize") + + def get_data_tensors(self): + row_tensors = () + col_tensors = () + if self._rowwise_storage is not None: + result = self._rowwise_storage.get_data_tensors() + row_tensors = result if isinstance(result, tuple) else (result,) + if self._columnwise_storage is not None: + result = self._columnwise_storage.get_data_tensors() + col_tensors = result if isinstance(result, tuple) else (result,) + return row_tensors + col_tensors + + def size(self, *args, **kwargs): + if self._rowwise_storage is not None: + return self._rowwise_storage.size(*args, **kwargs) + if self._columnwise_storage is not None: + return self._columnwise_storage.size(*args, **kwargs) + raise RuntimeError("HybridQuantizedTensorStorage has no data") + + @property + def device(self): + if self._rowwise_storage is not None: + return self._rowwise_storage.device + if self._columnwise_storage is not None: + return self._columnwise_storage.device + raise RuntimeError("HybridQuantizedTensorStorage has no data") + + def view(self, shape: torch.Size): + raise NotImplementedError( + "HybridQuantizedTensorStorage does not support view operations" + ) + + def get_metadata(self) -> Dict[str, Any]: + return { + "rowwise_storage": self._rowwise_storage, + "columnwise_storage": self._columnwise_storage, + "rowwise_quantizer": self._rowwise_quantizer, + "columnwise_quantizer": self._columnwise_quantizer, + "quantizer": self._quantizer, + "fake_dtype": self._dtype, + } + + def __repr__(self): + return ( + f"HybridQuantizedTensorStorage(" + f"rowwise={type(self._rowwise_storage).__name__}, " + f"columnwise={type(self._columnwise_storage).__name__}, " + f"dtype={self._dtype})" + ) From 19acc5ead252d2998c07c62a21076e312d764916 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 31 Mar 2026 14:55:26 +0000 Subject: [PATCH 02/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 349 ++++++++++-------- .../quantize_transpose_square_blockwise.cu | 3 +- .../pytorch/module/grouped_linear.py | 15 +- .../pytorch/tensor/hybrid_tensor.py | 25 +- .../tensor/storage/hybrid_tensor_storage.py | 6 +- 5 files changed, 235 insertions(+), 163 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 96a28744f33..5967c0816a6 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -192,9 +192,7 @@ def test_dequantize_reversed_close_to_original(self, input_tensor): hq = _make_hybrid_quantizer_fp4_row_fp8_col() result = hq.quantize(input_tensor) dequantized = result.dequantize() - torch.testing.assert_close( - dequantized.float(), input_tensor.float(), rtol=0.5, atol=1.0 - ) + torch.testing.assert_close(dequantized.float(), input_tensor.float(), rtol=0.5, atol=1.0) def test_storage_dequantize(self, input_tensor): hq = _make_hybrid_quantizer_fp8_row_fp4_col() @@ -412,18 +410,22 @@ def hybrid_fp8_factory(role): if role in ("linear_input", "linear_weight", "linear_output"): return HybridQuantizer( rowwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ), columnwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ), ) if role in ("linear_grad_output", "linear_grad_input"): return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, device="cuda", + tex.DType.kFloat8E5M2, + device="cuda", ) return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ) hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_fp8_factory) @@ -433,25 +435,24 @@ def hybrid_fp8_factory(role): loss_hybrid.backward() # Forward outputs must be bitwise identical - assert torch.equal(out_ref, out_hybrid), ( - f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" - ) + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" # Input gradients must be bitwise identical assert inp_ref.grad is not None and inp_hybrid.grad is not None - assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( - f"Input grad mismatch: max diff = " - f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" - ) + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" # Parameter gradients must be bitwise identical ref_params = dict(model_ref.named_parameters()) hybrid_params = dict(model_hybrid.named_parameters()) for name, p_ref in ref_params.items(): p_hyb = hybrid_params[name] - assert p_ref.grad is not None and p_hyb.grad is not None, ( - f"Missing gradient for param '{name}'" - ) + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" assert torch.equal(p_ref.grad, p_hyb.grad), ( f"Param '{name}' grad mismatch: max diff = " f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" @@ -494,18 +495,17 @@ def hybrid_mxfp8_factory(role): out_hybrid = model_hybrid(inp_hybrid) out_hybrid.float().sum().backward() - assert torch.equal(out_ref, out_hybrid), ( - f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" - ) - assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( - f"Input grad mismatch: max diff = " - f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" - ) + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" for name, p_ref in dict(model_ref.named_parameters()).items(): p_hyb = dict(model_hybrid.named_parameters())[name] - assert p_ref.grad is not None and p_hyb.grad is not None, ( - f"Missing gradient for param '{name}'" - ) + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" assert torch.equal(p_ref.grad, p_hyb.grad), ( f"Param '{name}' grad mismatch: max diff = " f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" @@ -540,18 +540,21 @@ def hybrid_block_fp8_factory(role): if role in ("linear_grad_output", "linear_grad_input"): return Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, columnwise=True, + rowwise=True, + columnwise=True, block_scaling_dim=dim, ) return HybridQuantizer( rowwise_quantizer=Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, columnwise=True, + rowwise=True, + columnwise=True, block_scaling_dim=dim, ), columnwise_quantizer=Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, columnwise=True, + rowwise=True, + columnwise=True, block_scaling_dim=dim, ), ) @@ -561,18 +564,17 @@ def hybrid_block_fp8_factory(role): out_hybrid = model_hybrid(inp_hybrid) out_hybrid.float().sum().backward() - assert torch.equal(out_ref, out_hybrid), ( - f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" - ) - assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( - f"Input grad mismatch: max diff = " - f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" - ) + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" for name, p_ref in dict(model_ref.named_parameters()).items(): p_hyb = dict(model_hybrid.named_parameters())[name] - assert p_ref.grad is not None and p_hyb.grad is not None, ( - f"Missing gradient for param '{name}'" - ) + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" assert torch.equal(p_ref.grad, p_hyb.grad), ( f"Param '{name}' grad mismatch: max diff = " f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" @@ -627,18 +629,17 @@ def hybrid_nvfp4_factory(role): out_hybrid = model_hybrid(inp_hybrid) out_hybrid.float().sum().backward() - assert torch.equal(out_ref, out_hybrid), ( - f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" - ) - assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( - f"Input grad mismatch: max diff = " - f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" - ) + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" for name, p_ref in dict(model_ref.named_parameters()).items(): p_hyb = dict(model_hybrid.named_parameters())[name] - assert p_ref.grad is not None and p_hyb.grad is not None, ( - f"Missing gradient for param '{name}'" - ) + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" assert torch.equal(p_ref.grad, p_hyb.grad), ( f"Param '{name}' grad mismatch: max diff = " f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" @@ -680,18 +681,17 @@ def hybrid_nvfp4_all_roles_factory(role): out_hybrid = model_hybrid(inp_hybrid) out_hybrid.float().sum().backward() - assert torch.equal(out_ref, out_hybrid), ( - f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" - ) - assert torch.equal(inp_ref.grad, inp_hybrid.grad), ( - f"Input grad mismatch: max diff = " - f"{(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" - ) + assert torch.equal( + out_ref, out_hybrid + ), f"Forward mismatch: max diff = {(out_ref - out_hybrid).abs().max().item()}" + assert torch.equal( + inp_ref.grad, inp_hybrid.grad + ), f"Input grad mismatch: max diff = {(inp_ref.grad - inp_hybrid.grad).abs().max().item()}" for name, p_ref in dict(model_ref.named_parameters()).items(): p_hyb = dict(model_hybrid.named_parameters())[name] - assert p_ref.grad is not None and p_hyb.grad is not None, ( - f"Missing gradient for param '{name}'" - ) + assert ( + p_ref.grad is not None and p_hyb.grad is not None + ), f"Missing gradient for param '{name}'" assert torch.equal(p_ref.grad, p_hyb.grad), ( f"Param '{name}' grad mismatch: max diff = " f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" @@ -711,7 +711,11 @@ def test_linear_fwd_bwd_fp8_row_nvfp4_col(self): model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() inp = torch.randn( - batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True, + batch, + in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, ) def mixed_factory(role): @@ -778,7 +782,10 @@ def mixed_factory(role): # FP8/FP4 quantization introduces error, but the result should be # in the same ballpark as BF16 torch.testing.assert_close( - out_mixed.float(), out_bf16.float(), rtol=0.25, atol=0.5, + out_mixed.float(), + out_bf16.float(), + rtol=0.25, + atol=0.5, ) @@ -817,12 +824,14 @@ def test_B_nt_returns_columnwise(self, hybrid_tensor): def test_tn_sub_storage_type(self, hybrid_tensor): assert isinstance( - _unwrap_hybrid_A(hybrid_tensor, "TN"), (Float8TensorStorage, Float8Tensor), + _unwrap_hybrid_A(hybrid_tensor, "TN"), + (Float8TensorStorage, Float8Tensor), ) def test_nt_sub_storage_type(self, hybrid_tensor): assert isinstance( - _unwrap_hybrid_B(hybrid_tensor, "NT"), (NVFP4TensorStorage, NVFP4Tensor), + _unwrap_hybrid_B(hybrid_tensor, "NT"), + (NVFP4TensorStorage, NVFP4Tensor), ) def test_non_hybrid_passthrough(self): @@ -852,26 +861,30 @@ def _make_uniform_hybrid_factory(self): def factory(role): if role in ("linear_grad_output", "linear_grad_input"): return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, device="cuda", + tex.DType.kFloat8E5M2, + device="cuda", ) return HybridQuantizer( rowwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ), columnwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ), ) + return factory def test_bias_grad_matches_vanilla_fp8(self): torch.manual_seed(456) in_features, out_features, batch = 64, 64, 16 - model_ref = Linear(in_features, out_features, bias=True, - params_dtype=torch.bfloat16).cuda() - model_hybrid = Linear(in_features, out_features, bias=True, - params_dtype=torch.bfloat16).cuda() + model_ref = Linear(in_features, out_features, bias=True, params_dtype=torch.bfloat16).cuda() + model_hybrid = Linear( + in_features, out_features, bias=True, params_dtype=torch.bfloat16 + ).cuda() model_hybrid.load_state_dict(model_ref.state_dict()) base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) @@ -884,31 +897,32 @@ def test_bias_grad_matches_vanilla_fp8(self): # Hybrid inp_hyb = base_inp.clone().detach().requires_grad_(True) - with autocast(enabled=True, - recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory())): + with autocast( + enabled=True, recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory()) + ): out_hyb = model_hybrid(inp_hyb) out_hyb.float().sum().backward() ref_bias_grad = dict(model_ref.named_parameters())["bias"].grad hyb_bias_grad = dict(model_hybrid.named_parameters())["bias"].grad assert ref_bias_grad is not None and hyb_bias_grad is not None - assert torch.equal(ref_bias_grad, hyb_bias_grad), ( - f"Bias grad mismatch: max diff = " - f"{(ref_bias_grad - hyb_bias_grad).abs().max().item()}" - ) + assert torch.equal( + ref_bias_grad, hyb_bias_grad + ), f"Bias grad mismatch: max diff = {(ref_bias_grad - hyb_bias_grad).abs().max().item()}" def test_no_bias_fwd_bwd(self): """Linear with bias=False skips bgrad_quantize entirely.""" torch.manual_seed(42) in_features, out_features, batch = 64, 64, 16 - model = Linear(in_features, out_features, bias=False, - params_dtype=torch.bfloat16).cuda() - inp = torch.randn(batch, in_features, device="cuda", - dtype=torch.bfloat16, requires_grad=True) + model = Linear(in_features, out_features, bias=False, params_dtype=torch.bfloat16).cuda() + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) - with autocast(enabled=True, - recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory())): + with autocast( + enabled=True, recipe=recipe.CustomRecipe(qfactory=self._make_uniform_hybrid_factory()) + ): out = model(inp) out.float().sum().backward() @@ -932,8 +946,7 @@ def test_matching_columnwise_formats_succeed(self): torch.manual_seed(42) # NVFP4 GEMM requires dimensions ≥ 128 for cuBLAS support. model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, - requires_grad=True) + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) def factory(role): if role in ("linear_input", "linear_weight"): @@ -954,8 +967,7 @@ def test_mismatched_columnwise_formats_raise(self): """NVFP4 input × FP8 grad_output columnwise → cuBLAS rejects.""" torch.manual_seed(42) model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, - requires_grad=True) + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) def factory(role): if role in ("linear_input", "linear_weight"): @@ -965,7 +977,8 @@ def factory(role): ) if role in ("linear_grad_output", "linear_grad_input"): return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, device="cuda", + tex.DType.kFloat8E5M2, + device="cuda", ) return None @@ -1015,8 +1028,9 @@ def test_nvfp4_row_fp8_col_full_fwd_bwd(self): in_features, out_features, batch = 128, 128, 32 model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(batch, in_features, device="cuda", - dtype=torch.bfloat16, requires_grad=True) + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) def factory(role): if role in ("linear_input", "linear_weight"): @@ -1065,8 +1079,9 @@ def test_hybrid_input_plain_weight_fwd_bwd(self): in_features, out_features, batch = 128, 128, 32 model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(batch, in_features, device="cuda", - dtype=torch.bfloat16, requires_grad=True) + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) def factory(role): if role == "linear_input": @@ -1076,11 +1091,13 @@ def factory(role): ) if role == "linear_weight": return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ) if role in ("linear_grad_output", "linear_grad_input"): return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, device="cuda", + tex.DType.kFloat8E5M2, + device="cuda", ) return None @@ -1104,13 +1121,15 @@ def test_plain_input_hybrid_weight_fwd_bwd(self): in_features, out_features, batch = 128, 128, 32 model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(batch, in_features, device="cuda", - dtype=torch.bfloat16, requires_grad=True) + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) def factory(role): if role == "linear_input": return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ) if role == "linear_weight": return HybridQuantizer( @@ -1119,7 +1138,8 @@ def factory(role): ) if role in ("linear_grad_output", "linear_grad_input"): return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, device="cuda", + tex.DType.kFloat8E5M2, + device="cuda", ) return None @@ -1142,6 +1162,7 @@ def factory(role): # Parametrized cross-format tests (stateless quantizers) # --------------------------------------------------------------------------- + def _make_mxfp8_quantizer(*, rowwise=True, columnwise=True): return MXFP8Quantizer( fp8_dtype=tex.DType.kFloat8E4M3, @@ -1209,25 +1230,26 @@ def _build_cross_format_params(): ("fp8_current", "mxfp8"), ("fp8_current", "nvfp4"), ("fp8_current", "block_fp8"), - ("mxfp8", "fp8_current"), - ("mxfp8", "mxfp8"), - ("mxfp8", "nvfp4"), - ("mxfp8", "block_fp8"), - ("block_fp8", "fp8_current"), - ("block_fp8", "mxfp8"), - ("block_fp8", "nvfp4"), - ("block_fp8", "block_fp8"), - ("nvfp4", "fp8_current"), - ("nvfp4", "mxfp8"), - ("nvfp4", "block_fp8"), + ("mxfp8", "fp8_current"), + ("mxfp8", "mxfp8"), + ("mxfp8", "nvfp4"), + ("mxfp8", "block_fp8"), + ("block_fp8", "fp8_current"), + ("block_fp8", "mxfp8"), + ("block_fp8", "nvfp4"), + ("block_fp8", "block_fp8"), + ("nvfp4", "fp8_current"), + ("nvfp4", "mxfp8"), + ("nvfp4", "block_fp8"), ] params = [] for row, col in combos: row_cfg = _QUANTIZER_CONFIGS[row] col_cfg = _QUANTIZER_CONFIGS[col] hw_skip = row_cfg[2] or col_cfg[2] - hw_reason = "; ".join(filter(None, [row_cfg[3] if row_cfg[2] else "", - col_cfg[3] if col_cfg[2] else ""])) + hw_reason = "; ".join( + filter(None, [row_cfg[3] if row_cfg[2] else "", col_cfg[3] if col_cfg[2] else ""]) + ) marks = [] if hw_skip: marks.append(pytest.mark.skipif(True, reason=hw_reason or "N/A")) @@ -1244,8 +1266,9 @@ def test_fwd_bwd(self, row_name, col_name): in_features, out_features, batch = 128, 128, 32 model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(batch, in_features, device="cuda", - dtype=torch.bfloat16, requires_grad=True) + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) row_cfg = _QUANTIZER_CONFIGS[row_name] col_cfg = _QUANTIZER_CONFIGS[col_name] @@ -1275,14 +1298,12 @@ def factory(role): loss.backward() assert inp.grad is not None, "Input gradient is None" - assert not torch.isnan(inp.grad).any(), ( - f"Input grad NaN ({row_name} row × {col_name} col)" - ) + assert not torch.isnan(inp.grad).any(), f"Input grad NaN ({row_name} row × {col_name} col)" for name, p in model.named_parameters(): assert p.grad is not None, f"Gradient for '{name}' is None" - assert not torch.isnan(p.grad).any(), ( - f"Gradient for '{name}' NaN ({row_name} row × {col_name} col)" - ) + assert not torch.isnan( + p.grad + ).any(), f"Gradient for '{name}' NaN ({row_name} row × {col_name} col)" # --------------------------------------------------------------------------- @@ -1311,8 +1332,9 @@ def test_fp8_fprop_mxfp8_dgrad_nvfp4_wgrad(self): in_features, out_features, batch = 128, 128, 32 model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(batch, in_features, device="cuda", - dtype=torch.bfloat16, requires_grad=True) + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) def factory(role): if role == "linear_weight": @@ -1353,8 +1375,9 @@ def test_nvfp4_fprop_fp8_dgrad_mxfp8_wgrad(self): in_features, out_features, batch = 128, 128, 32 model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(batch, in_features, device="cuda", - dtype=torch.bfloat16, requires_grad=True) + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) def factory(role): if role == "linear_weight": @@ -1395,8 +1418,9 @@ def test_same_dgrad_wgrad_reduces_to_plain_grad(self): in_features, out_features, batch = 128, 128, 32 model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn(batch, in_features, device="cuda", - dtype=torch.bfloat16, requires_grad=True) + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) def factory(role): if role == "linear_weight": @@ -1433,23 +1457,29 @@ def factory(role): def _make_hybrid_fp8_factory(): """Factory returning HybridQuantizer(FP8 row + FP8 col) for fwd roles, plain FP8 E5M2 for bwd roles.""" + def factory(role): if role in ("linear_input", "linear_weight", "linear_output"): return HybridQuantizer( rowwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ), columnwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ), ) if role in ("linear_grad_output", "linear_grad_input"): return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, device="cuda", + tex.DType.kFloat8E5M2, + device="cuda", ) return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda", + tex.DType.kFloat8E4M3, + device="cuda", ) + return factory @@ -1486,34 +1516,48 @@ def _run_fwd_bwd(self, model, inp): def test_linear(self): torch.manual_seed(500) model = Linear( - self.hidden_size, self.ffn_hidden_size, params_dtype=torch.bfloat16, + self.hidden_size, + self.ffn_hidden_size, + params_dtype=torch.bfloat16, ).cuda() inp = torch.randn( - self.batch, self.hidden_size, device="cuda", - dtype=torch.bfloat16, requires_grad=True, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, ) self._run_fwd_bwd(model, inp) def test_layernorm_linear(self): torch.manual_seed(501) model = LayerNormLinear( - self.hidden_size, self.ffn_hidden_size, params_dtype=torch.bfloat16, + self.hidden_size, + self.ffn_hidden_size, + params_dtype=torch.bfloat16, ).cuda() inp = torch.randn( - self.batch, self.hidden_size, device="cuda", - dtype=torch.bfloat16, requires_grad=True, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, ) self._run_fwd_bwd(model, inp) def test_layernorm_mlp(self): torch.manual_seed(502) model = LayerNormMLP( - hidden_size=self.hidden_size, ffn_hidden_size=self.ffn_hidden_size, + hidden_size=self.hidden_size, + ffn_hidden_size=self.ffn_hidden_size, params_dtype=torch.bfloat16, ).cuda() inp = torch.randn( - self.batch, self.hidden_size, device="cuda", - dtype=torch.bfloat16, requires_grad=True, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, ) self._run_fwd_bwd(model, inp) @@ -1521,12 +1565,17 @@ def test_grouped_linear(self): torch.manual_seed(504) num_gemms = 3 model = GroupedLinear( - num_gemms, self.hidden_size, self.ffn_hidden_size, + num_gemms, + self.hidden_size, + self.ffn_hidden_size, params_dtype=torch.bfloat16, ).cuda() inp = torch.randn( - self.batch, self.hidden_size, device="cuda", - dtype=torch.bfloat16, requires_grad=True, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, ) base = self.batch // num_gemms rem = self.batch % num_gemms @@ -1550,12 +1599,20 @@ def test_grouped_linear(self): def test_transformer_layer(self): torch.manual_seed(503) model = TransformerLayer( - self.hidden_size, self.ffn_hidden_size, self.num_heads, - hidden_dropout=0.0, attention_dropout=0.0, - fuse_qkv_params=True, params_dtype=torch.bfloat16, + self.hidden_size, + self.ffn_hidden_size, + self.num_heads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, ).cuda() inp = torch.randn( - self.seq_len, self.batch, self.hidden_size, device="cuda", - dtype=torch.bfloat16, requires_grad=True, + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, ) self._run_fwd_bwd(model, inp) diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu index 9c5cdb91efe..02d64bcfffa 100644 --- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu @@ -538,8 +538,7 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor } } if (return_identity) { - NVTE_CHECK(output.dtype == output_t.dtype, - "output and output_t need to have the same type."); + NVTE_CHECK(output.dtype == output_t.dtype, "output and output_t need to have the same type."); } NVTE_CHECK(scale_inv_t.shape.size() == 2, "scale_inv_t must have 2 dimensions."); diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 26c8e65af03..4f35d1859f3 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -86,10 +86,15 @@ def _hybrid_split_quantize(tensor, m_splits, quantizers): fake_dtype=tensor.dtype, ) for row, col, rq, cq, q in zip( - row_results, col_results, row_quantizers, col_quantizers, quantizers, + row_results, + col_results, + row_quantizers, + col_quantizers, + quantizers, ) ] + __all__ = ["GroupedLinear"] @@ -411,7 +416,9 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], for i in range(ctx.num_gemms): grad_biases[i] = grad_output_mats[i].sum(dim=0) grad_output = _hybrid_split_quantize( - grad_output_view, ctx.m_splits, ctx.grad_output_quantizers, + grad_output_view, + ctx.m_splits, + ctx.grad_output_quantizers, ) elif ctx.debug: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) @@ -503,7 +510,9 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], if ctx.fp8 and not ctx.debug and not input_hybrid: inputmats = tex.split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) elif ctx.fp8 and input_hybrid: - inputmats = _hybrid_split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) + inputmats = _hybrid_split_quantize( + inp_view, ctx.m_splits, ctx.input_quantizers + ) elif ctx.debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index c47cd925750..6073b5d1083 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -83,13 +83,19 @@ def make_empty( ) -> HybridQuantizedTensor: self.rowwise_quantizer.internal = True rowwise_empty = self.rowwise_quantizer.make_empty( - shape, dtype=dtype, device=device, pin_memory=pin_memory, + shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, ) self.rowwise_quantizer.internal = False self.columnwise_quantizer.internal = True columnwise_empty = self.columnwise_quantizer.make_empty( - shape, dtype=dtype, device=device, pin_memory=pin_memory, + shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, ) self.columnwise_quantizer.internal = False @@ -165,13 +171,16 @@ def __new__( return instance def __repr__(self, *, tensor_contents=None): - row_type = type(self._rowwise_storage).__name__ if self._rowwise_storage is not None else "None" - col_type = type(self._columnwise_storage).__name__ if self._columnwise_storage is not None else "None" + row_type = ( + type(self._rowwise_storage).__name__ if self._rowwise_storage is not None else "None" + ) + col_type = ( + type(self._columnwise_storage).__name__ + if self._columnwise_storage is not None + else "None" + ) return ( - f"HybridQuantizedTensor(" - f"rowwise={row_type}, " - f"columnwise={col_type}, " - f"dtype={self.dtype})" + f"HybridQuantizedTensor(rowwise={row_type}, columnwise={col_type}, dtype={self.dtype})" ) def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index 3668809e73e..e407d252ef9 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -134,9 +134,7 @@ def device(self): raise RuntimeError("HybridQuantizedTensorStorage has no data") def view(self, shape: torch.Size): - raise NotImplementedError( - "HybridQuantizedTensorStorage does not support view operations" - ) + raise NotImplementedError("HybridQuantizedTensorStorage does not support view operations") def get_metadata(self) -> Dict[str, Any]: return { @@ -150,7 +148,7 @@ def get_metadata(self) -> Dict[str, Any]: def __repr__(self): return ( - f"HybridQuantizedTensorStorage(" + "HybridQuantizedTensorStorage(" f"rowwise={type(self._rowwise_storage).__name__}, " f"columnwise={type(self._columnwise_storage).__name__}, " f"dtype={self._dtype})" From f80f5d0c6e48a88ffd33c21f229b06b559d14b2a Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 16 Apr 2026 08:21:01 +0000 Subject: [PATCH 03/97] Enable quantized_model_init Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 1154 ++++++++++++++++- .../pytorch/tensor/hybrid_tensor.py | 41 +- 2 files changed, 1192 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 5967c0816a6..008a9a5b9e1 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -13,6 +13,7 @@ from transformer_engine.common import recipe from transformer_engine.pytorch import ( autocast, + quantized_model_init, Linear, LayerNormLinear, LayerNormMLP, @@ -29,6 +30,7 @@ Float8TensorStorage, NVFP4Tensor, NVFP4TensorStorage, + QuantizedTensor, ) from transformer_engine.pytorch.cpp_extensions.gemm import ( _unwrap_hybrid_A, @@ -98,9 +100,9 @@ def test_creation(self): assert isinstance(hq.rowwise_quantizer, Float8CurrentScalingQuantizer) assert isinstance(hq.columnwise_quantizer, NVFP4Quantizer) - def test_compatible_recipe_is_none(self): + def test_compatible_recipe_is_custom_recipe(self): hq = _make_hybrid_quantizer_fp8_row_fp4_col() - assert hq._get_compatible_recipe() is None + assert hq._get_compatible_recipe() is recipe.CustomRecipe @requires_fp8_and_nvfp4 @@ -1616,3 +1618,1151 @@ def test_transformer_layer(self): requires_grad=True, ) self._run_fwd_bwd(model, inp) + + +# =========================================================================== +# Quantized Parameters (quantized_model_init) tests for hybrid quantization +# =========================================================================== + + +def _hybrid_custom_recipe(row_factory, col_factory, grad_factory=None): + """Build a CustomRecipe where forward roles use HybridQuantizer and + backward roles use a plain quantizer (or hybrid if grad_factory builds one). + + Parameters + ---------- + row_factory : callable() -> Quantizer + Creates the rowwise sub-quantizer for forward roles. + col_factory : callable() -> Quantizer + Creates the columnwise sub-quantizer for forward roles. + grad_factory : callable() -> Quantizer, optional + Creates the quantizer for grad_output/grad_input roles. + If None, uses col_factory (matching columnwise format for wgrad compatibility). + """ + if grad_factory is None: + grad_factory = col_factory + + def qfactory(role): + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=row_factory(), + columnwise_quantizer=col_factory(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return grad_factory() + return row_factory() + + return recipe.CustomRecipe(qfactory=qfactory) + + +# --------------------------------------------------------------------------- +# 1. quantized_model_init: model creation and parameter type verification +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridQuantizedModelInit: + """Verify that quantized_model_init with a hybrid CustomRecipe produces + HybridQuantizedTensor parameters.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_linear_weight_is_hybrid_quantized_tensor(self): + """model.weight should be a HybridQuantizedTensor after quantized_model_init.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert isinstance(weight, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(weight).__name__}" + ) + assert isinstance(weight, QuantizedTensor), ( + "HybridQuantizedTensor should be a QuantizedTensor subclass" + ) + + def test_linear_weight_has_both_sub_storages(self): + """Quantized param should have rowwise and columnwise sub-storages.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert weight.rowwise_sub_storage is not None, "Missing rowwise sub-storage" + assert weight.columnwise_sub_storage is not None, "Missing columnwise sub-storage" + + def test_linear_weight_shape_preserved(self): + """Quantized param should retain its logical shape.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 256, params_dtype=torch.bfloat16).cuda() + + assert model.weight.shape == torch.Size([256, 128]) + + def test_linear_bias_stays_bf16(self): + """Bias should remain BF16 (not quantized).""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, bias=True, params_dtype=torch.bfloat16).cuda() + + assert not isinstance(model.bias, QuantizedTensor), ( + "Bias should not be a QuantizedTensor" + ) + assert model.bias.dtype == torch.bfloat16 + + def test_layernorm_linear_weight_is_hybrid(self): + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = LayerNormLinear(128, 128, params_dtype=torch.bfloat16).cuda() + + assert isinstance(model.weight, HybridQuantizedTensor) + + def test_dequantize_close_to_original(self): + """Dequantized hybrid param should be close to the BF16 init values.""" + hybrid_recipe = self._hybrid_fp8_recipe() + + # Create a non-quantized reference + torch.manual_seed(42) + ref_model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + ref_weight = ref_model.weight.detach().clone() + + # Create quantized model with the same seed + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + dq_weight = model.weight.dequantize() + torch.testing.assert_close(dq_weight.float(), ref_weight.float(), rtol=0.125, atol=0.1) + + def test_preserve_high_precision_init_val(self): + """preserve_high_precision_init_val should store original BF16 on CPU.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init( + enabled=True, + recipe=hybrid_recipe, + preserve_high_precision_init_val=True, + ): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + weight = model.weight + assert isinstance(weight, HybridQuantizedTensor) + assert hasattr(weight, "get_high_precision_init_val") + hp_val = weight.get_high_precision_init_val() + assert hp_val is not None, "High-precision init val should be stored" + assert hp_val.device.type == "cpu" + assert hp_val.shape == weight.shape + + +# --------------------------------------------------------------------------- +# 2. get_weight_workspace cache invalidation for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridWeightWorkspaceCache: + """Test that get_weight_workspace handles HybridQuantizedTensorStorage + correctly for the quantized-params early-return path and the BF16 cache path.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_quantized_param_skips_workspace(self): + """When weight is already a HybridQuantizedTensor (quantized params), + get_weight_workspace should return it directly without creating a workspace.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + + assert out.shape == (32, 128) + assert not torch.isnan(out).any() + + def test_bf16_weight_creates_hybrid_workspace(self): + """When weight is BF16 and recipe produces HybridQuantizer, the workspace + should be a HybridQuantizedTensor.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + + assert out.shape == (32, 128) + assert not torch.isnan(out).any() + + def test_workspace_cache_reuse_across_microbatches(self): + """Cached hybrid workspace should be reused on 2nd+ microbatches.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with autocast(enabled=True, recipe=hybrid_recipe): + with torch.no_grad(): + out1 = model(inp, is_first_microbatch=True) + out2 = model(inp, is_first_microbatch=False) + + # Both should produce valid, identical results (same weight, same input) + assert not torch.isnan(out1).any() + assert not torch.isnan(out2).any() + torch.testing.assert_close(out1, out2) + + def test_workspace_cache_invalidation_on_usage_change(self): + """If usage requirements change (e.g. inference→training), the cache + should be invalidated and a fresh workspace created.""" + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + hybrid_recipe = self._hybrid_fp8_recipe() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + # First pass: inference (no columnwise needed) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_infer = model(inp, is_first_microbatch=True) + assert not torch.isnan(out_infer).any() + + # Second pass: training (columnwise now needed for backward) + with autocast(enabled=True, recipe=hybrid_recipe): + out_train = model(inp, is_first_microbatch=True) + loss = out_train.float().sum() + loss.backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + + +# --------------------------------------------------------------------------- +# 3. _update_weight_quantizers for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridUpdateWeightQuantizers: + """Test that quantizer refresh propagates correctly to hybrid sub-quantizers.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_quantized_param_survives_multiple_forward_passes(self): + """Weight should remain a HybridQuantizedTensor across multiple forward passes, + each of which triggers init_fp8_metadata → potential quantizer updates.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + for i in range(3): + inp_i = inp.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp_i) + out.float().sum().backward() + assert not torch.isnan(out).any(), f"NaN at iteration {i}" + assert inp_i.grad is not None, f"No input grad at iteration {i}" + + assert isinstance(model.weight, HybridQuantizedTensor), ( + "Weight lost HybridQuantizedTensor type after multiple passes" + ) + + +# --------------------------------------------------------------------------- +# 4. Recipe correspondence validation +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridRecipeCorrespondence: + """Test _check_weight_tensor_recipe_correspondence with hybrid params.""" + + def _hybrid_fp8_recipe(self): + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + def test_hybrid_param_with_matching_recipe_does_not_raise(self): + """Forward pass with matching recipe should not raise.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + assert not torch.isnan(out).any() + + def test_hybrid_param_with_mismatched_recipe_raises(self): + """Forward pass with a non-CustomRecipe on a hybrid param should raise.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + mismatch_recipe = recipe.Float8CurrentScaling() + with pytest.raises(RuntimeError, match="Recipe mismatch"): + with torch.no_grad(): + with autocast(enabled=True, recipe=mismatch_recipe): + model(inp) + + +# --------------------------------------------------------------------------- +# 5. quantize_ in-place update for hybrid +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridQuantizeInPlace: + """Test in-place re-quantization (quantize_) for HybridQuantizedTensor. + + This is needed for the optimizer writeback path (param.quantize_(master_weight)) + and the workspace cache update path (out.quantize_(new_bf16_weight)). + """ + + def test_quantize_inplace_updates_data(self): + """quantize_() should re-quantize both sub-storages from new BF16 data.""" + torch.manual_seed(42) + hq = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + original = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + tensor = hq.quantize(original) + + dq_before = tensor.dequantize().clone() + + # Update with different data + new_data = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + tensor.quantize_(new_data) + + dq_after = tensor.dequantize() + + # Should be close to new data, not old data + diff_new = (dq_after.float() - new_data.float()).abs().mean() + diff_old = (dq_after.float() - original.float()).abs().mean() + assert diff_new < diff_old, ( + f"After quantize_(), data is closer to old ({diff_old:.4f}) " + f"than new ({diff_new:.4f})" + ) + + def test_quantize_inplace_preserves_tensor_identity(self): + """quantize_() should update in-place, not create a new tensor.""" + hq = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + original = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + tensor = hq.quantize(original) + + tensor_id = id(tensor) + new_data = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + result = tensor.quantize_(new_data) + + assert id(tensor) == tensor_id, "quantize_() should return same object" + + # noop_flag is a delayed-scaling feature; not tested here since + # delayed scaling is out of scope for hybrid quantization. + + +# --------------------------------------------------------------------------- +# 6. FusedAdam with hybrid quantized params +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFusedAdam: + """Test FusedAdam optimizer with HybridQuantizedTensor parameters.""" + + def _build_hybrid_model(self): + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model, hybrid_recipe + + def test_fused_adam_accepts_hybrid_params(self): + """FusedAdam should not crash when given HybridQuantizedTensor params.""" + model, _ = self._build_hybrid_model() + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + assert optimizer is not None + + def test_fused_adam_master_weights_track_reference(self): + """FP32 master weights should closely track a reference Adam optimizer. + + Small divergence is expected because HybridQuantizedTensor.float() + may take a slightly different dequantization path than + detach().clone().float() through __torch_dispatch__. + """ + model, _ = self._build_hybrid_model() + + ref_params = [p.detach().clone().float() for p in model.parameters()] + + options = {"lr": 5e-4, "betas": (0.9, 0.999), "eps": 1e-8, "weight_decay": 0} + ref_optim = torch.optim.Adam(ref_params, **options) + tst_optim = te.optimizers.FusedAdam( + list(model.parameters()), + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + **options, + ) + + for _ in range(5): + for p_ref, p in zip(ref_params, model.parameters()): + p_ref.grad = torch.rand_like(p_ref) + p.decoupled_grad = p_ref.grad.clone() + ref_optim.step() + tst_optim.step() + + master_params = [ + tst_optim.get_unscaled_state(p, "master_param") for p in model.parameters() + ] + torch.testing.assert_close(ref_params, master_params, rtol=1e-3, atol=1e-3) + + def test_fused_adam_param_remains_hybrid_after_step(self): + """Weight params should still be HybridQuantizedTensors after optimizer step.""" + model, _ = self._build_hybrid_model() + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + ) + + for _ in range(3): + for p in model.parameters(): + p.decoupled_grad = torch.rand_like(p.float()) + optimizer.step() + + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance(p, HybridQuantizedTensor), ( + f"{name} lost HybridQuantizedTensor type: {type(p).__name__}" + ) + + def test_fused_adam_requires_master_weights(self): + """FusedAdam without master_weights should raise for hybrid quantized params.""" + model, _ = self._build_hybrid_model() + + with pytest.raises(RuntimeError, match="master_weights"): + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=False, + ) + for p in model.parameters(): + p.grad = torch.rand_like(p.float()).to(p.dtype) + optimizer.step() + + +# --------------------------------------------------------------------------- +# 7. End-to-end training loop: fwd + bwd + optimizer step +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridQuantizedParamsEndToEnd: + """Full training loop: quantized_model_init + autocast fwd + bwd + FusedAdam.step().""" + + def _build_model_and_recipe(self): + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model, hybrid_recipe + + def test_training_loop_loss_decreases(self): + """Loss should decrease over a few training steps.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses = [] + for i in range(7): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + + for name, p in model.named_parameters(): + assert p.grad is not None, f"Step {i}: {name} has no gradient" + assert torch.isfinite(p.grad).all(), f"Step {i}: {name} has non-finite grad" + + optimizer.step() + + assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" + + def test_training_loop_params_remain_quantized(self): + """Params should remain HybridQuantizedTensors after training.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + for _ in range(3): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + optimizer.step() + + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance(p, HybridQuantizedTensor), ( + f"{name} is {type(p).__name__}, not HybridQuantizedTensor" + ) + + def test_training_loop_optimizer_states_are_fp32(self): + """Optimizer states should be FP32.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_model_and_recipe() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + for _ in range(2): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + output.float().sum().backward() + optimizer.step() + + for name, p in model.named_parameters(): + state = optimizer.state[p] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + if "bias" not in name: + assert state["master_param"].dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# 8. Mixed-format quantized params (e.g. MXFP8 row + NVFP4 col) +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (mxfp8_available and nvfp4_available), + reason=f"MXFP8: {reason_for_no_mxfp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestHybridMixedFormatQuantizedParams: + """Quantized params with genuinely different formats per direction.""" + + def _build_mixed_model(self, in_features=256, out_features=256): + """MXFP8 rowwise (for fprop TN) + NVFP4 columnwise (for wgrad NT).""" + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + col_factory=lambda: NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + grad_factory=lambda: NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear( + in_features, out_features, params_dtype=torch.bfloat16 + ).cuda() + return model, hybrid_recipe + + def test_mixed_format_param_creation(self): + """Model init with mixed MXFP8/NVFP4 hybrid should produce a + HybridQuantizedTensor parameter.""" + model, _ = self._build_mixed_model() + assert isinstance(model.weight, HybridQuantizedTensor) + + def test_mixed_format_forward_only(self): + """Forward pass with mixed-format quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + inp = torch.randn(32, 256, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + + assert out.shape == (32, 256) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + + def test_mixed_format_forward_backward(self): + """Full fwd+bwd with mixed-format quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + inp = torch.randn(32, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None + assert not torch.isnan(inp.grad).any() + for name, p in model.named_parameters(): + assert p.grad is not None, f"No gradient for {name}" + assert not torch.isnan(p.grad).any(), f"NaN gradient for {name}" + + def test_mixed_format_training_loop(self): + """End-to-end training loop with mixed-format hybrid quantized params.""" + torch.manual_seed(42) + model, hybrid_recipe = self._build_mixed_model() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses = [] + for i in range(5): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance(p, HybridQuantizedTensor), ( + f"{name} is {type(p).__name__}" + ) + + def test_mixed_format_sub_storage_types(self): + """Verify that sub-storages have the correct types (MXFP8 vs NVFP4).""" + model, _ = self._build_mixed_model() + weight = model.weight + from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( + MXFP8TensorStorage, + ) + + row = weight.rowwise_sub_storage + col = weight.columnwise_sub_storage + assert isinstance(row, MXFP8TensorStorage) or hasattr(row, "_rowwise_data"), ( + f"Expected MXFP8 rowwise sub-storage, got {type(row).__name__}" + ) + assert isinstance(col, NVFP4TensorStorage) or hasattr(col, "_rowwise_data"), ( + f"Expected NVFP4 columnwise sub-storage, got {type(col).__name__}" + ) + + +# --------------------------------------------------------------------------- +# 9. Quantized params equivalence: vanilla vs hybrid (same format both dirs) +# --------------------------------------------------------------------------- + + +def _hybrid_fp8_current_qfactory(role): + """Hybrid FP8 current scaling (E4M3 both dirs, E5M2 for grad).""" + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + +def _hybrid_mxfp8_qfactory(role): + """Hybrid MXFP8 (E4M3 both dirs).""" + if role in ("linear_grad_output", "linear_grad_input"): + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + + +def _hybrid_block_fp8_qfactory(role): + """Hybrid block FP8 (E4M3 both dirs).""" + dim = 2 if role == "linear_weight" else 1 + if role in ("linear_grad_output", "linear_grad_input"): + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, columnwise=True, block_scaling_dim=dim, + ) + return HybridQuantizer( + rowwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, columnwise=True, block_scaling_dim=dim, + ), + columnwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, columnwise=True, block_scaling_dim=dim, + ), + ) + + +def _hybrid_nvfp4_qfactory(role): + """Hybrid NVFP4 (E2M1 both dirs).""" + if role in ("linear_grad_output", "linear_grad_input"): + return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + return HybridQuantizer( + rowwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + columnwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + ) + + +class _QuantizedParamsEquivalenceBase: + """Base for comparing vanilla vs hybrid quantized params training. + + When the hybrid quantizer uses the same format in both directions, + the full quantized_model_init + training loop should produce + equivalent results to the vanilla (non-hybrid) quantized params path. + """ + + hidden_size = 256 + num_steps = 5 + + def _vanilla_recipe(self): + raise NotImplementedError + + def _hybrid_recipe(self): + raise NotImplementedError + + def _build_models(self): + """Create two models with identical init: one vanilla, one hybrid.""" + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=self._vanilla_recipe()): + model_ref = Linear( + self.hidden_size, self.hidden_size, params_dtype=torch.bfloat16, + ).cuda() + + torch.manual_seed(42) + with quantized_model_init(enabled=True, recipe=self._hybrid_recipe()): + model_hyb = Linear( + self.hidden_size, self.hidden_size, params_dtype=torch.bfloat16, + ).cuda() + + return model_ref, model_hyb + + def _run_training_loop(self, model, train_recipe, x, target, num_steps): + optimizer = te.optimizers.FusedAdam( + model.parameters(), lr=1e-3, + master_weights=True, master_weight_dtype=torch.float32, + ) + losses = [] + for _ in range(num_steps): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=train_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + master_params = [ + optimizer.get_unscaled_state(p, "master_param") + for p in model.parameters() + if p.requires_grad + ] + return losses, master_params + + def _test_equivalence(self): + model_ref, model_hyb = self._build_models() + + torch.manual_seed(99) + x = torch.randn(4, 32, self.hidden_size, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses_ref, masters_ref = self._run_training_loop( + model_ref, self._vanilla_recipe(), x, target, self.num_steps, + ) + losses_hyb, masters_hyb = self._run_training_loop( + model_hyb, self._hybrid_recipe(), x, target, self.num_steps, + ) + + # Losses should be very close (same quantization, same training dynamics) + for i, (lr, lh) in enumerate(zip(losses_ref, losses_hyb)): + assert abs(lr - lh) < 0.1 * max(abs(lr), 1e-6), ( + f"Step {i}: loss diverged — vanilla={lr:.6f}, hybrid={lh:.6f}" + ) + + # Master weights should be close after training + for i, (mr, mh) in enumerate(zip(masters_ref, masters_hyb)): + torch.testing.assert_close(mr, mh, rtol=1e-3, atol=1e-3, msg=( + f"Master weight {i} diverged after {self.num_steps} steps" + )) + + +@requires_fp8 +class TestQuantizedParamsEquivalenceFP8CurrentScaling(_QuantizedParamsEquivalenceBase): + """Vanilla Float8CurrentScaling vs hybrid FP8 current (same format both dirs). + + Note: vanilla Float8Tensor params use the fused multi_tensor_adam_fp8 + kernel in FusedAdam, while HybridQuantizedTensor falls to the FP32 + master + quantize_() writeback path. These are numerically different + codepaths, so we use relaxed tolerances. + """ + + def _vanilla_recipe(self): + return recipe.Float8CurrentScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_fp8_current_qfactory) + + def test_equivalence(self): + model_ref, model_hyb = self._build_models() + + torch.manual_seed(99) + x = torch.randn(4, 32, self.hidden_size, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses_ref, _ = self._run_training_loop( + model_ref, self._vanilla_recipe(), x, target, self.num_steps, + ) + losses_hyb, _ = self._run_training_loop( + model_hyb, self._hybrid_recipe(), x, target, self.num_steps, + ) + + # Both should decrease (training works in both paths) + assert losses_ref[-1] < losses_ref[0], f"Vanilla loss didn't decrease: {losses_ref}" + assert losses_hyb[-1] < losses_hyb[0], f"Hybrid loss didn't decrease: {losses_hyb}" + + # Losses should be in the same ballpark (different optimizer kernels + # cause small divergence that compounds over steps) + for i, (lr, lh) in enumerate(zip(losses_ref, losses_hyb)): + assert abs(lr - lh) / max(abs(lr), 1e-6) < 0.5, ( + f"Step {i}: losses diverged too much — vanilla={lr:.6f}, hybrid={lh:.6f}" + ) + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestQuantizedParamsEquivalenceMXFP8(_QuantizedParamsEquivalenceBase): + """Vanilla MXFP8BlockScaling vs hybrid MXFP8 (same format both dirs).""" + + def _vanilla_recipe(self): + return recipe.MXFP8BlockScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_mxfp8_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +@pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) +class TestQuantizedParamsEquivalenceBlockFP8(_QuantizedParamsEquivalenceBase): + """Vanilla Float8BlockScaling vs hybrid block FP8 (same format both dirs).""" + + def _vanilla_recipe(self): + return recipe.Float8BlockScaling() + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +@pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) +class TestQuantizedParamsEquivalenceNVFP4(_QuantizedParamsEquivalenceBase): + """Vanilla NVFP4BlockScaling vs hybrid NVFP4 (same format both dirs). + + RHT, stochastic rounding, and 2D quantization disabled for determinism. + """ + + def _vanilla_recipe(self): + return recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + ) + + def _hybrid_recipe(self): + return recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) + + def test_equivalence(self): + self._test_equivalence() + + +# --------------------------------------------------------------------------- +# 10. State dict save/load (checkpointing) for hybrid quantized params +# --------------------------------------------------------------------------- + + +# Module-level qfactories (picklable, required for checkpoint serialization). + + +def _checkpoint_hybrid_fp8_qfactory(role): + """Module-level qfactory (picklable) for checkpoint tests.""" + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + columnwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + +@requires_fp8 +class TestHybridCheckpoint: + """Test state_dict save/load round-trips for models with hybrid quantized params.""" + + def _hybrid_fp8_recipe(self): + return recipe.CustomRecipe(qfactory=_checkpoint_hybrid_fp8_qfactory) + + def test_state_dict_save_load_roundtrip(self): + """state_dict → save → load → same model should produce identical outputs.""" + torch.manual_seed(42) + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_before = model(inp) + + state_dict = model.state_dict() + + # Create a fresh model and load + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model2 = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + model2.load_state_dict(state_dict) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_after = model2(inp) + + torch.testing.assert_close(out_before, out_after) + + def test_state_dict_contains_weight(self): + """state_dict should contain the weight key.""" + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + sd = model.state_dict() + assert "weight" in sd, f"state_dict keys: {list(sd.keys())}" + + def test_load_bf16_state_dict_into_hybrid_model(self): + """Loading a BF16 state_dict into a hybrid quantized model should work. + + This is the common scenario: pretrained BF16 weights loaded into a + model initialized with quantized_model_init. + """ + torch.manual_seed(42) + hybrid_recipe = self._hybrid_fp8_recipe() + + # Create BF16 model and get its state_dict + ref_model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + bf16_state_dict = ref_model.state_dict() + + # Create hybrid quantized model + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + # Load BF16 weights into hybrid model + model.load_state_dict(bf16_state_dict) + + # Verify model produces valid output + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out = model(inp) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + + def test_state_dict_torch_save_load(self): + """Full round-trip through torch.save/torch.load (file-based).""" + import tempfile + import os + + torch.manual_seed(42) + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_before = model(inp) + + with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as f: + torch.save(model.state_dict(), f.name) + tmp_path = f.name + + try: + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model2 = Linear(128, 128, params_dtype=torch.bfloat16).cuda() + state_dict = torch.load(tmp_path, weights_only=False) + model2.load_state_dict(state_dict) + + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + out_after = model2(inp) + + torch.testing.assert_close(out_before, out_after) + finally: + os.unlink(tmp_path) + + def test_checkpoint_resume_training(self): + """Save mid-training, load into new model+optimizer, verify training continues.""" + import tempfile + import os + + torch.manual_seed(42) + hybrid_recipe = self._hybrid_fp8_recipe() + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), lr=1e-3, + master_weights=True, master_weight_dtype=torch.float32, + ) + + x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + # Train for a few steps + for _ in range(3): + optimizer.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + optimizer.step() + + loss_before_save = loss.item() + + # Save checkpoint + with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as f: + torch.save({ + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + }, f.name) + tmp_path = f.name + + try: + # Load into fresh model + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model2 = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), lr=1e-3, + master_weights=True, master_weight_dtype=torch.float32, + ) + + checkpoint = torch.load(tmp_path, weights_only=False) + model2.load_state_dict(checkpoint["model"]) + optimizer2.load_state_dict(checkpoint["optimizer"]) + + # Continue training -- loss should not spike + optimizer2.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=hybrid_recipe): + output2 = model2(x) + loss_after_load = torch.nn.functional.mse_loss(output2, target).item() + + assert loss_after_load <= loss_before_save * 1.5, ( + f"Loss spiked after checkpoint resume: {loss_before_save:.4f} → {loss_after_load:.4f}" + ) + finally: + os.unlink(tmp_path) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 6073b5d1083..4d1a5ec7adc 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -111,13 +111,52 @@ def make_empty( quantizer=self, ) + def update_quantized( + self, + src: torch.Tensor, + dst: QuantizedTensorStorage, + *, + noop_flag: Optional[torch.Tensor] = None, + ) -> QuantizedTensorStorage: + """Re-quantize both sub-storages of a hybrid tensor in-place. + + Delegates to each sub-quantizer's update_quantized, which writes + new quantized data + scales into the existing sub-storage buffers. + """ + if not isinstance(dst, HybridQuantizedTensorStorage): + raise ValueError( + f"HybridQuantizer can only update HybridQuantizedTensorStorage, got {type(dst).__name__}" + ) + if dst._rowwise_storage is not None: + self.rowwise_quantizer.update_quantized( + src, dst._rowwise_storage, noop_flag=noop_flag + ) + if dst._columnwise_storage is not None: + self.columnwise_quantizer.update_quantized( + src, dst._columnwise_storage, noop_flag=noop_flag + ) + return dst + def set_usage( self, *, rowwise: Optional[bool] = None, columnwise: Optional[bool] = None ) -> None: super().set_usage(rowwise=rowwise, columnwise=columnwise) def _get_compatible_recipe(self): - return None + # HybridQuantizer is only reachable via CustomRecipe (the qfactory + # returns HybridQuantizer per role). Checking that the autocast recipe + # is also CustomRecipe catches the obvious mismatch (e.g. hybrid + # quantized_model_init + built-in MXFP8BlockScaling autocast). + # We trust that users who write a CustomRecipe know what they're doing + # with regard to per-operand scaling mode compatibility. + # TODO(negvet): improve to validate that the autocast recipe's + # sub-quantizer scaling modes are compatible with each sub-storage's + # scaling mode (e.g. rowwise MXFP8 weight requires MXFP8 input for + # fprop TN, columnwise NVFP4 weight requires NVFP4 grad_output for + # wgrad NT). + from transformer_engine.common.recipe import CustomRecipe # avoid circular import + + return CustomRecipe class HybridQuantizedTensor(HybridQuantizedTensorStorage, QuantizedTensor): From 2185b30987c95a0c443281a73b2538e1152337d2 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 17 Apr 2026 12:59:14 +0000 Subject: [PATCH 04/97] FSDP support Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/conftest.py | 22 + .../distributed/fsdp2_tests/fsdp2_utils.py | 60 ++ .../fsdp2_tests/run_fsdp2_fused_adam.py | 493 ++++++++++++++- .../fsdp2_tests/run_fsdp2_mem_leak.py | 150 +++++ .../fsdp2_tests/run_fsdp2_model.py | 111 ++++ tests/pytorch/distributed/test_torch_fsdp2.py | 51 ++ tests/pytorch/test_hybrid_quantization.py | 564 ++++++++++++++++++ .../pytorch/quantized_tensor.py | 59 ++ .../pytorch/tensor/float8_tensor.py | 45 +- .../pytorch/tensor/hybrid_tensor.py | 353 ++++++++++- .../pytorch/tensor/mxfp8_tensor.py | 85 ++- .../tensor/storage/float8_tensor_storage.py | 12 + .../tensor/storage/hybrid_tensor_storage.py | 14 +- .../tensor/storage/mxfp8_tensor_storage.py | 76 ++- 14 files changed, 2026 insertions(+), 69 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index bf9db094d2e..94bcc43f060 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -45,6 +45,13 @@ def _check_nvfp4_support(): ("NVFP4BlockScaling", _check_nvfp4_support), ] +_HYBRID_RECIPE_CONFIGS = [ + ("HybridFP8CurrentScaling", fp8.check_fp8_support), + ("HybridMXFP8", fp8.check_mxfp8_support), + ("HybridFloat8BlockScaling", fp8.check_fp8_block_scaling_support), + ("HybridMixed_MXFP8_FP8", fp8.check_mxfp8_support), +] + def _parametrize_recipes(): params = [] @@ -56,6 +63,16 @@ def _parametrize_recipes(): return params +def _parametrize_hybrid_recipes(): + params = [] + for name, check_fn in _HYBRID_RECIPE_CONFIGS: + supported, reason = check_fn() + params.append( + pytest.param(name, id=name, marks=pytest.mark.skipif(not supported, reason=reason)) + ) + return params + + # ── Session / per-test fixtures ────────────────────────────────────── @pytest.fixture(scope="session", autouse=True) def dist_init(): @@ -83,3 +100,8 @@ def _cleanup(): @pytest.fixture(params=_parametrize_recipes()) def recipe_name(request): return request.param + + +@pytest.fixture(params=_parametrize_hybrid_recipes()) +def hybrid_recipe_name(request): + return request.param diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 178ce62375d..1cf81b1db07 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -12,6 +12,66 @@ def get_recipe_from_string(recipe): return getattr(transformer_engine.common.recipe, recipe)() +def get_hybrid_recipe_from_string(recipe): + """Build a CustomRecipe that uses HybridQuantizer with the given base format. + + Supported values: + "HybridFP8CurrentScaling" — FP8 current for both directions + "HybridMXFP8" — MXFP8 for both directions + "HybridMixed_MXFP8_FP8" — MXFP8 rowwise + FP8 current columnwise + """ + import transformer_engine_torch as tex + from transformer_engine.pytorch import ( + Float8CurrentScalingQuantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + HybridQuantizer, + ) + + _BUILDERS = { + "HybridFP8CurrentScaling": lambda: dict( + row=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), + ), + "HybridMXFP8": lambda: dict( + row=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + col=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + grad=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + ), + "HybridFloat8BlockScaling": lambda: dict( + row=lambda: Float8BlockQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True), + col=lambda: Float8BlockQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True), + grad=lambda: Float8BlockQuantizer(fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True), + ), + "HybridMixed_MXFP8_FP8": lambda: dict( + row=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + col=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), + ), + } + + if recipe not in _BUILDERS: + raise ValueError( + f"Unknown hybrid recipe '{recipe}'. Supported: {sorted(_BUILDERS.keys())}" + ) + + builders = _BUILDERS[recipe]() + row_fn, col_fn, grad_fn = builders["row"], builders["col"], builders["grad"] + + def qfactory(role): + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=row_fn(), + columnwise_quantizer=col_fn(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return grad_fn() + return row_fn() + + return transformer_engine.common.recipe.CustomRecipe(qfactory=qfactory) + + def save_custom_attrs(module): custom_attrs = {} for name, param in module.named_parameters(): diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 42df06ed7fc..36d25b0d7af 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -103,7 +103,7 @@ def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None, use_meta_ return model -def _shard_model(model, world_size): +def _shard_model(model, world_size, reshard_after_forward=None): """Apply FSDP2 sharding with save/restore custom attrs. If the model was created on the meta device (e.g. for FP8 init), @@ -112,12 +112,24 @@ def _shard_model(model, world_size): restore_custom_attrs is called last so it applies to the final parameter objects. For meta-device models, reset_parameters() replaces params via module_setattr (base.py:1336-1339), so attrs must be restored afterward. + + Parameters + ---------- + reshard_after_forward : bool, optional + Passed through to ``fully_shard``. ``None`` (default) keeps FSDP2's + own default: ``True`` for child modules, ``False`` for the root. + ``False`` on child modules keeps the full-precision gathered weight + alive through backward, exercising the iter-2+ buffer-reuse path + inside the same forward/backward rather than across training steps. """ has_meta_params = any(p.is_meta for p in model.parameters()) custom_attrs = save_custom_attrs(model) mesh = DeviceMesh("cuda", list(range(world_size))) + shard_kwargs = {"mesh": mesh} + if reshard_after_forward is not None: + shard_kwargs["reshard_after_forward"] = reshard_after_forward for child in model.children(): - fully_shard(child, mesh=mesh) + fully_shard(child, **shard_kwargs) fully_shard(model, mesh=mesh) if has_meta_params: for module in model.modules(): @@ -915,6 +927,462 @@ def test_dcp_resharding_load(recipe_name): os.remove(ref_output_path) +# --------------------------------------------------------------------------- +# Hybrid quantization + FSDP2 tests +# --------------------------------------------------------------------------- + + +def _build_hybrid_model(hybrid_recipe, use_meta_device=True): + """Build a model with quantized_model_init using a hybrid CustomRecipe.""" + ctx = te.quantized_model_init(enabled=True, recipe=hybrid_recipe) + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + if use_meta_device: + kwargs["device"] = "meta" + with ctx: + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + **kwargs, + ) + for _ in range(NUM_LAYERS) + ] + ) + return model + + +def test_fused_adam_hybrid_master_weights(hybrid_recipe_name): + """FusedAdam + master_weights + FSDP2 + hybrid quantized_model_init. + + Verifies: + - Params are DTensors wrapping HybridQuantizedTensor local shards + - Training loop completes without error + - Optimizer states are FP32 + - Loss decreases over training steps + """ + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1) in reset_sharded_param. " + "Same root cause as vanilla Float8BlockScaling + quantized_model_init." + ) + + from transformer_engine.pytorch import HybridQuantizedTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + hybrid_count = sum( + 1 + for _, p in model.named_parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + for param in model.parameters(): + state = optimizer.state[param] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + if "master_param" in state: + assert state["master_param"].dtype == torch.float32 + + assert losses[-1] < losses[0], ( + f"Loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" + ) + + +@pytest.mark.parametrize("reshard_after_forward", [True, False]) +def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name, reshard_after_forward): + """Hybrid FusedAdam loop under both ``reshard_after_forward`` settings. + + ``reshard_after_forward=True`` is FSDP2's default: the gathered weight is + dropped after forward and a second all-gather happens in backward — + meaning ``fsdp_post_all_gather(out=...)`` is invoked twice per training + step (once per pass) on the same gathered buffer. ``False`` keeps the + gathered weight alive through backward — only one gather per step, and + the gathered copy persists across forward/backward within the same step. + + Both modes must complete cleanly and produce a decreasing loss. This + locks in that the hybrid hooks handle both FSDP2 schedules, and forms a + regression harness for a future bandwidth optimization (P1.1) that would + split forward-only / backward-only buffers. + """ + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1) in reset_sharded_param." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size, reshard_after_forward=reshard_after_forward) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + assert losses[-1] < losses[0], ( + f"[reshard_after_forward={reshard_after_forward}] " + f"loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" + ) + + +def test_fused_adam_hybrid_bf16_vs_hybrid_parity(hybrid_recipe_name): + """Compare hybrid+FSDP2 loss trajectory against BF16+FSDP2 within tolerance. + + This is a sanity check that hybrid quantized training converges similarly + to BF16 training, not a bitwise-exact comparison. + """ + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1)." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + def run_training(model, recipe_for_autocast): + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + losses = [] + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=(recipe_for_autocast is not None), recipe=recipe_for_autocast): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + return losses + + # BF16 baseline + torch.manual_seed(42) + torch.cuda.manual_seed(42) + bf16_model = _build_model(fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_losses = run_training(bf16_model, None) + + # Hybrid + torch.manual_seed(42) + torch.cuda.manual_seed(42) + hybrid_model = _build_hybrid_model(hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_losses = run_training(hybrid_model, hybrid_recipe) + + assert hybrid_losses[-1] < hybrid_losses[0], ( + f"Hybrid loss did not decrease: {hybrid_losses}" + ) + assert bf16_losses[-1] < bf16_losses[0], f"BF16 loss did not decrease: {bf16_losses}" + + # Verify hybrid and bf16 loss trajectories are within the same order of magnitude. + # Quantized training may diverge from bf16, but should not be wildly different. + for step, (h_loss, b_loss) in enumerate(zip(hybrid_losses, bf16_losses)): + ratio = h_loss / max(b_loss, 1e-10) + assert 0.1 < ratio < 10.0, ( + f"Step {step}: hybrid loss ({h_loss:.4f}) and bf16 loss ({b_loss:.4f}) " + f"differ by more than 10x (ratio={ratio:.2f})" + ) + + +def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): + """Validate that FSDP2 all-gather + post-gather reconstruction produces + correct results by comparing ``unshard(param).dequantize()`` with a manual + all-gather of dequantized local shards. + + For the stateless formats supported here (per-tensor FP8, MXFP8), FSDP2's + all-gather concatenates rowwise bytes along dim-0 and the per-block / + per-tensor scales follow the same concatenation. Dequantizing the gathered + bytes is therefore bitwise-identical to concatenating the dequantized + shards — the tolerance is effectively ``assert_equal``. + """ + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1)." + ) + + from transformer_engine.pytorch import HybridQuantizedTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + with te.autocast(enabled=True, recipe=hybrid_recipe): + _ = model(x) + + # Stateless formats gather identical bytes; dequantize must match exactly. + _TIGHT_TOLERANCE = { + "HybridFP8CurrentScaling": dict(rtol=0.0, atol=0.0), + "HybridMXFP8": dict(rtol=0.0, atol=0.0), + "HybridMixed_MXFP8_FP8": dict(rtol=0.0, atol=0.0), + } + tolerance = _TIGHT_TOLERANCE.get(hybrid_recipe_name, dict(rtol=1e-6, atol=1e-6)) + + checked = 0 + for name, param in model.named_parameters(): + if not (isinstance(param, DTensor) and isinstance(param._local_tensor, QuantizedTensor)): + continue + local_shard = param._local_tensor + + local_deq = local_shard.dequantize().contiguous() + gathered_list = [torch.zeros_like(local_deq) for _ in range(world_size)] + dist.all_gather(gathered_list, local_deq) + manual_full = torch.cat(gathered_list, dim=0) + + full_param = param.full_tensor() + if isinstance(full_param, QuantizedTensor): + fsdp_full_deq = full_param.dequantize() + else: + fsdp_full_deq = full_param.float() + + torch.testing.assert_close( + manual_full.float(), + fsdp_full_deq[: manual_full.shape[0]].float(), + msg=lambda m, n=name: f"Allgather mismatch for {n}: {m}", + **tolerance, + ) + checked += 1 + + assert checked > 0, "No quantized DTensor params found to validate" + + +def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): + """Exercise MXFP8 block-scale unpad/pad on a sharded Linear whose shard + dim-0 is block-aligned (divisible by 32) but NOT divisible by 128. + + MXFP8 block scales are stored with ``[128, 4]`` / ``[4, 128]`` alignment + padding, which must be stripped before FSDP2's dim-0 all-gather and + re-applied after. With ``HIDDEN_SIZE`` and ``FFN_HIDDEN_SIZE`` both + divisible by 128, the default model never forces this code path, so this + test uses a hand-picked Linear size. + + Regression test for the "pre-fix" bug where + ``HybridQuantizedTensor.fsdp_pre_all_gather`` pulled raw tensor fields via + ``get_metadata()`` without unpadding the scale — the padded bytes would + have been interleaved at every rank boundary in the gather output. + """ + from fsdp2_utils import get_hybrid_recipe_from_string + + world_size, device = _get_dist_info() + + # FSDP2 shards a Linear weight of shape (out_features, in_features) along + # dim-0, so each rank holds `out_features / world_size` rows. Pick + # per-rank shard dim-0 = 96: divisible by MXFP8_BLOCK_SCALING_SIZE (32) + # so data alignment holds, but NOT divisible by 128 so the rowwise + # scale-inv needs alignment padding on the sharded copy. This is the + # shape that exercises the unpad-before-gather / pad-after-gather + # behaviour in MXFP8TensorStorage.fsdp_{extract,assign}_buffers. + per_rank_out = 96 + out_features = per_rank_out * world_size + in_features = 128 # arbitrary, divisible by 32; not sharded by FSDP2 here + assert per_rank_out % 32 == 0, "MXFP8 data alignment precondition" + assert per_rank_out % 128 != 0, "Test precondition: shard must need scale padding" + + for recipe_name in ("HybridMXFP8", "HybridMixed_MXFP8_FP8"): + hybrid_recipe = get_hybrid_recipe_from_string(recipe_name) + + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + te.Linear( + in_features, + out_features, + params_dtype=torch.bfloat16, + device="meta", + ), + ) + model = _shard_model(model, world_size) + + # Batch (leading) dim must be divisible by MXFP8_BLOCK_SCALING_SIZE (32). + x = torch.randn(32, in_features, dtype=torch.bfloat16, device=device) + with te.autocast(enabled=True, recipe=hybrid_recipe): + out = model(x) + out.sum().backward() + + # Compare dim-0 all-gather (bytes) with FSDP2's reconstruction. + for name, param in model.named_parameters(): + if not ( + isinstance(param, DTensor) + and isinstance(param._local_tensor, QuantizedTensor) + ): + continue + local_shard = param._local_tensor + local_deq = local_shard.dequantize().contiguous() + gathered_list = [torch.zeros_like(local_deq) for _ in range(world_size)] + dist.all_gather(gathered_list, local_deq) + manual_full = torch.cat(gathered_list, dim=0) + + full_param = param.full_tensor() + fsdp_full_deq = ( + full_param.dequantize() + if isinstance(full_param, QuantizedTensor) + else full_param.float() + ) + + torch.testing.assert_close( + manual_full.float(), + fsdp_full_deq[: manual_full.shape[0]].float(), + rtol=0.0, + atol=0.0, + msg=lambda m, n=name, r=recipe_name: ( + f"[{r}] Allgather mismatch for {n} at awkward shard shape: {m}" + ), + ) + + +def test_hybrid_dcp_output_parity(hybrid_recipe_name): + """DCP save+load roundtrip: output after load matches output before save. + + Trains a hybrid model, saves with DCP, loads into fresh model, + and asserts forward output parity. + """ + import torch.distributed.checkpoint as dcp + + pytest.xfail( + "CustomRecipe with closure-based qfactory cannot be pickled by DCP. " + "Requires module-level picklable factory functions for DCP compatibility." + ) + + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1)." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + rank = int(os.environ["LOCAL_RANK"]) + checkpoint_dir = os.path.join("/tmp", f"hybrid_dcp_test_{os.getpid()}") + + try: + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + optimizer = te.optimizers.FusedAdam( + model.parameters(), lr=1e-3, + master_weights=True, master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + F.mse_loss(output, target).backward() + optimizer.step() + + with torch.no_grad(): + with te.autocast(enabled=True, recipe=hybrid_recipe): + ref_output = model(x).clone() + + save_state = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + } + dcp.save(save_state, checkpoint_id=checkpoint_dir) + + model2 = _build_hybrid_model(hybrid_recipe) + model2 = _shard_model(model2, world_size) + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), lr=1e-3, + master_weights=True, master_weight_dtype=torch.float32, + ) + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + out_tmp = model2(x) + F.mse_loss(out_tmp, target).backward() + optimizer2.step() + + state_to_load = { + "model": model2.state_dict(), + "optimizer": optimizer2.state_dict(), + } + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model2.load_state_dict(state_to_load["model"]) + optimizer2.load_state_dict(state_to_load["optimizer"]) + + with torch.no_grad(): + with te.autocast(enabled=True, recipe=hybrid_recipe): + loaded_output = model2(x) + + torch.testing.assert_close( + loaded_output, ref_output, rtol=0, atol=0, + msg=lambda m: f"DCP roundtrip output mismatch: {m}", + ) + finally: + dist.barrier() + if rank == 0: + import shutil + shutil.rmtree(checkpoint_dir, ignore_errors=True) + + TESTS = { "fused_adam_fp8_master_weights": test_fused_adam_fp8_master_weights, "fused_adam_fp8_master_weights_no_meta": test_fused_adam_fp8_master_weights_no_meta, @@ -927,6 +1395,18 @@ def test_dcp_resharding_load(recipe_name): "dcp_resharding_save": test_dcp_resharding_save, "dcp_resharding_load": test_dcp_resharding_load, "safetensors_fp32_export": test_safetensors_fp32_export, + "fused_adam_hybrid_master_weights": test_fused_adam_hybrid_master_weights, + "fused_adam_hybrid_bf16_vs_hybrid_parity": test_fused_adam_hybrid_bf16_vs_hybrid_parity, + "fused_adam_hybrid_allgather_correctness": test_fused_adam_hybrid_allgather_correctness, + "fused_adam_hybrid_mxfp8_awkward_shard_shape": ( + test_fused_adam_hybrid_mxfp8_awkward_shard_shape + ), + "hybrid_dcp_output_parity": test_hybrid_dcp_output_parity, +} + +# Hybrid tests that are NOT parametrized by recipe (they sweep internally). +_HYBRID_NON_PARAMETRIZED_TESTS = { + "fused_adam_hybrid_mxfp8_awkward_shard_shape", } @@ -944,6 +1424,10 @@ def test_dcp_resharding_load(recipe_name): "Float8BlockScaling", "MXFP8BlockScaling", "NVFP4BlockScaling", + "HybridFP8CurrentScaling", + "HybridMXFP8", + "HybridFloat8BlockScaling", + "HybridMixed_MXFP8_FP8", ], ) args = parser.parse_args() @@ -953,7 +1437,10 @@ def test_dcp_resharding_load(recipe_name): torch.manual_seed(42) torch.cuda.manual_seed(42) try: - TESTS[args.test](args.recipe) + if args.test in _HYBRID_NON_PARAMETRIZED_TESTS: + TESTS[args.test]() + else: + TESTS[args.test](args.recipe) finally: if dist.is_initialized(): dist.destroy_process_group() diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py index 387d3a96441..25212d4db0f 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py @@ -468,12 +468,152 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in ) +# ── Hybrid quantization memory tests ───────────────────────────────── + +def _build_hybrid_model(num_layers, hybrid_recipe, use_meta_device=True): + """Build a model with quantized_model_init using a hybrid CustomRecipe.""" + ctx = te.quantized_model_init(enabled=True, recipe=hybrid_recipe) + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + if use_meta_device: + kwargs["device"] = "meta" + with ctx: + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + **kwargs, + ) + for _ in range(num_layers) + ] + ) + return model + + +def test_hybrid_no_excess_forward_memory(hybrid_recipe_name): + """Hybrid quantized weights should not accumulate across layers during forward. + + Same methodology as test_fp8_temp_accumulation_across_layers but for + hybrid quantized tensors. + """ + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1)." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # bf16 baseline + bf16_model = _build_model(NUM_LAYERS, fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_optimizer = te.optimizers.FusedAdam( + bf16_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(bf16_model, bf16_optimizer, None, x, target) + bf16_increments = _measure_forward_increments(bf16_model, bf16_optimizer, None, x, target) + bf16_avg = sum(bf16_increments) / len(bf16_increments) + + del bf16_model, bf16_optimizer + gc.collect() + torch.cuda.empty_cache() + + # Hybrid model + hybrid_model = _build_hybrid_model(NUM_LAYERS, hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_optimizer = te.optimizers.FusedAdam( + hybrid_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(hybrid_model, hybrid_optimizer, hybrid_recipe, x, target) + hybrid_increments = _measure_forward_increments( + hybrid_model, hybrid_optimizer, hybrid_recipe, x, target, + ) + hybrid_avg = sum(hybrid_increments) / len(hybrid_increments) + + excess_per_layer = hybrid_avg - bf16_avg + tolerance_per_layer = 50 * 1024 # 50 KiB + + assert excess_per_layer <= tolerance_per_layer, ( + "Hybrid per-layer forward memory increment exceeds bf16 baseline by " + f"{excess_per_layer/1024:.1f} KiB/layer (tolerance: {tolerance_per_layer/1024:.1f} KiB). " + f"bf16 avg: {bf16_avg/1024:.1f} KiB/layer, hybrid avg: {hybrid_avg/1024:.1f} KiB/layer." + ) + + +def test_hybrid_transpose_cache_after_backward(hybrid_recipe_name): + """Detect transpose caches from hybrid sub-storages persisting after backward.""" + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1)." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # bf16 baseline + bf16_model = _build_model(NUM_LAYERS, fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_optimizer = te.optimizers.FusedAdam( + bf16_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(bf16_model, bf16_optimizer, None, x, target) + bf16_bwd_delta = _measure_backward_memory_delta(bf16_model, bf16_optimizer, None, x, target) + + del bf16_model, bf16_optimizer + gc.collect() + torch.cuda.empty_cache() + + # Hybrid model + hybrid_model = _build_hybrid_model(NUM_LAYERS, hybrid_recipe) + hybrid_model = _shard_model(hybrid_model, world_size) + hybrid_optimizer = te.optimizers.FusedAdam( + hybrid_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(hybrid_model, hybrid_optimizer, hybrid_recipe, x, target) + hybrid_bwd_delta = _measure_backward_memory_delta( + hybrid_model, hybrid_optimizer, hybrid_recipe, x, target, + ) + + excess = hybrid_bwd_delta - bf16_bwd_delta + tolerance = 256 * 1024 # 256 KiB + + assert excess <= tolerance, ( + f"Hybrid backward retains {excess/1024**2:.2f} MiB more than bf16 baseline. " + f"bf16 backward delta: {bf16_bwd_delta/1024**2:.2f} MiB, " + f"hybrid backward delta: {hybrid_bwd_delta/1024**2:.2f} MiB." + ) + + # ── Standalone runner ──────────────────────────────────────────────── TESTS = { "bf16_no_excess_forward_memory": test_bf16_no_excess_forward_memory, "bf16_no_excess_backward_memory": test_bf16_no_excess_backward_memory, "fp8_temp_accumulation_across_layers": test_fp8_temp_accumulation_across_layers, "transpose_cache_retained_after_backward": test_transpose_cache_retained_after_backward, + "hybrid_no_excess_forward_memory": test_hybrid_no_excess_forward_memory, + "hybrid_transpose_cache_after_backward": test_hybrid_transpose_cache_after_backward, } if __name__ == "__main__": @@ -489,6 +629,10 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in "Float8BlockScaling", "MXFP8BlockScaling", "NVFP4BlockScaling", + "HybridFP8CurrentScaling", + "HybridMXFP8", + "HybridFloat8BlockScaling", + "HybridMixed_MXFP8_FP8", ], ) parser.add_argument("--quantized-model-init", action="store_true", default=False) @@ -504,11 +648,17 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in "fp8_temp_accumulation_across_layers", "transpose_cache_retained_after_backward", } + _HYBRID_PARAMETRIZED_TESTS = { + "hybrid_no_excess_forward_memory", + "hybrid_transpose_cache_after_backward", + } try: test_fn = TESTS[args.test] if args.test in _PARAMETRIZED_TESTS: test_fn(args.recipe, args.quantized_model_init) + elif args.test in _HYBRID_PARAMETRIZED_TESTS: + test_fn(args.recipe) else: test_fn() finally: diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index fce565ed9a1..1ab4fa1fd59 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -387,5 +387,116 @@ def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): _run_training(args) +def test_distributed_hybrid(hybrid_recipe_name): + """FSDP2 training with hybrid quantized_model_init. + + Uses quantized_model_init with a hybrid CustomRecipe and verifies that + training completes without error with a TransformerLayer model. + """ + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1)." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + device="meta", + ) + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + *[te.TransformerLayer(512, 2048, 8, **kwargs) for _ in range(2)] + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + model = shard_model_with_fsdp2(model, mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + optimizer = optim.Adam(model.parameters(), lr=1e-3) + + inp_shape = (128, 16, 512) + out_shape = (128, 16, 512) + + for iteration in range(3): + optimizer.zero_grad() + input_data = torch.randn(inp_shape, device=device, dtype=torch.bfloat16) + target = torch.randn(out_shape, device=device, dtype=torch.bfloat16) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(input_data) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + dist_print(f"Hybrid iteration {iteration} completed with loss {loss.item()}") + + +def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): + """FSDP2 training with hybrid params and reshard_after_forward=True. + + A single LayerNormLinear as the root module gets reshard_after_forward=True + from FSDP2. This exercises the forward-reshard-backward-reshard cycle where + split/as_strided/slice dispatch ops fire every iteration. + """ + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1)." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + in_features = 512 + out_features = in_features * 3 + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = te.LayerNormLinear( + in_features, out_features, + params_dtype=torch.bfloat16, + device="meta", + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + fully_shard(model, mesh=mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + optimizer = optim.Adam(model.parameters(), lr=1e-3) + + for iteration in range(5): + optimizer.zero_grad() + x = torch.randn(128, 16, in_features, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, out_features, device=device, dtype=torch.bfloat16) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + dist_print(f"Hybrid reshard_after_fwd iter {iteration}, loss {loss.item():.4f}") + + if __name__ == "__main__": sys.exit(_train(_parse_args())) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index f386659b6cd..e3c33b0b36d 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -165,6 +165,57 @@ def test_fsdp2_fused_adam_dcp_resharding(recipe): assert result.returncode == 0, f"DCP resharding load phase failed: {result.returncode}" +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_hybrid_fused_adam_tests(): + """FSDP2 FusedAdam tests with hybrid quantized params (parametrized by hybrid recipe).""" + test_path = _FSDP2_DIR / "run_fsdp2_fused_adam.py" + nproc = min(NUM_PROCS, 2) + run_distributed( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + "-k", + "hybrid", + ], + valid_returncodes=(0, 5), + env=os.environ, + timeout=600, + ) + + +@pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_hybrid_model_tests(): + """FSDP2 model tests with hybrid quantized params (parametrized by hybrid recipe).""" + test_path = _FSDP2_DIR / "run_fsdp2_model.py" + run_distributed( + [ + "torchrun", + f"--nproc_per_node={NUM_PROCS}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + "-k", + "hybrid", + ], + valid_returncodes=(0, 5), + env=os.environ, + timeout=600, + ) + + def test_dummy() -> None: """Dummy test diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 008a9a5b9e1..bb5298bbf88 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -2766,3 +2766,567 @@ def test_checkpoint_resume_training(self): ) finally: os.unlink(tmp_path) + + +# --------------------------------------------------------------------------- +# 11. FSDP2 prerequisites: __torch_dispatch__ ops that FSDP2 relies on +# --------------------------------------------------------------------------- + +aten = torch.ops.aten + + +def _make_hybrid_param_for_dispatch(row_factory, col_factory, grad_factory=None, + in_features=256, out_features=256): + """Create a HybridQuantizedTensor weight via quantized_model_init for dispatch tests.""" + hybrid_recipe = _hybrid_custom_recipe(row_factory, col_factory, grad_factory) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + return model.weight + + +def _fp8_row_factory(): + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + +def _fp8_col_factory(): + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + +def _fp8_grad_factory(): + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") + + +def _mxfp8_factory(): + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + + +_dispatch_configs = [ + pytest.param("fp8_fp8", id="same-format-fp8"), +] +if mxfp8_available: + _dispatch_configs.append(pytest.param("mxfp8_mxfp8", id="same-format-mxfp8")) + + +def _get_dispatch_hybrid_param(config_name): + """Return a HybridQuantizedTensor weight for the given config.""" + if config_name == "fp8_fp8": + return _make_hybrid_param_for_dispatch( + _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + ) + elif config_name == "mxfp8_mxfp8": + return _make_hybrid_param_for_dispatch( + _mxfp8_factory, _mxfp8_factory, + grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + ) + else: + raise ValueError(f"Unknown config: {config_name}") + + +@requires_fp8 +class TestHybridTorchDispatchFSDP2Ops: + """Test aten ops that FSDP2 relies on to preserve the HybridQuantizedTensor type. + + Each op is called directly via torch.ops.aten and the result is verified to + still be HybridQuantizedTensor with valid sub-storages. + """ + + @pytest.fixture(params=_dispatch_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _get_dispatch_hybrid_param(request.param) + + def test_split_preserves_hybrid_type(self, hybrid_param): + """torch.split must return a list of HybridQuantizedTensor pieces.""" + dim0 = hybrid_param.shape[0] + chunk_size = dim0 // 2 + pieces = torch.split(hybrid_param, chunk_size, dim=0) + assert len(pieces) >= 2 + for piece in pieces: + assert isinstance(piece, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(piece).__name__}" + ) + assert piece.rowwise_sub_storage is not None + assert piece.columnwise_sub_storage is not None + + total_rows = sum(p.shape[0] for p in pieces) + assert total_rows == dim0 + + orig_deq = hybrid_param.dequantize() + reassembled = torch.cat([p.dequantize() for p in pieces], dim=0) + torch.testing.assert_close(orig_deq, reassembled) + + def test_split_sub_storage_types_preserved(self, hybrid_param): + """After split, sub-storage types must match the original.""" + orig_row_type = type(hybrid_param.rowwise_sub_storage) + orig_col_type = type(hybrid_param.columnwise_sub_storage) + + chunk_size = hybrid_param.shape[0] // 2 + pieces = torch.split(hybrid_param, chunk_size, dim=0) + for piece in pieces: + assert type(piece.rowwise_sub_storage) is orig_row_type + assert type(piece.columnwise_sub_storage) is orig_col_type + + def test_view_preserves_hybrid_type(self, hybrid_param): + """view must return a HybridQuantizedTensor (used by FSDP2 reset_sharded_param).""" + shape_2d = hybrid_param.shape + result = aten.view.default(hybrid_param, list(shape_2d)) + assert isinstance(result, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(result).__name__}" + ) + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_view_same_shape_preserves_hybrid(self, hybrid_param): + """view with same shape must return HybridQuantizedTensor.""" + shape_2d = list(hybrid_param.shape) + result = aten.view.default(hybrid_param, shape_2d) + assert isinstance(result, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(result).__name__}" + ) + + def test_as_strided_noop_preserves_hybrid(self, hybrid_param): + """as_strided with matching shape/strides is a no-op that preserves type.""" + shape = tuple(hybrid_param.size()) + strides = (shape[-1], 1) + result = aten.as_strided.default(hybrid_param, list(shape), list(strides)) + assert isinstance(result, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(result).__name__}" + ) + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_slice_noop_preserves_hybrid(self, hybrid_param): + """slice with full range is a no-op that preserves type.""" + result = aten.slice.Tensor(hybrid_param, 0, 0, hybrid_param.size(0)) + assert isinstance(result, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(result).__name__}" + ) + assert result.rowwise_sub_storage is not None + + def test_copy_between_hybrid_tensors(self, hybrid_param): + """copy_ between compatible HybridQuantizedTensors copies quantized data directly.""" + src_deq = hybrid_param.dequantize().clone() + dst = hybrid_param._quantizer.make_empty( + shape=hybrid_param.shape, dtype=hybrid_param.dtype, device=hybrid_param.device, + ) + assert isinstance(dst, HybridQuantizedTensor) + + aten.copy_.default(dst, hybrid_param) + dst_deq = dst.dequantize() + torch.testing.assert_close(src_deq, dst_deq) + + def test_copy_from_bf16_to_hybrid(self, hybrid_param): + """copy_ from BF16 into HybridQuantizedTensor triggers quantize_.""" + param = hybrid_param.detach() + bf16_data = torch.randn_like(param.dequantize()) + aten.copy_.default(param, bf16_data) + result_deq = param.dequantize() + assert isinstance(param, HybridQuantizedTensor) + assert result_deq.shape == bf16_data.shape + + def test_new_zeros_returns_hybrid(self, hybrid_param): + """new_zeros must return a usable HybridQuantizedTensor container. + + FSDP2 calls ``new_zeros`` only to allocate an all-gather destination + buffer that is immediately overwritten by ``copy_``; the initial + contents are never observed. The hybrid dispatch therefore delegates + to ``HybridQuantizer.make_empty`` (uninitialized bytes) rather than + quantizing a BF16 zeros temporary. This test asserts the contract + we actually depend on — correct container type / shape / sub-storage + presence, and the ability to copy into it and read back — NOT that + the raw dequantize value happens to be zero. + """ + new_shape = list(hybrid_param.shape) + result = aten.new_zeros.default(hybrid_param, new_shape) + + # Structural contract: FSDP2 needs a HybridQuantizedTensor with both + # sub-storages populated so the gathered buffers have a destination. + assert isinstance(result, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(result).__name__}" + ) + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + assert type(result.rowwise_sub_storage) is type(hybrid_param.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(hybrid_param.columnwise_sub_storage) + + # Functional contract: the container must be writable via copy_ from + # another hybrid (how FSDP2 populates the buffer post-gather). + aten.copy_.default(result, hybrid_param) + torch.testing.assert_close(result.dequantize(), hybrid_param.dequantize()) + + def test_empty_like_returns_hybrid(self, hybrid_param): + """empty_like must return a HybridQuantizedTensor.""" + result = aten.empty_like.default(hybrid_param) + assert isinstance(result, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(result).__name__}" + ) + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + + def test_clone_returns_hybrid(self, hybrid_param): + """clone must return an independent HybridQuantizedTensor with same data.""" + result = aten.clone.default(hybrid_param) + assert isinstance(result, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(result).__name__}" + ) + assert result is not hybrid_param + torch.testing.assert_close(result.dequantize(), hybrid_param.dequantize()) + + +# --------------------------------------------------------------------------- +# 12. FSDP2 prerequisites: fsdp_pre_all_gather protocol +# --------------------------------------------------------------------------- + + +def _make_fsdp_protocol_param(config_name): + """Create a HybridQuantizedTensor weight for FSDP protocol tests.""" + if config_name == "fp8_fp8": + r = _hybrid_custom_recipe(_fp8_row_factory, _fp8_col_factory, _fp8_grad_factory) + elif config_name == "mxfp8_fp8": + r = _hybrid_custom_recipe(_mxfp8_factory, _fp8_col_factory, _fp8_grad_factory) + else: + raise ValueError(f"Unknown config: {config_name}") + with quantized_model_init(enabled=True, recipe=r): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model.weight + + +_fsdp_protocol_configs = [pytest.param("fp8_fp8", id="same-format")] +if mxfp8_available: + _fsdp_protocol_configs.append(pytest.param("mxfp8_fp8", id="mixed-mxfp8-fp8")) + + +@requires_fp8 +class TestHybridFsdpPreAllGatherProtocol: + """Test the fsdp_pre_all_gather method on HybridQuantizedTensor. + + These tests call the method directly (no actual all-gather communication) + to verify the protocol contract: returns (sharded_tensors, metadata) where + sharded_tensors is a tuple of plain torch.Tensor. + """ + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_pre_all_gather_returns_tuple_pair(self, hybrid_param): + """fsdp_pre_all_gather returns (sharded_tensors, metadata).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + assert isinstance(sharded_tensors, tuple), ( + f"sharded_tensors should be tuple, got {type(sharded_tensors).__name__}" + ) + assert len(sharded_tensors) > 0, "sharded_tensors should not be empty" + assert isinstance(metadata, tuple), ( + f"metadata should be tuple, got {type(metadata).__name__}" + ) + + def test_pre_all_gather_buffers_are_plain_tensors(self, hybrid_param): + """Every element in sharded_tensors must be a plain torch.Tensor.""" + sharded_tensors, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + for i, t in enumerate(sharded_tensors): + assert isinstance(t, torch.Tensor), ( + f"sharded_tensors[{i}] should be torch.Tensor, got {type(t).__name__}" + ) + assert not isinstance(t, QuantizedTensor), ( + f"sharded_tensors[{i}] should NOT be QuantizedTensor subclass" + ) + + def test_pre_all_gather_buffer_count_consistent(self, hybrid_param): + """Buffer count must be the same across repeated calls (FSDP2 buffer reuse).""" + sharded_1, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + sharded_2, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + assert len(sharded_1) == len(sharded_2), ( + f"Buffer count changed: {len(sharded_1)} vs {len(sharded_2)}" + ) + + def test_pre_all_gather_metadata_sufficient_for_reconstruction(self, hybrid_param): + """Metadata must contain enough info to reconstruct the tensor.""" + _, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + assert metadata is not None + assert len(metadata) > 0, "metadata should not be empty" + + +# --------------------------------------------------------------------------- +# 13. FSDP2 prerequisites: fsdp_post_all_gather protocol +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFsdpPostAllGatherProtocol: + """Test the fsdp_post_all_gather method on HybridQuantizedTensor. + + Simulates the post-all-gather phase by passing the sharded_tensors + from pre_all_gather directly (mimicking a single-rank all-gather). + """ + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_post_all_gather_first_call_returns_hybrid_tensor(self, hybrid_param): + """With out=None, post_all_gather returns (HybridQuantizedTensor, outputs).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + result, ag_outputs = hybrid_param.fsdp_post_all_gather( + sharded_tensors, metadata, hybrid_param.dtype, out=None, + ) + assert isinstance(result, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(result).__name__}" + ) + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_post_all_gather_buffer_reuse(self, hybrid_param): + """On second call with out=previous, the same object is returned (buffer reuse).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + first_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, metadata, hybrid_param.dtype, out=None, + ) + + second_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, metadata, hybrid_param.dtype, out=first_result, + ) + assert second_result is first_result, ( + "Buffer reuse: post_all_gather(out=prev) should return the same object" + ) + + def test_post_all_gather_dequantize_matches_original(self, hybrid_param): + """Reconstructed tensor should dequantize close to the original.""" + orig_deq = hybrid_param.dequantize() + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, metadata, hybrid_param.dtype, out=None, + ) + result_deq = result.dequantize() + torch.testing.assert_close(orig_deq, result_deq) + + def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): + """Reconstructed tensor's sub-storages match the original types.""" + orig_row_type = type(hybrid_param.rowwise_sub_storage) + orig_col_type = type(hybrid_param.columnwise_sub_storage) + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, metadata, hybrid_param.dtype, out=None, + ) + assert type(result.rowwise_sub_storage) is orig_row_type + assert type(result.columnwise_sub_storage) is orig_col_type + + +# --------------------------------------------------------------------------- +# 14. FSDP2 prerequisites: pre/post roundtrip +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFsdpRoundtrip: + """End-to-end single-process roundtrip (pre -> post) without communication.""" + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_pre_post_roundtrip_preserves_data(self, hybrid_param): + """pre_all_gather -> post_all_gather(out=None) -> dequantize matches original.""" + orig_deq = hybrid_param.dequantize() + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, metadata, hybrid_param.dtype, out=None, + ) + torch.testing.assert_close(orig_deq, result.dequantize()) + + def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): + """Second roundtrip with out=previous preserves data (iteration 2+ simulation).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + first_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, metadata, hybrid_param.dtype, out=None, + ) + + sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + second_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_2, metadata_2, hybrid_param.dtype, out=first_result, + ) + assert second_result is first_result + torch.testing.assert_close(hybrid_param.dequantize(), second_result.dequantize()) + + def test_scale_refresh_across_iterations(self): + """After a sharded optimizer-style requantize, iter-2+ gathers see the new scale. + + Per-tensor FP8 does NOT include ``_scale_inv`` in ``fsdp_buffer_fields`` + (only ``_data`` is gathered; the scalar scale travels via iter-1 + metadata). This relies on the invariant that the sharded and gathered + ``Float8Tensor`` s share the same ``_scale_inv`` tensor object, and + that ``Float8CurrentScalingQuantizer.update_quantized`` writes the new + scale in place rather than replacing the tensor reference. If either + invariant broke, the gathered copy would carry a stale scale on + iter-2+ and silently apply the wrong dequantization. + + This test locks the invariant down by forcing a radically different + scale between iterations and asserting the gathered tensor's + dequantization tracks the sharded one. + """ + torch.manual_seed(42) + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + hybrid_param = model.weight + + # Iter-1 gather with the initial (small-magnitude) weights + sharded_tensors_1, metadata_1 = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + gathered, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_1, metadata_1, hybrid_param.dtype, out=None, + ) + + # Simulate an optimizer writeback that produces a much larger weight; + # Float8CurrentScalingQuantizer.update_quantized must recompute + # _scale_inv for this range. If the gathered copy didn't see the new + # scale, the dequantize below would disagree with the sharded copy. + huge_master = torch.randn_like(hybrid_param.dequantize()) * 100.0 + hybrid_param._quantizer.update_quantized(huge_master, hybrid_param) + + # Iter-2+ path: reuse the gathered buffer + sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( + mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + gathered_refreshed, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_2, metadata_2, hybrid_param.dtype, out=gathered, + ) + assert gathered_refreshed is gathered + + # The gathered copy must now reflect the new sharded scale, not the + # tiny original scale. + torch.testing.assert_close( + hybrid_param.dequantize(), + gathered_refreshed.dequantize(), + ) + # And the magnitude really did change (sanity: this test would pass + # vacuously if update_quantized didn't actually change anything). + assert gathered_refreshed.dequantize().abs().max() > 10.0, ( + "update_quantized did not produce a sufficiently different " + "weight; the scale-refresh invariant is not being exercised" + ) + + def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): + """Hybrid FSDP2 with an NVFP4 sub-storage must raise a clear error. + + Per the hybrid FSDP2 design (see ``hybrid_quantization_fsdp.md`` §9 + Gap 5), NVFP4 FSDP2 support is not implemented yet — packed FP4 data + alignment for dim-0 splitting, columnwise dequant, and RHT cache + handling all need work. Until that lands, hybrid pre-all-gather must + refuse an NVFP4 sub-storage cleanly via the ``fsdp_buffer_fields`` + protocol rather than silently producing wrong data. + + This test pins that contract: any hybrid whose sub-storage does not + implement ``fsdp_buffer_fields`` raises ``NotImplementedError`` at + ``fsdp_pre_all_gather`` time. The prior version of this test + inadvertently asserted the opposite when buffer extraction used + implicit ``get_metadata()``-based tensor scanning. + """ + if not (fp8_available and nvfp4_available): + pytest.skip("Requires FP8 + NVFP4 support") + + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: NVFP4Quantizer(), + col_factory=lambda: NVFP4Quantizer(), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + param = model.weight + + # Clean refusal: hybrid's pre_all_gather raises an NVFP4-specific + # message pointing to the design doc, not a generic + # "NVFP4Tensor does not implement fsdp_buffer_fields" from deep inside + # the base class. + with pytest.raises(NotImplementedError) as exc_info: + param.fsdp_pre_all_gather( + mesh=None, orig_size=param.shape, contiguous_orig_stride=None, + module=None, mp_policy=None, + ) + msg = str(exc_info.value) + assert "NVFP4Tensor" in msg + assert "hybrid_quantization_fsdp.md" in msg + assert "fsdp_buffer_fields" in msg + + +# --------------------------------------------------------------------------- +# 15. FSDP2 prerequisites: make_like correctness +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridMakeLike: + """Test that make_like produces correct copies for __torch_dispatch__ usage.""" + + def _make_hybrid_param(self): + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model.weight + + def test_make_like_preserves_sub_storages(self): + """make_like result has the same sub-storage types, quantizers, and dtype.""" + param = self._make_hybrid_param() + copy = HybridQuantizedTensor.make_like(param) + + assert isinstance(copy, HybridQuantizedTensor) + assert copy.dtype == param.dtype + assert copy.shape == param.shape + assert type(copy.rowwise_sub_storage) is type(param.rowwise_sub_storage) + assert type(copy.columnwise_sub_storage) is type(param.columnwise_sub_storage) + torch.testing.assert_close(copy.dequantize(), param.dequantize()) + + def test_make_like_is_independent(self): + """make_like result should not share the same tensor identity.""" + param = self._make_hybrid_param() + copy = HybridQuantizedTensor.make_like(param) + assert copy is not param diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index a7722f777e3..b378fb7d1ec 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -130,6 +130,65 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: f"{self.__class__.__name__} class does not implement copy_from_storage function" ) + # ── FSDP2 buffer protocol ─────────────────────────────────────── + # + # These three methods decouple FSDP2 all-gather buffer extraction from + # format-specific padding/layout tricks. `HybridQuantizedTensor` uses them + # to aggregate buffers from its two sub-storages without knowing each + # sub-storage's internal field layout. + # + # Contract: + # * ``fsdp_buffer_fields`` returns an ordered tuple of attribute names + # on *self* that hold the tensor buffers that must be all-gathered. + # Scalars/metadata that only need broadcasting (e.g. per-tensor FP8 + # ``_scale_inv``) are NOT listed here — they travel via the hook's + # metadata tuple instead. + # * ``fsdp_extract_buffers`` returns ``(buffers, reassembly_meta)``. + # The default implementation reads the fields as-is. Sub-storages with + # gather-time padding (MXFP8 block scales) override this to strip the + # padding before gather. + # * ``fsdp_assign_gathered`` writes the gathered buffers back into the + # storage's fields. Sub-storages with gather-time padding override + # this to re-apply the padding before assignment. + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Ordered attribute names holding tensor buffers gathered by FSDP2.""" + raise NotImplementedError( + f"{self.__class__.__name__} class does not implement fsdp_buffer_fields" + ) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Return ``(buffers, reassembly_meta)`` for FSDP2 all-gather. + + Default implementation reads the fields named by ``fsdp_buffer_fields`` + verbatim. Override when the on-disk layout differs from the + gather-ready layout (e.g. MXFP8 block scales carry alignment padding). + """ + names = self.fsdp_buffer_fields() + buffers = tuple(getattr(self, name) for name in names) + return buffers, {"field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered buffers into the fields named in ``meta``. + + Override when the gather-ready layout needs a format-specific transform + (e.g. MXFP8 scales must be padded back to ``[128, 4]`` / ``[4, 128]``). + """ + names = meta["field_names"] + if len(names) != len(gathered): + raise RuntimeError( + f"{type(self).__name__}.fsdp_assign_gathered got " + f"{len(gathered)} buffers for {len(names)} fields" + ) + for name, buf in zip(names, gathered): + setattr(self, name, buf) + def prepare_for_saving( *tensors: Union[torch.Tensor, QuantizedTensorStorage], diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 256250ff641..3beb90926a6 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -547,8 +547,12 @@ def detach(self) -> Float8Tensor: def clone(self) -> Float8Tensor: # pylint: disable=missing-function-docstring - assert self._data is not None - data = self._data.detach().clone() + # ``_data`` may be None for columnwise-only sub-storages of a + # HybridQuantizedTensor on architectures without native non-TN FP8 + # GEMM (Hopper / L40), where columnwise-only Float8 allocates + # ``_transpose`` instead of ``_data``. On Blackwell+ the C++ + # override keeps ``_data`` populated even in columnwise-only mode. + data = self._data.detach().clone() if self._data is not None else None data_transpose = None if self._transpose is not None: data_transpose = self._transpose.detach().clone() @@ -710,23 +714,25 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.split.Tensor: tensor = args[0] data = tensor._data - func_out = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - t_func_out = [None] * len(func_out) - # Compute corresponding split of the transpose cache if available + # _data may be None for columnwise-only sub-storages (hybrid quantization) + if data is not None: + func_out = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) + else: + func_out = None + + t_func_out = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose - ndim = data.dim() - # Figure out the original split dim + ndim = tensor.dim() if "dim" in kwargs: dim_to_split = kwargs["dim"] else: dim_to_split = args[2] if len(args) > 2 else 0 - # Dimension along which transpose needs to be split t_dim = 0 if dim_to_split == ndim - 1 else dim_to_split + 1 t_func_out = transpose.__torch_dispatch__( func, @@ -734,12 +740,23 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [transpose, args[1], t_dim], kwargs, ) + + ref_out = func_out if func_out is not None else t_func_out + if ref_out is None: + return super().__torch_dispatch__(func, types, args, kwargs) + + num_splits = len(ref_out) + if func_out is None: + func_out = [None] * num_splits + if t_func_out is None: + t_func_out = [None] * num_splits + outs = [ Float8Tensor.make_like( tensor, data=split_tensor, data_transpose=split_transpose_tensor, - shape=split_tensor.shape, + shape=(split_tensor.shape if split_tensor is not None else split_transpose_tensor.shape), ) for split_tensor, split_transpose_tensor in zip(func_out, t_func_out) ] diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 4d1a5ec7adc..b132cab10b5 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -149,11 +149,27 @@ def _get_compatible_recipe(self): # quantized_model_init + built-in MXFP8BlockScaling autocast). # We trust that users who write a CustomRecipe know what they're doing # with regard to per-operand scaling mode compatibility. - # TODO(negvet): improve to validate that the autocast recipe's - # sub-quantizer scaling modes are compatible with each sub-storage's - # scaling mode (e.g. rowwise MXFP8 weight requires MXFP8 input for - # fprop TN, columnwise NVFP4 weight requires NVFP4 grad_output for - # wgrad NT). + # + # TODO(negvet): validate per-operand scaling-mode compatibility at + # recipe-build time instead of at cuBLAS-dispatch time. Concretely: + # 1. Walk the qfactory outputs for a given module_type (``linear``, + # ``grouped_linear``, ``dpa``) — call the factory for each + # ``QuantizerRole.tensor_type`` the module uses. + # 2. Extract the scaling_mode of each sub-quantizer: + # weight_row, weight_col (from HybridQuantizer) + # input_row, input_col (from HybridQuantizer) + # grad_output_row, grad_output_col (plain quantizer OR + # HybridQuantizer) + # 3. Assert the three GEMM pairs share a scaling_mode each: + # fprop TN: weight_row == input_row (FormatA) + # dgrad NN: weight_col == grad_output_row (FormatB) + # wgrad NT: input_col == grad_output_col (FormatC) + # Mismatches raise ``ValueError`` naming the offending slots, e.g. + # "dgrad GEMM: weight columnwise format (MXFP8) does not match + # grad_output rowwise format (NVFP4)". + # 4. Blocked on `semantic_quantizer_roles` / PR #2620 for the + # ``QuantizerRole`` dataclass — the factory signature is role- + # aware only on that branch. from transformer_engine.common.recipe import CustomRecipe # avoid circular import return CustomRecipe @@ -233,9 +249,336 @@ def detach(self) -> HybridQuantizedTensor: def get_metadata(self) -> Dict[str, Any]: return HybridQuantizedTensorStorage.get_metadata(self) + # ── FSDP2 protocol ────────────────────────────────────────────── + + def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): + """Extract plain tensor buffers from both sub-storages for FSDP2 all-gather. + + Always send both directions. This gives a stable buffer count/shape + across forward and backward, at the cost of gathering the unused + direction each pass. No requantization, no BF16 fallback. + + Buffer extraction is delegated to each sub-storage's + :meth:`QuantizedTensorStorage.fsdp_extract_buffers`, which strips any + format-specific padding (e.g. MXFP8 block-scale alignment) before the + gather so concatenation along dim-0 is well-defined. + + TODO(negvet): bandwidth optimization — pack both directions into a + single flat buffer sized ``max(row_bytes, col_bytes)`` (not + ``row_bytes + col_bytes``) to halve comm volume for asymmetric format + pairs. Planned implementation: a new per-sub-storage + ``fsdp_pack_into(flat_buffer, offset, meta)`` helper that layouts + both directions back-to-back with offsets stored in the metadata + tuple; ``fsdp_post_all_gather`` would slice the gathered flat buffer + using those offsets. + """ + # Quick, targeted error for sub-storages whose FSDP2 support isn't + # implemented yet (e.g. NVFP4). Without this, users hit + # NotImplementedError from deep inside fsdp_extract_buffers with a + # generic message. + for role, sub in ( + ("rowwise", self._rowwise_storage), + ("columnwise", self._columnwise_storage), + ): + if sub is None: + continue + try: + sub.fsdp_buffer_fields() + except NotImplementedError as err: + raise NotImplementedError( + f"Hybrid FSDP2 all-gather is not supported for a " + f"{type(sub).__name__} {role} sub-storage: it does not " + f"implement fsdp_buffer_fields. " + f"See hybrid_quantization_fsdp.md section 9 (Gap 5) — " + f"NVFP4 sub-storages need packed-FP4 dim-0 alignment, " + f"columnwise dequantization and RHT-cache handling before " + f"they can be gathered. Use a supported sub-quantizer " + f"(Float8CurrentScaling, MXFP8, Float8Block) or run without " + f"FSDP2." + ) from err + + row_buffers: Tuple[Optional[torch.Tensor], ...] = () + col_buffers: Tuple[Optional[torch.Tensor], ...] = () + row_meta: Optional[Dict[str, Any]] = None + col_meta: Optional[Dict[str, Any]] = None + if self._rowwise_storage is not None: + row_buffers, row_meta = self._rowwise_storage.fsdp_extract_buffers() + if self._columnwise_storage is not None: + col_buffers, col_meta = self._columnwise_storage.fsdp_extract_buffers() + + sharded_tensors = row_buffers + col_buffers + + metadata = ( + len(row_buffers), + row_meta, + col_meta, + self._rowwise_storage, # original sharded sub-storage (for make_like on iter-1) + self._columnwise_storage, + self._rowwise_quantizer, + self._columnwise_quantizer, + self._quantizer, + ) + return sharded_tensors, metadata + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[HybridQuantizedTensor] = None, + ): + """Reconstruct HybridQuantizedTensor from all-gathered buffers. + + On iteration 1 (``out=None``): clone each sub-storage via + :meth:`make_like` from the sharded original, then delegate the + gathered-buffer writeback (and any format-specific re-padding) to + :meth:`QuantizedTensorStorage.fsdp_assign_gathered`. + On iteration 2+ (``out=prev``): delegate directly to the existing + sub-storages' ``fsdp_assign_gathered``. + """ + ( + n_row_buffers, + row_meta, + col_meta, + orig_row_sub, + orig_col_sub, + row_quantizer, + col_quantizer, + hybrid_quantizer, + ) = metadata + + row_gathered = all_gather_outputs[:n_row_buffers] + col_gathered = all_gather_outputs[n_row_buffers:] + + def _infer_shape(gathered_buffers): + for buf in gathered_buffers: + if buf is not None: + return buf.shape + return None + + if out is not None: + # Iteration 2+: in-place field update on existing sub-storages + if out._rowwise_storage is not None and row_meta is not None: + out._rowwise_storage.fsdp_assign_gathered(row_gathered, row_meta) + if out._columnwise_storage is not None and col_meta is not None: + out._columnwise_storage.fsdp_assign_gathered(col_gathered, col_meta) + else: + # First iteration: clone the original sharded sub-storages via make_like, + # then write gathered (full-size) buffers via each sub-storage's own + # fsdp_assign_gathered so padding is re-applied where applicable. + row_sub = None + if orig_row_sub is not None and isinstance(orig_row_sub, QuantizedTensor): + gathered_shape = _infer_shape(row_gathered) + row_sub = type(orig_row_sub).make_like(orig_row_sub, shape=gathered_shape) + if row_meta is not None: + row_sub.fsdp_assign_gathered(row_gathered, row_meta) + + col_sub = None + if orig_col_sub is not None and isinstance(orig_col_sub, QuantizedTensor): + gathered_shape = _infer_shape(col_gathered) + col_sub = type(orig_col_sub).make_like(orig_col_sub, shape=gathered_shape) + if col_meta is not None: + col_sub.fsdp_assign_gathered(col_gathered, col_meta) + + ref_sub = row_sub if row_sub is not None else col_sub + out = HybridQuantizedTensor( + shape=( + ref_sub.shape + if ref_sub is not None + else _infer_shape(row_gathered + col_gathered) + ), + dtype=param_dtype, + rowwise_storage=row_sub, + columnwise_storage=col_sub, + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + quantizer=hybrid_quantizer, + ) + + return out, all_gather_outputs + + @classmethod + def _delegate_reshape_op(cls, func, tensor, args, kwargs): + """Delegate a shape-altering op (slice, as_strided) to each sub-storage. + + Returns a new ``HybridQuantizedTensor`` when every non-None sub-storage + returns a ``QuantizedTensorStorage`` of the same kind (i.e. real + op support, as Float8Tensor provides via its own + ``__torch_dispatch__``). Returns ``None`` when any sub-storage + dequantized to a plain ``torch.Tensor`` (i.e. the sub-storage does not + support this op — MXFP8Tensor / Float8BlockwiseQTensor fall through + that way for real slicing today). On ``None`` the caller should defer + to ``super().__torch_dispatch__`` for a consistent BF16 fallback. + """ + def _delegate(sub): + if sub is None: + return None + return func(sub, *args[1:], **kwargs) + + row_out = _delegate(tensor._rowwise_storage) + col_out = _delegate(tensor._columnwise_storage) + + row_ok = row_out is None or isinstance(row_out, QuantizedTensorStorage) + col_ok = col_out is None or isinstance(col_out, QuantizedTensorStorage) + if not (row_ok and col_ok): + return None + if row_out is None and col_out is None: + return None + + ref = row_out if row_out is not None else col_out + return HybridQuantizedTensor( + shape=ref.shape, + dtype=tensor.dtype, + rowwise_storage=row_out, + columnwise_storage=col_out, + rowwise_quantizer=tensor._rowwise_quantizer, + columnwise_quantizer=tensor._columnwise_quantizer, + quantizer=tensor._quantizer, + ) + @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): + if kwargs is None: + kwargs = {} + if func == aten.detach.default: return args[0].detach() + # ── FSDP2: view ────────────────────────────────────────────── + if func == aten.view.default: + tensor = args[0] + shape = args[1] + row_view = None + col_view = None + if tensor._rowwise_storage is not None: + row_view = tensor._rowwise_storage.view(shape) + if tensor._columnwise_storage is not None: + col_view = tensor._columnwise_storage.view(shape) + return HybridQuantizedTensor( + shape=shape, + dtype=tensor.dtype, + rowwise_storage=row_view, + columnwise_storage=col_view, + rowwise_quantizer=tensor._rowwise_quantizer, + columnwise_quantizer=tensor._columnwise_quantizer, + quantizer=tensor._quantizer, + ) + + # ── FSDP2: split ───────────────────────────────────────────── + if func == aten.split.Tensor: + tensor = args[0] + split_size = args[1] + dim = kwargs.get("dim", args[2] if len(args) > 2 else 0) + + if dim != 0: + return super().__torch_dispatch__(func, types, args, kwargs) + + row_pieces = ( + torch.split(tensor._rowwise_storage, split_size, dim=dim) + if tensor._rowwise_storage is not None else None + ) + col_pieces = ( + torch.split(tensor._columnwise_storage, split_size, dim=dim) + if tensor._columnwise_storage is not None else None + ) + + if row_pieces is None and col_pieces is None: + return super().__torch_dispatch__(func, types, args, kwargs) + + num_pieces = len(row_pieces) if row_pieces is not None else len(col_pieces) + return [ + HybridQuantizedTensor( + shape=(row_pieces[i].shape if row_pieces is not None else col_pieces[i].shape), + dtype=tensor.dtype, + rowwise_storage=row_pieces[i] if row_pieces is not None else None, + columnwise_storage=col_pieces[i] if col_pieces is not None else None, + rowwise_quantizer=tensor._rowwise_quantizer, + columnwise_quantizer=tensor._columnwise_quantizer, + quantizer=tensor._quantizer, + ) + for i in range(num_pieces) + ] + + # ── FSDP2: as_strided / slice ──────────────────────────────── + # Fast path for no-op (common during FSDP2 reset_sharded_param); + # otherwise delegate per sub-storage so we inherit each sub-storage's + # own support level. Float8Tensor implements real slicing/as_strided + # via `_data.__torch_dispatch__`; MXFP8Tensor and Float8BlockwiseQTensor + # handle only the no-op case and fall through to dequantize for real + # ops (matching their vanilla FSDP2 behaviour). If any sub-storage + # returns a plain torch.Tensor (dequantized), we can't rewrap into a + # hybrid so we fall through to super() for a consistent BF16 fallback. + if func == aten.as_strided.default: + tensor = args[0] + shape = args[1] + strides = args[2] + if ( + len(shape) == len(strides) == 2 + and tuple(strides) == (shape[-1], 1) + and tuple(shape) == tuple(tensor.size()) + ): + return HybridQuantizedTensor.make_like(tensor) + return cls._delegate_reshape_op(func, tensor, args, kwargs) or \ + super().__torch_dispatch__(func, types, args, kwargs) + + if func == aten.slice.Tensor: + tensor = args[0] + dim = args[1] + start = args[2] + length = args[3] + if start == 0 and length == tensor.size(dim): + return HybridQuantizedTensor.make_like(tensor) + return cls._delegate_reshape_op(func, tensor, args, kwargs) or \ + super().__torch_dispatch__(func, types, args, kwargs) + + # ── FSDP2: copy_ ───────────────────────────────────────────── + # Fast path for hybrid-to-hybrid (FSDP2 fills buffer allocated via + # new_zeros/make_empty). Other src types (e.g. a BF16 master weight + # during checkpoint load) fall through to QuantizedTensor's base + # dispatch which routes to ``dst.quantize_(src)``. + if func == aten.copy_.default: + dst, src = args[0], args[1] + if isinstance(dst, HybridQuantizedTensor) and isinstance(src, HybridQuantizedTensor): + if dst._rowwise_storage is not None and src._rowwise_storage is not None: + aten.copy_.default(dst._rowwise_storage, src._rowwise_storage) + if dst._columnwise_storage is not None and src._columnwise_storage is not None: + aten.copy_.default(dst._columnwise_storage, src._columnwise_storage) + return dst + + # ── FSDP2: new_zeros ───────────────────────────────────────── + if func == aten.new_zeros.default: + tensor = args[0] + new_shape = args[1] + if tensor._quantizer is not None: + # FSDP2 allocates new_zeros buffers as all-gather destinations + # that are immediately overwritten by copy_. Use make_empty + # (uninitialized storage with the right container shape/fields). + return tensor._quantizer.make_empty( + new_shape, + dtype=tensor.dtype, + device=tensor.device, + ) + + # ── FSDP2: clone ───────────────────────────────────────────── + if func == aten.clone.default: + tensor = args[0] + row_clone = ( + torch.clone(tensor._rowwise_storage) + if tensor._rowwise_storage is not None else None + ) + col_clone = ( + torch.clone(tensor._columnwise_storage) + if tensor._columnwise_storage is not None else None + ) + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=row_clone, + columnwise_storage=col_clone, + rowwise_quantizer=tensor._rowwise_quantizer, + columnwise_quantizer=tensor._columnwise_quantizer, + quantizer=tensor._quantizer, + ) + return super().__torch_dispatch__(func, types, args, kwargs) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 965f59b320e..3c960e653aa 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -344,8 +344,8 @@ def detach(self) -> MXFP8Tensor: def clone(self) -> MXFP8Tensor: # pylint: disable=missing-function-docstring - assert self._rowwise_data is not None - rowwise_data = self._rowwise_data.detach().clone() + # _rowwise_data may be None for columnwise-only sub-storages (hybrid quantization) + rowwise_data = self._rowwise_data.detach().clone() if self._rowwise_data is not None else None columnwise_data = None if self._columnwise_data is not None: columnwise_data = self._columnwise_data.detach().clone() @@ -458,66 +458,60 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ): return super().__torch_dispatch__(func, types, args, kwargs) - out_data = [] - for data in [tensor._rowwise_data, tensor._columnwise_data]: - func_out = ( - data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - if data is not None - else None + def _split_data(data): + if data is None: + return None + return data.__torch_dispatch__( + func, types, [data] + list(args[1:]), kwargs, ) - out_data.append(func_out) + + row_data_splits = _split_data(tensor._rowwise_data) + col_data_splits = _split_data(tensor._columnwise_data) scale_invs = [tensor._rowwise_scale_inv, tensor._columnwise_scale_inv] split_sizes_for_scale = [split_size, split_size // MXFP8_BLOCK_SCALING_SIZE] - # Padding requirements: rowwise dim0 should be divisble by 128, columnwise dim0 should be divisble by 4 padding_multiples = [128, 4] + scale_splits = [] for scale_inv, scale_split_size, pad_multiple in zip( scale_invs, split_sizes_for_scale, padding_multiples ): - scale_inv_out = ( - scale_inv.__torch_dispatch__( - func, - types, - [scale_inv, scale_split_size] + list(args[2:]), - kwargs, - ) - if scale_inv is not None - else None - ) - scale_inv_out = list(scale_inv_out) if scale_inv_out is not None else None - # Pad scale_inv_out to be a multiple of pad_multiple - if scale_inv_out is not None: - for idx, split_scale_inv_out in enumerate(scale_inv_out): - current_shape = split_scale_inv_out.shape - pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple - if pad_dim0 > 0: - scale_inv_out[idx] = torch.nn.functional.pad( - split_scale_inv_out, (0, 0, 0, pad_dim0) - ) - out_data.append(scale_inv_out) + if scale_inv is None: + scale_splits.append(None) + continue + scale_inv_out = list(scale_inv.__torch_dispatch__( + func, types, + [scale_inv, scale_split_size] + list(args[2:]), kwargs, + )) + for idx, split_scale_inv_out in enumerate(scale_inv_out): + current_shape = split_scale_inv_out.shape + pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple + if pad_dim0 > 0: + scale_inv_out[idx] = torch.nn.functional.pad( + split_scale_inv_out, (0, 0, 0, pad_dim0) + ) + scale_splits.append(scale_inv_out) + row_scale_splits, col_scale_splits = scale_splits + + ref_splits = row_data_splits if row_data_splits is not None else col_data_splits + num_splits = len(ref_splits) return [ MXFP8Tensor( shape=( - splitted_tensor_data[0].size() - if splitted_tensor_data[0] is not None - else splitted_tensor_data[1].size() + row_data_splits[i].size() + if row_data_splits is not None + else col_data_splits[i].size() ), dtype=tensor.dtype, - rowwise_data=splitted_tensor_data[0], - rowwise_scale_inv=splitted_tensor_data[2], - columnwise_data=splitted_tensor_data[1], - columnwise_scale_inv=splitted_tensor_data[3], + rowwise_data=row_data_splits[i] if row_data_splits is not None else None, + rowwise_scale_inv=row_scale_splits[i] if row_scale_splits is not None else None, + columnwise_data=col_data_splits[i] if col_data_splits is not None else None, + columnwise_scale_inv=col_scale_splits[i] if col_scale_splits is not None else None, quantizer=tensor._quantizer, requires_grad=False, fp8_dtype=tensor._fp8_dtype, with_gemm_swizzled_scales=False, ) - for splitted_tensor_data in zip(*out_data) + for i in range(num_splits) ] if func == torch.ops.aten.as_strided.default: @@ -607,6 +601,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, ) + if func == torch.ops.aten.clone.default: + return cls.clone(args[0]) + # Default case return super().__torch_dispatch__(func, types, args, kwargs) diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index de7f8f58e24..dceee83cfdc 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -226,6 +226,10 @@ def __repr__(self): def _create_transpose(self): """Update FP8 transpose cache""" data = self._data + # Columnwise-only Float8Tensors (e.g. hybrid quantization sub-storages) + # have _data=None — nothing to transpose. + if data is None: + return if not data.is_contiguous(): data = data.contiguous() self._transpose = tex.fp8_transpose(data, self._fp8_dtype, out=self._transpose) @@ -279,3 +283,11 @@ def get_usages(self) -> Dict[str, bool]: else: usages["columnwise"] = self._transpose is not None and not self._transpose_invalid return usages + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for per-tensor FP8. + + ``_scale_inv`` is a per-tensor scalar; it travels through the hook's + metadata tuple (mirroring :meth:`Float8Tensor.fsdp_pre_all_gather`). + """ + return ("_data",) diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index e407d252ef9..3fb224d6405 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -133,8 +133,18 @@ def device(self): return self._columnwise_storage.device raise RuntimeError("HybridQuantizedTensorStorage has no data") - def view(self, shape: torch.Size): - raise NotImplementedError("HybridQuantizedTensorStorage does not support view operations") + def view(self, *shape): + """View delegates to each sub-storage. Used by FSDP2 reset_sharded_param.""" + row_view = self._rowwise_storage.view(*shape) if self._rowwise_storage is not None else None + col_view = self._columnwise_storage.view(*shape) if self._columnwise_storage is not None else None + return HybridQuantizedTensorStorage( + rowwise_storage=row_view, + columnwise_storage=col_view, + rowwise_quantizer=self._rowwise_quantizer, + columnwise_quantizer=self._columnwise_quantizer, + quantizer=self._quantizer, + fake_dtype=self._dtype, + ) def get_metadata(self) -> Dict[str, Any]: return { diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 7bbe809c9d6..5da14ba0a4d 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -15,7 +15,7 @@ from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from ...constants import TE_DType as torch_to_transformer_engine_dtype +from ...constants import TE_DType as torch_to_transformer_engine_dtype, MXFP8_BLOCK_SCALING_SIZE from ...utils import _empty_tensor @@ -308,3 +308,77 @@ def get_usages(self) -> Dict[str, bool]: "rowwise": self._rowwise_data is not None, "columnwise": self._columnwise_data is not None, } + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for MXFP8. + + Block scales are per-block and direction-specific — each direction + gathers both its data buffer and its scale-inv buffer. ``None``-valued + directions (e.g. a columnwise-only sub-storage in hybrid quantization) + are excluded so the gather tuple only contains real tensors. + """ + fields = [] + if self._rowwise_data is not None: + fields.extend(("_rowwise_data", "_rowwise_scale_inv")) + if self._columnwise_data is not None: + fields.extend(("_columnwise_data", "_columnwise_scale_inv")) + return tuple(fields) + + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract MXFP8 buffers, unpadding block-scale alignment before gather. + + MXFP8 kernels require scale-inv tensors aligned to ``[128, 4]`` + (rowwise) and ``[4, 128]`` (columnwise). That padding is attached to + the local shard but would produce misaligned concatenation under + FSDP2's dim-0 all-gather. Strip it here and re-apply in + :meth:`fsdp_assign_gathered`. + """ + if self._with_gemm_swizzled_scales: + raise NotImplementedError( + "FSDP2 is only supported for MXFP8Tensors with compact scales" + ) + names = self.fsdp_buffer_fields() + buffers = [] + shape = self.size() + flattened_in_shape0 = math.prod(shape[:-1]) + for name in names: + t = getattr(self, name) + if name == "_rowwise_scale_inv" and t is not None: + if t.size(0) != flattened_in_shape0: + t = t[:flattened_in_shape0] + elif name == "_columnwise_scale_inv" and t is not None: + expected = flattened_in_shape0 // MXFP8_BLOCK_SCALING_SIZE + if t.size(0) != expected: + t = t[:expected] + buffers.append(t) + return tuple(buffers), {"field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered MXFP8 buffers back, re-padding block scales. + + Inverse of :meth:`fsdp_extract_buffers`: the gathered scale-inv tensors + are padded back up to ``[128, 4]`` / ``[4, 128]`` alignment before + being assigned to the storage. + """ + names = meta["field_names"] + if len(names) != len(gathered): + raise RuntimeError( + f"MXFP8TensorStorage.fsdp_assign_gathered got " + f"{len(gathered)} buffers for {len(names)} fields" + ) + for name, buf in zip(names, gathered): + if buf is not None and name == "_rowwise_scale_inv": + pad = (128 - buf.size(0) % 128) % 128 + if pad > 0: + buf = torch.nn.functional.pad(buf, (0, 0, 0, pad)) + elif buf is not None and name == "_columnwise_scale_inv": + pad = (4 - buf.size(0) % 4) % 4 + if pad > 0: + buf = torch.nn.functional.pad(buf, (0, 0, 0, pad)) + setattr(self, name, buf) From f22a395477fe511ba94a884e7d9297631d391b22 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Apr 2026 13:00:41 +0000 Subject: [PATCH 05/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../distributed/fsdp2_tests/fsdp2_utils.py | 16 +- .../fsdp2_tests/run_fsdp2_fused_adam.py | 33 +- .../fsdp2_tests/run_fsdp2_mem_leak.py | 33 +- .../fsdp2_tests/run_fsdp2_model.py | 3 +- tests/pytorch/test_hybrid_quantization.py | 492 ++++++++++-------- .../pytorch/tensor/float8_tensor.py | 6 +- .../pytorch/tensor/hybrid_tensor.py | 48 +- .../pytorch/tensor/mxfp8_tensor.py | 25 +- .../tensor/storage/hybrid_tensor_storage.py | 4 +- .../tensor/storage/mxfp8_tensor_storage.py | 2 +- 10 files changed, 399 insertions(+), 263 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 1cf81b1db07..e3702e91044 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -40,9 +40,15 @@ def get_hybrid_recipe_from_string(recipe): grad=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), ), "HybridFloat8BlockScaling": lambda: dict( - row=lambda: Float8BlockQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True), - col=lambda: Float8BlockQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True), - grad=lambda: Float8BlockQuantizer(fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True), + row=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ), + col=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ), + grad=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True + ), ), "HybridMixed_MXFP8_FP8": lambda: dict( row=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), @@ -52,9 +58,7 @@ def get_hybrid_recipe_from_string(recipe): } if recipe not in _BUILDERS: - raise ValueError( - f"Unknown hybrid recipe '{recipe}'. Supported: {sorted(_BUILDERS.keys())}" - ) + raise ValueError(f"Unknown hybrid recipe '{recipe}'. Supported: {sorted(_BUILDERS.keys())}") builders = _BUILDERS[recipe]() row_fn, col_fn, grad_fn = builders["row"], builders["col"], builders["grad"] diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 36d25b0d7af..85c58d17f8d 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1017,9 +1017,7 @@ def test_fused_adam_hybrid_master_weights(hybrid_recipe_name): if "master_param" in state: assert state["master_param"].dtype == torch.float32 - assert losses[-1] < losses[0], ( - f"Loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" - ) + assert losses[-1] < losses[0], f"Loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" @pytest.mark.parametrize("reshard_after_forward", [True, False]) @@ -1130,9 +1128,7 @@ def run_training(model, recipe_for_autocast): hybrid_model = _shard_model(hybrid_model, world_size) hybrid_losses = run_training(hybrid_model, hybrid_recipe) - assert hybrid_losses[-1] < hybrid_losses[0], ( - f"Hybrid loss did not decrease: {hybrid_losses}" - ) + assert hybrid_losses[-1] < hybrid_losses[0], f"Hybrid loss did not decrease: {hybrid_losses}" assert bf16_losses[-1] < bf16_losses[0], f"BF16 loss did not decrease: {bf16_losses}" # Verify hybrid and bf16 loss trajectories are within the same order of magnitude. @@ -1266,8 +1262,7 @@ def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): # Compare dim-0 all-gather (bytes) with FSDP2's reconstruction. for name, param in model.named_parameters(): if not ( - isinstance(param, DTensor) - and isinstance(param._local_tensor, QuantizedTensor) + isinstance(param, DTensor) and isinstance(param._local_tensor, QuantizedTensor) ): continue local_shard = param._local_tensor @@ -1288,9 +1283,7 @@ def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): fsdp_full_deq[: manual_full.shape[0]].float(), rtol=0.0, atol=0.0, - msg=lambda m, n=name, r=recipe_name: ( - f"[{r}] Allgather mismatch for {n} at awkward shard shape: {m}" - ), + msg=lambda m, n=name, r=recipe_name: f"[{r}] Allgather mismatch for {n} at awkward shard shape: {m}", ) @@ -1324,8 +1317,10 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): model = _build_hybrid_model(hybrid_recipe) model = _shard_model(model, world_size) optimizer = te.optimizers.FusedAdam( - model.parameters(), lr=1e-3, - master_weights=True, master_weight_dtype=torch.float32, + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) @@ -1351,8 +1346,10 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): model2 = _build_hybrid_model(hybrid_recipe) model2 = _shard_model(model2, world_size) optimizer2 = te.optimizers.FusedAdam( - model2.parameters(), lr=1e-3, - master_weights=True, master_weight_dtype=torch.float32, + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) optimizer2.zero_grad(set_to_none=True) with te.autocast(enabled=True, recipe=hybrid_recipe): @@ -1373,13 +1370,17 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): loaded_output = model2(x) torch.testing.assert_close( - loaded_output, ref_output, rtol=0, atol=0, + loaded_output, + ref_output, + rtol=0, + atol=0, msg=lambda m: f"DCP roundtrip output mismatch: {m}", ) finally: dist.barrier() if rank == 0: import shutil + shutil.rmtree(checkpoint_dir, ignore_errors=True) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py index 25212d4db0f..4ddfb02b5f2 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py @@ -470,6 +470,7 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in # ── Hybrid quantization memory tests ───────────────────────────────── + def _build_hybrid_model(num_layers, hybrid_recipe, use_meta_device=True): """Build a model with quantized_model_init using a hybrid CustomRecipe.""" ctx = te.quantized_model_init(enabled=True, recipe=hybrid_recipe) @@ -520,7 +521,10 @@ def test_hybrid_no_excess_forward_memory(hybrid_recipe_name): bf16_model = _build_model(NUM_LAYERS, fp8_init=False) bf16_model = _shard_model(bf16_model, world_size) bf16_optimizer = te.optimizers.FusedAdam( - bf16_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, + bf16_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) for _ in range(WARMUP_STEPS): _run_training_step(bf16_model, bf16_optimizer, None, x, target) @@ -535,12 +539,19 @@ def test_hybrid_no_excess_forward_memory(hybrid_recipe_name): hybrid_model = _build_hybrid_model(NUM_LAYERS, hybrid_recipe) hybrid_model = _shard_model(hybrid_model, world_size) hybrid_optimizer = te.optimizers.FusedAdam( - hybrid_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, + hybrid_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) for _ in range(WARMUP_STEPS): _run_training_step(hybrid_model, hybrid_optimizer, hybrid_recipe, x, target) hybrid_increments = _measure_forward_increments( - hybrid_model, hybrid_optimizer, hybrid_recipe, x, target, + hybrid_model, + hybrid_optimizer, + hybrid_recipe, + x, + target, ) hybrid_avg = sum(hybrid_increments) / len(hybrid_increments) @@ -574,7 +585,10 @@ def test_hybrid_transpose_cache_after_backward(hybrid_recipe_name): bf16_model = _build_model(NUM_LAYERS, fp8_init=False) bf16_model = _shard_model(bf16_model, world_size) bf16_optimizer = te.optimizers.FusedAdam( - bf16_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, + bf16_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) for _ in range(WARMUP_STEPS): _run_training_step(bf16_model, bf16_optimizer, None, x, target) @@ -588,12 +602,19 @@ def test_hybrid_transpose_cache_after_backward(hybrid_recipe_name): hybrid_model = _build_hybrid_model(NUM_LAYERS, hybrid_recipe) hybrid_model = _shard_model(hybrid_model, world_size) hybrid_optimizer = te.optimizers.FusedAdam( - hybrid_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, + hybrid_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) for _ in range(WARMUP_STEPS): _run_training_step(hybrid_model, hybrid_optimizer, hybrid_recipe, x, target) hybrid_bwd_delta = _measure_backward_memory_delta( - hybrid_model, hybrid_optimizer, hybrid_recipe, x, target, + hybrid_model, + hybrid_optimizer, + hybrid_recipe, + x, + target, ) excess = hybrid_bwd_delta - bf16_bwd_delta diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 1ab4fa1fd59..0f73c5a582c 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -471,7 +471,8 @@ def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): out_features = in_features * 3 with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): model = te.LayerNormLinear( - in_features, out_features, + in_features, + out_features, params_dtype=torch.bfloat16, device="meta", ) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index bb5298bbf88..0d81da14764 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -1667,12 +1667,8 @@ class TestHybridQuantizedModelInit: def _hybrid_fp8_recipe(self): return _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), @@ -1685,12 +1681,12 @@ def test_linear_weight_is_hybrid_quantized_tensor(self): model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() weight = model.weight - assert isinstance(weight, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(weight).__name__}" - ) - assert isinstance(weight, QuantizedTensor), ( - "HybridQuantizedTensor should be a QuantizedTensor subclass" - ) + assert isinstance( + weight, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(weight).__name__}" + assert isinstance( + weight, QuantizedTensor + ), "HybridQuantizedTensor should be a QuantizedTensor subclass" def test_linear_weight_has_both_sub_storages(self): """Quantized param should have rowwise and columnwise sub-storages.""" @@ -1716,9 +1712,7 @@ def test_linear_bias_stays_bf16(self): with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(128, 128, bias=True, params_dtype=torch.bfloat16).cuda() - assert not isinstance(model.bias, QuantizedTensor), ( - "Bias should not be a QuantizedTensor" - ) + assert not isinstance(model.bias, QuantizedTensor), "Bias should not be a QuantizedTensor" assert model.bias.dtype == torch.bfloat16 def test_layernorm_linear_weight_is_hybrid(self): @@ -1776,12 +1770,8 @@ class TestHybridWeightWorkspaceCache: def _hybrid_fp8_recipe(self): return _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), @@ -1865,12 +1855,8 @@ class TestHybridUpdateWeightQuantizers: def _hybrid_fp8_recipe(self): return _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), @@ -1892,9 +1878,9 @@ def test_quantized_param_survives_multiple_forward_passes(self): assert not torch.isnan(out).any(), f"NaN at iteration {i}" assert inp_i.grad is not None, f"No input grad at iteration {i}" - assert isinstance(model.weight, HybridQuantizedTensor), ( - "Weight lost HybridQuantizedTensor type after multiple passes" - ) + assert isinstance( + model.weight, HybridQuantizedTensor + ), "Weight lost HybridQuantizedTensor type after multiple passes" # --------------------------------------------------------------------------- @@ -1908,12 +1894,8 @@ class TestHybridRecipeCorrespondence: def _hybrid_fp8_recipe(self): return _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), @@ -1962,9 +1944,7 @@ def test_quantize_inplace_updates_data(self): """quantize_() should re-quantize both sub-storages from new BF16 data.""" torch.manual_seed(42) hq = HybridQuantizer( - rowwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), columnwise_quantizer=Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda" ), @@ -1983,17 +1963,14 @@ def test_quantize_inplace_updates_data(self): # Should be close to new data, not old data diff_new = (dq_after.float() - new_data.float()).abs().mean() diff_old = (dq_after.float() - original.float()).abs().mean() - assert diff_new < diff_old, ( - f"After quantize_(), data is closer to old ({diff_old:.4f}) " - f"than new ({diff_new:.4f})" - ) + assert ( + diff_new < diff_old + ), f"After quantize_(), data is closer to old ({diff_old:.4f}) than new ({diff_new:.4f})" def test_quantize_inplace_preserves_tensor_identity(self): """quantize_() should update in-place, not create a new tensor.""" hq = HybridQuantizer( - rowwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), columnwise_quantizer=Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda" ), @@ -2022,12 +1999,8 @@ class TestHybridFusedAdam: def _build_hybrid_model(self): hybrid_recipe = _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), @@ -2098,9 +2071,9 @@ def test_fused_adam_param_remains_hybrid_after_step(self): for name, p in model.named_parameters(): if "bias" not in name: - assert isinstance(p, HybridQuantizedTensor), ( - f"{name} lost HybridQuantizedTensor type: {type(p).__name__}" - ) + assert isinstance( + p, HybridQuantizedTensor + ), f"{name} lost HybridQuantizedTensor type: {type(p).__name__}" def test_fused_adam_requires_master_weights(self): """FusedAdam without master_weights should raise for hybrid quantized params.""" @@ -2128,12 +2101,8 @@ class TestHybridQuantizedParamsEndToEnd: def _build_model_and_recipe(self): hybrid_recipe = _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), @@ -2199,9 +2168,9 @@ def test_training_loop_params_remain_quantized(self): for name, p in model.named_parameters(): if "bias" not in name: - assert isinstance(p, HybridQuantizedTensor), ( - f"{name} is {type(p).__name__}, not HybridQuantizedTensor" - ) + assert isinstance( + p, HybridQuantizedTensor + ), f"{name} is {type(p).__name__}, not HybridQuantizedTensor" def test_training_loop_optimizer_states_are_fp32(self): """Optimizer states should be FP32.""" @@ -2251,9 +2220,7 @@ def _build_mixed_model(self, in_features=256, out_features=256): grad_factory=lambda: NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): - model = Linear( - in_features, out_features, params_dtype=torch.bfloat16 - ).cuda() + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() return model, hybrid_recipe def test_mixed_format_param_creation(self): @@ -2321,9 +2288,7 @@ def test_mixed_format_training_loop(self): assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" for name, p in model.named_parameters(): if "bias" not in name: - assert isinstance(p, HybridQuantizedTensor), ( - f"{name} is {type(p).__name__}" - ) + assert isinstance(p, HybridQuantizedTensor), f"{name} is {type(p).__name__}" def test_mixed_format_sub_storage_types(self): """Verify that sub-storages have the correct types (MXFP8 vs NVFP4).""" @@ -2335,12 +2300,12 @@ def test_mixed_format_sub_storage_types(self): row = weight.rowwise_sub_storage col = weight.columnwise_sub_storage - assert isinstance(row, MXFP8TensorStorage) or hasattr(row, "_rowwise_data"), ( - f"Expected MXFP8 rowwise sub-storage, got {type(row).__name__}" - ) - assert isinstance(col, NVFP4TensorStorage) or hasattr(col, "_rowwise_data"), ( - f"Expected NVFP4 columnwise sub-storage, got {type(col).__name__}" - ) + assert isinstance(row, MXFP8TensorStorage) or hasattr( + row, "_rowwise_data" + ), f"Expected MXFP8 rowwise sub-storage, got {type(row).__name__}" + assert isinstance(col, NVFP4TensorStorage) or hasattr( + col, "_rowwise_data" + ), f"Expected NVFP4 columnwise sub-storage, got {type(col).__name__}" # --------------------------------------------------------------------------- @@ -2352,9 +2317,7 @@ def _hybrid_fp8_current_qfactory(role): """Hybrid FP8 current scaling (E4M3 both dirs, E5M2 for grad).""" if role in ("linear_input", "linear_weight", "linear_output"): return HybridQuantizer( - rowwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), columnwise_quantizer=Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda" ), @@ -2380,16 +2343,22 @@ def _hybrid_block_fp8_qfactory(role): if role in ("linear_grad_output", "linear_grad_input"): return Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, columnwise=True, block_scaling_dim=dim, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, ) return HybridQuantizer( rowwise_quantizer=Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, columnwise=True, block_scaling_dim=dim, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, ), columnwise_quantizer=Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, columnwise=True, block_scaling_dim=dim, + rowwise=True, + columnwise=True, + block_scaling_dim=dim, ), ) @@ -2426,21 +2395,27 @@ def _build_models(self): torch.manual_seed(42) with quantized_model_init(enabled=True, recipe=self._vanilla_recipe()): model_ref = Linear( - self.hidden_size, self.hidden_size, params_dtype=torch.bfloat16, + self.hidden_size, + self.hidden_size, + params_dtype=torch.bfloat16, ).cuda() torch.manual_seed(42) with quantized_model_init(enabled=True, recipe=self._hybrid_recipe()): model_hyb = Linear( - self.hidden_size, self.hidden_size, params_dtype=torch.bfloat16, + self.hidden_size, + self.hidden_size, + params_dtype=torch.bfloat16, ).cuda() return model_ref, model_hyb def _run_training_loop(self, model, train_recipe, x, target, num_steps): optimizer = te.optimizers.FusedAdam( - model.parameters(), lr=1e-3, - master_weights=True, master_weight_dtype=torch.float32, + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) losses = [] for _ in range(num_steps): @@ -2466,23 +2441,35 @@ def _test_equivalence(self): target = torch.randn_like(x) losses_ref, masters_ref = self._run_training_loop( - model_ref, self._vanilla_recipe(), x, target, self.num_steps, + model_ref, + self._vanilla_recipe(), + x, + target, + self.num_steps, ) losses_hyb, masters_hyb = self._run_training_loop( - model_hyb, self._hybrid_recipe(), x, target, self.num_steps, + model_hyb, + self._hybrid_recipe(), + x, + target, + self.num_steps, ) # Losses should be very close (same quantization, same training dynamics) for i, (lr, lh) in enumerate(zip(losses_ref, losses_hyb)): - assert abs(lr - lh) < 0.1 * max(abs(lr), 1e-6), ( - f"Step {i}: loss diverged — vanilla={lr:.6f}, hybrid={lh:.6f}" - ) + assert abs(lr - lh) < 0.1 * max( + abs(lr), 1e-6 + ), f"Step {i}: loss diverged — vanilla={lr:.6f}, hybrid={lh:.6f}" # Master weights should be close after training for i, (mr, mh) in enumerate(zip(masters_ref, masters_hyb)): - torch.testing.assert_close(mr, mh, rtol=1e-3, atol=1e-3, msg=( - f"Master weight {i} diverged after {self.num_steps} steps" - )) + torch.testing.assert_close( + mr, + mh, + rtol=1e-3, + atol=1e-3, + msg=f"Master weight {i} diverged after {self.num_steps} steps", + ) @requires_fp8 @@ -2509,10 +2496,18 @@ def test_equivalence(self): target = torch.randn_like(x) losses_ref, _ = self._run_training_loop( - model_ref, self._vanilla_recipe(), x, target, self.num_steps, + model_ref, + self._vanilla_recipe(), + x, + target, + self.num_steps, ) losses_hyb, _ = self._run_training_loop( - model_hyb, self._hybrid_recipe(), x, target, self.num_steps, + model_hyb, + self._hybrid_recipe(), + x, + target, + self.num_steps, ) # Both should decrease (training works in both paths) @@ -2522,9 +2517,9 @@ def test_equivalence(self): # Losses should be in the same ballpark (different optimizer kernels # cause small divergence that compounds over steps) for i, (lr, lh) in enumerate(zip(losses_ref, losses_hyb)): - assert abs(lr - lh) / max(abs(lr), 1e-6) < 0.5, ( - f"Step {i}: losses diverged too much — vanilla={lr:.6f}, hybrid={lh:.6f}" - ) + assert ( + abs(lr - lh) / max(abs(lr), 1e-6) < 0.5 + ), f"Step {i}: losses diverged too much — vanilla={lr:.6f}, hybrid={lh:.6f}" @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") @@ -2591,9 +2586,7 @@ def _checkpoint_hybrid_fp8_qfactory(role): """Module-level qfactory (picklable) for checkpoint tests.""" if role in ("linear_input", "linear_weight", "linear_output"): return HybridQuantizer( - rowwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), columnwise_quantizer=Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda" ), @@ -2716,8 +2709,10 @@ def test_checkpoint_resume_training(self): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() optimizer = te.optimizers.FusedAdam( - model.parameters(), lr=1e-3, - master_weights=True, master_weight_dtype=torch.float32, + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") @@ -2736,10 +2731,13 @@ def test_checkpoint_resume_training(self): # Save checkpoint with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as f: - torch.save({ - "model": model.state_dict(), - "optimizer": optimizer.state_dict(), - }, f.name) + torch.save( + { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + }, + f.name, + ) tmp_path = f.name try: @@ -2747,8 +2745,10 @@ def test_checkpoint_resume_training(self): with quantized_model_init(enabled=True, recipe=hybrid_recipe): model2 = Linear(256, 256, params_dtype=torch.bfloat16).cuda() optimizer2 = te.optimizers.FusedAdam( - model2.parameters(), lr=1e-3, - master_weights=True, master_weight_dtype=torch.float32, + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) checkpoint = torch.load(tmp_path, weights_only=False) @@ -2762,7 +2762,8 @@ def test_checkpoint_resume_training(self): loss_after_load = torch.nn.functional.mse_loss(output2, target).item() assert loss_after_load <= loss_before_save * 1.5, ( - f"Loss spiked after checkpoint resume: {loss_before_save:.4f} → {loss_after_load:.4f}" + f"Loss spiked after checkpoint resume: {loss_before_save:.4f} →" + f" {loss_after_load:.4f}" ) finally: os.unlink(tmp_path) @@ -2775,8 +2776,9 @@ def test_checkpoint_resume_training(self): aten = torch.ops.aten -def _make_hybrid_param_for_dispatch(row_factory, col_factory, grad_factory=None, - in_features=256, out_features=256): +def _make_hybrid_param_for_dispatch( + row_factory, col_factory, grad_factory=None, in_features=256, out_features=256 +): """Create a HybridQuantizedTensor weight via quantized_model_init for dispatch tests.""" hybrid_recipe = _hybrid_custom_recipe(row_factory, col_factory, grad_factory) with quantized_model_init(enabled=True, recipe=hybrid_recipe): @@ -2811,11 +2813,14 @@ def _get_dispatch_hybrid_param(config_name): """Return a HybridQuantizedTensor weight for the given config.""" if config_name == "fp8_fp8": return _make_hybrid_param_for_dispatch( - _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, ) elif config_name == "mxfp8_mxfp8": return _make_hybrid_param_for_dispatch( - _mxfp8_factory, _mxfp8_factory, + _mxfp8_factory, + _mxfp8_factory, grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), ) else: @@ -2842,9 +2847,9 @@ def test_split_preserves_hybrid_type(self, hybrid_param): pieces = torch.split(hybrid_param, chunk_size, dim=0) assert len(pieces) >= 2 for piece in pieces: - assert isinstance(piece, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(piece).__name__}" - ) + assert isinstance( + piece, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(piece).__name__}" assert piece.rowwise_sub_storage is not None assert piece.columnwise_sub_storage is not None @@ -2870,9 +2875,9 @@ def test_view_preserves_hybrid_type(self, hybrid_param): """view must return a HybridQuantizedTensor (used by FSDP2 reset_sharded_param).""" shape_2d = hybrid_param.shape result = aten.view.default(hybrid_param, list(shape_2d)) - assert isinstance(result, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(result).__name__}" - ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" assert result.rowwise_sub_storage is not None assert result.columnwise_sub_storage is not None @@ -2880,34 +2885,36 @@ def test_view_same_shape_preserves_hybrid(self, hybrid_param): """view with same shape must return HybridQuantizedTensor.""" shape_2d = list(hybrid_param.shape) result = aten.view.default(hybrid_param, shape_2d) - assert isinstance(result, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(result).__name__}" - ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" def test_as_strided_noop_preserves_hybrid(self, hybrid_param): """as_strided with matching shape/strides is a no-op that preserves type.""" shape = tuple(hybrid_param.size()) strides = (shape[-1], 1) result = aten.as_strided.default(hybrid_param, list(shape), list(strides)) - assert isinstance(result, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(result).__name__}" - ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" assert result.rowwise_sub_storage is not None assert result.columnwise_sub_storage is not None def test_slice_noop_preserves_hybrid(self, hybrid_param): """slice with full range is a no-op that preserves type.""" result = aten.slice.Tensor(hybrid_param, 0, 0, hybrid_param.size(0)) - assert isinstance(result, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(result).__name__}" - ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" assert result.rowwise_sub_storage is not None def test_copy_between_hybrid_tensors(self, hybrid_param): """copy_ between compatible HybridQuantizedTensors copies quantized data directly.""" src_deq = hybrid_param.dequantize().clone() dst = hybrid_param._quantizer.make_empty( - shape=hybrid_param.shape, dtype=hybrid_param.dtype, device=hybrid_param.device, + shape=hybrid_param.shape, + dtype=hybrid_param.dtype, + device=hybrid_param.device, ) assert isinstance(dst, HybridQuantizedTensor) @@ -2941,9 +2948,9 @@ def test_new_zeros_returns_hybrid(self, hybrid_param): # Structural contract: FSDP2 needs a HybridQuantizedTensor with both # sub-storages populated so the gathered buffers have a destination. - assert isinstance(result, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(result).__name__}" - ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" assert result.shape == hybrid_param.shape assert result.rowwise_sub_storage is not None assert result.columnwise_sub_storage is not None @@ -2958,18 +2965,18 @@ def test_new_zeros_returns_hybrid(self, hybrid_param): def test_empty_like_returns_hybrid(self, hybrid_param): """empty_like must return a HybridQuantizedTensor.""" result = aten.empty_like.default(hybrid_param) - assert isinstance(result, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(result).__name__}" - ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" assert result.shape == hybrid_param.shape assert result.rowwise_sub_storage is not None def test_clone_returns_hybrid(self, hybrid_param): """clone must return an independent HybridQuantizedTensor with same data.""" result = aten.clone.default(hybrid_param) - assert isinstance(result, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(result).__name__}" - ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" assert result is not hybrid_param torch.testing.assert_close(result.dequantize(), hybrid_param.dequantize()) @@ -3014,50 +3021,65 @@ def hybrid_param(self, request): def test_pre_all_gather_returns_tuple_pair(self, hybrid_param): """fsdp_pre_all_gather returns (sharded_tensors, metadata).""" sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, - ) - assert isinstance(sharded_tensors, tuple), ( - f"sharded_tensors should be tuple, got {type(sharded_tensors).__name__}" + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) + assert isinstance( + sharded_tensors, tuple + ), f"sharded_tensors should be tuple, got {type(sharded_tensors).__name__}" assert len(sharded_tensors) > 0, "sharded_tensors should not be empty" - assert isinstance(metadata, tuple), ( - f"metadata should be tuple, got {type(metadata).__name__}" - ) + assert isinstance( + metadata, tuple + ), f"metadata should be tuple, got {type(metadata).__name__}" def test_pre_all_gather_buffers_are_plain_tensors(self, hybrid_param): """Every element in sharded_tensors must be a plain torch.Tensor.""" sharded_tensors, _ = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) for i, t in enumerate(sharded_tensors): - assert isinstance(t, torch.Tensor), ( - f"sharded_tensors[{i}] should be torch.Tensor, got {type(t).__name__}" - ) - assert not isinstance(t, QuantizedTensor), ( - f"sharded_tensors[{i}] should NOT be QuantizedTensor subclass" - ) + assert isinstance( + t, torch.Tensor + ), f"sharded_tensors[{i}] should be torch.Tensor, got {type(t).__name__}" + assert not isinstance( + t, QuantizedTensor + ), f"sharded_tensors[{i}] should NOT be QuantizedTensor subclass" def test_pre_all_gather_buffer_count_consistent(self, hybrid_param): """Buffer count must be the same across repeated calls (FSDP2 buffer reuse).""" sharded_1, _ = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) sharded_2, _ = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, - ) - assert len(sharded_1) == len(sharded_2), ( - f"Buffer count changed: {len(sharded_1)} vs {len(sharded_2)}" + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) + assert len(sharded_1) == len( + sharded_2 + ), f"Buffer count changed: {len(sharded_1)} vs {len(sharded_2)}" def test_pre_all_gather_metadata_sufficient_for_reconstruction(self, hybrid_param): """Metadata must contain enough info to reconstruct the tensor.""" _, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) assert metadata is not None assert len(metadata) > 0, "metadata should not be empty" @@ -3084,15 +3106,21 @@ def hybrid_param(self, request): def test_post_all_gather_first_call_returns_hybrid_tensor(self, hybrid_param): """With out=None, post_all_gather returns (HybridQuantizedTensor, outputs).""" sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) result, ag_outputs = hybrid_param.fsdp_post_all_gather( - sharded_tensors, metadata, hybrid_param.dtype, out=None, - ) - assert isinstance(result, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(result).__name__}" + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" assert result.shape == hybrid_param.shape assert result.rowwise_sub_storage is not None assert result.columnwise_sub_storage is not None @@ -3100,30 +3128,45 @@ def test_post_all_gather_first_call_returns_hybrid_tensor(self, hybrid_param): def test_post_all_gather_buffer_reuse(self, hybrid_param): """On second call with out=previous, the same object is returned (buffer reuse).""" sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) first_result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, metadata, hybrid_param.dtype, out=None, + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, ) second_result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, metadata, hybrid_param.dtype, out=first_result, - ) - assert second_result is first_result, ( - "Buffer reuse: post_all_gather(out=prev) should return the same object" + sharded_tensors, + metadata, + hybrid_param.dtype, + out=first_result, ) + assert ( + second_result is first_result + ), "Buffer reuse: post_all_gather(out=prev) should return the same object" def test_post_all_gather_dequantize_matches_original(self, hybrid_param): """Reconstructed tensor should dequantize close to the original.""" orig_deq = hybrid_param.dequantize() sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, metadata, hybrid_param.dtype, out=None, + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, ) result_deq = result.dequantize() torch.testing.assert_close(orig_deq, result_deq) @@ -3134,11 +3177,17 @@ def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): orig_col_type = type(hybrid_param.columnwise_sub_storage) sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, metadata, hybrid_param.dtype, out=None, + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, ) assert type(result.rowwise_sub_storage) is orig_row_type assert type(result.columnwise_sub_storage) is orig_col_type @@ -3163,30 +3212,48 @@ def test_pre_post_roundtrip_preserves_data(self, hybrid_param): orig_deq = hybrid_param.dequantize() sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, metadata, hybrid_param.dtype, out=None, + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, ) torch.testing.assert_close(orig_deq, result.dequantize()) def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): """Second roundtrip with out=previous preserves data (iteration 2+ simulation).""" sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) first_result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, metadata, hybrid_param.dtype, out=None, + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, ) sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) second_result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors_2, metadata_2, hybrid_param.dtype, out=first_result, + sharded_tensors_2, + metadata_2, + hybrid_param.dtype, + out=first_result, ) assert second_result is first_result torch.testing.assert_close(hybrid_param.dequantize(), second_result.dequantize()) @@ -3209,7 +3276,9 @@ def test_scale_refresh_across_iterations(self): """ torch.manual_seed(42) hybrid_recipe = _hybrid_custom_recipe( - _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -3217,11 +3286,17 @@ def test_scale_refresh_across_iterations(self): # Iter-1 gather with the initial (small-magnitude) weights sharded_tensors_1, metadata_1 = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) gathered, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors_1, metadata_1, hybrid_param.dtype, out=None, + sharded_tensors_1, + metadata_1, + hybrid_param.dtype, + out=None, ) # Simulate an optimizer writeback that produces a much larger weight; @@ -3233,11 +3308,17 @@ def test_scale_refresh_across_iterations(self): # Iter-2+ path: reuse the gathered buffer sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( - mesh=None, orig_size=hybrid_param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) gathered_refreshed, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors_2, metadata_2, hybrid_param.dtype, out=gathered, + sharded_tensors_2, + metadata_2, + hybrid_param.dtype, + out=gathered, ) assert gathered_refreshed is gathered @@ -3287,8 +3368,11 @@ def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): # the base class. with pytest.raises(NotImplementedError) as exc_info: param.fsdp_pre_all_gather( - mesh=None, orig_size=param.shape, contiguous_orig_stride=None, - module=None, mp_policy=None, + mesh=None, + orig_size=param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, ) msg = str(exc_info.value) assert "NVFP4Tensor" in msg @@ -3307,7 +3391,9 @@ class TestHybridMakeLike: def _make_hybrid_param(self): hybrid_recipe = _hybrid_custom_recipe( - _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 3beb90926a6..88ee7900af5 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -756,7 +756,11 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): tensor, data=split_tensor, data_transpose=split_transpose_tensor, - shape=(split_tensor.shape if split_tensor is not None else split_transpose_tensor.shape), + shape=( + split_tensor.shape + if split_tensor is not None + else split_transpose_tensor.shape + ), ) for split_tensor, split_transpose_tensor in zip(func_out, t_func_out) ] diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index b132cab10b5..1c80193f407 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -125,12 +125,11 @@ def update_quantized( """ if not isinstance(dst, HybridQuantizedTensorStorage): raise ValueError( - f"HybridQuantizer can only update HybridQuantizedTensorStorage, got {type(dst).__name__}" + "HybridQuantizer can only update HybridQuantizedTensorStorage, got" + f" {type(dst).__name__}" ) if dst._rowwise_storage is not None: - self.rowwise_quantizer.update_quantized( - src, dst._rowwise_storage, noop_flag=noop_flag - ) + self.rowwise_quantizer.update_quantized(src, dst._rowwise_storage, noop_flag=noop_flag) if dst._columnwise_storage is not None: self.columnwise_quantizer.update_quantized( src, dst._columnwise_storage, noop_flag=noop_flag @@ -286,15 +285,15 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m sub.fsdp_buffer_fields() except NotImplementedError as err: raise NotImplementedError( - f"Hybrid FSDP2 all-gather is not supported for a " + "Hybrid FSDP2 all-gather is not supported for a " f"{type(sub).__name__} {role} sub-storage: it does not " - f"implement fsdp_buffer_fields. " - f"See hybrid_quantization_fsdp.md section 9 (Gap 5) — " - f"NVFP4 sub-storages need packed-FP4 dim-0 alignment, " - f"columnwise dequantization and RHT-cache handling before " - f"they can be gathered. Use a supported sub-quantizer " - f"(Float8CurrentScaling, MXFP8, Float8Block) or run without " - f"FSDP2." + "implement fsdp_buffer_fields. " + "See hybrid_quantization_fsdp.md section 9 (Gap 5) — " + "NVFP4 sub-storages need packed-FP4 dim-0 alignment, " + "columnwise dequantization and RHT-cache handling before " + "they can be gathered. Use a supported sub-quantizer " + "(Float8CurrentScaling, MXFP8, Float8Block) or run without " + "FSDP2." ) from err row_buffers: Tuple[Optional[torch.Tensor], ...] = () @@ -312,7 +311,7 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m len(row_buffers), row_meta, col_meta, - self._rowwise_storage, # original sharded sub-storage (for make_like on iter-1) + self._rowwise_storage, # original sharded sub-storage (for make_like on iter-1) self._columnwise_storage, self._rowwise_quantizer, self._columnwise_quantizer, @@ -411,6 +410,7 @@ def _delegate_reshape_op(cls, func, tensor, args, kwargs): that way for real slicing today). On ``None`` the caller should defer to ``super().__torch_dispatch__`` for a consistent BF16 fallback. """ + def _delegate(sub): if sub is None: return None @@ -476,11 +476,13 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): row_pieces = ( torch.split(tensor._rowwise_storage, split_size, dim=dim) - if tensor._rowwise_storage is not None else None + if tensor._rowwise_storage is not None + else None ) col_pieces = ( torch.split(tensor._columnwise_storage, split_size, dim=dim) - if tensor._columnwise_storage is not None else None + if tensor._columnwise_storage is not None + else None ) if row_pieces is None and col_pieces is None: @@ -519,8 +521,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): and tuple(shape) == tuple(tensor.size()) ): return HybridQuantizedTensor.make_like(tensor) - return cls._delegate_reshape_op(func, tensor, args, kwargs) or \ - super().__torch_dispatch__(func, types, args, kwargs) + return cls._delegate_reshape_op( + func, tensor, args, kwargs + ) or super().__torch_dispatch__(func, types, args, kwargs) if func == aten.slice.Tensor: tensor = args[0] @@ -529,8 +532,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): length = args[3] if start == 0 and length == tensor.size(dim): return HybridQuantizedTensor.make_like(tensor) - return cls._delegate_reshape_op(func, tensor, args, kwargs) or \ - super().__torch_dispatch__(func, types, args, kwargs) + return cls._delegate_reshape_op( + func, tensor, args, kwargs + ) or super().__torch_dispatch__(func, types, args, kwargs) # ── FSDP2: copy_ ───────────────────────────────────────────── # Fast path for hybrid-to-hybrid (FSDP2 fills buffer allocated via @@ -565,11 +569,13 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): tensor = args[0] row_clone = ( torch.clone(tensor._rowwise_storage) - if tensor._rowwise_storage is not None else None + if tensor._rowwise_storage is not None + else None ) col_clone = ( torch.clone(tensor._columnwise_storage) - if tensor._columnwise_storage is not None else None + if tensor._columnwise_storage is not None + else None ) return HybridQuantizedTensor( shape=tensor.shape, diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 3c960e653aa..2f7867e8254 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -345,7 +345,9 @@ def detach(self) -> MXFP8Tensor: def clone(self) -> MXFP8Tensor: # pylint: disable=missing-function-docstring # _rowwise_data may be None for columnwise-only sub-storages (hybrid quantization) - rowwise_data = self._rowwise_data.detach().clone() if self._rowwise_data is not None else None + rowwise_data = ( + self._rowwise_data.detach().clone() if self._rowwise_data is not None else None + ) columnwise_data = None if self._columnwise_data is not None: columnwise_data = self._columnwise_data.detach().clone() @@ -462,7 +464,10 @@ def _split_data(data): if data is None: return None return data.__torch_dispatch__( - func, types, [data] + list(args[1:]), kwargs, + func, + types, + [data] + list(args[1:]), + kwargs, ) row_data_splits = _split_data(tensor._rowwise_data) @@ -478,10 +483,14 @@ def _split_data(data): if scale_inv is None: scale_splits.append(None) continue - scale_inv_out = list(scale_inv.__torch_dispatch__( - func, types, - [scale_inv, scale_split_size] + list(args[2:]), kwargs, - )) + scale_inv_out = list( + scale_inv.__torch_dispatch__( + func, + types, + [scale_inv, scale_split_size] + list(args[2:]), + kwargs, + ) + ) for idx, split_scale_inv_out in enumerate(scale_inv_out): current_shape = split_scale_inv_out.shape pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple @@ -505,7 +514,9 @@ def _split_data(data): rowwise_data=row_data_splits[i] if row_data_splits is not None else None, rowwise_scale_inv=row_scale_splits[i] if row_scale_splits is not None else None, columnwise_data=col_data_splits[i] if col_data_splits is not None else None, - columnwise_scale_inv=col_scale_splits[i] if col_scale_splits is not None else None, + columnwise_scale_inv=( + col_scale_splits[i] if col_scale_splits is not None else None + ), quantizer=tensor._quantizer, requires_grad=False, fp8_dtype=tensor._fp8_dtype, diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index 3fb224d6405..a2da94f614e 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -136,7 +136,9 @@ def device(self): def view(self, *shape): """View delegates to each sub-storage. Used by FSDP2 reset_sharded_param.""" row_view = self._rowwise_storage.view(*shape) if self._rowwise_storage is not None else None - col_view = self._columnwise_storage.view(*shape) if self._columnwise_storage is not None else None + col_view = ( + self._columnwise_storage.view(*shape) if self._columnwise_storage is not None else None + ) return HybridQuantizedTensorStorage( rowwise_storage=row_view, columnwise_storage=col_view, diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 5da14ba0a4d..b6ef0e79440 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -369,7 +369,7 @@ def fsdp_assign_gathered( names = meta["field_names"] if len(names) != len(gathered): raise RuntimeError( - f"MXFP8TensorStorage.fsdp_assign_gathered got " + "MXFP8TensorStorage.fsdp_assign_gathered got " f"{len(gathered)} buffers for {len(names)} fields" ) for name, buf in zip(names, gathered): From 103fffe393a8be73b8448a3baf68892089c4afb3 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 22 Apr 2026 14:59:37 +0000 Subject: [PATCH 06/97] Enable CPU offloading Signed-off-by: Evgeny --- tests/pytorch/test_cpu_offloading.py | 96 +++++ tests/pytorch/test_cpu_offloading_v1.py | 74 +++- tests/pytorch/test_hybrid_quantization.py | 333 ++++++++++++++++++ .../pytorch/tensor/hybrid_tensor.py | 46 ++- .../tensor/storage/hybrid_tensor_storage.py | 15 + 5 files changed, 562 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 50196782f2b..9ad83156f3e 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -28,6 +28,43 @@ mxfp8_available, _ = FP8GlobalStateManager.is_mxfp8_available() nvfp4_available, _ = FP8GlobalStateManager.is_nvfp4_available() +def _hybrid_fp8_mxfp8_qfactory(role): + """Hybrid CustomRecipe factory: FP8 current-scaling rowwise + MXFP8 columnwise. + + Forward roles -> HybridQuantizer; backward roles -> plain MXFP8 so + dgrad/wgrad operand pairs share a single scaling mode. Catch-all + returns plain FP8 for non-``linear_*`` roles used by layernorm_linear, + layernorm_mlp, multihead_attention, and transformer_layer. + """ + if role in ("linear_input", "linear_weight", "linear_output"): + return te.HybridQuantizer( + rowwise_quantizer=te.Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2) + return te.Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + +def _hybrid_mxfp8_nvfp4_qfactory(role): + """Hybrid CustomRecipe factory: MXFP8 rowwise + NVFP4 columnwise. + + Mirrors ``mxfp8_fwd_nvfp4_bwd_quantizer_factory`` from + ``custom_recipes/quantization_nvfp4.py``. grad_output uses plain NVFP4 + (both directions) so wgrad's columnwise operand matches. + """ + if role in ("linear_input", "linear_weight", "linear_output"): + return te.HybridQuantizer( + rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + + quantization_recipes: List[Optional[recipe.Recipe]] = [None] if fp8_available: quantization_recipes.extend((recipe.Float8CurrentScaling(), recipe.DelayedScaling())) @@ -37,6 +74,10 @@ quantization_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: quantization_recipes.append(recipe.NVFP4BlockScaling()) +if fp8_available and mxfp8_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_fp8_mxfp8_qfactory)) +if mxfp8_available and nvfp4_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_mxfp8_nvfp4_qfactory)) model_config = { @@ -178,6 +219,15 @@ def create_tensor(recipe: Optional[recipe.Recipe], requires_grad: bool = False) elif recipe.nvfp4(): quantizer = te.tensor.nvfp4_tensor.NVFP4Quantizer() return quantizer(tensor) + elif recipe.custom(): + # CustomRecipe: invoke the qfactory for the ``linear_weight`` role + # as a representative quantizer (returns a HybridQuantizer for the + # hybrid factories registered at module scope). + quantizer = recipe.qfactory("linear_weight") + if quantizer is None: + # Fallback: factory did not supply a weight quantizer. + return tensor.requires_grad_() if requires_grad else tensor + return quantizer(tensor) @staticmethod def create_recipe_ctx(recipe: Optional[recipe.Recipe]): @@ -432,6 +482,22 @@ def test_sanity(self, layer_type, recipe, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Skip hybrid (CustomRecipe) on ops-based LayerNormMLP: the ops-based + # LayerNorm passes the quantizer directly to the fused C++ kernel which + # does not recognize HybridQuantizer (cf. design-doc TODO; the regular + # layernorm_mlp.py has an unfused fallback but the ops path does not + # yet). Unrelated to CPU offload. + # grouped_linear is NOT skipped here — it passes test_sanity with + # hybrid; only memory-accounting assertions trip it in test_memory / + # test_manual_synchronization. + if ( + layer_type in ("layernorm_mlp_ops",) + and recipe is not None + and recipe.custom() + ): + pytest.skip( + f"Hybrid CustomRecipe + {layer_type} integration is not yet complete" + ) recipe_ctx = Utils.create_recipe_ctx(recipe) init_cuda_memory = Utils.get_cuda_memory_mb() @@ -480,6 +546,19 @@ def test_memory(self, layer_type, recipe, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Memory-accounting checks fail for grouped_linear with hybrid because + # `_hybrid_split_quantize` produces per-group HybridQuantizedTensorStorage + # whose individual sub-buffers don't all cross the 256K-element offload + # threshold — the net GPU memory drop after offload is smaller than the + # analytical estimate. Correctness (test_sanity, test_numerics) passes. + if ( + layer_type in ("layernorm_mlp_ops", "grouped_linear") + and recipe is not None + and recipe.custom() + ): + pytest.skip( + f"Hybrid CustomRecipe + {layer_type} integration is not yet complete" + ) offload_ctx, sync_function = get_cpu_offload_context( enabled=True, @@ -571,6 +650,15 @@ def test_manual_synchronization(self, recipe, layer_type, backward_override): and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + # Same memory-accounting caveat as test_memory (see comment there). + if ( + layer_type in ("layernorm_mlp_ops", "grouped_linear") + and recipe is not None + and recipe.custom() + ): + pytest.skip( + f"Hybrid CustomRecipe + {layer_type} integration is not yet complete" + ) offload_ctx, sync_function, manual_controller = get_cpu_offload_context( enabled=True, @@ -650,6 +738,14 @@ def test_numerics( and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") + if ( + layer_type in ("layernorm_mlp_ops",) + and recipe is not None + and recipe.custom() + ): + pytest.skip( + f"Hybrid CustomRecipe + {layer_type} integration is not yet complete" + ) recipe_ctx = Utils.create_recipe_ctx(recipe) diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index 153bceca7d5..aa128d258aa 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -11,6 +11,7 @@ import torch import transformer_engine.pytorch as te +import transformer_engine_torch as tex from transformer_engine.common import recipe from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported @@ -19,10 +20,53 @@ # Check supported quantization schemes fp8_available = te.is_fp8_available() mxfp8_available = te.is_mxfp8_available() +nvfp4_available = te.is_nvfp4_available() + + +def _hybrid_fp8_mxfp8_qfactory(role): + """Hybrid CustomRecipe factory: FP8 current-scaling rowwise + MXFP8 columnwise. + + Forward roles get a HybridQuantizer; backward/grad roles get a plain + MXFP8 quantizer so dgrad/wgrad GEMMs see a single scaling mode per + operand pair. Catch-all returns plain FP8 for non-``linear_*`` roles + (layernorm_linear, layernorm_mlp, multihead_attention, transformer_layer). + """ + if role in ("linear_input", "linear_weight", "linear_output"): + return te.HybridQuantizer( + rowwise_quantizer=te.Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2) + return te.Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + +def _hybrid_mxfp8_nvfp4_qfactory(role): + """Hybrid CustomRecipe factory: MXFP8 rowwise + NVFP4 columnwise. + + Mirrors the ``mxfp8_fwd_nvfp4_bwd_quantizer_factory`` headline recipe + from ``custom_recipes/quantization_nvfp4.py``. grad_output uses plain + NVFP4 (both directions) so wgrad's columnwise operand matches. + """ + if role in ("linear_input", "linear_weight", "linear_output"): + return te.HybridQuantizer( + rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantization_recipes: Optional[recipe.Recipe] = [None] if fp8_available: quantization_recipes.extend((recipe.Float8CurrentScaling(), recipe.DelayedScaling())) +if fp8_available and mxfp8_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_fp8_mxfp8_qfactory)) +if mxfp8_available and nvfp4_available: + quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_mxfp8_nvfp4_qfactory)) model_config = { "small": ModelConfig(8, 512, 8, 64, num_layers=5, eps=0.1), @@ -100,6 +144,15 @@ def _estimate_cached_weight_size( if quantization_recipe is None: return 0 + # Hybrid (CustomRecipe) caches two sub-storages per weight with + # potentially different formats. Returning ``None`` signals the caller + # to skip the exact memory-accounting assertion — the ``memory_with_offload + # < memory_without_offload`` check still applies. Deriving an analytical + # estimate here is blocked on the FSDP2-style packing optimization still + # being a TODO in hybrid_quantization_design.md. + if quantization_recipe.custom(): + return None + # Count number of weight param elements param_elements = 0 for module in modules: @@ -184,6 +237,21 @@ def _measure_cached_memory( def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: str) -> None: """Check that CPU offloading runs and has expected memory usage.""" + # Skip hybrid (CustomRecipe) on module types whose integration with hybrid + # is not yet complete (preexisting, independent of CPU offload): + # - layernorm_mlp_ops: the ops-based LayerNorm passes the quantizer + # directly to the fused C++ kernel which does not recognize + # HybridQuantizer (cf. design doc; the regular layernorm_mlp.py has + # an unfused fallback but the ops path does not yet). + if ( + model_name in ("layernorm_mlp_ops",) + and quantization_recipe is not None + and quantization_recipe.custom() + ): + pytest.skip( + f"Hybrid CustomRecipe + {model_name} integration is not yet complete" + ) + # Construct model modules_list = [model_types[model_name]() for _ in range(NUM_LAYERS)] if model_name in ["multihead_attention", "transformer_layer"]: @@ -212,4 +280,8 @@ def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: s modules_list, quantization_recipe, ) - assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON + # ``_estimate_cached_weight_size`` returns ``None`` for recipes whose + # analytical cached-weight size is not worked out (CustomRecipe / hybrid); + # in that case the memory-savings assertion above is the only check. + if memory_from_cached_weights is not None: + assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 0d81da14764..442f7de0bd5 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -262,6 +262,176 @@ def test_repr_after_drop(self, hybrid_tensor): assert "columnwise=None" in r +requires_mxfp8 = pytest.mark.skipif( + not mxfp8_available, + reason=f"MXFP8: {reason_for_no_mxfp8}", +) + + +@requires_fp8_and_nvfp4 +class TestHybridClear: + """Test HybridQuantizedTensorStorage.clear() — buffer deallocation. + + ``clear()`` is invoked by cpu_offload_v1 after the offloader has taken + its own reference to the extracted buffers, to free the GPU-resident + originals. HybridQuantizedTensorStorage delegates to each sub-storage's + own clear(), which replaces primary data buffers with empty tensors. + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + @staticmethod + def _primary_data_numels(sub_storage): + """Collect numel() of primary data buffers on a sub-storage. + + After ``clear()`` every entry should be 0 (native sub-storages set + ``t.data = _empty_tensor()`` on the primary buffers). + """ + if sub_storage is None: + return [] + data = sub_storage.get_data_tensors() + if not isinstance(data, tuple): + data = (data,) + return [t.numel() for t in data if t is not None] + + def test_clear_delegates_to_both_sub_storages(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + + row_before = self._primary_data_numels(ht.rowwise_sub_storage) + col_before = self._primary_data_numels(ht.columnwise_sub_storage) + assert row_before and all(n > 0 for n in row_before) + assert col_before and all(n > 0 for n in col_before) + + ht.clear() + + row_after = self._primary_data_numels(ht.rowwise_sub_storage) + col_after = self._primary_data_numels(ht.columnwise_sub_storage) + assert all(n == 0 for n in row_after) + assert all(n == 0 for n in col_after) + + @requires_mxfp8 + def test_clear_delegates_mxfp8_nvfp4(self, input_tensor): + """Per-block sub-storage path (MXFP8 rowwise + NVFP4 columnwise).""" + hq = HybridQuantizer( + rowwise_quantizer=_make_mxfp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + ht = hq.quantize(input_tensor) + ht.clear() + for sub in (ht.rowwise_sub_storage, ht.columnwise_sub_storage): + for n in self._primary_data_numels(sub): + assert n == 0 + + def test_clear_with_rowwise_only(self, input_tensor): + """columnwise sub-storage is None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(columnwise_usage=False) + assert ht.columnwise_sub_storage is None + + ht.clear() + + assert all(n == 0 for n in self._primary_data_numels(ht.rowwise_sub_storage)) + + def test_clear_with_columnwise_only(self, input_tensor): + """rowwise sub-storage is None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(rowwise_usage=False) + assert ht.rowwise_sub_storage is None + + ht.clear() + + assert all(n == 0 for n in self._primary_data_numels(ht.columnwise_sub_storage)) + + def test_clear_with_both_sub_storages_dropped(self, input_tensor): + """Both sub-storages are None — clear() must not crash.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.update_usage(rowwise_usage=False, columnwise_usage=False) + assert ht.rowwise_sub_storage is None + assert ht.columnwise_sub_storage is None + + ht.clear() # must not raise + + def test_clear_is_idempotent(self, input_tensor): + """Calling clear() twice must not raise and leaves buffers empty.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + ht.clear() + ht.clear() + for sub in (ht.rowwise_sub_storage, ht.columnwise_sub_storage): + for n in self._primary_data_numels(sub): + assert n == 0 + + +@requires_fp8_and_nvfp4 +class TestHybridDetachIsolation: + """``HybridQuantizedTensor.detach()`` must produce a hybrid whose + sub-storage wrappers are NOT shared with ``self`` — so that + ``detached.prepare_for_saving()`` does not null out fields on the + original. + + This is the property cpu_offload_v2 relies on at + ``cpu_offload.py:378-382``: + + tensor_copy = tensor.detach() + saved_tensors, _ = tensor_copy.prepare_for_saving() + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + def test_detach_produces_distinct_sub_storage_wrappers(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + detached = ht.detach() + + assert detached is not ht + assert detached._rowwise_storage is not ht._rowwise_storage + assert detached._columnwise_storage is not ht._columnwise_storage + + def test_detach_prepare_for_saving_does_not_affect_original(self, input_tensor): + """prepare_for_saving on the detach() copy must not null the original. + + This is the specific invariant the cpu_offload_v2 push/reload cycle + depends on. Without it, a second push on the same tensor — or even + a bare ``.device`` read during offload eligibility checks — hits + `` has no data!``. + """ + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + + detached = ht.detach() + _ = detached.prepare_for_saving() + + # Original must still be usable: dequantize, .device, repeated clone + _ = ht.dequantize() + _ = ht.device + + def test_detach_shares_underlying_buffers(self, input_tensor): + """Buffer tensors are shared (no GPU allocation) — only wrappers differ.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + ht = hq.quantize(input_tensor) + detached = ht.detach() + + orig_row_buffers = ht._rowwise_storage.get_data_tensors() + new_row_buffers = detached._rowwise_storage.get_data_tensors() + if not isinstance(orig_row_buffers, tuple): + orig_row_buffers = (orig_row_buffers,) + new_row_buffers = (new_row_buffers,) + for a, b in zip(orig_row_buffers, new_row_buffers): + if a is None and b is None: + continue + assert a is b, "detach() must share buffer tensors, not copy them" + + @requires_fp8_and_nvfp4 class TestHybridSaveRestore: """Test prepare_for_saving / restore_from_saved round-trip.""" @@ -1308,6 +1478,169 @@ def factory(role): ).any(), f"Gradient for '{name}' NaN ({row_name} row × {col_name} col)" +# --------------------------------------------------------------------------- +# CPU offload push/pop protocol (v2 OffloadableLayerState path) +# --------------------------------------------------------------------------- + + +class TestHybridCpuOffloadPushPop: + """Exercise the cpu_offload_v2 push/pop protocol on HybridQuantizedTensor. + + Uses :class:`OffloadableLayerState` directly — same pattern as + ``test_cpu_offloading.py::TestsOffloadableLayerState::test_general``. + Each test runs the full cycle: + + push → start_offload → release_activation_forward_gpu_memory + → start_reload → pop → release_all_memory + + The push path decomposes the hybrid via ``prepare_for_saving`` + (HybridQuantizedTensorStorage), recursively pushes each sub-storage + buffer, then reconstructs on pop via ``restore_from_saved``. Sub-buffers + below the 256K-element offload threshold (e.g. small block scales) are + returned unchanged; large data buffers round-trip through CPU. + """ + + # Hybrid tensor shape — each sub-storage primary buffer must exceed the + # cpu_offload _check_if_offload threshold (256K elements) so the path is + # actually exercised end-to-end. + _SHAPE = (1024, 1024) + + def _run_roundtrip(self, hybrid_tensor): + """Push → offload → release → reload → pop one hybrid tensor. + + Returns the reloaded tensor (a new HybridQuantizedTensor instance + reconstructed from the gathered-back buffers). + """ + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + + tid = state.push_tensor(hybrid_tensor) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + reloaded = state.pop_tensor(tid) + torch.cuda.synchronize() + + try: + return reloaded + finally: + state.release_all_memory() + + @pytest.mark.parametrize("row_name,col_name", _build_cross_format_params()) + def test_push_pop_roundtrip(self, row_name, col_name): + """Dequantize-equivalence round-trip across the full 14-pair matrix.""" + torch.manual_seed(42) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + row_cfg = _QUANTIZER_CONFIGS[row_name] + col_cfg = _QUANTIZER_CONFIGS[col_name] + hq = HybridQuantizer( + rowwise_quantizer=row_cfg[0](), + columnwise_quantizer=col_cfg[0](), + ) + hybrid = hq.quantize(inp) + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + torch.testing.assert_close(reloaded.dequantize(), expected) + + @requires_fp8_and_nvfp4 + def test_push_pop_preserves_sub_storage_types(self): + """Reconstructed hybrid preserves each sub-storage's concrete type.""" + torch.manual_seed(7) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + row_type = type(hybrid.rowwise_sub_storage) + col_type = type(hybrid.columnwise_sub_storage) + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded.rowwise_sub_storage, row_type) + assert isinstance(reloaded.columnwise_sub_storage, col_type) + + @requires_fp8_and_nvfp4 + def test_push_pop_with_rowwise_only(self): + """Columnwise sub-storage dropped pre-push — roundtrip still works.""" + torch.manual_seed(11) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + hybrid.update_usage(columnwise_usage=False) + assert hybrid.columnwise_sub_storage is None + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert reloaded.columnwise_sub_storage is None + assert reloaded.rowwise_sub_storage is not None + torch.testing.assert_close(reloaded.dequantize(), expected) + + @requires_fp8_and_nvfp4 + def test_push_pop_with_columnwise_only(self): + """Rowwise sub-storage dropped pre-push — roundtrip still works. + + Uses the reversed hybrid (NVFP4 rowwise + FP8 columnwise) so that + ``hybrid.dequantize()`` can fall through to the columnwise sub-storage. + ``HybridQuantizedTensorStorage.dequantize`` prefers rowwise and only + falls back to columnwise when rowwise is ``None``; NVFP4 does not yet + support columnwise-only dequantize, but Float8 does. + """ + torch.manual_seed(13) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp4_row_fp8_col() + hybrid = hq.quantize(inp) + hybrid.update_usage(rowwise_usage=False) + assert hybrid.rowwise_sub_storage is None + expected = hybrid.dequantize() + + reloaded = self._run_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert reloaded.rowwise_sub_storage is None + assert reloaded.columnwise_sub_storage is not None + torch.testing.assert_close(reloaded.dequantize(), expected) + + @requires_fp8_and_nvfp4 + def test_push_pop_roundtrip_does_not_leak_intermediate_buffers(self): + """After release_all_memory the offloader holds no hybrid buffers. + + Sanity check that the v2 cycle completes cleanly — no dangling CPU + pinned buffers left behind on a one-shot push/pop. + """ + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + torch.manual_seed(17) + inp = torch.randn(*self._SHAPE, dtype=torch.bfloat16, device="cuda") + + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hybrid = hq.quantize(inp) + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + + tid = state.push_tensor(hybrid) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + _ = state.pop_tensor(tid) + torch.cuda.synchronize() + state.release_all_memory() + + assert len(state.fwd_gpu_tensor_group.tensor_list) == 0 + assert len(state.cpu_tensor_group.tensor_list) == 0 + assert len(state.bwd_gpu_tensor_group.tensor_list) == 0 + assert state.state == "not_offloaded" + + # --------------------------------------------------------------------------- # 3-format hybrid: different quantization for fprop, dgrad, wgrad # --------------------------------------------------------------------------- diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 1c80193f407..c607036e9f1 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -243,7 +243,51 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: return HybridQuantizedTensorStorage.dequantize(self, dtype=dtype) def detach(self) -> HybridQuantizedTensor: - return HybridQuantizedTensor.make_like(self) + """Return a new HybridQuantizedTensor with cloned sub-storage wrappers. + + Each sub-storage is re-wrapped via its own ``make_like`` so the + new hybrid tensor has independent sub-storage objects that share + the *underlying* buffer tensors with ``self``. This is required for + the cpu_offload_v2 pattern at ``cpu_offload.py:378-382``:: + + tensor_copy = tensor.detach() + saved_tensors, _ = tensor_copy.prepare_for_saving() # nulls fields + + If ``detach()`` merely shared sub-storage objects, the + ``prepare_for_saving`` call above would null out fields on the + original ``tensor`` too (since both hybrids would point at the same + sub-storage Python objects), and subsequent operations — even a + bare ``.device`` read during ``_check_if_offload`` for a follow-up + ``push_tensor`` on the same original — would crash with + `` has no data!``. + """ + row = None + if self._rowwise_storage is not None: + row_cls = type(self._rowwise_storage) + if hasattr(row_cls, "make_like"): + row = row_cls.make_like(self._rowwise_storage) + else: + # Storage-only sub-storages (HybridQuantizer.internal=True + # path) don't have make_like; the cpu_offload_v2 path does + # not hit this branch, but keep the behaviour safe by + # sharing the reference as before. + row = self._rowwise_storage + col = None + if self._columnwise_storage is not None: + col_cls = type(self._columnwise_storage) + if hasattr(col_cls, "make_like"): + col = col_cls.make_like(self._columnwise_storage) + else: + col = self._columnwise_storage + return HybridQuantizedTensor( + shape=self.shape, + dtype=self.dtype, + rowwise_storage=row, + columnwise_storage=col, + rowwise_quantizer=self._rowwise_quantizer, + columnwise_quantizer=self._columnwise_quantizer, + quantizer=self._quantizer, + ) def get_metadata(self) -> Dict[str, Any]: return HybridQuantizedTensorStorage.get_metadata(self) diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index a2da94f614e..f92d596934d 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -71,6 +71,21 @@ def update_usage( if columnwise_usage is not None and not columnwise_usage: self._columnwise_storage = None + def clear(self): + """Deallocate both sub-storages' buffers. + + Delegates to each sub-storage's own ``clear()``; no-op when a + sub-storage is ``None`` (columnwise-only or rowwise-only hybrid). + + Used by ``cpu_offload_v1`` after the offloader has taken its own + reference to the extracted buffers, to release the GPU-resident + originals. + """ + if self._rowwise_storage is not None: + self._rowwise_storage.clear() + if self._columnwise_storage is not None: + self._columnwise_storage.clear() + def get_usages(self) -> Dict[str, bool]: return { "rowwise": self._rowwise_storage is not None, From 16fb371b2b42e24a3e74a5735d7fa0e6db6bcf69 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 24 Apr 2026 12:44:09 +0000 Subject: [PATCH 07/97] Activation recomputation Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 686 ++++++++++++++++++++++ 1 file changed, 686 insertions(+) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 442f7de0bd5..e93111adfa6 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -3749,3 +3749,689 @@ def test_make_like_is_independent(self): param = self._make_hybrid_param() copy = HybridQuantizedTensor.make_like(param) assert copy is not param + + +# --------------------------------------------------------------------------- +# 16. Activation recomputation (torch.utils.checkpoint / te.checkpoint) +# --------------------------------------------------------------------------- + + +def _reset_rng(seed: int = 1234): + """Reset deterministic RNG for reproducible forward/backward comparisons. + + Activation recompute relies on RNG equality between the first forward + and the recomputed forward. These tests use dropout-free modules, so + RNG advancement doesn't affect numerics, but we still reset between + runs so the reference (no-recompute) and checkpointed paths see + identical weight init, input, and grad_output seeds. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _collect_outputs(out, inp, model): + """Gather forward output, input grad, and parameter grads into a flat list. + + Mirrors ``test_numerics.py::_test_e2e_*_recompute`` conventions so the + comparison against a non-recomputed baseline is a simple zip. + """ + results = [out.detach().clone()] + if inp.grad is not None: + results.append(inp.grad.detach().clone()) + for _, p in model.named_parameters(): + if p.requires_grad and p.grad is not None: + results.append(p.grad.detach().clone()) + return results + + +def _assert_outputs_bitwise_equal(ref, test, label): + """All stateless same-format hybrid recipes should be bitwise-identical + under activation recompute: same input bytes → same quantized bytes → + same GEMM result. Any drift means the recompute path silently diverged + (e.g. fell back to a different quantization path).""" + assert len(ref) == len(test), f"{label}: output count mismatch" + for i, (r, t) in enumerate(zip(ref, test)): + torch.testing.assert_close( + t, r, rtol=0, atol=0, msg=f"{label}: bitwise mismatch at output {i}" + ) + + +@requires_fp8 +class TestHybridActivationRecompute: + """Activation recomputation around TE modules with a hybrid CustomRecipe. + + Probes the interaction between ``HybridQuantizedTensor`` / + ``HybridQuantizedTensorStorage`` and the three activation-checkpoint + paths in use today: + + * ``te.checkpoint(fn, ..., use_reentrant=True)`` — reentrant path; wraps + ``torch.autograd.Function`` that re-runs the forward under + ``activation_recompute_forward(recompute_phase=True)``. This is the + Megatron-style path. + * ``te.checkpoint(fn, ..., use_reentrant=False)`` — non-reentrant path; + uses ``_checkpoint_hook`` (torch saved-tensors hooks) to discard + saved tensors on the first forward and recompute them on unpack. + * ``torch.utils.checkpoint.checkpoint(fn, ..., use_reentrant=False)`` + — vanilla PyTorch path without TE wrapper. Exercised because users + (and some Megatron configs) invoke it directly around TE modules. + + Failure modes it catches: + + * Silent BF16 fallback during recompute (would break bitwise parity + but pass loose tolerance — hence the bitwise assertion for + same-format stateless recipes). + * ``HybridQuantizedTensorStorage.prepare_for_saving`` / + ``restore_from_saved`` chain losing a sub-storage across the + save-for-backward boundary. + * ``HybridQuantizedTensor`` subclass being stripped by the autograd + engine (would manifest as ``AttributeError`` on the recomputed + tensor). + """ + + in_features = 128 + out_features = 128 + batch = 32 + + # ----- helpers --------------------------------------------------- + + def _same_format_fp8_recipe(self): + """Same-format FP8 current scaling both directions → bitwise-safe + baseline. Matches + :class:`TestHybridGemmBitwiseIdentical` construction so + recompute parity can be asserted bitwise-equal.""" + return _hybrid_custom_recipe( + row_factory=_fp8_row_factory, + col_factory=_fp8_col_factory, + grad_factory=_fp8_grad_factory, + ) + + def _same_format_mxfp8_recipe(self): + """Same-format MXFP8 both directions — stateless, per-block scales + computed from the tensor content; bitwise-stable under recompute.""" + return _hybrid_custom_recipe( + row_factory=_mxfp8_factory, + col_factory=_mxfp8_factory, + grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + ) + + def _cross_format_fp8_mxfp8_recipe(self): + """Cross-format FP8 row + MXFP8 col — the canonical hybrid + scenario. Numerical parity is not bitwise because the wgrad GEMM + uses MXFP8 scaling modes on both operands (so grad_output must be + MXFP8 columnwise), pairing differently from the fprop path.""" + return _hybrid_custom_recipe( + row_factory=_fp8_row_factory, + col_factory=_mxfp8_factory, + grad_factory=_mxfp8_factory, + ) + + def _run_linear(self, recipe_obj, *, checkpoint_fn=None): + """Build a fresh Linear, run forward+backward, return collected + outputs. ``checkpoint_fn`` is an optional callable of the form + ``fn(model, inp) -> output`` that wraps the forward in an + activation-checkpoint implementation; ``None`` is the reference + (non-recompute) baseline. + """ + _reset_rng(seed=4242) + model = Linear( + self.in_features, self.out_features, params_dtype=torch.bfloat16 + ).cuda() + inp = torch.randn( + self.batch, + self.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + inp.retain_grad() + + with autocast(enabled=True, recipe=recipe_obj): + out = checkpoint_fn(model, inp) if checkpoint_fn is not None else model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + def _run_transformer_layer(self, recipe_obj, *, checkpoint_fn=None): + """Small TransformerLayer (no dropout, fuse_qkv) with optional + activation checkpointing around the whole block.""" + _reset_rng(seed=5151) + hidden = 128 + ffn = 128 + nheads = 4 + seq = 8 + bs = 4 + + model = TransformerLayer( + hidden, + ffn, + nheads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + ).cuda() + + inp = torch.randn( + seq, bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + inp.retain_grad() + + with autocast(enabled=True, recipe=recipe_obj): + out = checkpoint_fn(model, inp) if checkpoint_fn is not None else model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + # ----- te.checkpoint, reentrant --------------------------------- + + def test_te_checkpoint_reentrant_linear_fp8_bitwise(self): + """te.checkpoint(use_reentrant=True) around te.Linear with + same-format FP8 hybrid → bitwise parity with non-recompute. + + This is the Megatron-style activation-recompute path. Bitwise + parity catches silent BF16 fallback (would pass loose tolerance). + """ + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_reentrant_linear_mxfp8_bitwise(self): + """Same as FP8 but MXFP8 hybrid — per-block scales must recompute + identically. Asserts that the MXFP8 path does not get disabled + during recompute.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) MXFP8") + + # ----- te.checkpoint, non-reentrant ----------------------------- + + def test_te_checkpoint_non_reentrant_linear_fp8_bitwise(self): + """te.checkpoint(use_reentrant=False) — the saved-tensors-hooks + path. Different recompute infra (``_checkpoint_hook``) than the + reentrant path; validates the hybrid activation survives the + pack/unpack transport.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_non_reentrant_linear_mxfp8_bitwise(self): + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) MXFP8") + + # ----- torch.utils.checkpoint (vanilla, non-reentrant) ---------- + # + # These tests document a *known* TE-level incompatibility between + # vanilla ``torch.utils.checkpoint.checkpoint(..., use_reentrant=False)`` + # and TE's weight-workspace cache (``_linear_forward_impl`` in + # ``module/linear.py``). The mechanism: + # + # * First forward: ``quantize_weight`` takes the cache-miss path, + # creating a fresh hybrid workspace and threading it into + # ``prepare_for_saving`` → ``ctx.save_for_backward``. + # * Recompute forward: the workspace is already populated on the + # module, so ``quantize_weight`` takes the cache-hit path and + # saves a different tensor-count. + # + # Vanilla ``torch.utils.checkpoint`` (``use_reentrant=False``) + # enforces a strict count match between original-forward and + # recompute-forward ``save_for_backward`` calls, and rejects the + # discrepancy with ``CheckpointError: A different number of tensors + # was saved``. The 2:1 count ratio (``8`` forward vs ``4`` recompute) + # is a hybrid signature — both sub-storages are saved on cache-miss + # and only the remaining one on cache-hit. + # + # ``te.checkpoint`` avoids this by threading ``is_first_microbatch`` + # / ``skip_fp8_weight_update`` correctly across the recompute phase, + # which is why the ``te.checkpoint`` tests above pass bitwise. + # + # Keeping the xfail'd tests here: + # 1. pins the boundary — users hitting this failure get a clear + # diagnosis and pointer to ``te.checkpoint``; + # 2. becomes a regression signal if the underlying cache-vs- + # checkpoint interaction is ever resolved (the xfail flips to + # an unexpected pass). + # + # Not hybrid-specific *in nature* (any quantized TE module with + # weight-workspace caching hits it under vanilla torch checkpoint), + # but hybrid amplifies and surfaces it via the 2x sub-storage count. + + _TORCH_CHECKPOINT_CACHE_XFAIL = pytest.mark.xfail( + raises=torch.utils.checkpoint.CheckpointError, + strict=True, + reason=( + "Vanilla torch.utils.checkpoint(use_reentrant=False) is" + " incompatible with TE's weight-workspace cache: cache-miss" + " on the first forward saves a different tensor count than" + " cache-hit on recompute. Use te.checkpoint instead (tested" + " above)." + ), + ) + + @_TORCH_CHECKPOINT_CACHE_XFAIL + def test_torch_checkpoint_non_reentrant_linear_fp8_bitwise(self): + """Vanilla ``torch.utils.checkpoint.checkpoint`` without TE wrapper + around a hybrid-quantized te.Linear. + + Users invoke ``torch.utils.checkpoint`` directly in many Megatron + branches and custom recomputation schemes. Currently fails due to + the weight-workspace cache interaction documented above; pins the + boundary so a future fix would flip this to an unexpected pass. + """ + def fn(model, inp): + return torch.utils.checkpoint.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "torch.utils.checkpoint FP8") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + @_TORCH_CHECKPOINT_CACHE_XFAIL + def test_torch_checkpoint_non_reentrant_linear_mxfp8_bitwise(self): + def fn(model, inp): + return torch.utils.checkpoint.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._same_format_mxfp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "torch.utils.checkpoint MXFP8") + + # ----- cross-format + recompute (functional, loose tolerance) --- + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_te_checkpoint_reentrant_linear_cross_format(self): + """Cross-format hybrid (FP8 row + MXFP8 col) under activation + recompute. Numerics are allowed to drift from non-recompute only + through paths recompute is allowed to affect; in practice they + should still match tightly because the recipe is stateless. Loose + tolerance catches only catastrophic silent fallbacks.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_linear(self._cross_format_fp8_mxfp8_recipe(), checkpoint_fn=None) + test = self._run_linear(self._cross_format_fp8_mxfp8_recipe(), checkpoint_fn=fn) + # Expected to match bitwise since both quantizers are stateless + # and the input bytes are identical between the two runs. Use a + # strict tolerance; if this ever drifts it's a real bug. + _assert_outputs_bitwise_equal( + ref, test, "te.checkpoint(reentrant) FP8xMXFP8 cross-format" + ) + + # ----- TransformerLayer ----------------------------------------- + + def test_te_checkpoint_reentrant_transformer_layer_fp8(self): + """te.checkpoint(reentrant) around a full TransformerLayer under + hybrid FP8. Exercises LayerNormLinear + DPA + LayerNormMLP in one + shot — the ``with_quantized_norm=False`` unfused path for hybrid + in ``layernorm_linear.py`` / ``layernorm_mlp.py`` must produce + the same result when recomputed. + + Asserted bitwise: the module uses ``hidden_dropout=0.0``, + ``attention_dropout=0.0``, and ``te.checkpoint`` restores RNG + state before recompute, so every kernel sees identical inputs and + there are no stochastic ops. Non-determinism at this level would + indicate a real regression (e.g. a kernel quietly taking a + non-deterministic code path) — not measurement noise.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=True) + + ref = self._run_transformer_layer( + self._same_format_fp8_recipe(), checkpoint_fn=None + ) + test = self._run_transformer_layer( + self._same_format_fp8_recipe(), checkpoint_fn=fn + ) + _assert_outputs_bitwise_equal( + ref, test, "te.checkpoint(reentrant) TransformerLayer FP8" + ) + + def test_te_checkpoint_non_reentrant_transformer_layer_fp8(self): + """Same TransformerLayer setup but through the non-reentrant + saved-tensors-hooks recompute path. Same bitwise-equality + rationale as the reentrant variant above.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=False) + + ref = self._run_transformer_layer( + self._same_format_fp8_recipe(), checkpoint_fn=None + ) + test = self._run_transformer_layer( + self._same_format_fp8_recipe(), checkpoint_fn=fn + ) + _assert_outputs_bitwise_equal( + ref, test, "te.checkpoint(non-reentrant) TransformerLayer FP8" + ) + + # ----- quantized_model_init + recompute ------------------------- + + def test_te_checkpoint_reentrant_quantized_model_init_fp8_bitwise(self): + """Combine ``quantized_model_init`` (persistent + HybridQuantizedTensor weights) with activation recompute — + verifies the recompute path doesn't try to re-quantize an already- + quantized weight incorrectly, and the HybridQuantizer workspace + caching stays consistent across first-forward + recomputed-forward.""" + import transformer_engine.pytorch as te_pytorch + + hybrid_recipe = self._same_format_fp8_recipe() + + def _build_and_run(use_checkpoint): + _reset_rng(seed=7777) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear( + self.in_features, self.out_features, params_dtype=torch.bfloat16 + ).cuda() + inp = torch.randn( + self.batch, + self.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + inp.retain_grad() + with autocast(enabled=True, recipe=hybrid_recipe): + if use_checkpoint: + out = te_pytorch.checkpoint(model, inp, use_reentrant=True) + else: + out = model(inp) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + ref = _build_and_run(use_checkpoint=False) + test = _build_and_run(use_checkpoint=True) + _assert_outputs_bitwise_equal( + ref, test, "quantized_model_init + te.checkpoint(reentrant) FP8" + ) + + # ----- GroupedLinear + recompute -------------------------------- + + def _run_grouped_linear(self, recipe_obj, *, checkpoint_fn=None): + """Build a GroupedLinear, run forward+backward with optional + activation checkpointing around the module. Exercises the + ``_hybrid_split_quantize`` code path under recompute. + + GroupedLinear is the MoE token-dispatch kernel: a single batch + is split along dim-0 into ``num_gemms`` chunks and each chunk + goes through its own weight matrix. Under hybrid quantization, + ``_hybrid_split_quantize`` (``module/grouped_linear.py``) runs + ``tex.split_quantize`` twice (once per sub-quantizer direction) + and zips the results into a list of ``HybridQuantizedTensor`` + chunks — save-for-backward then receives a *list* of hybrid + tensors, not a single one, so the ``prepare_for_saving`` chain + has to handle an extended tensor-object list. + """ + _reset_rng(seed=9090) + num_gemms = 3 + hidden = 128 + ffn = 128 + bs = 24 + + model = GroupedLinear( + num_gemms, hidden, ffn, params_dtype=torch.bfloat16 + ).cuda() + inp = torch.randn( + bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + inp.retain_grad() + base = bs // num_gemms + rem = bs % num_gemms + m_splits = [base + (1 if i < rem else 0) for i in range(num_gemms)] + + with autocast(enabled=True, recipe=recipe_obj): + if checkpoint_fn is not None: + out = checkpoint_fn(model, inp, m_splits) + else: + out = model(inp, m_splits) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + def test_te_checkpoint_reentrant_grouped_linear_fp8_bitwise(self): + """GroupedLinear + te.checkpoint(reentrant) under same-format FP8 + hybrid. Exercises the MoE ``_hybrid_split_quantize`` + list-of- + hybrid-tensors save-for-backward path under recompute.""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp, m_splits): + return te_pytorch.checkpoint( + model, inp, m_splits, use_reentrant=True + ) + + ref = self._run_grouped_linear( + self._same_format_fp8_recipe(), checkpoint_fn=None + ) + test = self._run_grouped_linear( + self._same_format_fp8_recipe(), checkpoint_fn=fn + ) + _assert_outputs_bitwise_equal( + ref, test, "te.checkpoint(reentrant) GroupedLinear FP8" + ) + + def test_te_checkpoint_non_reentrant_grouped_linear_fp8_bitwise(self): + """Same GroupedLinear recompute setup but through the non- + reentrant saved-tensors-hooks path — verifies that the list of + hybrid activations survives the pack/unpack transport (one hook + invocation per split × per sub-storage buffer, not just one).""" + import transformer_engine.pytorch as te_pytorch + + def fn(model, inp, m_splits): + return te_pytorch.checkpoint( + model, inp, m_splits, use_reentrant=False + ) + + ref = self._run_grouped_linear( + self._same_format_fp8_recipe(), checkpoint_fn=None + ) + test = self._run_grouped_linear( + self._same_format_fp8_recipe(), checkpoint_fn=fn + ) + _assert_outputs_bitwise_equal( + ref, test, "te.checkpoint(non-reentrant) GroupedLinear FP8" + ) + + # ----- Selective attention recompute ---------------------------- + + def test_selective_attention_recompute_transformer_layer_fp8_bitwise(self): + """``TransformerLayer(..., checkpoint_core_attention=True)`` — + the Megatron default memory-savings pattern. + + Unlike full-layer recompute (``te.checkpoint(layer, inp)``), + selective attention recompute is a TransformerLayer-internal + option: only the DPA (dot-product attention) block is wrapped + in a checkpoint, everything else runs normally. This is a + *different* code path in ``transformer.py`` from the + ``te.checkpoint(...)`` tests above — DPA internally invokes its + own checkpoint context around the attention kernel. + + For hybrid, the question is whether a hybrid activation produced + by LayerNormLinear (QKV projection) survives the DPA-internal + recompute boundary (which saves it for backward) and is + consumable by the backward GEMM unchanged. + + Bitwise because the model uses ``hidden_dropout=0.0``, + ``attention_dropout=0.0``, and the DPA checkpoint restores RNG + state — so reference and recomputed paths should be identical + to the last bit.""" + _reset_rng(seed=5151) + hidden = 128 + ffn = 128 + nheads = 4 + seq = 8 + bs = 4 + + def _run(checkpoint_core_attention): + _reset_rng(seed=5151) + model = TransformerLayer( + hidden, + ffn, + nheads, + hidden_dropout=0.0, + attention_dropout=0.0, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + ).cuda() + inp = torch.randn( + seq, bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + inp.retain_grad() + with autocast(enabled=True, recipe=self._same_format_fp8_recipe()): + out = model(inp, checkpoint_core_attention=checkpoint_core_attention) + out.float().sum().backward() + return _collect_outputs(out, inp, model) + + ref = _run(checkpoint_core_attention=False) + test = _run(checkpoint_core_attention=True) + _assert_outputs_bitwise_equal( + ref, test, "checkpoint_core_attention TransformerLayer FP8" + ) + + # ----- Linear bitwise parametrized across all 4 stateless formats ----- + + @pytest.mark.parametrize( + "format_name,reentrant", + [ + pytest.param("fp8_current", True, id="fp8_current-reentrant"), + pytest.param("fp8_current", False, id="fp8_current-nonreentrant"), + pytest.param( + "mxfp8", + True, + id="mxfp8-reentrant", + marks=pytest.mark.skipif( + not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}" + ), + ), + pytest.param( + "mxfp8", + False, + id="mxfp8-nonreentrant", + marks=pytest.mark.skipif( + not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}" + ), + ), + pytest.param( + "block_fp8", + True, + id="block_fp8-reentrant", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"BlockFP8: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "block_fp8", + False, + id="block_fp8-nonreentrant", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"BlockFP8: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "nvfp4", + True, + id="nvfp4-reentrant", + marks=pytest.mark.skipif( + not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}" + ), + ), + pytest.param( + "nvfp4", + False, + id="nvfp4-nonreentrant", + marks=pytest.mark.skipif( + not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}" + ), + ), + ], + ) + def test_te_checkpoint_linear_all_stateless_formats_bitwise( + self, format_name, reentrant + ): + """Bitwise parity of Linear + te.checkpoint across all four + stateless hybrid formats (FP8 current, MXFP8, BlockFP8, NVFP4), + both reentrant and non-reentrant. + + Each format has a distinct history of columnwise-only kernel + support — BlockFP8 required C++ null-check patches before + columnwise-only mode worked, NVFP4 has packed FP4 layout plus + optional RHT cache, MXFP8 has [128,4]/[4,128] scale padding. + The recompute path exercises columnwise-only sub-quantizers + (rowwise is freed after fprop and only recreated on backward), + so format-specific columnwise-only handling is on the critical + path. + + A regression in any of these would silently fall back to BF16 + during recompute; bitwise equality catches that immediately.""" + import transformer_engine.pytorch as te_pytorch + + row_factory, col_factory_for_grad, hw_skip, hw_reason = _QUANTIZER_CONFIGS[ + format_name + ] + # Most formats have a distinct E5M2 variant for grad; NVFP4 has + # only one format (col_factory_for_grad is None → reuse + # row_factory, which is what the existing hybrid NVFP4 tests do). + grad_factory = col_factory_for_grad if col_factory_for_grad is not None else row_factory + + hybrid_recipe = _hybrid_custom_recipe( + row_factory=row_factory, + col_factory=row_factory, + grad_factory=grad_factory, + ) + + def fn(model, inp): + return te_pytorch.checkpoint(model, inp, use_reentrant=reentrant) + + ref = self._run_linear(hybrid_recipe, checkpoint_fn=None) + test = self._run_linear(hybrid_recipe, checkpoint_fn=fn) + label = f"te.checkpoint({'reentrant' if reentrant else 'non-reentrant'}) Linear {format_name}" + _assert_outputs_bitwise_equal(ref, test, label) + + # ----- save_for_backward round-trip (unit-level) ---------------- + + def test_prepare_restore_roundtrip_is_identity(self): + """Unit-level guarantee: the + ``prepare_for_saving`` / ``restore_from_saved`` chain used by + activation-recompute ``ctx.save_for_backward`` preserves both + sub-storages bitwise. + + This is the primitive the recompute path is built on; pinning it + here gives a focused failure signal independent of the module- + level recompute tests above.""" + torch.manual_seed(0) + inp = torch.randn(256, 256, dtype=torch.bfloat16, device="cuda") + hq = HybridQuantizer( + rowwise_quantizer=_fp8_row_factory(), + columnwise_quantizer=_fp8_col_factory(), + ) + hybrid = hq.quantize(inp) + expected = hybrid.dequantize() + + saved_tensors, saved_obj = hybrid.prepare_for_saving() + # Mimic the autograd ctx round-trip: all saved tensors pass + # through ``ctx.save_for_backward`` (a no-op for semantics). + leftover = saved_obj.restore_from_saved(list(saved_tensors)) + assert leftover == [], "restore_from_saved should consume every element" + torch.testing.assert_close(saved_obj.dequantize(), expected, rtol=0, atol=0) From a50fd63d533a3562b7f92c67c129b11eb2bf8af9 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 24 Apr 2026 14:10:48 +0000 Subject: [PATCH 08/97] TP/SP Signed-off-by: Evgeny --- tests/pytorch/distributed/run_hybrid_tp_sp.py | 543 ++++++++++++++++++ .../pytorch/distributed/test_hybrid_tp_sp.py | 153 +++++ tests/pytorch/test_hybrid_quantization.py | 46 ++ .../pytorch/tensor/hybrid_tensor.py | 73 +++ 4 files changed, 815 insertions(+) create mode 100644 tests/pytorch/distributed/run_hybrid_tp_sp.py create mode 100644 tests/pytorch/distributed/test_hybrid_tp_sp.py diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py new file mode 100644 index 00000000000..7353cd5b6ff --- /dev/null +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -0,0 +1,543 @@ +#!/usr/bin/python3 + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Distributed TP/SP numerics tests for hybrid quantization. + +Launched via torchrun from ``test_hybrid_tp_sp.py``. Each test compares a +tensor-parallel (optionally sequence-parallel) TE module running a hybrid +``CustomRecipe`` against a single-node reference module running the same +recipe. Weights are synchronized via ``_copy_params`` (shared with +``run_numerics.py``), so any drift between the two paths is a hybrid- +specific TP/SP issue rather than an initialization artifact. + +Test surface: + * ``te.Linear`` column-parallel and row-parallel, with and without + sequence parallelism. + * ``te.LayerNormLinear`` column-parallel with sequence parallelism — + the quantized-AG path that currently unfuses LN+quantize for + ``HybridQuantizer``. + * ``te.TransformerLayer`` with ``set_parallel_mode=True`` and SP on — + integration test hitting LayerNormLinear + DPA + LayerNormMLP + row- + parallel output projection in one shot. + +Only same-format hybrid recipes (FP8 current rowwise + FP8 current +columnwise; MXFP8 rowwise + MXFP8 columnwise) are exercised here so the +numerical signal is clean. Cross-format hybrid adds independent +numerical variation unrelated to TP/SP and is covered by single-GPU +tests already. + +Tolerances match upstream ``run_numerics.py`` per-format settings (see +``_get_tolerances``); they're loose enough to absorb the amax-reduction +and stochastic numerical asymmetries inherent to distributed FP8, tight +enough to catch a silent BF16 fallback on the hybrid sub-storage path. +""" + +import argparse +import datetime +import os +import sys +from pathlib import Path + +import torch +import torch.distributed as dist +from torch import nn + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common import recipe as te_recipe +from transformer_engine.pytorch import ( + Float8CurrentScalingQuantizer, + HybridQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) + +# Reuse helpers from run_numerics.py (sibling import — same pattern as +# run_numerics.py's own `from run_layer_with_overlap import _compare_tensors`). +TEST_ROOT = Path(__file__).parent.resolve() +sys.path.insert(0, str(TEST_ROOT)) +from run_layer_with_overlap import _compare_tensors # noqa: E402 + + +# ── Global state ───────────────────────────────────────────────────── + +SEQ_LEN = 32 +BATCH_SIZE = 32 +HIDDEN_SIZE = 128 +FFN_HIDDEN_SIZE = 128 +NR_HEADS = 4 + +WORLD_RANK = None +WORLD_SIZE = None +NCCL_WORLD = None +QUANTIZATION = None + +LOSS_FN = nn.MSELoss() + + +# ── Hybrid recipe factories ────────────────────────────────────────── +# +# Both rowwise and columnwise sub-quantizers use the same format so the +# observed distributed numerics only reflect TP/SP interactions and not +# cross-format composition noise. For comparison against vanilla built-in +# recipes that have well-understood TP/SP tolerances, see upstream +# ``run_numerics.py``. + + +def _make_fp8_current_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + +def _make_mxfp8_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=fp8_dtype) + + +def _hybrid_fp8_qfactory(role): + """FP8 current scaling in both directions for fwd roles; E5M2 for + grad roles (standard Hybrid:HYBRID format pairing).""" + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_current_quantizer(), + columnwise_quantizer=_make_fp8_current_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return _make_fp8_current_quantizer(fp8_dtype=tex.DType.kFloat8E5M2) + return _make_fp8_current_quantizer() + + +def _hybrid_mxfp8_qfactory(role): + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=_make_mxfp8_quantizer(), + columnwise_quantizer=_make_mxfp8_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return _make_mxfp8_quantizer(fp8_dtype=tex.DType.kFloat8E5M2) + return _make_mxfp8_quantizer() + + +def _make_nvfp4_quantizer(): + """Default NVFP4Quantizer: no RHT, no stochastic rounding, no 2D + scaling — matches upstream ``run_numerics.py::nvfp4_vanilla()`` which + strips the recipe to bare 1D block scaling for distributed TP + fairness. Same-format hybrid NVFP4 has no E5M2 variant (NVFP4 is a + single format family — E2M1 only), so grad roles reuse the same + NVFP4 quantizer.""" + return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + + +def _hybrid_nvfp4_qfactory(role): + if role in ("linear_input", "linear_weight", "linear_output"): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + if role in ("linear_grad_output", "linear_grad_input"): + return _make_nvfp4_quantizer() + return _make_nvfp4_quantizer() + + +def hybrid_recipe(): + """Fresh CustomRecipe instance per call (mirrors + ``run_numerics.quantization_recipe`` lifetime contract).""" + if QUANTIZATION == "hybrid_fp8": + return te_recipe.CustomRecipe(qfactory=_hybrid_fp8_qfactory) + if QUANTIZATION == "hybrid_mxfp8": + return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_qfactory) + if QUANTIZATION == "hybrid_nvfp4": + return te_recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) + raise ValueError(f"Unknown hybrid QUANTIZATION={QUANTIZATION!r}") + + +# ── Tolerances ─────────────────────────────────────────────────────── +# +# Upstream ``run_numerics.py::_get_tolerances`` uses (0.4, 0.25) for +# fp8_cs (loose because of sequence parallel & amax reduction) and +# (0.125, 0.0625) for other FP8 recipes. Hybrid with same-format +# sub-quantizers should inherit the underlying format's distributed +# behaviour — with slightly looser bounds to absorb the two-pass +# quantization (rowwise and columnwise quantizers run independently, so +# their outputs may differ by ~1 ULP from a single fused-quantize path +# in edge cases). + + +def _get_tolerances(): + if QUANTIZATION == "hybrid_fp8": + return {"rtol": 0.4, "atol": 0.25} + if QUANTIZATION == "hybrid_mxfp8": + return {"rtol": 0.2, "atol": 0.1} + if QUANTIZATION == "hybrid_nvfp4": + # Upstream ``run_numerics.py`` uses (0.125, 0.12) for vanilla + # NVFP4 with an open TODO to investigate why the tolerance is so + # large. Hybrid NVFP4 runs the same block-scaled kernel in each + # direction independently; bump atol modestly to absorb the + # two-pass asymmetry without hiding a real regression. + return {"rtol": 0.2, "atol": 0.15} + raise ValueError(f"No tolerances for QUANTIZATION={QUANTIZATION!r}") + + +# ── Distributed helpers ────────────────────────────────────────────── + + +def dist_print(msg, src=None, error=False): + stream = sys.stderr if error else sys.stdout + if WORLD_RANK == (0 if src is None else src): + stream.write(f"[rank{WORLD_RANK}] {msg}\n") + stream.flush() + + +def _gather(tensor, dim=0): + """All-gather with gradient scaling, matching + ``run_numerics.py::_gather``. Required because + ``torch.distributed.nn.functional.all_gather`` multiplies gradients + by WORLD_SIZE on the backward pass — so gradients in the + ``output_distributed`` backward would be WORLD_SIZE× too large + compared to ``output_single_node``.""" + + class HalfGradient(torch.autograd.Function): + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad_output): + return grad_output / WORLD_SIZE + + tensor = HalfGradient.apply(tensor) + gathered = torch.distributed.nn.functional.all_gather(tensor, group=NCCL_WORLD) + return torch.cat(gathered, dim=dim) + + +def _copy_params(model_distributed, model_single): + """Shard the single-node parameters into the TP-split distributed + model. Same algorithm as ``run_numerics.py::_copy_params``: for each + dim where shapes differ between the two params, slice the single- + node param along that dim using ``WORLD_RANK``.""" + for dp, sp in zip(model_distributed.parameters(), model_single.parameters()): + with torch.no_grad(): + to_copy = sp + for dim, _ in enumerate(dp.shape): + if dp.shape[dim] != sp.shape[dim]: + start = WORLD_RANK * dp.shape[dim] + end = (WORLD_RANK + 1) * dp.shape[dim] + indices = [slice(None)] * max(min(dim, len(dp.shape) - 1), 0) + indices.append(slice(start, end)) + if dim < len(dp.shape) - 1: + indices.append(slice(None)) + to_copy = sp[tuple(indices)] + dp.copy_(to_copy) + + +def _match_param_sizes(dist_param, single_param): + indices = [slice(None)] * len(single_param.shape) + for i in range(len(dist_param.shape)): + if dist_param.shape[i] != single_param.shape[i]: + start = WORLD_RANK * dist_param.shape[i] + end = (WORLD_RANK + 1) * dist_param.shape[i] + indices[i] = slice(start, end) + return single_param[tuple(indices)] + + +def _check_outputs(output_single, output_dist, label="outputs"): + failed = torch.tensor([0], dtype=torch.uint8, device="cuda") + f, info = _compare_tensors( + label, output_dist, output_single, **_get_tolerances() + ) + if f: + dist_print(info, src=WORLD_RANK, error=True) + failed[0] = int(f) + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: numerical check failed on at least one rank" + + +def _check_gradients(model_dist, model_single): + for i, ((name, pd), ps) in enumerate( + zip(model_dist.named_parameters(), model_single.parameters()) + ): + if pd.grad is None or ps.grad is None: + continue + failed = torch.tensor([0], dtype=torch.uint8, device="cuda") + ps_grad = _match_param_sizes(pd.grad, ps.grad) + f, info = _compare_tensors( + f"grad[{i}].{name}", pd.grad, ps_grad, **_get_tolerances() + ) + if f: + dist_print(info, src=WORLD_RANK, error=True) + failed[0] = int(f) + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"grad[{i}].{name}: failed on at least one rank" + + +def _apply_models(model_single, model_dist, inp_single, inp_dist, **kwargs): + """Run both models under te.autocast with a fresh hybrid recipe each + time. Both models see the same recipe instance-shape (CustomRecipe + with the same qfactory), but get independently-constructed + quantizers — matching how real training would instantiate them.""" + inp_single.requires_grad_() + inp_dist.requires_grad_() + with te.autocast(enabled=True, recipe=hybrid_recipe()): + out_single = model_single(inp_single, **kwargs) + with te.autocast(enabled=True, recipe=hybrid_recipe()): + out_dist = model_dist(inp_dist, **kwargs) + return out_single, out_dist + + +def _loss_backward(out_single, out_dist): + target = torch.randn_like(out_single) + LOSS_FN(out_single, target).backward() + LOSS_FN(out_dist, target).backward() + + +# ── Test 1: te.Linear TP (column + row) × SP (on/off) ──────────────── + + +def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): + dist_print( + f"linear: parallel_mode={parallel_mode} sequence_parallel={sequence_parallel}" + f" dtype={params_dtype}" + ) + + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + + model_single = te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=params_dtype).cuda() + model_dist = te.Linear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode=parallel_mode, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + # Prepare inputs matching run_numerics._test_linear's conventions. + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if parallel_mode == "row": + split = HIDDEN_SIZE // WORLD_SIZE + inp_dist = inp_single[:, WORLD_RANK * split : (WORLD_RANK + 1) * split].clone() + elif parallel_mode == "column": + if sequence_parallel: + # SP column: input is sharded along batch/sequence dim 0. + inp_single = torch.empty( + (WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE) + ).cuda().to(params_dtype) + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_dist = inp_single.clone() + else: + raise ValueError(parallel_mode) + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # For column-parallel: output is split along feature dim 1; gather. + # For row-parallel + SP: output is split along seq dim 0; gather. + if parallel_mode == "column" or (sequence_parallel and parallel_mode == "row"): + gather_dim = 1 if parallel_mode == "column" else 0 + out_dist = _gather(out_dist, dim=gather_dim) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"linear[{parallel_mode},sp={sequence_parallel}]") + + # Gradient check is only well-defined in these configurations (the + # others need cross-rank synchronization that the test doesn't + # perform — see run_numerics.py::_test_linear line 725 for the + # matching gate). + if parallel_mode == "column" or not sequence_parallel: + _check_gradients(model_dist, model_single) + + +def test_linear(): + for parallel_mode in ["column", "row"]: + for sequence_parallel in [False, True]: + _test_linear(parallel_mode, sequence_parallel) + + +# ── Test 2: te.LayerNormLinear column + SP ────────────────────────── + + +def _test_layernorm_linear(sequence_parallel, params_dtype=torch.bfloat16): + """Column-parallel LayerNormLinear. Exercises the SP all-gather path + that runs BEFORE quantization for hybrid (since + ``with_quantized_norm=False`` for HybridQuantizer — see + ``layernorm_linear.py:220``).""" + dist_print( + f"layernorm_linear: parallel_mode=column sequence_parallel={sequence_parallel}" + ) + + torch.manual_seed(23456) + torch.cuda.manual_seed(23456) + + model_single = te.LayerNormLinear( + HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=params_dtype + ).cuda() + model_dist = te.LayerNormLinear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode="column", + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + if sequence_parallel: + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # Column-parallel output: gather along dim 1. + out_dist = _gather(out_dist, dim=1) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"layernorm_linear[sp={sequence_parallel}]") + + +def test_layernorm_linear(): + for sequence_parallel in [False, True]: + _test_layernorm_linear(sequence_parallel) + + +# ── Test 3: te.TransformerLayer + TP + SP ─────────────────────────── + + +def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): + """Integration test: full TransformerLayer with TP and optional SP. + Hits LayerNormLinear(QKV), DPA, and LayerNormMLP all with hybrid + quantizers. If any of the unfused/hybrid code paths break something + visible to the backward graph, this catches it with a concrete + forward-output mismatch.""" + dist_print( + f"transformer_layer: parallel_mode=set sequence_parallel={sequence_parallel}" + ) + + torch.manual_seed(34567) + torch.cuda.manual_seed(34567) + + model_single = te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NR_HEADS, + attention_dropout=0.0, + hidden_dropout=0.0, + fuse_qkv_params=True, + params_dtype=params_dtype, + ).cuda() + model_dist = te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NR_HEADS, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + seq_length=WORLD_SIZE * SEQ_LEN if sequence_parallel else None, + attention_dropout=0.0, + hidden_dropout=0.0, + fuse_qkv_params=True, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + inp_single = ( + torch.randn((WORLD_SIZE * SEQ_LEN, BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + ) + if sequence_parallel: + inp_dist = inp_single[ + WORLD_RANK * SEQ_LEN : (WORLD_RANK + 1) * SEQ_LEN, :, : + ].contiguous() + else: + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + if sequence_parallel: + out_dist = _gather(out_dist, dim=0) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"transformer_layer[sp={sequence_parallel}]") + + +def test_transformer_layer(): + for sequence_parallel in [False, True]: + _test_transformer_layer(sequence_parallel) + + +# ── Driver ─────────────────────────────────────────────────────────── + + +def main(argv=None): + global WORLD_RANK, WORLD_SIZE, NCCL_WORLD, QUANTIZATION + + WORLD_RANK = int(os.getenv("RANK", "0")) + WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) + LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) + LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) + + assert WORLD_SIZE == LOCAL_SIZE, "This test is single-node only" + assert LOCAL_SIZE <= torch.cuda.device_count() + + torch.cuda.set_device(LOCAL_RANK) + dist.init_process_group( + backend="nccl", + rank=WORLD_RANK, + world_size=WORLD_SIZE, + timeout=datetime.timedelta(seconds=60), + init_method="env://", + device_id=torch.device(f"cuda:{LOCAL_RANK}"), + ) + NCCL_WORLD = dist.new_group(backend="nccl") + + parser = argparse.ArgumentParser() + parser.add_argument( + "--quantization", + type=str, + required=True, + choices=["hybrid_fp8", "hybrid_mxfp8", "hybrid_nvfp4"], + ) + parser.add_argument( + "--test", + type=str, + default="all", + choices=["all", "linear", "layernorm_linear", "transformer_layer"], + help="Run only the named test (speeds up iterative debugging)", + ) + args = parser.parse_args(argv) + QUANTIZATION = args.quantization + + test_map = { + "linear": test_linear, + "layernorm_linear": test_layernorm_linear, + "transformer_layer": test_transformer_layer, + } + if args.test == "all": + tests_to_run = list(test_map.values()) + else: + tests_to_run = [test_map[args.test]] + + for test_fn in tests_to_run: + dist_print(f"=== Starting {test_fn.__name__} ===") + test_fn() + dist.barrier() + dist_print(f"=== Passed {test_fn.__name__} ===") + + dist.destroy_process_group() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py new file mode 100644 index 00000000000..87e21255eb2 --- /dev/null +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -0,0 +1,153 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pytest driver for hybrid quantization TP/SP distributed tests. + +Launches ``run_hybrid_tp_sp.py`` via ``torchrun --nproc_per_node=N`` +and asserts a zero exit code. Rank-level numerical checks are performed +inside ``run_hybrid_tp_sp.py`` and propagated via ``dist.all_reduce`` +with ``ReduceOp.MAX`` on a failure flag, so a failure on any rank +fails the whole subprocess (and thus the pytest assertion). + +Mirrors the ``test_numerics.py`` ↔ ``run_numerics.py`` split pattern but +scoped to hybrid recipes only. Isolated from the main ``run_numerics.py`` +harness so that adding hybrid-specific cases here doesn't perturb the +larger vanilla-recipe test matrix. +""" + +import os +import subprocess +from pathlib import Path + +import pytest +import torch +import transformer_engine.pytorch as te + + +if torch.cuda.device_count() < 2: + pytest.skip( + "Distributed TP/SP tests need at least 2 GPUs.", + allow_module_level=True, + ) + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + +TEST_ROOT = Path(__file__).parent.resolve() +NUM_PROCS = min(2, torch.cuda.device_count()) +LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] + + +def _run_test(quantization: str, test: str = "all"): + script = TEST_ROOT / "run_hybrid_tp_sp.py" + cmd = LAUNCH_CMD + [str(script), "--quantization", quantization, "--test", test] + result = subprocess.run(cmd, env=os.environ, check=False) + assert result.returncode == 0, ( + f"run_hybrid_tp_sp.py (quantization={quantization}, test={test})" + f" exited with code {result.returncode}" + ) + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid FP8 current scaling (rowwise + columnwise same format) +# ────────────────────────────────────────────────────────────────────── +# +# FP8 current scaling is stateless per-tensor (amax computed from the +# live tensor) and therefore uses amax reduction across TP ranks when +# sequence parallelism is on. This is the tightest integration of +# hybrid with the distributed codepath: each rank computes an +# independent amax, reduces it across TP group, then both sub- +# quantizers (rowwise + columnwise) use the same reduced scale. If +# ``HybridQuantizer`` mis-plumbs the amax reduction through its two +# inner quantizers, we'd see numerical drift vs single-node. + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_hybrid_fp8_linear(): + """TP column + row × SP on/off for ``te.Linear`` under hybrid FP8 + current-scaling. Fine-grained: this runs first (cheapest, most + likely to surface TP-path hybrid bugs) so a failure here tells us + to stop before the more-expensive TransformerLayer run.""" + _run_test("hybrid_fp8", "linear") + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_hybrid_fp8_layernorm_linear(): + """Column-parallel ``te.LayerNormLinear`` with and without SP. + Probes the all-gather-before-quantize path that + ``layernorm_linear.py`` disables the fused norm for when + ``isinstance(input_quantizer, HybridQuantizer)``.""" + _run_test("hybrid_fp8", "layernorm_linear") + + +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_hybrid_fp8_transformer_layer(): + """Full ``te.TransformerLayer`` with ``set_parallel_mode=True`` and + SP on/off. Integration check hitting LayerNormLinear(QKV) → DPA → + LayerNormMLP → row-parallel output projection all under hybrid + FP8.""" + _run_test("hybrid_fp8", "transformer_layer") + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid MXFP8 (rowwise + columnwise same format) +# ────────────────────────────────────────────────────────────────────── +# +# MXFP8 is per-block (32-element microblocks), stateless, no amax +# reduction. Simpler distributed behaviour than FP8 current scaling, +# but exercises the ``[128, 4]`` / ``[4, 128]`` scale alignment padding +# through TP shards (each rank sees its own dim-0 slice which may not +# be a multiple of 128). + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8_linear(): + _run_test("hybrid_mxfp8", "linear") + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8_layernorm_linear(): + _run_test("hybrid_mxfp8", "layernorm_linear") + + +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8_transformer_layer(): + _run_test("hybrid_mxfp8", "transformer_layer") + + +# ────────────────────────────────────────────────────────────────────── +# Hybrid NVFP4 (rowwise + columnwise same format, 1D block scaling) +# ────────────────────────────────────────────────────────────────────── +# +# NVFP4 is the Rubin-era target recipe: 4-bit data (E2M1) with FP8 block +# scales on 16-element microblocks. The default ``NVFP4Quantizer()`` is +# 1D block scaling only — no RHT, no stochastic rounding, no 2D block +# scaling — matching upstream ``run_numerics.py::nvfp4_vanilla()``. +# Those more-sophisticated knobs are orthogonal to hybrid composition +# and can be layered in separately once baseline distributed NVFP4 +# hybrid is stable. +# +# Unlike FP8 current scaling, NVFP4 does not reduce amax across TP ranks +# (block-level scales are computed per-microblock locally), so SP +# amax-reduction issues don't apply. The tight interaction to watch is +# the packed FP4 dim-0 alignment in the TP shard — each rank sees a +# weight slice that may not be naturally aligned to the NVFP4 block +# boundary, and hybrid quantizes twice (rowwise + columnwise) on that +# shard. + + +@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") +def test_hybrid_nvfp4_linear(): + _run_test("hybrid_nvfp4", "linear") + + +@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") +def test_hybrid_nvfp4_layernorm_linear(): + _run_test("hybrid_nvfp4", "layernorm_linear") + + +@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") +def test_hybrid_nvfp4_transformer_layer(): + _run_test("hybrid_nvfp4", "transformer_layer") diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index e93111adfa6..5a0304b6166 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -104,6 +104,52 @@ def test_compatible_recipe_is_custom_recipe(self): hq = _make_hybrid_quantizer_fp8_row_fp4_col() assert hq._get_compatible_recipe() is recipe.CustomRecipe + def test_supports_only_rowwise_all_gather_nvfp4_columnwise(self): + """NVFP4 columnwise sub-quantizer forces rowwise-only AG. + + ``NVFP4Tensor.dequantize()`` raises ``NotImplementedError`` for + columnwise-only data, so the BF16 fallback in + ``gather_along_first_dim`` cannot operate on a columnwise-only + NVFP4 hybrid sub-storage. ``HybridQuantizer.supports_only_rowwise_all_gather`` + must return True in this case so ``_linear_forward_impl`` / + ``_linear_backward`` preserve rowwise data (which NVFP4 can + dequantize) instead. + """ + hq = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + + def test_supports_only_rowwise_all_gather_mxfp8_both(self): + """MXFP8 in both directions → columnwise dequant works → default + False so the save-columnwise (for wgrad) path stays active.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + hq = HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + assert hq.supports_only_rowwise_all_gather() is False + + def test_supports_only_rowwise_all_gather_fp8_current_propagates(self): + """Float8CurrentScalingQuantizer returns True for its own flag; + hybrid must propagate (not swallow) that semantics.""" + hq = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + + def test_supports_only_rowwise_all_gather_nvfp4_both(self): + """NVFP4 in both directions → columnwise sub-quantizer is NVFP4 + → forces rowwise-only AG regardless of rowwise flag.""" + hq = HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + ) + assert hq.supports_only_rowwise_all_gather() is True + @requires_fp8_and_nvfp4 class TestHybridQuantize: diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index c607036e9f1..9be5d9c3f94 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -141,6 +141,79 @@ def set_usage( ) -> None: super().set_usage(rowwise=rowwise, columnwise=columnwise) + def supports_only_rowwise_all_gather(self) -> bool: + """Whether TP activation all-gather must preserve rowwise data. + + Used by ``_linear_forward_impl`` / ``_linear_backward`` to decide + which direction of the saved activation to keep for the backward + input-AG: ``True`` keeps rowwise (drops columnwise), + ``False`` keeps columnwise (drops rowwise, default for block- + scaled formats whose columnwise is directly consumable by wgrad). + + Why hybrid needs a custom rule + ------------------------------ + ``gather_along_first_dim`` has no hybrid-specific dispatch, so + hybrid falls through to the generic BF16 fallback:: + + inp.dequantize() → all_gather BF16 → quantizer(out) + + The direction we preserve must therefore be one the hybrid can + dequantize. Two cases force rowwise preservation: + + 1. The rowwise sub-quantizer itself declares rowwise-only AG + (e.g. Float8 delayed / current scaling). Propagating keeps + hybrid consistent with its component semantics. + 2. The columnwise sub-quantizer is :class:`NVFP4Quantizer`: + ``NVFP4TensorStorage`` has no columnwise dequantize + (``_FromNVFP4Func.forward`` raises for ``is_colwise=True``), + so a columnwise-only NVFP4 sub-storage cannot traverse the + BF16 fallback. Rowwise preservation routes the fallback + through NVFP4's working rowwise dequantize instead. + + For MXFP8 / Float8Block / Float8CurrentScaling columnwise sub- + quantizers, columnwise dequantize works and the default + (``False``) keeps the smaller, wgrad-ready columnwise shard + saved — which is the more efficient memory choice. + + TODO(negvet): Add native hybrid dispatch to + ``gather_along_first_dim`` to remove the BF16 detour. + + * **Scope.** Branch at the top of ``gather_along_first_dim`` that + detects ``HybridQuantizedTensorStorage`` / ``HybridQuantizer``, + extracts ``rowwise_sub_storage`` and ``columnwise_sub_storage`` + with their sub-quantizers, dispatches each to its native + ``_all_gather_{fp8,mxfp8,nvfp4,fp8_blockwise}`` path, and wraps + the gathered sub-storages back into a ``HybridQuantizedTensor``. + Each per-format AG routine already supports rowwise-only or + columnwise-only input natively (including NVFP4 columnwise — + it gathers packed FP4 bytes without dequantize). + + * **Impact.** Replaces 2×–4× BF16 bandwidth cost with native + quantized AG. Mirrors the FSDP2 native-AG pattern we already + ship on ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather``. + Once it lands, the ``NVFP4Quantizer`` branch in this method + can be removed (columnwise NVFP4 AG works natively), leaving + only the rowwise-sub-quantizer propagation. + + * **Implementation notes.** Compose async handles across the two + per-direction AG calls into a single handle object with a + ``.wait()`` that waits on both. Pass ``out_shape=None`` to the + recursive calls so each format computes its own packed shape. + Preserve FP8 current / delayed rowwise-only semantics on + Hopper / L40 (``_all_gather_fp8`` reads ``inp._data`` which + may be ``None`` for a columnwise-only FP8 sub-storage on + those architectures). + """ + if self.rowwise_quantizer.supports_only_rowwise_all_gather(): + return True + # Local import avoids a circular dependency chain + # (nvfp4_tensor → quantized_tensor → hybrid_tensor at module import). + from .nvfp4_tensor import NVFP4Quantizer # noqa: PLC0415 + + if isinstance(self.columnwise_quantizer, NVFP4Quantizer): + return True + return False + def _get_compatible_recipe(self): # HybridQuantizer is only reachable via CustomRecipe (the qfactory # returns HybridQuantizer per role). Checking that the autocast recipe From 2214843f8e6c0dfac9981b9d409bbee00101182f Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 24 Apr 2026 14:54:54 +0000 Subject: [PATCH 09/97] Resolve comments: hybrid uniform list, make_empty try, __repr__, etc. Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 103 ++++++++++++++++++ .../pytorch/module/grouped_linear.py | 68 ++++++++++-- .../pytorch/tensor/hybrid_tensor.py | 42 +++---- .../tensor/storage/hybrid_tensor_storage.py | 12 +- 4 files changed, 192 insertions(+), 33 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 5a0304b6166..45557466c9f 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -1999,6 +1999,109 @@ def test_transformer_layer(self): self._run_fwd_bwd(model, inp) +@requires_fp8 +class TestHybridGroupedLinearClassifier: + """Unit tests for ``grouped_linear._is_hybrid_quantizer_list``. + + ``GroupedLinear`` dispatches its split-quantize between two mutually- + exclusive backends: ``tex.split_quantize`` (plain) and + ``_hybrid_split_quantize`` (all-hybrid). Neither can consume a mixed + list — ``tex.split_quantize`` doesn't recognise ``HybridQuantizer``, + and ``_hybrid_split_quantize`` calls ``q.rowwise_quantizer`` on every + element. Before the classifier was tightened, ``_has_hybrid_quantizer`` + used ``any(...)`` semantics: a single hybrid entry in a mixed list + would route to ``_hybrid_split_quantize`` and raise ``AttributeError`` + deep inside a grouped C++ call. These tests pin the new strict + classifier contract.""" + + def test_all_hybrid_returns_true(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _is_hybrid_quantizer_list, + ) + + quantizers = [ + _make_hybrid_quantizer_fp8_row_fp4_col() for _ in range(3) + ] + assert _is_hybrid_quantizer_list(quantizers) is True + + def test_all_plain_returns_false(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _is_hybrid_quantizer_list, + ) + + quantizers = [_make_fp8_quantizer() for _ in range(3)] + assert _is_hybrid_quantizer_list(quantizers) is False + + def test_all_none_returns_false(self): + """No quantizers at all (BF16 path) — classifier returns False so + the caller takes the non-fp8 branch.""" + from transformer_engine.pytorch.module.grouped_linear import ( + _is_hybrid_quantizer_list, + ) + + assert _is_hybrid_quantizer_list([None, None, None]) is False + + def test_mixed_hybrid_and_plain_raises(self): + """The actual bug: a mixed list used to silently route to + ``_hybrid_split_quantize`` and crash with ``AttributeError`` on + ``plain_q.rowwise_quantizer``. Now it fails fast at the + classifier with a user-actionable error.""" + from transformer_engine.pytorch.module.grouped_linear import ( + _is_hybrid_quantizer_list, + ) + + quantizers = [ + _make_hybrid_quantizer_fp8_row_fp4_col(), + _make_fp8_quantizer(), + _make_hybrid_quantizer_fp8_row_fp4_col(), + ] + with pytest.raises(ValueError) as exc_info: + _is_hybrid_quantizer_list(quantizers) + msg = str(exc_info.value) + # Error names both counts and points at the root cause so users + # can fix their ``qfactory`` without digging. + assert "mixes HybridQuantizer" in msg + assert "2 hybrid" in msg + assert "1 non-hybrid" in msg + assert "CustomRecipe" in msg and "qfactory" in msg + + def test_none_entries_ignored_when_remainder_is_uniform(self): + """None entries are filtered before uniformity check — a list + of hybrids plus a None must still classify as hybrid (not + mixed).""" + from transformer_engine.pytorch.module.grouped_linear import ( + _is_hybrid_quantizer_list, + ) + + quantizers = [ + _make_hybrid_quantizer_fp8_row_fp4_col(), + None, + _make_hybrid_quantizer_fp8_row_fp4_col(), + ] + assert _is_hybrid_quantizer_list(quantizers) is True + + def test_hybrid_split_quantize_rejects_plain_element(self): + """Defense-in-depth: even if a caller bypasses the classifier, + ``_hybrid_split_quantize`` itself asserts uniformity and raises + ``TypeError`` with a list of received types, rather than the + opaque ``AttributeError: 'Float8CurrentScalingQuantizer' object + has no attribute 'rowwise_quantizer'`` from the old code.""" + from transformer_engine.pytorch.module.grouped_linear import ( + _hybrid_split_quantize, + ) + + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + _make_hybrid_quantizer_fp8_row_fp4_col(), + _make_fp8_quantizer(), # Not hybrid — should trigger TypeError + ] + with pytest.raises(TypeError) as exc_info: + _hybrid_split_quantize(tensor, [16, 16], quantizers) + msg = str(exc_info.value) + assert "HybridQuantizer" in msg + assert "Float8CurrentScalingQuantizer" in msg + + # =========================================================================== # Quantized Parameters (quantized_model_init) tests for hybrid quantization # =========================================================================== diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 295bbaa6ff3..a182f922f7c 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -58,20 +58,68 @@ from ...debug.pytorch.debug_state import TEDebugState -def _has_hybrid_quantizer(quantizers): - """Check if any quantizer in the list is a HybridQuantizer.""" - return any(isinstance(q, HybridQuantizer) for q in quantizers if q is not None) +def _is_hybrid_quantizer_list(quantizers): + """Classify a GroupedLinear quantizer list as hybrid-uniform or plain-uniform. + + Returns ``True`` when every non-``None`` entry is a ``HybridQuantizer``, + ``False`` when none are. Raises :class:`ValueError` on a mixed list. + + Why it's a hard "or": neither dispatch branch at the call sites + supports a mixed list. ``tex.split_quantize`` (the plain path) does + not recognize ``HybridQuantizer``. :func:`_hybrid_split_quantize` + (the hybrid path) treats every entry as hybrid — it calls + ``q.rowwise_quantizer`` / ``q.columnwise_quantizer`` unconditionally, + which would ``AttributeError`` on a plain quantizer. Rejecting at + the classifier gives a clear actionable error instead of failing + deep inside a grouped C++ call. + + Supporting mixed would require a per-element grouped kernel that + accepts a heterogeneous quantizer vector; no such kernel exists + today. If that becomes a real requirement (e.g. per-expert hybrid + recipes in MoE), implement element-wise fallback here rather than + silently masking the mismatch. + """ + non_none = [q for q in quantizers if q is not None] + if not non_none: + return False + hybrid_count = sum(1 for q in non_none if isinstance(q, HybridQuantizer)) + if hybrid_count == 0: + return False + if hybrid_count == len(non_none): + return True + raise ValueError( + "GroupedLinear quantizer list mixes HybridQuantizer and non-hybrid" + f" quantizers ({hybrid_count} hybrid, {len(non_none) - hybrid_count}" + " non-hybrid). This combination is not supported: neither" + " `tex.split_quantize` nor `_hybrid_split_quantize` can consume a" + " heterogeneous list. Make the CustomRecipe `qfactory` return a" + " consistent type (all HybridQuantizer or all non-hybrid) across" + " every GEMM for the same role." + ) def _hybrid_split_quantize(tensor, m_splits, quantizers): - """Grouped split+quantize for HybridQuantizer lists. + """Grouped split+quantize for an **all-hybrid** quantizer list. + + Precondition: every ``q`` in ``quantizers`` is a ``HybridQuantizer``. + Enforce via :func:`_is_hybrid_quantizer_list` at the call site (or + the explicit assert below as a defense-in-depth). - Runs tex.split_quantize twice (once per direction with the native - sub-quantizers), then zips the results into HybridQuantizedTensorStorage. - Non-hybrid quantizers in the list fall back to per-split Python quantize. + Runs ``tex.split_quantize`` twice — once over the rowwise sub- + quantizers, once over the columnwise sub-quantizers — then zips + the two per-split results back into ``HybridQuantizedTensorStorage`` + per GEMM. Two grouped C++ calls instead of ``2 * num_gemms`` + ungrouped Python calls. """ from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage + if not all(isinstance(q, HybridQuantizer) for q in quantizers): + raise TypeError( + "_hybrid_split_quantize requires every quantizer to be a" + " HybridQuantizer; callers must gate on _is_hybrid_quantizer_list." + f" Got types: {[type(q).__name__ for q in quantizers]}" + ) + row_quantizers = [q.rowwise_quantizer for q in quantizers] col_quantizers = [q.columnwise_quantizer for q in quantizers] @@ -199,7 +247,7 @@ def forward( ) inp_view = inp.reshape(-1, in_features) inputmats: list - hybrid = _has_hybrid_quantizer(input_quantizers) + hybrid = _is_hybrid_quantizer_list(input_quantizers) if fp8 and not debug and not hybrid: # Disable bulk allocation when CPU offloading is active: offloading skips small # tensors (like scales), but bulk allocation shares storage across all tensors, @@ -415,7 +463,7 @@ def backward( grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) grad_output = [None] * ctx.num_gemms grad_biases = [None] * ctx.num_gemms - grad_output_hybrid = _has_hybrid_quantizer(ctx.grad_output_quantizers) + grad_output_hybrid = _is_hybrid_quantizer_list(ctx.grad_output_quantizers) if ctx.fp8 and not ctx.debug and not grad_output_hybrid: if ctx.use_bias: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) @@ -558,7 +606,7 @@ def backward( else: input_quantizer.set_usage(rowwise=False, columnwise=True) inputmats: list - input_hybrid = _has_hybrid_quantizer(ctx.input_quantizers) + input_hybrid = _is_hybrid_quantizer_list(ctx.input_quantizers) if ctx.fp8 and not ctx.debug and not input_hybrid: inputmats = tex.split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) elif ctx.fp8 and input_hybrid: diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 9be5d9c3f94..5feeefd0e3c 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -81,23 +81,28 @@ def make_empty( requires_grad: bool = False, pin_memory: bool = False, ) -> HybridQuantizedTensor: - self.rowwise_quantizer.internal = True - rowwise_empty = self.rowwise_quantizer.make_empty( - shape, - dtype=dtype, - device=device, - pin_memory=pin_memory, - ) - self.rowwise_quantizer.internal = False + # The ``internal`` flag toggles a shared-state mode on each sub- + # quantizer (returns bare ``*TensorStorage`` rather than a + # ``QuantizedTensor`` subclass from ``make_empty``). It is + # global to the quantizer instance, so a raise between + # ``internal=True`` and ``internal=False`` would leak the flag + # and corrupt subsequent non-hybrid quantize calls on the same + # sub-quantizer. ``try/finally`` guarantees the reset. + def _make_empty_internal(sub_quantizer): + prev_internal = sub_quantizer.internal + sub_quantizer.internal = True + try: + return sub_quantizer.make_empty( + shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + finally: + sub_quantizer.internal = prev_internal - self.columnwise_quantizer.internal = True - columnwise_empty = self.columnwise_quantizer.make_empty( - shape, - dtype=dtype, - device=device, - pin_memory=pin_memory, - ) - self.columnwise_quantizer.internal = False + rowwise_empty = _make_empty_internal(self.rowwise_quantizer) + columnwise_empty = _make_empty_internal(self.columnwise_quantizer) return HybridQuantizedTensor( shape=shape, @@ -136,11 +141,6 @@ def update_quantized( ) return dst - def set_usage( - self, *, rowwise: Optional[bool] = None, columnwise: Optional[bool] = None - ) -> None: - super().set_usage(rowwise=rowwise, columnwise=columnwise) - def supports_only_rowwise_all_gather(self) -> bool: """Whether TP activation all-gather must preserve rowwise data. diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index f92d596934d..d9f58734509 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -174,9 +174,17 @@ def get_metadata(self) -> Dict[str, Any]: } def __repr__(self): + row_type = ( + type(self._rowwise_storage).__name__ if self._rowwise_storage is not None else "None" + ) + col_type = ( + type(self._columnwise_storage).__name__ + if self._columnwise_storage is not None + else "None" + ) return ( "HybridQuantizedTensorStorage(" - f"rowwise={type(self._rowwise_storage).__name__}, " - f"columnwise={type(self._columnwise_storage).__name__}, " + f"rowwise={row_type}, " + f"columnwise={col_type}, " f"dtype={self._dtype})" ) From 88fe46731d1a48f07baa11894913d3181ec833d3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Apr 2026 14:56:03 +0000 Subject: [PATCH 10/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/run_hybrid_tp_sp.py | 28 ++---- tests/pytorch/test_cpu_offloading.py | 29 ++---- tests/pytorch/test_cpu_offloading_v1.py | 4 +- tests/pytorch/test_hybrid_quantization.py | 93 ++++++------------- 4 files changed, 41 insertions(+), 113 deletions(-) diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index 7353cd5b6ff..679f443c99a 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -244,9 +244,7 @@ def _match_param_sizes(dist_param, single_param): def _check_outputs(output_single, output_dist, label="outputs"): failed = torch.tensor([0], dtype=torch.uint8, device="cuda") - f, info = _compare_tensors( - label, output_dist, output_single, **_get_tolerances() - ) + f, info = _compare_tensors(label, output_dist, output_single, **_get_tolerances()) if f: dist_print(info, src=WORLD_RANK, error=True) failed[0] = int(f) @@ -262,9 +260,7 @@ def _check_gradients(model_dist, model_single): continue failed = torch.tensor([0], dtype=torch.uint8, device="cuda") ps_grad = _match_param_sizes(pd.grad, ps.grad) - f, info = _compare_tensors( - f"grad[{i}].{name}", pd.grad, ps_grad, **_get_tolerances() - ) + f, info = _compare_tensors(f"grad[{i}].{name}", pd.grad, ps_grad, **_get_tolerances()) if f: dist_print(info, src=WORLD_RANK, error=True) failed[0] = int(f) @@ -325,9 +321,7 @@ def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): elif parallel_mode == "column": if sequence_parallel: # SP column: input is sharded along batch/sequence dim 0. - inp_single = torch.empty( - (WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE) - ).cuda().to(params_dtype) + inp_single = torch.empty((WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) inp_single = _gather(inp_dist, dim=0).detach() else: @@ -368,16 +362,12 @@ def _test_layernorm_linear(sequence_parallel, params_dtype=torch.bfloat16): that runs BEFORE quantization for hybrid (since ``with_quantized_norm=False`` for HybridQuantizer — see ``layernorm_linear.py:220``).""" - dist_print( - f"layernorm_linear: parallel_mode=column sequence_parallel={sequence_parallel}" - ) + dist_print(f"layernorm_linear: parallel_mode=column sequence_parallel={sequence_parallel}") torch.manual_seed(23456) torch.cuda.manual_seed(23456) - model_single = te.LayerNormLinear( - HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=params_dtype - ).cuda() + model_single = te.LayerNormLinear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=params_dtype).cuda() model_dist = te.LayerNormLinear( HIDDEN_SIZE, HIDDEN_SIZE, @@ -420,9 +410,7 @@ def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): quantizers. If any of the unfused/hybrid code paths break something visible to the backward graph, this catches it with a concrete forward-output mismatch.""" - dist_print( - f"transformer_layer: parallel_mode=set sequence_parallel={sequence_parallel}" - ) + dist_print(f"transformer_layer: parallel_mode=set sequence_parallel={sequence_parallel}") torch.manual_seed(34567) torch.cuda.manual_seed(34567) @@ -457,9 +445,7 @@ def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): torch.randn((WORLD_SIZE * SEQ_LEN, BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) ) if sequence_parallel: - inp_dist = inp_single[ - WORLD_RANK * SEQ_LEN : (WORLD_RANK + 1) * SEQ_LEN, :, : - ].contiguous() + inp_dist = inp_single[WORLD_RANK * SEQ_LEN : (WORLD_RANK + 1) * SEQ_LEN, :, :].contiguous() else: inp_dist = inp_single.clone() diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 9ad83156f3e..a23e32646b4 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -28,6 +28,7 @@ mxfp8_available, _ = FP8GlobalStateManager.is_mxfp8_available() nvfp4_available, _ = FP8GlobalStateManager.is_nvfp4_available() + def _hybrid_fp8_mxfp8_qfactory(role): """Hybrid CustomRecipe factory: FP8 current-scaling rowwise + MXFP8 columnwise. @@ -490,14 +491,8 @@ def test_sanity(self, layer_type, recipe, backward_override): # grouped_linear is NOT skipped here — it passes test_sanity with # hybrid; only memory-accounting assertions trip it in test_memory / # test_manual_synchronization. - if ( - layer_type in ("layernorm_mlp_ops",) - and recipe is not None - and recipe.custom() - ): - pytest.skip( - f"Hybrid CustomRecipe + {layer_type} integration is not yet complete" - ) + if layer_type in ("layernorm_mlp_ops",) and recipe is not None and recipe.custom(): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") recipe_ctx = Utils.create_recipe_ctx(recipe) init_cuda_memory = Utils.get_cuda_memory_mb() @@ -556,9 +551,7 @@ def test_memory(self, layer_type, recipe, backward_override): and recipe is not None and recipe.custom() ): - pytest.skip( - f"Hybrid CustomRecipe + {layer_type} integration is not yet complete" - ) + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") offload_ctx, sync_function = get_cpu_offload_context( enabled=True, @@ -656,9 +649,7 @@ def test_manual_synchronization(self, recipe, layer_type, backward_override): and recipe is not None and recipe.custom() ): - pytest.skip( - f"Hybrid CustomRecipe + {layer_type} integration is not yet complete" - ) + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") offload_ctx, sync_function, manual_controller = get_cpu_offload_context( enabled=True, @@ -738,14 +729,8 @@ def test_numerics( and recipe.float8_block_scaling() ): pytest.skip("Fusible operations do not support FP8 block scaling recipe") - if ( - layer_type in ("layernorm_mlp_ops",) - and recipe is not None - and recipe.custom() - ): - pytest.skip( - f"Hybrid CustomRecipe + {layer_type} integration is not yet complete" - ) + if layer_type in ("layernorm_mlp_ops",) and recipe is not None and recipe.custom(): + pytest.skip(f"Hybrid CustomRecipe + {layer_type} integration is not yet complete") recipe_ctx = Utils.create_recipe_ctx(recipe) diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index aa128d258aa..c98f6098d9c 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -248,9 +248,7 @@ def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: s and quantization_recipe is not None and quantization_recipe.custom() ): - pytest.skip( - f"Hybrid CustomRecipe + {model_name} integration is not yet complete" - ) + pytest.skip(f"Hybrid CustomRecipe + {model_name} integration is not yet complete") # Construct model modules_list = [model_types[model_name]() for _ in range(NUM_LAYERS)] diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 45557466c9f..2ca341092ef 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -2019,9 +2019,7 @@ def test_all_hybrid_returns_true(self): _is_hybrid_quantizer_list, ) - quantizers = [ - _make_hybrid_quantizer_fp8_row_fp4_col() for _ in range(3) - ] + quantizers = [_make_hybrid_quantizer_fp8_row_fp4_col() for _ in range(3)] assert _is_hybrid_quantizer_list(quantizers) is True def test_all_plain_returns_false(self): @@ -4022,9 +4020,7 @@ def _run_linear(self, recipe_obj, *, checkpoint_fn=None): (non-recompute) baseline. """ _reset_rng(seed=4242) - model = Linear( - self.in_features, self.out_features, params_dtype=torch.bfloat16 - ).cuda() + model = Linear(self.in_features, self.out_features, params_dtype=torch.bfloat16).cuda() inp = torch.randn( self.batch, self.in_features, @@ -4059,9 +4055,7 @@ def _run_transformer_layer(self, recipe_obj, *, checkpoint_fn=None): params_dtype=torch.bfloat16, ).cuda() - inp = torch.randn( - seq, bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True - ) + inp = torch.randn(seq, bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) inp.retain_grad() with autocast(enabled=True, recipe=recipe_obj): @@ -4187,6 +4181,7 @@ def test_torch_checkpoint_non_reentrant_linear_fp8_bitwise(self): the weight-workspace cache interaction documented above; pins the boundary so a future fix would flip this to an unexpected pass. """ + def fn(model, inp): return torch.utils.checkpoint.checkpoint(model, inp, use_reentrant=False) @@ -4223,9 +4218,7 @@ def fn(model, inp): # Expected to match bitwise since both quantizers are stateless # and the input bytes are identical between the two runs. Use a # strict tolerance; if this ever drifts it's a real bug. - _assert_outputs_bitwise_equal( - ref, test, "te.checkpoint(reentrant) FP8xMXFP8 cross-format" - ) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) FP8xMXFP8 cross-format") # ----- TransformerLayer ----------------------------------------- @@ -4247,15 +4240,9 @@ def test_te_checkpoint_reentrant_transformer_layer_fp8(self): def fn(model, inp): return te_pytorch.checkpoint(model, inp, use_reentrant=True) - ref = self._run_transformer_layer( - self._same_format_fp8_recipe(), checkpoint_fn=None - ) - test = self._run_transformer_layer( - self._same_format_fp8_recipe(), checkpoint_fn=fn - ) - _assert_outputs_bitwise_equal( - ref, test, "te.checkpoint(reentrant) TransformerLayer FP8" - ) + ref = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) TransformerLayer FP8") def test_te_checkpoint_non_reentrant_transformer_layer_fp8(self): """Same TransformerLayer setup but through the non-reentrant @@ -4266,12 +4253,8 @@ def test_te_checkpoint_non_reentrant_transformer_layer_fp8(self): def fn(model, inp): return te_pytorch.checkpoint(model, inp, use_reentrant=False) - ref = self._run_transformer_layer( - self._same_format_fp8_recipe(), checkpoint_fn=None - ) - test = self._run_transformer_layer( - self._same_format_fp8_recipe(), checkpoint_fn=fn - ) + ref = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=fn) _assert_outputs_bitwise_equal( ref, test, "te.checkpoint(non-reentrant) TransformerLayer FP8" ) @@ -4339,12 +4322,8 @@ def _run_grouped_linear(self, recipe_obj, *, checkpoint_fn=None): ffn = 128 bs = 24 - model = GroupedLinear( - num_gemms, hidden, ffn, params_dtype=torch.bfloat16 - ).cuda() - inp = torch.randn( - bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True - ) + model = GroupedLinear(num_gemms, hidden, ffn, params_dtype=torch.bfloat16).cuda() + inp = torch.randn(bs, hidden, device="cuda", dtype=torch.bfloat16, requires_grad=True) inp.retain_grad() base = bs // num_gemms rem = bs % num_gemms @@ -4365,19 +4344,11 @@ def test_te_checkpoint_reentrant_grouped_linear_fp8_bitwise(self): import transformer_engine.pytorch as te_pytorch def fn(model, inp, m_splits): - return te_pytorch.checkpoint( - model, inp, m_splits, use_reentrant=True - ) + return te_pytorch.checkpoint(model, inp, m_splits, use_reentrant=True) - ref = self._run_grouped_linear( - self._same_format_fp8_recipe(), checkpoint_fn=None - ) - test = self._run_grouped_linear( - self._same_format_fp8_recipe(), checkpoint_fn=fn - ) - _assert_outputs_bitwise_equal( - ref, test, "te.checkpoint(reentrant) GroupedLinear FP8" - ) + ref = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) GroupedLinear FP8") def test_te_checkpoint_non_reentrant_grouped_linear_fp8_bitwise(self): """Same GroupedLinear recompute setup but through the non- @@ -4387,19 +4358,11 @@ def test_te_checkpoint_non_reentrant_grouped_linear_fp8_bitwise(self): import transformer_engine.pytorch as te_pytorch def fn(model, inp, m_splits): - return te_pytorch.checkpoint( - model, inp, m_splits, use_reentrant=False - ) + return te_pytorch.checkpoint(model, inp, m_splits, use_reentrant=False) - ref = self._run_grouped_linear( - self._same_format_fp8_recipe(), checkpoint_fn=None - ) - test = self._run_grouped_linear( - self._same_format_fp8_recipe(), checkpoint_fn=fn - ) - _assert_outputs_bitwise_equal( - ref, test, "te.checkpoint(non-reentrant) GroupedLinear FP8" - ) + ref = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=None) + test = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) + _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(non-reentrant) GroupedLinear FP8") # ----- Selective attention recompute ---------------------------- @@ -4453,9 +4416,7 @@ def _run(checkpoint_core_attention): ref = _run(checkpoint_core_attention=False) test = _run(checkpoint_core_attention=True) - _assert_outputs_bitwise_equal( - ref, test, "checkpoint_core_attention TransformerLayer FP8" - ) + _assert_outputs_bitwise_equal(ref, test, "checkpoint_core_attention TransformerLayer FP8") # ----- Linear bitwise parametrized across all 4 stateless formats ----- @@ -4516,9 +4477,7 @@ def _run(checkpoint_core_attention): ), ], ) - def test_te_checkpoint_linear_all_stateless_formats_bitwise( - self, format_name, reentrant - ): + def test_te_checkpoint_linear_all_stateless_formats_bitwise(self, format_name, reentrant): """Bitwise parity of Linear + te.checkpoint across all four stateless hybrid formats (FP8 current, MXFP8, BlockFP8, NVFP4), both reentrant and non-reentrant. @@ -4536,9 +4495,7 @@ def test_te_checkpoint_linear_all_stateless_formats_bitwise( during recompute; bitwise equality catches that immediately.""" import transformer_engine.pytorch as te_pytorch - row_factory, col_factory_for_grad, hw_skip, hw_reason = _QUANTIZER_CONFIGS[ - format_name - ] + row_factory, col_factory_for_grad, hw_skip, hw_reason = _QUANTIZER_CONFIGS[format_name] # Most formats have a distinct E5M2 variant for grad; NVFP4 has # only one format (col_factory_for_grad is None → reuse # row_factory, which is what the existing hybrid NVFP4 tests do). @@ -4555,7 +4512,9 @@ def fn(model, inp): ref = self._run_linear(hybrid_recipe, checkpoint_fn=None) test = self._run_linear(hybrid_recipe, checkpoint_fn=fn) - label = f"te.checkpoint({'reentrant' if reentrant else 'non-reentrant'}) Linear {format_name}" + label = ( + f"te.checkpoint({'reentrant' if reentrant else 'non-reentrant'}) Linear {format_name}" + ) _assert_outputs_bitwise_equal(ref, test, label) # ----- save_for_backward round-trip (unit-level) ---------------- From 485849193095221616d45109e60f3e50ac6b5042 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 29 Apr 2026 16:02:06 +0000 Subject: [PATCH 11/97] Respect usage Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 260 ++++++++++++++++++ transformer_engine/pytorch/distributed.py | 19 +- transformer_engine/pytorch/module/base.py | 9 + .../pytorch/module/layernorm_mlp.py | 2 +- .../pytorch/tensor/hybrid_tensor.py | 29 +- 5 files changed, 308 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 2ca341092ef..c4613476ad9 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -539,6 +539,266 @@ def test_make_empty_has_sub_storages(self): assert empty.columnwise_sub_storage is not None +@requires_fp8_and_nvfp4 +class TestHybridUsageFlagsRespected: + """``HybridQuantizer`` must skip directions whose parent usage flag is + False. Native quantizers honor ``rowwise_usage`` / ``columnwise_usage`` + inside the C++ kernel; hybrid sub-quantizers are pinned to one direction + in ``__init__``, so the parent's flags never reach C++ — the equivalent + skip lives in the Python composition layer. Modules call ``set_usage`` + extensively before each ``quantize`` (inference, output / grad_input + quantizers, AG paths), so honoring the flags avoids 2x quantization waste. + """ + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(42) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + # ── quantize_impl ──────────────────────────────────────────── + + def test_quantize_rowwise_only(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is None + + def test_quantize_columnwise_only(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=True) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is not None + + def test_quantize_both_false(self, input_tensor): + """``set_usage(False, False)`` mirrors ``update_usage(False, False)`` — + both produce an empty hybrid. No defensive assert (matches native).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=False) + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is None + + def test_quantize_both_true_default(self, input_tensor): + """Default state (both flags True) keeps both directions populated.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + out = hq.quantize(input_tensor) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is not None + + def test_quantize_internal_storage_rowwise_only(self, input_tensor): + """Internal storage path (used by FSDP2 / make_like flows) also + honors the gate.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + hq.internal = True + try: + out = hq.quantize(input_tensor) + assert isinstance(out, HybridQuantizedTensorStorage) + assert out.rowwise_sub_storage is not None + assert out.columnwise_sub_storage is None + finally: + hq.internal = False + + def test_quantize_flag_change_between_calls(self, input_tensor): + """A single quantizer can be re-used with different flags across + calls (which is exactly how modules use one ``input_quantizer`` / + ``weight_quantizer`` across forward / backward phases).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + + hq.set_usage(rowwise=True, columnwise=False) + out_row = hq.quantize(input_tensor) + assert out_row.rowwise_sub_storage is not None + assert out_row.columnwise_sub_storage is None + + hq.set_usage(rowwise=False, columnwise=True) + out_col = hq.quantize(input_tensor) + assert out_col.rowwise_sub_storage is None + assert out_col.columnwise_sub_storage is not None + + hq.set_usage(rowwise=True, columnwise=True) + out_both = hq.quantize(input_tensor) + assert out_both.rowwise_sub_storage is not None + assert out_both.columnwise_sub_storage is not None + + # ── make_empty ─────────────────────────────────────────────── + + def test_make_empty_rowwise_only(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=True, columnwise=False) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is not None + assert empty.columnwise_sub_storage is None + + def test_make_empty_columnwise_only(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=True) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is None + assert empty.columnwise_sub_storage is not None + + def test_make_empty_both_false(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.set_usage(rowwise=False, columnwise=False) + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + assert empty.rowwise_sub_storage is None + assert empty.columnwise_sub_storage is None + + # ── update_quantized ───────────────────────────────────────── + # + # Comparison strategy: snapshot raw data buffers via ``get_data_tensors()`` + # and compare bytes pre/post-update (same pattern as ``TestHybridClear``). + # Avoids per-format ``dequantize()`` limitations (NVFP4 columnwise raises + # NotImplementedError) and is a strictly stronger check — if the kernel + # writes, raw bytes differ regardless of whether dequant is reversible. + + @staticmethod + def _clone_data_tensors(sub_storage): + """Deep-clone the primary data buffers of a sub-storage.""" + if sub_storage is None: + return () + data = sub_storage.get_data_tensors() + if not isinstance(data, tuple): + data = (data,) + return tuple(t.clone() if t is not None else None for t in data) + + @staticmethod + def _assert_data_tensors_equal(snapshot, sub_storage): + """Assert sub-storage's current data buffers byte-match a prior snapshot.""" + assert sub_storage is not None + current = sub_storage.get_data_tensors() + if not isinstance(current, tuple): + current = (current,) + assert len(snapshot) == len(current), ( + f"Buffer count changed: {len(snapshot)} → {len(current)}" + ) + for before, after in zip(snapshot, current): + if before is None: + assert after is None + continue + assert after is not None + torch.testing.assert_close(before, after, rtol=0, atol=0) + + @staticmethod + def _assert_data_tensors_differ(snapshot, sub_storage): + """Assert at least one buffer changed bytes vs the prior snapshot.""" + assert sub_storage is not None + current = sub_storage.get_data_tensors() + if not isinstance(current, tuple): + current = (current,) + any_changed = False + for before, after in zip(snapshot, current): + if before is None or after is None: + continue + if not torch.equal(before, after): + any_changed = True + break + assert any_changed, "Expected at least one data buffer to change but none did" + + def test_update_quantized_rowwise_only_preserves_columnwise_data(self, input_tensor): + """``update_quantized`` must not refresh a direction whose parent flag + is False, even if the dst storage has that direction allocated. + + Mirrors how native ``tex.quantize(src, quantizer, dst, noop_flag)`` + skips a direction when ``quantizer.rowwise_usage=False`` even if the + dst storage has that direction allocated. + """ + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + # Fully populate both directions + dst = hq.quantize(input_tensor) + # Snapshot the columnwise raw buffers before the targeted rowwise-only update + col_before = self._clone_data_tensors(dst._columnwise_storage) + + # Switch to rowwise-only refresh and feed a substantially different src + hq.set_usage(rowwise=True, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + # Both sub-storage objects survive in-place; columnwise bytes untouched + self._assert_data_tensors_equal(col_before, dst._columnwise_storage) + + def test_update_quantized_columnwise_only_preserves_rowwise_data(self, input_tensor): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + + hq.set_usage(rowwise=False, columnwise=True) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + self._assert_data_tensors_equal(row_before, dst._rowwise_storage) + + def test_update_quantized_both_false_is_noop(self, input_tensor): + """``set_usage(False, False)`` then ``update_quantized`` must leave + both sub-storages' bytes untouched.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + col_before = self._clone_data_tensors(dst._columnwise_storage) + + hq.set_usage(rowwise=False, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + self._assert_data_tensors_equal(row_before, dst._rowwise_storage) + self._assert_data_tensors_equal(col_before, dst._columnwise_storage) + + def test_update_quantized_actually_refreshes_requested(self, input_tensor): + """Sanity check: when the parent flag is True, the corresponding + sub-storage IS refreshed (otherwise the previous tests would pass + vacuously by not refreshing anything).""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + dst = hq.quantize(input_tensor) + row_before = self._clone_data_tensors(dst._rowwise_storage) + + hq.set_usage(rowwise=True, columnwise=False) + new_src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") * 100 + hq.update_quantized(new_src, dst) + + # Rowwise bytes must differ — confirms update_quantized actually ran + self._assert_data_tensors_differ(row_before, dst._rowwise_storage) + + # ── te.Linear integration: inference path takes rowwise-only ─ + + def test_te_linear_inference_workspace_rowwise_only(self): + """``te.Linear`` forward under ``torch.no_grad()`` with a hybrid + ``CustomRecipe`` must produce a rowwise-only weight workspace. + ``linear.py:266-274`` sets ``weight_quantizer.set_usage(columnwise=False)`` + in inference; without the parent-flag gate, hybrid would still allocate + both directions. + """ + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + ) + torch.manual_seed(2026) + model = Linear(128, 256, bias=False, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 128, dtype=torch.bfloat16, device="cuda") + + # is_first_microbatch=True forces the cache_name="weight" path + # (see linear.py:1631) so the hybrid workspace persists in + # model._fp8_workspaces and we can inspect its sub-storages. + with torch.no_grad(): + with autocast(enabled=True, recipe=hybrid_recipe): + _ = model(x, is_first_microbatch=True) + + ws = model._fp8_workspaces.get("weight") + assert isinstance( + ws, HybridQuantizedTensorStorage + ), f"Expected hybrid weight workspace, got {type(ws).__name__}" + assert ws.rowwise_sub_storage is not None, "Rowwise sub-storage must be populated for fprop" + assert ws.columnwise_sub_storage is None, ( + "Inference forward must produce rowwise-only hybrid weight workspace; " + "columnwise quantization should have been skipped per " + "weight_quantizer.set_usage(rowwise=True, columnwise=False)." + ) + + @requires_fp8_and_nvfp4 class TestHybridTorchDispatch: """Test torch dispatch operations.""" diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index a0d4ac3530c..dc51e4635d8 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -43,6 +43,7 @@ from .tensor.mxfp8_tensor import MXFP8Quantizer from .tensor.nvfp4_tensor import NVFP4Quantizer from .tensor.float8_blockwise_tensor import Float8BlockQuantizer +from .tensor.hybrid_tensor import HybridQuantizer from .quantized_tensor import QuantizedTensorStorage, QuantizedTensor, Quantizer from .tensor.storage.float8_tensor_storage import Float8TensorStorage from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -1755,7 +1756,23 @@ def gather_along_first_dim( memory_format=torch.contiguous_format, ) torch.distributed.all_gather_into_tensor(out, inp, group=process_group) - out = quantizer(out) + # Hybrid override: callers drop columnwise before AG, expecting to + # synthesize it post-AG via ``update_usage(columnwise_usage=True)`` + # (native FP8's ``_create_transpose``). Hybrid has no synthesis path — + # that update_usage is a no-op — so re-quantize with both directions, + # mirroring what the planned native hybrid AG dispatch would produce + # (see the TODO in + # :meth:`HybridQuantizer.supports_only_rowwise_all_gather`); once + # native AG lands, hybrid won't reach this fallback. + if isinstance(quantizer, HybridQuantizer): + prev_row, prev_col = quantizer.rowwise_usage, quantizer.columnwise_usage + quantizer.set_usage(rowwise=True, columnwise=True) + try: + out = quantizer(out) + finally: + quantizer.set_usage(rowwise=prev_row, columnwise=prev_col) + else: + out = quantizer(out) return out, None # Dequantize quantized tensor if not supported diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 6236115b434..cf3fd5e52ec 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -46,6 +46,7 @@ from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage +from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage from ..utils import ( is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype, @@ -658,6 +659,14 @@ def _is_weight_workspace_valid( return False if quantizer.columnwise_usage and workspace._columnwise_data is None: return False + elif isinstance(workspace, HybridQuantizedTensorStorage): + # Workspace cached under one flag setting (e.g. inference with + # ``columnwise=False``) becomes stale when the next call needs the + # missing direction; invalidate so a fresh workspace is built. + if quantizer.rowwise_usage and workspace._rowwise_storage is None: + return False + if quantizer.columnwise_usage and workspace._columnwise_storage is None: + return False if isinstance(workspace, DebugQuantizedTensor) != isinstance(quantizer, DebugQuantizer): return False return True diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 23388baa758..0d9cef08917 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2362,7 +2362,7 @@ def _get_quantizers(self, fp8_output, is_grad_enabled): rowwise=True, columnwise=isinstance( fc2_input_quantizer, - (MXFP8Quantizer, Float8BlockQuantizer, NVFP4Quantizer), + (MXFP8Quantizer, Float8BlockQuantizer, NVFP4Quantizer, HybridQuantizer), ), ) fc2_input_quantizer.internal = True diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 5feeefd0e3c..ba05fd2ed1c 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -49,8 +49,15 @@ def __init__( self.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: - rowwise_result = self.rowwise_quantizer.quantize(tensor) - columnwise_result = self.columnwise_quantizer.quantize(tensor) + # Gate each sub-quantizer call on the parent usage flag. Sub-quantizers + # are pinned to one direction in ``__init__``; the parent flag decides + # whether to invoke them. + rowwise_result = ( + self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None + ) + columnwise_result = ( + self.columnwise_quantizer.quantize(tensor) if self.columnwise_usage else None + ) if self.internal: return HybridQuantizedTensorStorage( @@ -101,8 +108,12 @@ def _make_empty_internal(sub_quantizer): finally: sub_quantizer.internal = prev_internal - rowwise_empty = _make_empty_internal(self.rowwise_quantizer) - columnwise_empty = _make_empty_internal(self.columnwise_quantizer) + rowwise_empty = ( + _make_empty_internal(self.rowwise_quantizer) if self.rowwise_usage else None + ) + columnwise_empty = ( + _make_empty_internal(self.columnwise_quantizer) if self.columnwise_usage else None + ) return HybridQuantizedTensor( shape=shape, @@ -123,19 +134,19 @@ def update_quantized( *, noop_flag: Optional[torch.Tensor] = None, ) -> QuantizedTensorStorage: - """Re-quantize both sub-storages of a hybrid tensor in-place. + """Re-quantize sub-storages of a hybrid tensor in-place. - Delegates to each sub-quantizer's update_quantized, which writes - new quantized data + scales into the existing sub-storage buffers. + Each direction is refreshed only when the parent usage flag is set + **and** the corresponding sub-storage exists. """ if not isinstance(dst, HybridQuantizedTensorStorage): raise ValueError( "HybridQuantizer can only update HybridQuantizedTensorStorage, got" f" {type(dst).__name__}" ) - if dst._rowwise_storage is not None: + if self.rowwise_usage and dst._rowwise_storage is not None: self.rowwise_quantizer.update_quantized(src, dst._rowwise_storage, noop_flag=noop_flag) - if dst._columnwise_storage is not None: + if self.columnwise_usage and dst._columnwise_storage is not None: self.columnwise_quantizer.update_quantized( src, dst._columnwise_storage, noop_flag=noop_flag ) From ef31a9a622183877ad1d82aae6145ab0f87ee116 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 29 Apr 2026 16:04:30 +0000 Subject: [PATCH 12/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 14 +++++--------- transformer_engine/pytorch/tensor/hybrid_tensor.py | 8 ++------ 2 files changed, 7 insertions(+), 15 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index c4613476ad9..1e264ec4001 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -670,9 +670,9 @@ def _assert_data_tensors_equal(snapshot, sub_storage): current = sub_storage.get_data_tensors() if not isinstance(current, tuple): current = (current,) - assert len(snapshot) == len(current), ( - f"Buffer count changed: {len(snapshot)} → {len(current)}" - ) + assert len(snapshot) == len( + current + ), f"Buffer count changed: {len(snapshot)} → {len(current)}" for before, after in zip(snapshot, current): if before is None: assert after is None @@ -769,12 +769,8 @@ def test_te_linear_inference_workspace_rowwise_only(self): both directions. """ hybrid_recipe = _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), ) torch.manual_seed(2026) model = Linear(128, 256, bias=False, params_dtype=torch.bfloat16).cuda() diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index ba05fd2ed1c..0df86f93f5e 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -52,9 +52,7 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: # Gate each sub-quantizer call on the parent usage flag. Sub-quantizers # are pinned to one direction in ``__init__``; the parent flag decides # whether to invoke them. - rowwise_result = ( - self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None - ) + rowwise_result = self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None columnwise_result = ( self.columnwise_quantizer.quantize(tensor) if self.columnwise_usage else None ) @@ -108,9 +106,7 @@ def _make_empty_internal(sub_quantizer): finally: sub_quantizer.internal = prev_internal - rowwise_empty = ( - _make_empty_internal(self.rowwise_quantizer) if self.rowwise_usage else None - ) + rowwise_empty = _make_empty_internal(self.rowwise_quantizer) if self.rowwise_usage else None columnwise_empty = ( _make_empty_internal(self.columnwise_quantizer) if self.columnwise_usage else None ) From c7da5b234f382128dd465357db2b7f343f88bee8 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 20 May 2026 10:42:25 +0000 Subject: [PATCH 13/97] Misc minor fixes: comments, tests, etc. Signed-off-by: Evgeny --- qa/L0_pytorch_unittest/test.sh | 1 + qa/L1_pytorch_distributed_unittest/test.sh | 1 + .../distributed/fsdp2_tests/fsdp2_utils.py | 138 +++--- .../fsdp2_tests/run_fsdp2_fused_adam.py | 34 +- tests/pytorch/distributed/run_hybrid_tp_sp.py | 15 +- tests/pytorch/test_cpu_offloading.py | 18 +- tests/pytorch/test_cpu_offloading_v1.py | 12 +- tests/pytorch/test_hybrid_quantization.py | 393 +++++++++++++++--- .../pytorch/module/grouped_linear.py | 19 +- .../pytorch/tensor/hybrid_tensor.py | 110 ++++- .../tensor/storage/float8_tensor_storage.py | 38 ++ 11 files changed, 602 insertions(+), 177 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index c35dc4c063d..28518de9e42 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -48,6 +48,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index db13e9f1e0c..96e0803c749 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -27,6 +27,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PAT python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py || test_fail "test_torch_fsdp2.py" +python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_hybrid_tp_sp.xml $TE_PATH/tests/pytorch/distributed/test_hybrid_tp_sp.py || test_fail "test_hybrid_tp_sp.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_userbuffers.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py || test_fail "test_fusible_ops_with_userbuffers.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index e3702e91044..f8423f14ccf 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -5,75 +5,97 @@ """Shared utility functions for FSDP2 distributed tests.""" import transformer_engine.common.recipe -from transformer_engine.pytorch import QuantizedTensor +from transformer_engine.pytorch import HybridQuantizer, QuantizedTensor +from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + current_scaling_quantizer_factory, + float8_block_scaling_quantizer_factory, + mxfp8_quantizer_factory, +) def get_recipe_from_string(recipe): return getattr(transformer_engine.common.recipe, recipe)() +# ── Hybrid qfactories ───────────────────────────────────────────────── +# +# Module-level (picklable) qfactories used by ``get_hybrid_recipe_from_string``. +# Each factory composes one or two role-aware base factories from +# ``quantization_recipes_base`` per direction. Per-role behavior is delegated +# to the base factory — the hybrid layer only decides direction pairing. +# +# DCP serializes ``CustomRecipe`` via ``pickle``; closure-based qfactories +# (lambdas, inner functions referencing captured state) are not picklable, +# so the qfactory must live at module scope. See +# ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. + + +def _hybrid_fp8_current_qfactory(role): + """FP8 current-scaling rowwise + FP8 current-scaling columnwise.""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=current_scaling_quantizer_factory(role), + columnwise_quantizer=current_scaling_quantizer_factory(role), + ) + return current_scaling_quantizer_factory(role) + + +def _hybrid_mxfp8_qfactory(role): + """MXFP8 rowwise + MXFP8 columnwise.""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=mxfp8_quantizer_factory(role), + ) + return mxfp8_quantizer_factory(role) + + +def _hybrid_float8_block_qfactory(role): + """Float8 block-scaling rowwise + Float8 block-scaling columnwise.""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=float8_block_scaling_quantizer_factory(role), + columnwise_quantizer=float8_block_scaling_quantizer_factory(role), + ) + return float8_block_scaling_quantizer_factory(role) + + +def _hybrid_mixed_mxfp8_fp8_qfactory(role): + """MXFP8 rowwise + FP8 current columnwise (cross-format hybrid).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=current_scaling_quantizer_factory(role), + ) + return current_scaling_quantizer_factory(role) + + +_HYBRID_QFACTORIES = { + "HybridFP8CurrentScaling": _hybrid_fp8_current_qfactory, + "HybridMXFP8": _hybrid_mxfp8_qfactory, + "HybridFloat8BlockScaling": _hybrid_float8_block_qfactory, + "HybridMixed_MXFP8_FP8": _hybrid_mixed_mxfp8_fp8_qfactory, +} + + def get_hybrid_recipe_from_string(recipe): - """Build a CustomRecipe that uses HybridQuantizer with the given base format. + """Build a CustomRecipe wrapping a module-level (picklable) hybrid qfactory. Supported values: "HybridFP8CurrentScaling" — FP8 current for both directions - "HybridMXFP8" — MXFP8 for both directions - "HybridMixed_MXFP8_FP8" — MXFP8 rowwise + FP8 current columnwise + "HybridMXFP8" — MXFP8 for both directions + "HybridFloat8BlockScaling" — Float8 block scaling for both directions + "HybridMixed_MXFP8_FP8" — MXFP8 rowwise + FP8 current columnwise """ - import transformer_engine_torch as tex - from transformer_engine.pytorch import ( - Float8CurrentScalingQuantizer, - Float8BlockQuantizer, - MXFP8Quantizer, - HybridQuantizer, - ) - - _BUILDERS = { - "HybridFP8CurrentScaling": lambda: dict( - row=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), - col=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), - grad=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), - ), - "HybridMXFP8": lambda: dict( - row=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - col=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - grad=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), - ), - "HybridFloat8BlockScaling": lambda: dict( - row=lambda: Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True - ), - col=lambda: Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True - ), - grad=lambda: Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True - ), - ), - "HybridMixed_MXFP8_FP8": lambda: dict( - row=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - col=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), - grad=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), - ), - } - - if recipe not in _BUILDERS: - raise ValueError(f"Unknown hybrid recipe '{recipe}'. Supported: {sorted(_BUILDERS.keys())}") - - builders = _BUILDERS[recipe]() - row_fn, col_fn, grad_fn = builders["row"], builders["col"], builders["grad"] - - def qfactory(role): - if role in ("linear_input", "linear_weight", "linear_output"): - return HybridQuantizer( - rowwise_quantizer=row_fn(), - columnwise_quantizer=col_fn(), - ) - if role in ("linear_grad_output", "linear_grad_input"): - return grad_fn() - return row_fn() - - return transformer_engine.common.recipe.CustomRecipe(qfactory=qfactory) + if recipe not in _HYBRID_QFACTORIES: + raise ValueError( + f"Unknown hybrid recipe '{recipe}'. Supported: {sorted(_HYBRID_QFACTORIES.keys())}" + ) + return transformer_engine.common.recipe.CustomRecipe(qfactory=_HYBRID_QFACTORIES[recipe]) def save_custom_attrs(module): diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 8a6610f1ec4..68979ad20e4 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1459,23 +1459,41 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): """ import torch.distributed.checkpoint as dcp - pytest.xfail( - "CustomRecipe with closure-based qfactory cannot be pickled by DCP. " - "Requires module-level picklable factory functions for DCP compatibility." - ) - if hybrid_recipe_name == "HybridFloat8BlockScaling": pytest.xfail( "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " "quantized type through FSDP2 view(-1)." ) + if hybrid_recipe_name == "HybridFP8CurrentScaling": + pytest.xfail( + "HybridFP8CurrentScaling: per-tensor _scale_inv is not preserved " + "through DCP's tensor-storage-byte serialization path. " + "HybridQuantizedTensor.__reduce_ex__ correctly round-trips through " + "pickle (verified by torch.save/torch.load), but DCP bypasses " + "pickle and serializes the tensor's storage bytes — the scalar " + "_scale_inv is not enumerated as a separate tensor leaf and gets " + "lost. Vanilla Float8CurrentScaling avoids this because per-tensor " + "scale lives in module.fp8_meta (saved as extra_state), not on " + "the tensor; hybrid uses per-sub-storage scales without that " + "mirror. Fix path: implement __tensor_flatten__/__tensor_unflatten__ " + "across the quantized tensor stack so DCP can serialize the " + "per-leaf tensor fields directly. Loaded model output diverges by " + "~5e-2." + ) + from fsdp2_utils import get_hybrid_recipe_from_string hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) world_size, device = _get_dist_info() - rank = int(os.environ["LOCAL_RANK"]) - checkpoint_dir = os.path.join("/tmp", f"hybrid_dcp_test_{os.getpid()}") + rank = int(os.environ.get("RANK", "0")) + # Deterministic, rank-agnostic checkpoint dir so all ranks read/write + # the same DCP path. ``os.getpid()`` differs per rank under torchrun. + checkpoint_dir = f"/tmp/te_test_fsdp2_hybrid_dcp_parity_{hybrid_recipe_name}" + + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + dist.barrier() try: model = _build_hybrid_model(hybrid_recipe) @@ -1543,8 +1561,6 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): finally: dist.barrier() if rank == 0: - import shutil - shutil.rmtree(checkpoint_dir, ignore_errors=True) diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index 679f443c99a..e6e4d45e1c2 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -99,23 +99,25 @@ def _make_mxfp8_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): def _hybrid_fp8_qfactory(role): """FP8 current scaling in both directions for fwd roles; E5M2 for grad roles (standard Hybrid:HYBRID format pairing).""" - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( rowwise_quantizer=_make_fp8_current_quantizer(), columnwise_quantizer=_make_fp8_current_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return _make_fp8_current_quantizer(fp8_dtype=tex.DType.kFloat8E5M2) return _make_fp8_current_quantizer() def _hybrid_mxfp8_qfactory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( rowwise_quantizer=_make_mxfp8_quantizer(), columnwise_quantizer=_make_mxfp8_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return _make_mxfp8_quantizer(fp8_dtype=tex.DType.kFloat8E5M2) return _make_mxfp8_quantizer() @@ -131,12 +133,13 @@ def _make_nvfp4_quantizer(): def _hybrid_nvfp4_qfactory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( rowwise_quantizer=_make_nvfp4_quantizer(), columnwise_quantizer=_make_nvfp4_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return _make_nvfp4_quantizer() return _make_nvfp4_quantizer() diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index a23e32646b4..e54a80b6029 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -34,17 +34,18 @@ def _hybrid_fp8_mxfp8_qfactory(role): Forward roles -> HybridQuantizer; backward roles -> plain MXFP8 so dgrad/wgrad operand pairs share a single scaling mode. Catch-all - returns plain FP8 for non-``linear_*`` roles used by layernorm_linear, + returns plain FP8 for non-linear roles used by layernorm_linear, layernorm_mlp, multihead_attention, and transformer_layer. """ - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( rowwise_quantizer=te.Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda" ), columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2) return te.Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") @@ -56,12 +57,13 @@ def _hybrid_mxfp8_nvfp4_qfactory(role): ``custom_recipes/quantization_nvfp4.py``. grad_output uses plain NVFP4 (both directions) so wgrad's columnwise operand matches. """ - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), columnwise_quantizer=te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) @@ -221,10 +223,12 @@ def create_tensor(recipe: Optional[recipe.Recipe], requires_grad: bool = False) quantizer = te.tensor.nvfp4_tensor.NVFP4Quantizer() return quantizer(tensor) elif recipe.custom(): - # CustomRecipe: invoke the qfactory for the ``linear_weight`` role + # CustomRecipe: invoke the qfactory for the linear weight role # as a representative quantizer (returns a HybridQuantizer for the # hybrid factories registered at module scope). - quantizer = recipe.qfactory("linear_weight") + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = recipe.qfactory(QuantizerRole(module_type="linear", tensor_type="weight")) if quantizer is None: # Fallback: factory did not supply a weight quantizer. return tensor.requires_grad_() if requires_grad else tensor diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index c98f6098d9c..a5e5df8d200 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -28,17 +28,18 @@ def _hybrid_fp8_mxfp8_qfactory(role): Forward roles get a HybridQuantizer; backward/grad roles get a plain MXFP8 quantizer so dgrad/wgrad GEMMs see a single scaling mode per - operand pair. Catch-all returns plain FP8 for non-``linear_*`` roles + operand pair. Catch-all returns plain FP8 for non-linear roles (layernorm_linear, layernorm_mlp, multihead_attention, transformer_layer). """ - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( rowwise_quantizer=te.Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda" ), columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2) return te.Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") @@ -50,12 +51,13 @@ def _hybrid_mxfp8_nvfp4_qfactory(role): from ``custom_recipes/quantization_nvfp4.py``. grad_output uses plain NVFP4 (both directions) so wgrad's columnwise operand matches. """ - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), columnwise_quantizer=te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 1e264ec4001..86b7d4b9597 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -881,7 +881,11 @@ def test_linear_fwd_bwd_matches_vanilla_fp8(self): loss_ref.backward() def hybrid_fp8_factory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight", "output") + ): return HybridQuantizer( rowwise_quantizer=Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, @@ -892,7 +896,11 @@ def hybrid_fp8_factory(role): device="cuda", ), ) - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda", @@ -957,7 +965,11 @@ def test_linear_fwd_bwd_matches_vanilla_mxfp8(self): out_ref.float().sum().backward() def hybrid_mxfp8_factory(role): - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) return HybridQuantizer( rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), @@ -1010,8 +1022,10 @@ def test_linear_fwd_bwd_matches_vanilla_block_fp8(self): out_ref.float().sum().backward() def hybrid_block_fp8_factory(role): - dim = 2 if role == "linear_weight" else 1 - if role in ("linear_grad_output", "linear_grad_input"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + is_weight = is_linear and role.tensor_type == "weight" + dim = 2 if is_weight else 1 + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, @@ -1091,7 +1105,11 @@ def test_linear_fwd_bwd_matches_vanilla_nvfp4(self): out_ref.float().sum().backward() def hybrid_nvfp4_factory(role): - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) return HybridQuantizer( rowwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), @@ -1193,14 +1211,22 @@ def test_linear_fwd_bwd_fp8_row_nvfp4_col(self): ) def mixed_factory(role): - if role in ("linear_input", "linear_weight"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_nvfp4_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return _make_nvfp4_quantizer() - return None + return _make_fp8_quantizer() mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) @@ -1241,12 +1267,16 @@ def test_numerical_sanity_against_bf16(self): out_bf16 = model(inp) def mixed_factory(role): - if role in ("linear_input", "linear_weight"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_nvfp4_quantizer(), ) - return None + return _make_fp8_quantizer() mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) with torch.no_grad(): @@ -1333,7 +1363,11 @@ class TestHybridBiasGradient: def _make_uniform_hybrid_factory(self): def factory(role): - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda", @@ -1423,14 +1457,22 @@ def test_matching_columnwise_formats_succeed(self): inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) def factory(role): - if role in ("linear_input", "linear_weight"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_nvfp4_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return _make_nvfp4_quantizer() - return None + return _make_fp8_quantizer() with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): out = model(inp) @@ -1444,17 +1486,25 @@ def test_mismatched_columnwise_formats_raise(self): inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) def factory(role): - if role in ("linear_input", "linear_weight"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_nvfp4_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda", ) - return None + return _make_fp8_quantizer() with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): out = model(inp) @@ -1480,12 +1530,16 @@ def test_nvfp4_row_fp8_col_forward_only(self): inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) def factory(role): - if role in ("linear_input", "linear_weight"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): return HybridQuantizer( rowwise_quantizer=_make_nvfp4_quantizer(), columnwise_quantizer=_make_fp8_quantizer(), ) - return None + return _make_nvfp4_quantizer() mixed_recipe = recipe.CustomRecipe(qfactory=factory) with torch.no_grad(): @@ -1507,14 +1561,22 @@ def test_nvfp4_row_fp8_col_full_fwd_bwd(self): ) def factory(role): - if role in ("linear_input", "linear_weight"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): return HybridQuantizer( rowwise_quantizer=_make_nvfp4_quantizer(), columnwise_quantizer=_make_fp8_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return _make_fp8_quantizer() - return None + return _make_nvfp4_quantizer() mixed_recipe = recipe.CustomRecipe(qfactory=factory) with autocast(enabled=True, recipe=mixed_recipe): @@ -1558,22 +1620,23 @@ def test_hybrid_input_plain_weight_fwd_bwd(self): ) def factory(role): - if role == "linear_input": + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "input": return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_fp8_quantizer(), ) - if role == "linear_weight": + if is_linear and role.tensor_type == "weight": return Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda", ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda", ) - return None + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") mixed_recipe = recipe.CustomRecipe(qfactory=factory) with autocast(enabled=True, recipe=mixed_recipe): @@ -1600,22 +1663,23 @@ def test_plain_input_hybrid_weight_fwd_bwd(self): ) def factory(role): - if role == "linear_input": + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "input": return Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda", ) - if role == "linear_weight": + if is_linear and role.tensor_type == "weight": return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_fp8_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda", ) - return None + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") mixed_recipe = recipe.CustomRecipe(qfactory=factory) with autocast(enabled=True, recipe=mixed_recipe): @@ -1751,14 +1815,22 @@ def test_fwd_bwd(self, row_name, col_name): make_col_grad = col_cfg[1] if col_cfg[1] is not None else col_cfg[0] def factory(role): - if role in ("linear_input", "linear_weight"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("input", "weight") + ): return HybridQuantizer( rowwise_quantizer=make_row_e4m3(), columnwise_quantizer=make_col_e4m3(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type in ("grad_output", "grad_input") + ): return make_col_grad() - return None + return make_row_e4m3() mixed_recipe = recipe.CustomRecipe(qfactory=factory) with autocast(enabled=True, recipe=mixed_recipe): @@ -1974,22 +2046,23 @@ def test_fp8_fprop_mxfp8_dgrad_nvfp4_wgrad(self): ) def factory(role): - if role == "linear_weight": + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "weight": return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_mxfp8_quantizer(), ) - if role == "linear_input": + if is_linear and role.tensor_type == "input": return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_nvfp4_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return HybridQuantizer( rowwise_quantizer=_make_mxfp8_quantizer(), columnwise_quantizer=_make_nvfp4_quantizer(), ) - return None + return _make_fp8_quantizer() with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): out = model(inp) @@ -2017,22 +2090,23 @@ def test_nvfp4_fprop_fp8_dgrad_mxfp8_wgrad(self): ) def factory(role): - if role == "linear_weight": + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "weight": return HybridQuantizer( rowwise_quantizer=_make_nvfp4_quantizer(), columnwise_quantizer=_make_fp8_quantizer(), ) - if role == "linear_input": + if is_linear and role.tensor_type == "input": return HybridQuantizer( rowwise_quantizer=_make_nvfp4_quantizer(), columnwise_quantizer=_make_mxfp8_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_mxfp8_quantizer(), ) - return None + return _make_nvfp4_quantizer() with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): out = model(inp) @@ -2060,19 +2134,20 @@ def test_same_dgrad_wgrad_reduces_to_plain_grad(self): ) def factory(role): - if role == "linear_weight": + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "weight": return HybridQuantizer( rowwise_quantizer=_make_nvfp4_quantizer(), columnwise_quantizer=_make_mxfp8_quantizer(), ) - if role == "linear_input": + if is_linear and role.tensor_type == "input": return HybridQuantizer( rowwise_quantizer=_make_nvfp4_quantizer(), columnwise_quantizer=_make_mxfp8_quantizer(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return _make_mxfp8_quantizer() - return None + return _make_nvfp4_quantizer() with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): out = model(inp) @@ -2096,7 +2171,8 @@ def _make_hybrid_fp8_factory(): plain FP8 E5M2 for bwd roles.""" def factory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( rowwise_quantizer=Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, @@ -2107,7 +2183,7 @@ def factory(role): device="cuda", ), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda", @@ -2319,10 +2395,20 @@ def test_mixed_hybrid_and_plain_raises(self): assert "1 non-hybrid" in msg assert "CustomRecipe" in msg and "qfactory" in msg - def test_none_entries_ignored_when_remainder_is_uniform(self): - """None entries are filtered before uniformity check — a list - of hybrids plus a None must still classify as hybrid (not - mixed).""" + def test_none_plus_hybrid_raises(self): + """None entries mixed with HybridQuantizer must NOT classify as + all-hybrid: ``_hybrid_split_quantize`` would later iterate the + full list with ``isinstance(q, HybridQuantizer)`` and raise + ``TypeError`` on the ``None`` entry. The classifier rejects + upfront with a clear ValueError so users see a single, + actionable error. + + In current TE flows ``CustomRecipeState.make_quantizers`` + rejects ``None`` returns from ``qfactory``, so this combination + shouldn't actually arise — but if a future "intentional no-op" + ``IdentityQuantizer`` ever loosens that contract, this guard + prevents the silent crash. + """ from transformer_engine.pytorch.module.grouped_linear import ( _is_hybrid_quantizer_list, ) @@ -2332,7 +2418,12 @@ def test_none_entries_ignored_when_remainder_is_uniform(self): None, _make_hybrid_quantizer_fp8_row_fp4_col(), ] - assert _is_hybrid_quantizer_list(quantizers) is True + with pytest.raises(ValueError) as exc_info: + _is_hybrid_quantizer_list(quantizers) + msg = str(exc_info.value) + assert "mixes HybridQuantizer" in msg + assert "2 hybrid" in msg + assert "1 None" in msg def test_hybrid_split_quantize_rejects_plain_element(self): """Defense-in-depth: even if a caller bypasses the classifier, @@ -2379,12 +2470,13 @@ def _hybrid_custom_recipe(row_factory, col_factory, grad_factory=None): grad_factory = col_factory def qfactory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( rowwise_quantizer=row_factory(), columnwise_quantizer=col_factory(), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return grad_factory() return row_factory() @@ -3051,21 +3143,23 @@ def test_mixed_format_sub_storage_types(self): def _hybrid_fp8_current_qfactory(role): """Hybrid FP8 current scaling (E4M3 both dirs, E5M2 for grad).""" - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), columnwise_quantizer=Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda" ), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") def _hybrid_mxfp8_qfactory(role): """Hybrid MXFP8 (E4M3 both dirs).""" - if role in ("linear_grad_output", "linear_grad_input"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) return HybridQuantizer( rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), @@ -3075,8 +3169,10 @@ def _hybrid_mxfp8_qfactory(role): def _hybrid_block_fp8_qfactory(role): """Hybrid block FP8 (E4M3 both dirs).""" - dim = 2 if role == "linear_weight" else 1 - if role in ("linear_grad_output", "linear_grad_input"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + is_weight = is_linear and role.tensor_type == "weight" + dim = 2 if is_weight else 1 + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, @@ -3101,7 +3197,8 @@ def _hybrid_block_fp8_qfactory(role): def _hybrid_nvfp4_qfactory(role): """Hybrid NVFP4 (E2M1 both dirs).""" - if role in ("linear_grad_output", "linear_grad_input"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) return HybridQuantizer( rowwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), @@ -3320,14 +3417,15 @@ def test_equivalence(self): def _checkpoint_hybrid_fp8_qfactory(role): """Module-level qfactory (picklable) for checkpoint tests.""" - if role in ("linear_input", "linear_weight", "linear_output"): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), columnwise_quantizer=Float8CurrentScalingQuantizer( tex.DType.kFloat8E4M3, device="cuda" ), ) - if role in ("linear_grad_output", "linear_grad_input"): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") @@ -4154,6 +4252,175 @@ def test_make_like_is_independent(self): assert copy is not param +# --------------------------------------------------------------------------- +# 15b. Hopper-only paths: columnwise-only Float8 sub-storage +# --------------------------------------------------------------------------- +# +# On architectures where ``is_non_tn_fp8_gemm_supported()`` returns False +# (Hopper sm_90, L40 sm_89), per-tensor FP8 GEMM only supports the TN +# layout — non-TN layouts are simulated by feeding pre-transposed data. +# So a columnwise-only ``Float8TensorStorage`` (used as a hybrid sub- +# storage) holds its quantized data in ``_transpose`` instead of +# ``_data``, with ``_data = None``. +# +# This is the exact layout the FSDP2 buffer protocol must recognize +# when the sub-storage is part of a ``HybridQuantizedTensor`` parameter. +# These tests pin the contracts that would break if the buffer +# protocol regressed to the unconditional ``("_data",)`` field name +# (which would all-gather a ``None`` tensor on Hopper). +# +# Skip on Blackwell where the C++ kernel always populates ``_data`` and +# the columnwise-only Float8 path doesn't exercise ``_transpose``. + +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported # noqa: E402 + +requires_hopper_fp8 = pytest.mark.skipif( + is_non_tn_fp8_gemm_supported() or not fp8_available, + reason=( + "Hopper-only: requires per-tensor FP8 with non-TN GEMM unsupported " + "(Hopper sm_90 / L40 sm_89). On Blackwell the C++ kernel populates " + "_data even for columnwise-only mode, so the _transpose-only path " + "is not exercised." + ), +) + + +@requires_hopper_fp8 +class TestHybridFloat8ColumnwiseOnlyHopperPath: + """Float8TensorStorage columnwise-only sub-storage exercises the + ``_transpose`` field on Hopper. The FSDP2 buffer protocol must + recognize this layout. + """ + + def _make_columnwise_only_float8_storage(self): + """Build a Float8TensorStorage in the layout a columnwise-only + hybrid sub-storage would have on Hopper: ``_data=None`` and + the actual quantized bytes in ``_transpose``. + """ + q = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + src = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda") + out = q(src) + # Columnwise-only Float8 on Hopper: _data is None, _transpose holds data + assert out._data is None, ( + f"Test precondition failed: expected _data is None on Hopper, got {out._data}" + ) + assert out._transpose is not None, "Test precondition failed: _transpose is None" + return out + + def test_fsdp_buffer_fields_returns_transpose(self): + """``fsdp_buffer_fields`` must return ``("_transpose",)`` when + ``_data`` is ``None`` and ``_transpose`` is populated. The + unconditional ``("_data",)`` would have FSDP2 all-gather a + ``None`` tensor on Hopper hybrid + FSDP2. + """ + storage = self._make_columnwise_only_float8_storage() + assert storage.fsdp_buffer_fields() == ("_transpose",) + + def test_fsdp_extract_buffers_returns_transpose_data(self): + """``fsdp_extract_buffers`` (default impl, reads named fields) + must return the actual ``_transpose`` tensor, not ``None``. + """ + storage = self._make_columnwise_only_float8_storage() + buffers, meta = storage.fsdp_extract_buffers() + assert len(buffers) == 1 + assert buffers[0] is not None + assert buffers[0] is storage._transpose + assert meta["field_names"] == ("_transpose",) + + def test_fsdp_assign_gathered_resets_transpose_invalid(self): + """After the gathered transpose buffer is written back via + ``fsdp_assign_gathered``, ``_transpose_invalid`` must be False + — otherwise ``update_usage`` / ``get_usages`` would treat the + freshly gathered transpose as stale on first use. + """ + storage = self._make_columnwise_only_float8_storage() + # Simulate stale state pre-gather + storage._transpose_invalid = True + new_transpose = torch.zeros_like(storage._transpose) + storage.fsdp_assign_gathered((new_transpose,), {"field_names": ("_transpose",)}) + assert storage._transpose is new_transpose + assert storage._transpose_invalid is False + # And ``get_usages`` correctly reports columnwise-available + assert storage.get_usages()["columnwise"] is True + + def test_fsdp_buffer_fields_falls_back_to_data_when_both_present(self): + """A normally-constructed Float8TensorStorage has ``_data`` + populated; ``fsdp_buffer_fields`` should still prefer ``_data`` + — direction-aware logic must not regress the vanilla path. + """ + q = Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + src = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda") + out = q(src) + assert out._data is not None + assert out.fsdp_buffer_fields() == ("_data",) + + +@requires_hopper_fp8 +class TestHybridFsdpPostAllGatherUpdateUsage: + """``HybridQuantizedTensor.fsdp_post_all_gather`` must call + ``update_usage`` on each sub-storage after writing gathered data + (mirroring vanilla ``Float8Tensor.fsdp_post_all_gather:888``). + Without it, on Hopper a previously-cached ``_transpose`` from the + prior iteration is silently reused with the new ``_data``, producing + incorrect dgrad / wgrad GEMMs. + """ + + def _make_param(self): + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(64, 64, params_dtype=torch.bfloat16).cuda() + return model.weight + + def test_iter2_invalidates_stale_transpose_on_rowwise_substorage(self): + """Simulates iter-2+ buffer reuse: pre-existing ``out`` with a + possibly-stale ``_transpose`` cache; after ``fsdp_post_all_gather`` + the rowwise sub-storage's ``_transpose`` must be invalidated / + regenerated to match the freshly gathered ``_data``. + """ + param = self._make_param() + # Build a plausible iter-2+ "out" with stale state. + out = HybridQuantizedTensor.make_like(param) + # Rowwise sub-storage on Hopper has _data populated. Force a stale + # _transpose and invalidate flag to mimic the regression scenario. + if out._rowwise_storage._transpose is None: + # Set up a fake stale transpose (non-None, marked invalid by + # the mismatching shape would catch nothing, so just plant + # a tensor and clear the invalid flag to "valid"). + out._rowwise_storage._transpose = torch.zeros_like(out._rowwise_storage._data).t() + out._rowwise_storage._transpose_invalid = False + stale_transpose_id = id(out._rowwise_storage._transpose) + + # Drive a real all-gather round trip via the protocol + sharded_tensors, metadata = param.fsdp_pre_all_gather( + mesh=None, + orig_size=param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + out2, _ = param.fsdp_post_all_gather( + sharded_tensors, metadata, param.dtype, out=out + ) + + # After fsdp_post_all_gather, the rowwise sub-quantizer is pinned + # columnwise=False, so update_usage(rowwise=True, columnwise=False) + # must clear the stale _transpose (preventing the silent + # stale-cache regression on Hopper). + assert out2._rowwise_storage._transpose is None or ( + out2._rowwise_storage._transpose_invalid + and id(out2._rowwise_storage._transpose) != stale_transpose_id + ), "Stale _transpose was not invalidated after fsdp_post_all_gather" + + # --------------------------------------------------------------------------- # 16. Activation recomputation (torch.utils.checkpoint / te.checkpoint) # --------------------------------------------------------------------------- diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index d625903a175..e5fa724196b 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -86,16 +86,21 @@ def _is_hybrid_quantizer_list(quantizers): hybrid_count = sum(1 for q in non_none if isinstance(q, HybridQuantizer)) if hybrid_count == 0: return False - if hybrid_count == len(non_none): + # Reject a list mixing HybridQuantizer with None entries: ``_hybrid_split_quantize`` + # subsequently iterates the *full* list with ``isinstance(q, HybridQuantizer)`` which + # would raise ``TypeError`` on the ``None`` entries. Forcing the list to be entirely + # non-``None`` before claiming "all hybrid" matches the dispatch's actual capability. + if hybrid_count == len(quantizers): return True raise ValueError( "GroupedLinear quantizer list mixes HybridQuantizer and non-hybrid" - f" quantizers ({hybrid_count} hybrid, {len(non_none) - hybrid_count}" - " non-hybrid). This combination is not supported: neither" - " `tex.split_quantize` nor `_hybrid_split_quantize` can consume a" - " heterogeneous list. Make the CustomRecipe `qfactory` return a" - " consistent type (all HybridQuantizer or all non-hybrid) across" - " every GEMM for the same role." + f" quantizers ({hybrid_count} hybrid," + f" {len(non_none) - hybrid_count} non-hybrid," + f" {len(quantizers) - len(non_none)} None). This combination is not" + " supported: neither `tex.split_quantize` nor `_hybrid_split_quantize`" + " can consume a heterogeneous list. Make the CustomRecipe `qfactory`" + " return a consistent type (all HybridQuantizer or all non-hybrid)" + " across every GEMM for the same role." ) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 0df86f93f5e..5e8a1f87832 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -86,29 +86,23 @@ def make_empty( requires_grad: bool = False, pin_memory: bool = False, ) -> HybridQuantizedTensor: - # The ``internal`` flag toggles a shared-state mode on each sub- - # quantizer (returns bare ``*TensorStorage`` rather than a - # ``QuantizedTensor`` subclass from ``make_empty``). It is - # global to the quantizer instance, so a raise between - # ``internal=True`` and ``internal=False`` would leak the flag - # and corrupt subsequent non-hybrid quantize calls on the same - # sub-quantizer. ``try/finally`` guarantees the reset. - def _make_empty_internal(sub_quantizer): - prev_internal = sub_quantizer.internal - sub_quantizer.internal = True - try: - return sub_quantizer.make_empty( - shape, - dtype=dtype, - device=device, - pin_memory=pin_memory, - ) - finally: - sub_quantizer.internal = prev_internal - - rowwise_empty = _make_empty_internal(self.rowwise_quantizer) if self.rowwise_usage else None + # Mirror ``quantize_impl``: invoke each sub-quantizer with its own + # ``internal`` setting (no toggle), so the produced sub-storages have + # the same type that ``quantize_impl`` would produce via + # ``sub_quantizer.quantize(tensor)``. + rowwise_empty = ( + self.rowwise_quantizer.make_empty( + shape, dtype=dtype, device=device, pin_memory=pin_memory + ) + if self.rowwise_usage + else None + ) columnwise_empty = ( - _make_empty_internal(self.columnwise_quantizer) if self.columnwise_usage else None + self.columnwise_quantizer.make_empty( + shape, dtype=dtype, device=device, pin_memory=pin_memory + ) + if self.columnwise_usage + else None ) return HybridQuantizedTensor( @@ -372,6 +366,58 @@ def detach(self) -> HybridQuantizedTensor: def get_metadata(self) -> Dict[str, Any]: return HybridQuantizedTensorStorage.get_metadata(self) + @classmethod + def _make_in_reduce_ex( + cls, + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + rowwise_quantizer: Optional[Quantizer], + columnwise_quantizer: Optional[Quantizer], + quantizer: Optional[Quantizer], + dtype: torch.dtype, + shape: torch.Size, + ) -> HybridQuantizedTensor: + """Build HybridQuantizedTensor, for use in ``__reduce_ex__``.""" + return HybridQuantizedTensor( + shape=shape, + dtype=dtype, + rowwise_storage=rowwise_storage, + columnwise_storage=columnwise_storage, + rowwise_quantizer=rowwise_quantizer, + columnwise_quantizer=columnwise_quantizer, + quantizer=quantizer, + ) + + def __reduce_ex__(self, protocol: int) -> tuple: + """Custom pickling. + + Without this, the default ``torch.Tensor.__reduce_ex__`` rebuilds + the parameter as a plain ``torch.Tensor``, dropping the + sub-storages and per-tensor scale state. DCP then reloads the + parameter via ``aten.copy_(dst, plain_tensor)`` which routes to + ``dst.quantize_(plain_tensor)`` — re-quantizing dequantized data + loses precision. + + Mirrors the per-format ``__reduce_ex__`` on ``Float8Tensor``, + ``MXFP8Tensor``, ``NVFP4Tensor``, and ``Float8BlockwiseQTensor``. + Each sub-storage is itself pickled via its own ``__reduce_ex__`` + (preserving FP8 bytes + ``_scale_inv``); the quantizers travel as + regular Python objects and must therefore be picklable + themselves. + """ + return ( + HybridQuantizedTensor._make_in_reduce_ex, + ( + self._rowwise_storage, + self._columnwise_storage, + self._rowwise_quantizer, + self._columnwise_quantizer, + self._quantizer, + self.dtype, + self.shape, + ), + ) + # ── FSDP2 protocol ────────────────────────────────────────────── def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): @@ -480,12 +526,30 @@ def _infer_shape(gathered_buffers): return buf.shape return None + # ``update_usage`` after gathered writeback mirrors what vanilla + # ``Float8Tensor.fsdp_post_all_gather`` does + # — invalidates any stale ``_transpose`` cache on Float8 sub-storages + # and recreates the transpose on non-Hopper architectures where the + # FP8 cuBLAS path requires it. No-op on Blackwell. The flags come + # from the sub-quantizer's pinned direction (set by + # ``HybridQuantizer.__init__``), so we honor whatever the inner + # quantizer thinks its direction is. + def _sync_usage(sub_storage, sub_quantizer): + if sub_storage is None or sub_quantizer is None: + return + sub_storage.update_usage( + rowwise_usage=sub_quantizer.rowwise_usage, + columnwise_usage=sub_quantizer.columnwise_usage, + ) + if out is not None: # Iteration 2+: in-place field update on existing sub-storages if out._rowwise_storage is not None and row_meta is not None: out._rowwise_storage.fsdp_assign_gathered(row_gathered, row_meta) + _sync_usage(out._rowwise_storage, out._rowwise_quantizer) if out._columnwise_storage is not None and col_meta is not None: out._columnwise_storage.fsdp_assign_gathered(col_gathered, col_meta) + _sync_usage(out._columnwise_storage, out._columnwise_quantizer) else: # First iteration: clone the original sharded sub-storages via make_like, # then write gathered (full-size) buffers via each sub-storage's own @@ -496,6 +560,7 @@ def _infer_shape(gathered_buffers): row_sub = type(orig_row_sub).make_like(orig_row_sub, shape=gathered_shape) if row_meta is not None: row_sub.fsdp_assign_gathered(row_gathered, row_meta) + _sync_usage(row_sub, row_quantizer) col_sub = None if orig_col_sub is not None and isinstance(orig_col_sub, QuantizedTensor): @@ -503,6 +568,7 @@ def _infer_shape(gathered_buffers): col_sub = type(orig_col_sub).make_like(orig_col_sub, shape=gathered_shape) if col_meta is not None: col_sub.fsdp_assign_gathered(col_gathered, col_meta) + _sync_usage(col_sub, col_quantizer) ref_sub = row_sub if row_sub is not None else col_sub out = HybridQuantizedTensor( diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index dceee83cfdc..ac3a6182a19 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -289,5 +289,43 @@ def fsdp_buffer_fields(self) -> Tuple[str, ...]: ``_scale_inv`` is a per-tensor scalar; it travels through the hook's metadata tuple (mirroring :meth:`Float8Tensor.fsdp_pre_all_gather`). + + Direction-aware: a vanilla Float8Tensor parameter has ``_data`` + populated, but a columnwise-only sub-storage (used inside + ``HybridQuantizedTensor`` on Hopper / L40 where non-TN FP8 GEMM is + not natively supported) holds its quantized data in ``_transpose`` + instead. Returning ``("_data",)`` unconditionally would have + ``fsdp_extract_buffers`` produce ``(None,)`` and FSDP2 would + all-gather a ``None`` tensor. + + The per-sub-storage direction is fixed at construction (pinned by + ``HybridQuantizer.__init__`` via ``set_usage``), so this check is + stable across iterations even though it inspects the current + field state. """ + if self._data is not None: + return ("_data",) + if self._transpose is not None: + return ("_transpose",) + # Degenerate: fully empty storage. Fall back to ``_data`` so the + # base ``fsdp_extract_buffers`` returns ``(None,)`` — same surface + # the caller would have seen pre-direction-aware logic. return ("_data",) + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered Float8 buffers back, refreshing ``_transpose_invalid``. + + The base implementation just ``setattr``s the gathered tensors into + the named fields. For Float8 we additionally need to clear + ``_transpose_invalid`` when the gathered field is ``_transpose`` — + otherwise a freshly-gathered transpose buffer is treated as stale + on first use (see :attr:`_transpose_invalid` semantics in + ``update_usage`` / ``get_usages``). + """ + super().fsdp_assign_gathered(gathered, meta) + if "_transpose" in meta["field_names"]: + self._transpose_invalid = False From a164cd3610d673169fdbd1f6b5aa6cb7ab926ea4 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 21 May 2026 13:52:41 +0000 Subject: [PATCH 14/97] Towards MLM integration Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 607 +++++++++++++++++++++ transformer_engine/pytorch/tensor/utils.py | 192 +++++++ 2 files changed, 799 insertions(+) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 86b7d4b9597..c6caac786c7 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -19,6 +19,7 @@ LayerNormMLP, TransformerLayer, GroupedLinear, + Float8Quantizer, Float8CurrentScalingQuantizer, MXFP8Quantizer, Float8BlockQuantizer, @@ -2711,6 +2712,612 @@ def test_quantized_param_survives_multiple_forward_passes(self): ), "Weight lost HybridQuantizedTensor type after multiple passes" +# --------------------------------------------------------------------------- +# 3b. quantize_master_weights + post_all_gather_processing for hybrid params +# +# Covers the supported (same-format) cases and the rejected (cross-format, +# missing sub-storage, unsupported sub-quantizer) cases. The supported subset +# is the first incremental hybrid integration with the distributed-optimizer +# quantized-param all-gather flow. Cross-format support is deferred to a +# follow-up; the tests below pin the NotImplementedError contract so the +# rejection messaging stays clear as the feature evolves. +# --------------------------------------------------------------------------- + + +def _ensure_single_rank_dp_group(): + """Return a single-rank NCCL process group for hybrid quantize_master_weights + tests. Mirrors the local-pytest setup in + `tests/pytorch/distributed/test_cast_master_weights_to_fp8.py` so we can call + `torch.distributed.all_reduce` against a trivial group from inside the + per-format helpers. The group is created lazily on first call and reused + across tests within the same pytest process. + """ + # pylint: disable=import-outside-toplevel + import tempfile + import pathlib + + if not torch.distributed.is_initialized(): + torch.cuda.set_device(0) + with tempfile.NamedTemporaryFile(delete=False) as f: + rendezvous_file = pathlib.Path(f.name) + torch.distributed.init_process_group( + backend="nccl", + init_method=rendezvous_file.resolve().as_uri(), + rank=0, + world_size=1, + ) + return torch.distributed.GroupMember.WORLD + + +def _hybrid_recipe_fp8_current(): + """Same-format Float8CurrentScaling on both directions (supported).""" + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + +def _make_delayed_quantizer(fp8_dtype=None): + """Construct a ``Float8Quantizer`` (delayed scaling) with locally-allocated + scale/amax buffers for single-shot unit tests. + + The full delayed-scaling lifecycle (``FP8GlobalStateManager`` updating + ``amax_history`` -> ``scale`` across iterations) is out of scope here; for + ``quantize_master_weights`` we only need the helper to read/write + ``quantizer.amax`` / ``quantizer.scale`` / ``model_weight._scale_inv``, + which works with any pair of 1-element float32 tensors. Initial scale=1.0 + and amax=0.0 mirror the cold-start state ``FP8GlobalStateManager`` would + initialize for the first iteration. + """ + if fp8_dtype is None: + fp8_dtype = tex.DType.kFloat8E4M3 + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=fp8_dtype, + ) + + +def _hybrid_recipe_fp8_delayed(): + """Same-format Float8 delayed scaling on both directions (supported).""" + return _hybrid_custom_recipe( + row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_fp8_delayed_row_current_col(): + """Cross-format per-tensor Float8: delayed rowwise + current columnwise. + + Routed per-direction: row sub-storage -> delayed bucket, col sub-storage + -> current bucket. The two helpers run independently (no shared state), + so each direction's scale is computed via its own scaling lifecycle. + """ + return _hybrid_custom_recipe( + row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + # grad_factory matches the columnwise direction so the wgrad GEMM's + # grad_output sub-quantizer pairs with the input/weight col format. + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + + +def _hybrid_recipe_fp8_current_row_delayed_col(): + """Cross-format per-tensor Float8: current rowwise + delayed columnwise. + + Reversed variant of ``_hybrid_recipe_fp8_delayed_row_current_col``: row + sub-storage -> current bucket, col sub-storage -> delayed bucket. + """ + return _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_mxfp8(): + """Same-format MXFP8 on both directions (rejected today; TODO).""" + return _hybrid_custom_recipe( + row_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E5M2), + ) + + +def _hybrid_recipe_blockwise(): + """Same-format Float8Blockwise on both directions (rejected today; TODO).""" + return _hybrid_custom_recipe( + row_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ), + col_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ), + grad_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True + ), + ) + + +def _build_hybrid_linear_weight(out_features, in_features, hybrid_recipe): + """Build a `HybridQuantizedTensor` weight via `quantized_model_init`. + + Returns (weight, fp32_high_precision_init_val) where the high-precision + init val is on GPU so we can use it as the "master" weight in + quantize_master_weights tests. + """ + torch.manual_seed(42) + with quantized_model_init( + enabled=True, + recipe=hybrid_recipe, + preserve_high_precision_init_val=True, + ): + model = Linear( + in_features, out_features, bias=False, params_dtype=torch.bfloat16 + ).cuda() + + weight = model.weight + assert isinstance(weight, HybridQuantizedTensor), ( + f"Expected HybridQuantizedTensor, got {type(weight).__name__}" + ) + hp_init_cpu = weight.get_high_precision_init_val() + assert hp_init_cpu is not None, "preserve_high_precision_init_val should populate the cpu val" + hp_init = hp_init_cpu.to(weight.device).float() + return weight, hp_init + + +def _hybrid_param_for(out_features, in_features, hybrid_recipe): + """Same as `_build_hybrid_linear_weight` but discards the init val.""" + weight, _ = _build_hybrid_linear_weight(out_features, in_features, hybrid_recipe) + return weight + + +@requires_fp8 +class TestHybridQuantizeMasterWeights: + """`quantize_master_weights` + `post_all_gather_processing` for hybrid params. + + Dispatch is per-direction: each sub-storage is routed independently into the + per-format bucket matching its own sub-quantizer type. Currently-supported + sub-quantizer types can mix freely across directions (e.g. Float8 delayed + row + Float8 current col), single-direction hybrid (one sub-storage dropped + via ``update_usage``) routes the live direction(s) only; per-block sub- + quantizers (MXFP8, NVFP4, Float8Blockwise) raise NotImplementedError + regardless of which direction they appear in. + + Supported subset (per-tensor Float8) -- positive tests verify the present + sub-storage(s) dequantize close to the master weight after the cast: + + * Float8CurrentScaling on both directions (same-format, full master) + * Float8CurrentScaling on both directions (DP-sharded master, non-zero + start_offset) + * Float8 delayed scaling on both directions (same-format) + * Float8 delayed row + Float8 current col (cross-format; row -> delayed + bucket, col -> current bucket) + * Float8 current row + Float8 delayed col (cross-format, reversed) + * Single-direction (rowwise-only) hybrid via ``update_usage`` + * Single-direction (columnwise-only) hybrid via ``update_usage`` + + Rejected subset (NotImplementedError / ValueError) -- negative tests pin + the per-direction rejection contract and the both-None guardrail: + * MXFP8 as a hybrid sub-quantizer (rowwise OR columnwise) + * NVFP4 as a hybrid sub-quantizer (rowwise OR columnwise) + * Float8Blockwise as a hybrid sub-quantizer + * Both sub-storages dropped (caller bug: nothing left to cast) + """ + + # ---------- Positive tests (same-format) ---------- + + def test_fp8_current_same_format_full_master(self): + """Full master (start_offset=0) routes both sub-storages through the + existing per-format current-scaling helper. Verifies both directions + dequantize close to the master weight after the cast. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + # Distributed-optimizer convention: master weight is the flat FP32 shard + # owned by the current rank (or the full param for non-distributed cases). + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + # FP8 E4M3 round-trip; matches the loose tolerance the equivalent + # native-FP8-current test uses (e.g. test_dequantize_close_to_original). + torch.testing.assert_close( + dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + torch.testing.assert_close( + dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + + def test_fp8_current_nonzero_start_offset(self): + """Mimic DP-sharded master: master covers logical elements + [start_offset, start_offset + master.numel()) of the full model weight. + Verifies that the shared logical start_offset is honored by both + sub-storages' per-format routings. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master_full = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + + half = hp_master_full.numel() // 2 + hp_master_shard = hp_master_full.view(-1)[half:].contiguous() + start_offset = half + + quantize_master_weights([weight], [hp_master_shard], [start_offset], group=group) + post_all_gather_processing([weight]) + + # The second-half slice (which the master shard covered) should match. + # The first-half slice was already written at quantized_model_init time + # from the same high-precision init val, so it should also match (modulo + # the second amax all-reduce shifting the per-tensor scale). + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32).view(-1) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32).view(-1) + torch.testing.assert_close( + dq_row[start_offset:], hp_master_shard, rtol=0.125, atol=0.1 + ) + torch.testing.assert_close( + dq_col[start_offset:], hp_master_shard, rtol=0.125, atol=0.1 + ) + + def test_fp8_delayed_same_format_full_master(self): + """Same-format delayed scaling on both directions. Both sub-storages + route into the delayed-scaling bucket as independent entries; the + helper processes them with a single bucket-wide amax all-reduce. + Verifies each direction dequantizes close to the master weight. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_delayed() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close( + dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + torch.testing.assert_close( + dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + + def test_fp8_delayed_row_current_col_full_master(self): + """Cross-format per-tensor Float8: delayed row + current col. + + Pins the new per-direction routing: row sub-storage goes to the + delayed bucket, col sub-storage goes to the current bucket. Each + helper runs independently on its single-entry bucket, with no + cross-pollination between the two scaling lifecycles. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_delayed_row_current_col() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + assert isinstance(weight._rowwise_quantizer, Float8Quantizer) + assert isinstance(weight._columnwise_quantizer, Float8CurrentScalingQuantizer) + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close( + dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + torch.testing.assert_close( + dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + + def test_fp8_current_row_delayed_col_full_master(self): + """Cross-format per-tensor Float8: current row + delayed col. + + Reversed variant of the test above — pins that the per-direction + loop's second iteration (col) reaches the delayed dispatch arm + independently of what the rowwise iteration did. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current_row_delayed_col() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is not None + assert isinstance(weight._rowwise_quantizer, Float8CurrentScalingQuantizer) + assert isinstance(weight._columnwise_quantizer, Float8Quantizer) + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close( + dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + torch.testing.assert_close( + dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + + # NOTE: Per-block sub-quantizers (MXFP8, NVFP4, Float8Blockwise) are not + # supported as hybrid sub-quantizers by this initial integration, regardless + # of which direction they appear in. See the per-direction rejection tests + # below (``test_mxfp8_*_raises`` covers both rowwise and columnwise rejection + # of MXFP8; ``test_nvfp4_*_raises`` and ``test_blockwise_*_raises`` similarly). + # The TODO block above ``_route_hybrid_to_buckets`` in tensor/utils.py + # documents the upstream constraints (single-direction cast helper / kernel + # support) whose unblocker drops per-block format support in for free. + + # ---------- Negative tests (per-direction rejection contract) ---------- + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_mxfp8_rowwise_raises(self): + """MXFP8 in the rowwise sub-quantizer is rejected per-direction. + + ``_cast_master_weights_to_fp8_mxfp8_scaling`` assumes each entry's + ``model_weight`` has BOTH ``_rowwise_*`` and ``_columnwise_*`` populated + (the underlying partial-cast kernel is bidirectional), while a hybrid + sub-storage is single-direction by construction. See + ``TODO(hybrid-mxfp8-distopt)`` in tensor/utils.py for the unblocker shape. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_mxfp8() + # Shape must be a multiple of MXFP8 block size (32) on both axes. + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="MXFP8Quantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_mxfp8_columnwise_raises(self): + """MXFP8 in the columnwise sub-quantizer is rejected per-direction. + + Pairs FP8 current scaling in the rowwise slot (supported) with MXFP8 + in the columnwise slot (rejected). The rowwise iteration of + ``_route_hybrid_to_buckets`` routes the FP8 sub-storage into the + current-scaling bucket cleanly; the columnwise iteration then hits + MXFP8 and raises. Pins that per-direction dispatch visits and rejects + the columnwise sub-quantizer too — not just the rowwise one. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda" + ), + col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + # Shape must be a multiple of MXFP8 block size (32) on both axes. + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="MXFP8Quantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") + def test_nvfp4_rowwise_raises(self): + """NVFP4 in the rowwise sub-quantizer is rejected per-direction. + + The NVFP4 cast path is blocked on a pair of upstream constraints + documented in the TODO block above ``_route_hybrid_to_buckets`` in + tensor/utils.py. + + NOTE: this test exercises the rejection at the ``quantize_master_weights`` + entrypoint, but ``quantized_model_init`` with 2D NVFP4 hybrid already + fails earlier (single-direction 2D NVFP4 quantize is rejected by the + kernel). Use ``with_2d_quantization=False`` so the param can be + constructed and the rejection surfaces at the cast site we care about. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False + ), + col_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False + ), + grad_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False + ), + ) + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="NVFP4Quantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8 block scaling: {reason_for_no_fp8_block_scaling}", + ) + def test_blockwise_rowwise_raises(self): + """Float8BlockQuantizer in the rowwise sub-quantizer is rejected + per-direction (no e2e factory uses it; TODO marker in tensor/utils.py). + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_blockwise() + weight, hp_master = _build_hybrid_linear_weight(128, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="Float8BlockQuantizer rowwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + def test_rowwise_only_fp8_current_full_master(self): + """Single-direction hybrid: columnwise dropped via update_usage. + + Pins that the per-direction loop in `_route_hybrid_to_buckets` skips + the dropped direction silently and routes only the present (rowwise) + sub-storage. Useful for inference / memory-saving paths that + deliberately keep only the fprop-side direction. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=True, columnwise_usage=False) + assert weight._rowwise_storage is not None + assert weight._columnwise_storage is None + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + # Columnwise stays dropped (the cast must not silently revive it). + assert weight._columnwise_storage is None + # Rowwise is populated and dequantizes close to the master. + dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close( + dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + + def test_columnwise_only_fp8_current_full_master(self): + """Single-direction hybrid: rowwise dropped via update_usage. + + Reversed variant — verifies the column-only iteration of the per- + direction loop reaches the dispatch and routes correctly. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=False, columnwise_usage=True) + assert weight._rowwise_storage is None + assert weight._columnwise_storage is not None + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + + # Rowwise stays dropped (the cast must not silently revive it). + assert weight._rowwise_storage is None + # Columnwise is populated and dequantizes close to the master. + dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) + torch.testing.assert_close( + dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 + ) + + def test_both_sub_storages_none_raises(self): + """Both sub-storages dropped via update_usage — nothing left to cast. + + This is the only remaining sub-storage-presence guardrail after the + single-direction enablement: a fully-dropped hybrid weight reaching + `quantize_master_weights` is a caller bug, not a deferred feature, + so we surface it as a ValueError. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + weight.update_usage(rowwise_usage=False, columnwise_usage=False) + assert weight._rowwise_storage is None + assert weight._columnwise_storage is None + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(ValueError, match="both rowwise and columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + + +@requires_fp8 +class TestHybridPostAllGatherProcessing: + """Hybrid branch of `post_all_gather_processing` is exercised indirectly by + the positive `TestHybridQuantizeMasterWeights` tests; the case below pins + an additional invariant that the routing logic must preserve. + """ + + def test_post_ag_idempotent_for_fp8_current_hybrid(self): + """Calling post_all_gather_processing twice on a same-format Float8 + hybrid must not corrupt the sub-storages. + """ + from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + post_all_gather_processing, + ) + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_recipe_fp8_current() + weight, hp_master = _build_hybrid_linear_weight(64, 64, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + quantize_master_weights([weight], [master_flat], [0], group=group) + post_all_gather_processing([weight]) + dq_row_first = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col_first = weight._columnwise_storage.dequantize(dtype=torch.float32) + + post_all_gather_processing([weight]) + dq_row_second = weight._rowwise_storage.dequantize(dtype=torch.float32) + dq_col_second = weight._columnwise_storage.dequantize(dtype=torch.float32) + + torch.testing.assert_close(dq_row_first, dq_row_second, rtol=0.0, atol=0.0) + torch.testing.assert_close(dq_col_first, dq_col_second, rtol=0.0, atol=0.0) + + # --------------------------------------------------------------------------- # 4. Recipe correspondence validation # --------------------------------------------------------------------------- diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 8b22097f7e0..1b4ec8cae3a 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -19,6 +19,7 @@ from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer +from .hybrid_tensor import HybridQuantizedTensor, HybridQuantizer from ..optimizers.multi_tensor_apply import multi_tensor_applier from ..utils import is_non_tn_fp8_gemm_supported from ..constants import NVFP4_BLOCK_SCALING_SIZE @@ -67,6 +68,18 @@ def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): del old_rowwise elif isinstance(tensor, MXFP8Tensor): raise NotImplementedError("replace_raw_data for MXFP8Tensor is not supported yet") + elif isinstance(tensor, HybridQuantizedTensor): + # The distopt all-gather buffer routes at the rowwise sub-storage only; + # the columnwise sub-storage is refreshed each iteration via + # ``HybridQuantizer.update_quantized``. The underlying call delegates + # to the rowwise sub-storage's own ``replace_raw_data`` (which may + # raise for sub-storage types that don't implement it). + if tensor._rowwise_storage is None: + raise NotImplementedError( + "replace_raw_data for HybridQuantizedTensor without a rowwise " + "sub-storage is not supported." + ) + replace_raw_data(tensor._rowwise_storage, new_raw_data) else: raise ValueError(f"replace_raw_data for {type(tensor)} is not supported yet") @@ -165,6 +178,15 @@ def quantize_master_weights( mxfp8_scaling_params.append( (model_weight, master_weight, start_offset, fsdp_shard_model_weight) ) + elif isinstance(quantizer, HybridQuantizer): + _route_hybrid_to_buckets( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + delayed_scaling_params=delayed_scaling_params, + current_scaling_params=current_scaling_params, + ) else: raise ValueError(f"quantize_master_weights for {type(quantizer)} is not supported yet") @@ -933,6 +955,161 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( ) +# --------------------------------------------------------------------------------------------- +# HybridQuantizer helpers for `quantize_master_weights` / `post_all_gather_processing`. +# +# Dispatch is per-direction: `_route_hybrid_to_buckets` iterates over both sub-storages +# of a `HybridQuantizedTensor` and routes each one independently into the per-format +# bucket matching its own sub-quantizer type. Row and col make their own decisions and +# can mix any pair of currently-supported sub-quantizers. +# +# Supported (per-tensor Float8 sub-quantizers, in any per-direction combination): +# - Float8Quantizer (delayed scaling) +# - Float8CurrentScalingQuantizer (current scaling) +# +# Per-tensor Float8 works because `_cast_master_weights_to_fp8_{delayed,current}_scaling` +# accept any Float8Tensor (single direction is fine — each entry is one Float8Tensor +# with its own `_scale_inv` and the helper writes that one entry's `_data`). Each +# hybrid sub-storage IS a single-direction Float8Tensor, so we route them as two +# independent entries (into the same bucket for same-format, or into different +# buckets for cross-format Float8 — e.g. delayed row + current col). +# +# Single-direction hybrid (only one sub-storage populated, e.g. after +# `update_usage(columnwise=False)`) routes the present direction only — the +# per-direction loop skips dropped sub-storages. Both-None hybrids raise ValueError. +# Per-block sub-quantizers still hit their per-direction TODO regardless of single +# vs both direction. +# +# Not supported (raise NotImplementedError per-direction + TODO): +# +# - MXFP8Quantizer as a hybrid sub-quantizer (any direction) +# TODO(hybrid-mxfp8-distopt): the distopt cast kernels +# (`tex.mxfp8_scaling_compute_partial_amax`, `tex.mxfp8_scaling_partial_cast`) +# are bidirectional — both rowwise and colwise outputs required — so they +# cannot ingest a single-direction hybrid sub-storage. (Unrelated to the +# regular `tex.quantize` kernel used by forward/backward, which natively +# supports single-direction output.) Unblocker: add single-direction +# variants of the two distopt kernels, then route hybrid sub-storages +# per-direction into `mxfp8_scaling_params` matching the Float8 path above. +# Also unlocks cross-format MXFP8 row + col. +# +# - NVFP4Quantizer as a hybrid sub-quantizer (any direction) +# TODO(hybrid-nvfp4-distopt): load-bearing blocker is the kernel assertion +# `return_identity || !use_2d_quantization` in +# `quantize_transpose_vector_blockwise_fp4.cu`, which rejects exactly the +# columnwise-only 2D configuration that `HybridQuantizer.__init__` produces +# for the col sub-quantizer. Blocks hybrid 2D NVFP4 weight construction at +# `quantized_model_init` time. 1D NVFP4 is unaffected. The assertion is an +# explicitly-marked unwritten code path, not an algorithmic limit (see the +# kernel author's note above the early-return guard). +# +# Secondary blocker (gated on the kernel fix): the distopt helper +# `_cast_master_weights_to_nvfp4_2d` writes only `_rowwise_data` and relies +# on per-tensor post-AG `_create_columnwise()` — for hybrid, the columnwise +# data needs to land in a SEPARATE col sub-storage, so the post-AG branch +# must be made hybrid-aware (derive `col_sub._columnwise_data` from +# `row_sub`'s gathered rowwise). +# +# - Float8BlockQuantizer as a hybrid sub-quantizer +# TODO(hybrid-fp8-blockwise): same shape as the NVFP4 secondary blocker — +# `_cast_master_weights_to_fp8_blockwise_scaling` writes only `_rowwise_data` +# with per-tensor post-AG `_create_columnwise()` that doesn't reach hybrid's +# separate col sub-storage. Unlike NVFP4, there is no kernel-level +# construction blocker (the Block FP8 kernel natively supports +# columnwise-only mode), so hybrid Block FP8 weights construct fine via the +# non-distopt FusedAdam path today; only the sharded-master distopt cast +# path is blocked. Unblocker is a Python-side hybrid-aware post-AG branch; +# no C++ work needed. +# +# --------------------------------------------------------------------------------------------- + + +def _route_hybrid_to_buckets( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + *, + delayed_scaling_params, + current_scaling_params, +): + """Decompose a `HybridQuantizedTensor` into per-direction entries and route each + into the appropriate per-format bucket used by `quantize_master_weights`. + + Per-direction dispatch: each sub-storage routes independently based on its + own sub-quantizer type. Per-tensor Float8 sub-quantizers (delayed and/or + current scaling) are supported in any combination per direction; single- + direction hybrid (one sub-storage dropped via ``update_usage``) is also + supported. See the TODO block above this helper for the per-block-format + rejection rationale and unblocker shapes. + """ + row_sub = model_weight._rowwise_storage + col_sub = model_weight._columnwise_storage + sub_q_row = model_weight._rowwise_quantizer + sub_q_col = model_weight._columnwise_quantizer + + if row_sub is None and col_sub is None: + raise ValueError( + "quantize_master_weights called on HybridQuantizedTensor with both " + "rowwise and columnwise sub-storages dropped (via update_usage). " + "Nothing to cast — this is most likely a caller bug." + ) + + # Per-direction routing: each (sub_storage, sub_quantizer) pair selects its + # own bucket based on the sub-quantizer's type. Directions that have been + # dropped via ``update_usage`` are silently skipped. + for direction, sub_storage, sub_q in ( + ("rowwise", row_sub, sub_q_row), + ("columnwise", col_sub, sub_q_col), + ): + if sub_storage is None: + continue + entry = (sub_storage, master_weight, start_offset, fsdp_shard_model_weight) + if isinstance(sub_q, Float8Quantizer): + # Delayed scaling: the per-format helper iterates entries + # independently and does a per-DP amax all-reduce across the bucket. + delayed_scaling_params.append(entry) + elif isinstance(sub_q, Float8CurrentScalingQuantizer): + current_scaling_params.append(entry) + elif isinstance(sub_q, MXFP8Quantizer): + # TODO(hybrid-mxfp8-distopt): the distopt cast kernels are + # bidirectional, so a single-direction hybrid sub-storage cannot be + # fed in. See top-of-file TODO block for the unblocker (single- + # direction variants of the two distopt kernels). + raise NotImplementedError( + f"quantize_master_weights for HybridQuantizer with MXFP8Quantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for the unblocker shape." + ) + elif isinstance(sub_q, NVFP4Quantizer): + # TODO(hybrid-nvfp4-distopt): load-bearing blocker is the kernel + # assertion that rejects columnwise-only 2D NVFP4 — which is + # exactly what hybrid's col sub-quantizer pin produces. Secondary + # blocker (gated on the kernel fix) is the per-tensor post-AG + # `_create_columnwise()` not reaching hybrid's separate col + # sub-storage. See top-of-file TODO block for details. + raise NotImplementedError( + f"quantize_master_weights for HybridQuantizer with NVFP4Quantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for details." + ) + elif isinstance(sub_q, Float8BlockQuantizer): + # TODO(hybrid-fp8-blockwise): same shape as the NVFP4 secondary + # blocker (and only that one — no kernel-level construction + # blocker for Block FP8). Python-side post-AG fix. See top-of-file + # TODO block for details. + raise NotImplementedError( + f"quantize_master_weights for HybridQuantizer with Float8BlockQuantizer " + f"{direction} sub-quantizer is not supported yet. See the TODO " + "block above _route_hybrid_to_buckets for details." + ) + else: + raise NotImplementedError( + "quantize_master_weights for HybridQuantizer with " + f"{type(sub_q).__name__} {direction} sub-quantizer is not supported yet." + ) + + def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Tensor]]): """ Post-processing after all-gather for weights in distributed optimizer. @@ -941,6 +1118,13 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten - Plain pytorch tensor: noop. For NVFP4 tensors, uses batched multi-tensor processing to reduce CPU overhead. + + For `HybridQuantizedTensor`, recurses per-direction so that each + sub-storage's native post-processing runs (e.g. Float8 Hopper transpose-cache + pre-creation). Per-block sub-quantizers are rejected at + `quantize_master_weights` time, so by the time we reach here each present + sub-storage is a `Float8Tensor` and the recursive call hits the native + Float8 branch above. """ if not isinstance(model_weights, list): model_weights = [model_weights] @@ -963,6 +1147,14 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten elif isinstance(model_weight, MXFP8Tensor): # MXFP8 scaling: no need to do anything. pass + elif isinstance(model_weight, HybridQuantizedTensor): + # Per-direction post-processing: each Float8 sub-storage routes + # through the recursive call (None / other-type sub-storages are + # silently skipped by the isinstance filter — they would have been + # rejected upstream in `quantize_master_weights`). + for sub in (model_weight._rowwise_storage, model_weight._columnwise_storage): + if isinstance(sub, Float8Tensor): + post_all_gather_processing(sub) elif isinstance(model_weight, QuantizedTensor): raise ValueError(f"post_processing for {type(model_weight)} is not supported") From 62e766831dda7cbc50286b52cf3e63ca04d461c0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 21 May 2026 13:54:52 +0000 Subject: [PATCH 15/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 96 ++++++---------------- transformer_engine/pytorch/tensor/utils.py | 6 +- 2 files changed, 30 insertions(+), 72 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index c6caac786c7..bfca2dfaa4b 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -2752,15 +2752,9 @@ def _ensure_single_rank_dp_group(): def _hybrid_recipe_fp8_current(): """Same-format Float8CurrentScaling on both directions (supported).""" return _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - grad_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), ) @@ -2803,14 +2797,10 @@ def _hybrid_recipe_fp8_delayed_row_current_col(): """ return _hybrid_custom_recipe( row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), - col_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), # grad_factory matches the columnwise direction so the wgrad GEMM's # grad_output sub-quantizer pairs with the input/weight col format. - grad_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, device="cuda" - ), + grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), ) @@ -2821,9 +2811,7 @@ def _hybrid_recipe_fp8_current_row_delayed_col(): sub-storage -> current bucket, col sub-storage -> delayed bucket. """ return _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), ) @@ -2866,14 +2854,12 @@ def _build_hybrid_linear_weight(out_features, in_features, hybrid_recipe): recipe=hybrid_recipe, preserve_high_precision_init_val=True, ): - model = Linear( - in_features, out_features, bias=False, params_dtype=torch.bfloat16 - ).cuda() + model = Linear(in_features, out_features, bias=False, params_dtype=torch.bfloat16).cuda() weight = model.weight - assert isinstance(weight, HybridQuantizedTensor), ( - f"Expected HybridQuantizedTensor, got {type(weight).__name__}" - ) + assert isinstance( + weight, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(weight).__name__}" hp_init_cpu = weight.get_high_precision_init_val() assert hp_init_cpu is not None, "preserve_high_precision_init_val should populate the cpu val" hp_init = hp_init_cpu.to(weight.device).float() @@ -2947,12 +2933,8 @@ def test_fp8_current_same_format_full_master(self): dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) # FP8 E4M3 round-trip; matches the loose tolerance the equivalent # native-FP8-current test uses (e.g. test_dequantize_close_to_original). - torch.testing.assert_close( - dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) - torch.testing.assert_close( - dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) def test_fp8_current_nonzero_start_offset(self): """Mimic DP-sharded master: master covers logical elements @@ -2982,12 +2964,8 @@ def test_fp8_current_nonzero_start_offset(self): # the second amax all-reduce shifting the per-tensor scale). dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32).view(-1) dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32).view(-1) - torch.testing.assert_close( - dq_row[start_offset:], hp_master_shard, rtol=0.125, atol=0.1 - ) - torch.testing.assert_close( - dq_col[start_offset:], hp_master_shard, rtol=0.125, atol=0.1 - ) + torch.testing.assert_close(dq_row[start_offset:], hp_master_shard, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col[start_offset:], hp_master_shard, rtol=0.125, atol=0.1) def test_fp8_delayed_same_format_full_master(self): """Same-format delayed scaling on both directions. Both sub-storages @@ -3012,12 +2990,8 @@ def test_fp8_delayed_same_format_full_master(self): assert weight._columnwise_storage is not None dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) - torch.testing.assert_close( - dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) - torch.testing.assert_close( - dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) def test_fp8_delayed_row_current_col_full_master(self): """Cross-format per-tensor Float8: delayed row + current col. @@ -3046,12 +3020,8 @@ def test_fp8_delayed_row_current_col_full_master(self): assert isinstance(weight._columnwise_quantizer, Float8CurrentScalingQuantizer) dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) - torch.testing.assert_close( - dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) - torch.testing.assert_close( - dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) def test_fp8_current_row_delayed_col_full_master(self): """Cross-format per-tensor Float8: current row + delayed col. @@ -3079,12 +3049,8 @@ def test_fp8_current_row_delayed_col_full_master(self): assert isinstance(weight._columnwise_quantizer, Float8Quantizer) dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) - torch.testing.assert_close( - dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) - torch.testing.assert_close( - dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) # NOTE: Per-block sub-quantizers (MXFP8, NVFP4, Float8Blockwise) are not # supported as hybrid sub-quantizers by this initial integration, regardless @@ -3133,9 +3099,7 @@ def test_mxfp8_columnwise_raises(self): group = _ensure_single_rank_dp_group() hybrid_recipe = _hybrid_custom_recipe( - row_factory=lambda: Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" @@ -3228,9 +3192,7 @@ def test_rowwise_only_fp8_current_full_master(self): assert weight._columnwise_storage is None # Rowwise is populated and dequantizes close to the master. dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) - torch.testing.assert_close( - dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) + torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) def test_columnwise_only_fp8_current_full_master(self): """Single-direction hybrid: rowwise dropped via update_usage. @@ -3258,9 +3220,7 @@ def test_columnwise_only_fp8_current_full_master(self): assert weight._rowwise_storage is None # Columnwise is populated and dequantizes close to the master. dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) - torch.testing.assert_close( - dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1 - ) + torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) def test_both_sub_storages_none_raises(self): """Both sub-storages dropped via update_usage — nothing left to cast. @@ -4913,9 +4873,9 @@ def _make_columnwise_only_float8_storage(self): src = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda") out = q(src) # Columnwise-only Float8 on Hopper: _data is None, _transpose holds data - assert out._data is None, ( - f"Test precondition failed: expected _data is None on Hopper, got {out._data}" - ) + assert ( + out._data is None + ), f"Test precondition failed: expected _data is None on Hopper, got {out._data}" assert out._transpose is not None, "Test precondition failed: _transpose is None" return out @@ -5014,9 +4974,7 @@ def test_iter2_invalidates_stale_transpose_on_rowwise_substorage(self): module=None, mp_policy=None, ) - out2, _ = param.fsdp_post_all_gather( - sharded_tensors, metadata, param.dtype, out=out - ) + out2, _ = param.fsdp_post_all_gather(sharded_tensors, metadata, param.dtype, out=out) # After fsdp_post_all_gather, the rowwise sub-quantizer is pinned # columnwise=False, so update_usage(rowwise=True, columnwise=False) diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 1b4ec8cae3a..6a1cd57c7ab 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1077,7 +1077,7 @@ def _route_hybrid_to_buckets( # fed in. See top-of-file TODO block for the unblocker (single- # direction variants of the two distopt kernels). raise NotImplementedError( - f"quantize_master_weights for HybridQuantizer with MXFP8Quantizer " + "quantize_master_weights for HybridQuantizer with MXFP8Quantizer " f"{direction} sub-quantizer is not supported yet. See the TODO " "block above _route_hybrid_to_buckets for the unblocker shape." ) @@ -1089,7 +1089,7 @@ def _route_hybrid_to_buckets( # `_create_columnwise()` not reaching hybrid's separate col # sub-storage. See top-of-file TODO block for details. raise NotImplementedError( - f"quantize_master_weights for HybridQuantizer with NVFP4Quantizer " + "quantize_master_weights for HybridQuantizer with NVFP4Quantizer " f"{direction} sub-quantizer is not supported yet. See the TODO " "block above _route_hybrid_to_buckets for details." ) @@ -1099,7 +1099,7 @@ def _route_hybrid_to_buckets( # blocker for Block FP8). Python-side post-AG fix. See top-of-file # TODO block for details. raise NotImplementedError( - f"quantize_master_weights for HybridQuantizer with Float8BlockQuantizer " + "quantize_master_weights for HybridQuantizer with Float8BlockQuantizer " f"{direction} sub-quantizer is not supported yet. See the TODO " "block above _route_hybrid_to_buckets for details." ) From 731651610124843701a3e4d2f85543d0565493dd Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 3 Jun 2026 13:01:30 +0000 Subject: [PATCH 16/97] Resolve comments: improve fsdp/tp/sp tests + amax reduction fix Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/fsdp2_utils.py | 22 +- .../fsdp2_tests/run_fsdp2_fused_adam.py | 356 +++++++++++++---- .../fsdp2_tests/run_fsdp2_mem_leak.py | 12 +- .../fsdp2_tests/run_fsdp2_model.py | 127 +++++- tests/pytorch/distributed/run_hybrid_tp_sp.py | 371 +++++++++++++++--- .../pytorch/distributed/test_hybrid_tp_sp.py | 81 ++++ tests/pytorch/test_hybrid_quantization.py | 10 +- transformer_engine/pytorch/module/base.py | 5 +- .../pytorch/module/layernorm_linear.py | 9 + .../pytorch/module/layernorm_mlp.py | 9 + transformer_engine/pytorch/module/linear.py | 9 + .../pytorch/tensor/hybrid_tensor.py | 38 ++ 12 files changed, 899 insertions(+), 150 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index f8423f14ccf..7b71cb04c39 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -17,19 +17,6 @@ def get_recipe_from_string(recipe): return getattr(transformer_engine.common.recipe, recipe)() -# ── Hybrid qfactories ───────────────────────────────────────────────── -# -# Module-level (picklable) qfactories used by ``get_hybrid_recipe_from_string``. -# Each factory composes one or two role-aware base factories from -# ``quantization_recipes_base`` per direction. Per-role behavior is delegated -# to the base factory — the hybrid layer only decides direction pairing. -# -# DCP serializes ``CustomRecipe`` via ``pickle``; closure-based qfactories -# (lambdas, inner functions referencing captured state) are not picklable, -# so the qfactory must live at module scope. See -# ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. - - def _hybrid_fp8_current_qfactory(role): """FP8 current-scaling rowwise + FP8 current-scaling columnwise.""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -74,6 +61,11 @@ def _hybrid_mixed_mxfp8_fp8_qfactory(role): return current_scaling_quantizer_factory(role) +# The qfactories above are registered here as module-level functions (not +# lambdas or closures) on purpose: DCP serializes ``CustomRecipe`` via +# ``pickle``, and closure-based qfactories (or inner functions capturing state) +# are not picklable. Keeping them at module scope lets them pickle by reference. +# See ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. _HYBRID_QFACTORIES = { "HybridFP8CurrentScaling": _hybrid_fp8_current_qfactory, "HybridMXFP8": _hybrid_mxfp8_qfactory, @@ -85,6 +77,10 @@ def _hybrid_mixed_mxfp8_fp8_qfactory(role): def get_hybrid_recipe_from_string(recipe): """Build a CustomRecipe wrapping a module-level (picklable) hybrid qfactory. + Each hybrid qfactory composes one or two role-aware base factories from + ``quantization_recipes_base`` per direction; per-role behavior is delegated + to the base factory and the hybrid layer only decides the direction pairing. + Supported values: "HybridFP8CurrentScaling" — FP8 current for both directions "HybridMXFP8" — MXFP8 for both directions diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index d53dad0a721..591bd643867 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1059,7 +1059,6 @@ def test_dcp_resharding_load(recipe_name): def _build_hybrid_model(hybrid_recipe, use_meta_device=True): """Build a model with quantized_model_init using a hybrid CustomRecipe.""" - ctx = te.quantized_model_init(enabled=True, recipe=hybrid_recipe) kwargs = dict( fuse_qkv_params=True, params_dtype=torch.bfloat16, @@ -1068,7 +1067,7 @@ def _build_hybrid_model(hybrid_recipe, use_meta_device=True): ) if use_meta_device: kwargs["device"] = "meta" - with ctx: + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): model = torch.nn.Sequential( *[ te.TransformerLayer( @@ -1142,24 +1141,27 @@ def test_fused_adam_hybrid_master_weights(hybrid_recipe_name): if "master_param" in state: assert state["master_param"].dtype == torch.float32 - assert losses[-1] < losses[0], f"Loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" -@pytest.mark.parametrize("reshard_after_forward", [True, False]) -def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name, reshard_after_forward): - """Hybrid FusedAdam loop under both ``reshard_after_forward`` settings. +def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name): + """Hybrid FusedAdam training must be numerically invariant to FSDP2's + ``reshard_after_forward`` schedule. - ``reshard_after_forward=True`` is FSDP2's default: the gathered weight is - dropped after forward and a second all-gather happens in backward — - meaning ``fsdp_post_all_gather(out=...)`` is invoked twice per training - step (once per pass) on the same gathered buffer. ``False`` keeps the - gathered weight alive through backward — only one gather per step, and - the gathered copy persists across forward/backward within the same step. + ``reshard_after_forward`` only changes *when* the gathered weight is + materialized/freed, not the math: ``True`` (FSDP2's child-module default) + drops the gathered weight after forward and re-gathers it in backward -- + invoking ``fsdp_post_all_gather(out=...)`` twice per step -- while ``False`` + keeps the gathered copy alive through backward (one gather per step). The + gathered quantized bytes are identical either way, so both schedules must + produce **bitwise-identical** loss trajectories. - Both modes must complete cleanly and produce a decreasing loss. This - locks in that the hybrid hooks handle both FSDP2 schedules, and forms a - regression harness for a future bandwidth optimization (P1.1) that would - split forward-only / backward-only buffers. + Strictly stronger than "loss decreased": it locks in that the hybrid + all-gather hooks are schedule-invariant across both FSDP2 passes, and + regression-guards the future P1.1 buffer-split bandwidth optimization. """ if hybrid_recipe_name == "HybridFloat8BlockScaling": pytest.xfail( @@ -1172,32 +1174,47 @@ def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name, reshard_after_fo hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) world_size, device = _get_dist_info() - model = _build_hybrid_model(hybrid_recipe) - model = _shard_model(model, world_size, reshard_after_forward=reshard_after_forward) - - optimizer = te.optimizers.FusedAdam( - model.parameters(), - lr=1e-3, - master_weights=True, - master_weight_dtype=torch.float32, - ) - + # Shared, fixed input/target so the two schedules are compared on identical data. x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) target = torch.randn_like(x) - losses = [] - for _ in range(NUM_STEPS): - optimizer.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=hybrid_recipe): - output = model(x) - loss = F.mse_loss(output, target) - losses.append(loss.item()) - loss.backward() - optimizer.step() + def run(reshard_after_forward): + # Re-seed so both schedules get identical weight init from reset_parameters(). + torch.manual_seed(42) + torch.cuda.manual_seed(42) + model = _shard_model( + _build_hybrid_model(hybrid_recipe), + world_size, + reshard_after_forward=reshard_after_forward, + ) + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + losses = [] + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + return losses - assert losses[-1] < losses[0], ( - f"[reshard_after_forward={reshard_after_forward}] " - f"loss did not decrease: {losses[0]:.4f} -> {losses[-1]:.4f}" + losses_resharded = run(reshard_after_forward=True) # re-gather in backward + losses_kept = run(reshard_after_forward=False) # keep gathered weight through backward + + # Both schedules must train (strictly monotonic decrease) ... + assert all( + losses_resharded[i + 1] < losses_resharded[i] for i in range(NUM_STEPS - 1) + ), f"reshard_after_forward=True loss not strictly decreasing: {losses_resharded}" + # ... and be numerically identical, since reshard is a schedule, not a math change. + assert losses_resharded == losses_kept, ( + "reshard_after_forward changed numerics (must be schedule-invariant): " + f"True={losses_resharded} vs False={losses_kept}" ) @@ -1256,14 +1273,211 @@ def run_training(model, recipe_for_autocast): assert hybrid_losses[-1] < hybrid_losses[0], f"Hybrid loss did not decrease: {hybrid_losses}" assert bf16_losses[-1] < bf16_losses[0], f"BF16 loss did not decrease: {bf16_losses}" - # Verify hybrid and bf16 loss trajectories are within the same order of magnitude. - # Quantized training may diverge from bf16, but should not be wildly different. + # Hybrid stays within a few % of bf16 (seed-fixed). + rel_tol = 0.10 for step, (h_loss, b_loss) in enumerate(zip(hybrid_losses, bf16_losses)): - ratio = h_loss / max(b_loss, 1e-10) - assert 0.1 < ratio < 10.0, ( - f"Step {step}: hybrid loss ({h_loss:.4f}) and bf16 loss ({b_loss:.4f}) " - f"differ by more than 10x (ratio={ratio:.2f})" + rel_diff = abs(h_loss - b_loss) / max(abs(b_loss), 1e-10) + assert rel_diff < rel_tol, ( + f"Step {step}: hybrid loss ({h_loss:.4f}) vs bf16 ({b_loss:.4f}) " + f"differ by {rel_diff * 100:.2f}% (> {rel_tol * 100:.0f}%)" + ) + + +# Same-format hybrid -> the vanilla recipe it must match bitwise. Cross-format +# hybrids (e.g. HybridMixed_MXFP8_FP8) have no single-format vanilla equivalent. +_HYBRID_TO_BASE_RECIPE = { + "HybridFP8CurrentScaling": "Float8CurrentScaling", + "HybridMXFP8": "MXFP8BlockScaling", + "HybridFloat8BlockScaling": "Float8BlockScaling", +} + + +def _build_linear_parity_stack(recipe): + """Two bare ``te.Linear`` layers under ``quantized_model_init`` for + hybrid-vs-vanilla bitwise parity. + """ + with te.quantized_model_init(enabled=True, recipe=recipe): + return torch.nn.Sequential( + te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=torch.bfloat16, device="meta"), + te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, params_dtype=torch.bfloat16, device="meta"), + ) + + +def test_fused_adam_hybrid_vs_base_recipe_parity(hybrid_recipe_name): + """Same-format hybrid must match its vanilla recipe bitwise through the full + FSDP2 + FusedAdam loop. + + Forward output, weight gradients, and per-step loss are all asserted + bitwise-identical every iteration -- regression guard for the amax-reduction + fix. Uses a bare ``te.Linear`` stack (see ``_build_linear_parity_stack``) to + isolate GEMM-operand quantization. + """ + if hybrid_recipe_name == "HybridFloat8BlockScaling": + pytest.xfail( + "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " + "quantized type through FSDP2 view(-1) in reset_sharded_param." + ) + if hybrid_recipe_name not in _HYBRID_TO_BASE_RECIPE: + pytest.skip( + f"{hybrid_recipe_name} is cross-format; no single-format vanilla " + "recipe to compare against." + ) + + from fsdp2_utils import get_hybrid_recipe_from_string + + base_recipe_name = _HYBRID_TO_BASE_RECIPE[hybrid_recipe_name] + world_size, device = _get_dist_info() + + # Shared, fixed input/target; the comparison is per-rank (base vs hybrid on + # the same rank), so cross-rank input consistency does not matter. + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + def run_training(build_fn, recipe_for_autocast): + # Re-seed so both models get identical init from reset_parameters() (run + # after sharding); with same-format quantization and a dropout-free loop + # the full trajectory is then deterministic. + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + model = _shard_model(build_fn(), world_size) + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) + first_output = None + losses = [] + grads_per_step = [] + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe_for_autocast): + output = model(x) + if step == 0: + first_output = output.detach().clone() + loss = F.mse_loss(output, target) + losses.append(loss.detach().clone()) + loss.backward() + # Snapshot grad local shards before the optimizer consumes them + # (p.grad is a DTensor under FSDP2) to assert backward parity directly. + step_grads = [] + for p in model.parameters(): + g = p.grad + if g is None: + step_grads.append(None) + else: + g = g.to_local() if isinstance(g, DTensor) else g + step_grads.append(g.detach().clone()) + grads_per_step.append(step_grads) + optimizer.step() + return first_output, losses, grads_per_step + + base_recipe = get_recipe_from_string(base_recipe_name) + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + + base_first, base_losses, base_grads = run_training( + lambda: _build_linear_parity_stack(base_recipe), base_recipe + ) + hybrid_first, hybrid_losses, hybrid_grads = run_training( + lambda: _build_linear_parity_stack(hybrid_recipe), hybrid_recipe + ) + + # (1) First forward: bitwise-identical (the core operand-equivalence check). + torch.testing.assert_close( + hybrid_first, + base_first, + rtol=0.0, + atol=0.0, + msg=lambda m: ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] first forward output not" + f" bitwise-identical (a same-format hybrid must match its vanilla recipe" + f" before any optimizer step): {m}" + ), + ) + + # (2) Every per-step loss: bitwise-identical across the whole optimizer loop. + for step, (b_loss, h_loss) in enumerate(zip(base_losses, hybrid_losses)): + torch.testing.assert_close( + h_loss, + b_loss, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} loss not" + f" bitwise-identical to the vanilla recipe: {m}" + ), + ) + + # (3) Backward: every weight-gradient shard at every step bitwise-identical + # (implied by the loss trajectory, but asserted directly to be explicit). + for step, (b_step, h_step) in enumerate(zip(base_grads, hybrid_grads)): + for i, (b_grad, h_grad) in enumerate(zip(b_step, h_step)): + assert (b_grad is None) == (h_grad is None), ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} param {i}" + " gradient presence differs between hybrid and vanilla" + ) + if b_grad is None: + continue + torch.testing.assert_close( + h_grad, + b_grad, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step, i=i: ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i}" + f" gradient not bitwise-identical to the vanilla recipe: {m}" + ), + ) + + +def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): + """Per-tensor hybrid weights must share ONE amax-reduced scale across FSDP2 + shards -- tolerance-free regression guard for the amax-reduction fix. + + Without cross-shard reduction each rank quantizes its shard with a local amax + and the scales differ; with the fix they match. Checked directly on the + sharded weight (no forward). Block-scaled formats (MXFP8) are skipped. + """ + if hybrid_recipe_name != "HybridFP8CurrentScaling": + pytest.skip("scale-uniformity check applies to per-tensor current scaling only") + + from transformer_engine.pytorch import HybridQuantizedTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + world_size, device = _get_dist_info() + if world_size < 2: + pytest.skip("needs >=2 ranks to compare shard scales") + + hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) + model = _build_hybrid_model(hybrid_recipe) + model = _shard_model(model, world_size) + + checked = 0 + for name, param in model.named_parameters(): + if not ( + isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ): + continue + row_sub = param._local_tensor._rowwise_storage + scale_inv = getattr(row_sub, "_scale_inv", None) if row_sub is not None else None + if scale_inv is None: + continue + local_scale = scale_inv.detach().reshape(-1).clone() + gathered = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered, local_scale) + for r in range(1, world_size): + torch.testing.assert_close( + gathered[r], + gathered[0], + rtol=0.0, + atol=0.0, + msg=lambda m, n=name, r=r: ( + f"{n}: rank {r} rowwise _scale_inv differs from rank 0 -- cross-shard " + f"amax reduction was not applied to the hybrid current-scaling weight: {m}" + ), + ) + checked += 1 + assert checked > 0, "no hybrid current-scaling weights found to check" def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): @@ -1296,14 +1510,6 @@ def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): with te.autocast(enabled=True, recipe=hybrid_recipe): _ = model(x) - # Stateless formats gather identical bytes; dequantize must match exactly. - _TIGHT_TOLERANCE = { - "HybridFP8CurrentScaling": dict(rtol=0.0, atol=0.0), - "HybridMXFP8": dict(rtol=0.0, atol=0.0), - "HybridMixed_MXFP8_FP8": dict(rtol=0.0, atol=0.0), - } - tolerance = _TIGHT_TOLERANCE.get(hybrid_recipe_name, dict(rtol=1e-6, atol=1e-6)) - checked = 0 for name, param in model.named_parameters(): if not (isinstance(param, DTensor) and isinstance(param._local_tensor, QuantizedTensor)): @@ -1321,11 +1527,13 @@ def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): else: fsdp_full_deq = full_param.float() + # Stateless formats gather identical bytes -> exact match. torch.testing.assert_close( manual_full.float(), fsdp_full_deq[: manual_full.shape[0]].float(), msg=lambda m, n=name: f"Allgather mismatch for {n}: {m}", - **tolerance, + rtol=0.0, + atol=0.0, ) checked += 1 @@ -1361,8 +1569,17 @@ def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): per_rank_out = 96 out_features = per_rank_out * world_size in_features = 128 # arbitrary, divisible by 32; not sharded by FSDP2 here - assert per_rank_out % 32 == 0, "MXFP8 data alignment precondition" - assert per_rank_out % 128 != 0, "Test precondition: shard must need scale padding" + assert per_rank_out % 32 == 0, ( + f"Test setup error: per_rank_out={per_rank_out} (= out_features / world_size, " + f"world_size={world_size}) must be a multiple of the MXFP8 block size (32) so the " + f"sharded weight's data stays block-aligned. Pick a per_rank_out divisible by 32." + ) + assert per_rank_out % 128 != 0, ( + f"Test setup error: per_rank_out={per_rank_out} must NOT be a multiple of 128, or the " + f"rowwise scale-inv needs no alignment padding and this test stops exercising the MXFP8 " + f"unpad-before-gather / pad-after-gather path it exists to cover. Pick a per_rank_out " + f"divisible by 32 but not 128 (e.g. 96)." + ) for recipe_name in ("HybridMXFP8", "HybridMixed_MXFP8_FP8"): hybrid_recipe = get_hybrid_recipe_from_string(recipe_name) @@ -1427,20 +1644,19 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): ) if hybrid_recipe_name == "HybridFP8CurrentScaling": + # TODO: preserve hybrid current-scaling primary-weight scales across DCP + # by implementing __tensor_flatten__/__tensor_unflatten__ on the quantized + # tensor stack (HybridQuantizedTensor + its Float8Tensor sub-storages) so + # DCP serializes fp8 data + fp32 _scale_inv as explicit tensor leaves + # instead of round-tripping through a dequantized bf16 weight. pytest.xfail( - "HybridFP8CurrentScaling: per-tensor _scale_inv is not preserved " - "through DCP's tensor-storage-byte serialization path. " - "HybridQuantizedTensor.__reduce_ex__ correctly round-trips through " - "pickle (verified by torch.save/torch.load), but DCP bypasses " - "pickle and serializes the tensor's storage bytes — the scalar " - "_scale_inv is not enumerated as a separate tensor leaf and gets " - "lost. Vanilla Float8CurrentScaling avoids this because per-tensor " - "scale lives in module.fp8_meta (saved as extra_state), not on " - "the tensor; hybrid uses per-sub-storage scales without that " - "mirror. Fix path: implement __tensor_flatten__/__tensor_unflatten__ " - "across the quantized tensor stack so DCP can serialize the " - "per-leaf tensor fields directly. Loaded model output diverges by " - "~5e-2." + "HybridFP8CurrentScaling: hybrid current-scaling primary-weight " + "_scale_inv is not preserved across DCP. DCP stores each weight as a " + "dequantized bf16 leaf (no scale leaf) and re-quantizes on load; this " + "is idempotent for a single per-tensor Float8Tensor (vanilla " + "round-trips bitwise) but not for the hybrid's two sub-storages, so " + "the loaded output diverges ~5e-2. torch.save/load and the FP32 " + "master weight are unaffected. See the TODO above for the fix." ) from fsdp2_utils import get_hybrid_recipe_from_string @@ -1540,6 +1756,10 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): "safetensors_fp32_export": test_safetensors_fp32_export, "fused_adam_hybrid_master_weights": test_fused_adam_hybrid_master_weights, "fused_adam_hybrid_bf16_vs_hybrid_parity": test_fused_adam_hybrid_bf16_vs_hybrid_parity, + "fused_adam_hybrid_vs_base_recipe_parity": test_fused_adam_hybrid_vs_base_recipe_parity, + "fused_adam_hybrid_scale_uniform_across_shards": ( + test_fused_adam_hybrid_scale_uniform_across_shards + ), "fused_adam_hybrid_allgather_correctness": test_fused_adam_hybrid_allgather_correctness, "fused_adam_hybrid_mxfp8_awkward_shard_shape": ( test_fused_adam_hybrid_mxfp8_awkward_shard_shape diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py index f729dff5355..5aa09ae6c46 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py @@ -457,7 +457,6 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in def _build_hybrid_model(num_layers, hybrid_recipe, use_meta_device=True): """Build a model with quantized_model_init using a hybrid CustomRecipe.""" - ctx = te.quantized_model_init(enabled=True, recipe=hybrid_recipe) kwargs = dict( fuse_qkv_params=True, params_dtype=torch.bfloat16, @@ -466,7 +465,7 @@ def _build_hybrid_model(num_layers, hybrid_recipe, use_meta_device=True): ) if use_meta_device: kwargs["device"] = "meta" - with ctx: + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): model = torch.nn.Sequential( *[ te.TransformerLayer( @@ -540,6 +539,11 @@ def test_hybrid_no_excess_forward_memory(hybrid_recipe_name): hybrid_avg = sum(hybrid_increments) / len(hybrid_increments) excess_per_layer = hybrid_avg - bf16_avg + # Basis: forward growth is constant per layer (no accumulation) for both bf16 and + # hybrid; the excess is just hybrid's extra per-layer quantized buffers. Measured + # excess: ~3 KiB (FP8 current) / ~7 KiB (mixed MXFP8+FP8) / ~12 KiB (MXFP8). A + # leaked layer's quantized weights would be hundreds of KiB, so 50 KiB sits above + # the real per-layer overhead and well below a leak. tolerance_per_layer = 50 * 1024 # 50 KiB assert excess_per_layer <= tolerance_per_layer, ( @@ -602,6 +606,10 @@ def test_hybrid_transpose_cache_after_backward(hybrid_recipe_name): ) excess = hybrid_bwd_delta - bf16_bwd_delta + # Basis: hybrid retains no more than bf16 after backward+step — measured excess is + # slightly negative (~-0.02..-0.09 MiB vs a ~2 MiB bf16 delta). The tolerance only + # absorbs allocator/measurement noise; a genuinely retained gathered weight or + # transpose cache would be MiB-scale (>> 256 KiB). tolerance = 256 * 1024 # 256 KiB assert excess <= tolerance, ( diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 38a22915b02..6b3c2c2d2f0 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -27,6 +27,7 @@ """ import gc +import math import os import sys import argparse @@ -254,6 +255,51 @@ def _check_fp8_fsdp2_allgather(model): module.reshard() +@torch.no_grad() +def _check_hybrid_fsdp2_allgather(model): + """All-gather correctness for hybrid quantized params under FSDP2. + + Mirrors :func:`_check_fp8_fsdp2_allgather` but for hybrid params: compares + FSDP2's quantized all-gather (``unshard`` -> ``dequantize``) against a manual + fp32 dequantize-then-allgather of the local shards. This is self-referential + (no vanilla reference model), so it is robust to the norm-fusion difference + between hybrid and vanilla recipes -- hybrid auto-disables ``with_quantized_norm`` + while vanilla fuses, so a hybrid-vs-vanilla comparison would not match. + """ + # Manual fp32 weight allgather of the (possibly hybrid) local shards. + fp32_allgathered_params = {} + for name, param in model.named_parameters(): + assert isinstance(param, DTensor) + local_tensor = param._local_tensor + device_mesh = param.device_mesh + dist_group = ( + device_mesh.get_group(mesh_dim="shard") + if device_mesh.ndim > 1 + else device_mesh.get_group() + ) + # ``dequantize`` is a no-op on plain (non-quantized) bf16 params such as + # LayerNorm weights/biases, and dequantizes hybrid sub-storages otherwise. + local_hp = local_tensor.dequantize() + gathered_tensor = [ + torch.zeros_like(local_hp) for _ in range(dist.get_world_size(group=dist_group)) + ] + dist.all_gather(gathered_tensor, local_hp, group=dist_group) + fp32_allgathered_params[name] = torch.cat(gathered_tensor, dim=0) + # Quantized allgather via FSDP2. + for module in model.modules(): + if hasattr(module, "unshard"): + module.unshard() + # Hybrid sub-storages (e.g. MXFP8 scale unpad/repad through FSDP2) can introduce + # small numerical differences vs the manual dequantize-then-allgather path. + tols = dict(atol=5e-4, rtol=5e-3) + for name, param in model.named_parameters(): + torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **tols) + # Revert model to original sharded state. + for module in model.modules(): + if hasattr(module, "reshard"): + module.reshard() + + def _run_training(args): """Core training logic. Assumes dist is already initialized.""" device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") @@ -408,8 +454,12 @@ def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): def test_distributed_hybrid(hybrid_recipe_name): """FSDP2 training with hybrid quantized_model_init. - Uses quantized_model_init with a hybrid CustomRecipe and verifies that - training completes without error with a TransformerLayer model. + Uses quantized_model_init with a hybrid CustomRecipe on a TransformerLayer + model: + - params are DTensors wrapping HybridQuantizedTensor local shards, + - the loss is finite and strictly decreasing on fixed data (grads flow), + - params keep their hybrid quantized type across optimizer.step(), + - FSDP2's quantized all-gather matches a manual fp32 dequant-then-allgather. """ if hybrid_recipe_name == "HybridFloat8BlockScaling": pytest.xfail( @@ -446,21 +496,50 @@ def test_distributed_hybrid(hybrid_recipe_name): module.reset_parameters() restore_custom_attrs(model, custom_attrs) + from transformer_engine.pytorch import HybridQuantizedTensor + + def _hybrid_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + + # quantized_model_init must produce HybridQuantizedTensor local shards. + hybrid_count = _hybrid_param_count() + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + optimizer = optim.Adam(model.parameters(), lr=1e-3) - inp_shape = (128, 16, 512) - out_shape = (128, 16, 512) + input_data = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + losses = [] for iteration in range(3): optimizer.zero_grad() - input_data = torch.randn(inp_shape, device=device, dtype=torch.bfloat16) - target = torch.randn(out_shape, device=device, dtype=torch.bfloat16) with te.autocast(enabled=True, recipe=hybrid_recipe): output = model(input_data) loss = F.mse_loss(output, target) loss.backward() optimizer.step() - dist_print(f"Hybrid iteration {iteration} completed with loss {loss.item()}") + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Hybrid iteration {iteration} completed with loss {loss_val}") + + # Training must actually progress on fixed data: strictly monotonic decrease. + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + # Params must stay HybridQuantizedTensor after the optimizer step -- guards a + # silent dequantize-to-bf16 through the FSDP2 / optimizer path. + assert _hybrid_param_count() == hybrid_count, ( + "HybridQuantizedTensor params lost their quantized type after optimizer.step()" + ) + + # FSDP2 quantized all-gather must match a manual fp32 dequant-then-allgather. + _check_hybrid_fsdp2_allgather(model) def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): @@ -503,18 +582,46 @@ def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): module.reset_parameters() restore_custom_attrs(model, custom_attrs) + from transformer_engine.pytorch import HybridQuantizedTensor + + def _hybrid_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, HybridQuantizedTensor) + ) + + hybrid_count = _hybrid_param_count() + assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" + optimizer = optim.Adam(model.parameters(), lr=1e-3) + x = torch.randn(128, 16, in_features, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, out_features, device=device, dtype=torch.bfloat16) + + losses = [] for iteration in range(5): optimizer.zero_grad() - x = torch.randn(128, 16, in_features, device=device, dtype=torch.bfloat16) - target = torch.randn(128, 16, out_features, device=device, dtype=torch.bfloat16) with te.autocast(enabled=True, recipe=hybrid_recipe): output = model(x) loss = F.mse_loss(output, target) loss.backward() optimizer.step() - dist_print(f"Hybrid reshard_after_fwd iter {iteration}, loss {loss.item():.4f}") + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Hybrid reshard_after_fwd iter {iteration}, loss {loss_val:.4f}") + + # The forward-reshard-backward-reshard cycle must still train: strict decrease. + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" + + # Params must survive the split/as_strided/slice reshard dispatch ops with + # their hybrid quantized type intact. + assert _hybrid_param_count() == hybrid_count, ( + "HybridQuantizedTensor params lost their quantized type after optimizer.step()" + ) if __name__ == "__main__": diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index e6e4d45e1c2..12203ec39c2 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -14,26 +14,23 @@ ``run_numerics.py``), so any drift between the two paths is a hybrid- specific TP/SP issue rather than an initialization artifact. -Test surface: - * ``te.Linear`` column-parallel and row-parallel, with and without - sequence parallelism. - * ``te.LayerNormLinear`` column-parallel with sequence parallelism — - the quantized-AG path that currently unfuses LN+quantize for - ``HybridQuantizer``. - * ``te.TransformerLayer`` with ``set_parallel_mode=True`` and SP on — - integration test hitting LayerNormLinear + DPA + LayerNormMLP + row- - parallel output projection in one shot. - -Only same-format hybrid recipes (FP8 current rowwise + FP8 current -columnwise; MXFP8 rowwise + MXFP8 columnwise) are exercised here so the -numerical signal is clean. Cross-format hybrid adds independent -numerical variation unrelated to TP/SP and is covered by single-GPU -tests already. - -Tolerances match upstream ``run_numerics.py`` per-format settings (see -``_get_tolerances``); they're loose enough to absorb the amax-reduction -and stochastic numerical asymmetries inherent to distributed FP8, tight -enough to catch a silent BF16 fallback on the hybrid sub-storage path. +Test surface: ``te.Linear`` (column/row x SP on/off, plus a bitwise +hybrid-vs-vanilla operand-equivalence check), ``te.LayerNormLinear``, +``te.LayerNormMLP``, and ``te.TransformerLayer`` (all with SP on/off). The +non-attention tests also compare per-parameter gradients in the no-SP configs, +where grads align directly with the single-node reference. + +Recipes: same-format (FP8-current, MXFP8, NVFP4) for a clean signal and the +bitwise check, plus a cross-format one (MXFP8 fwd / NVFP4 bwd) that exercises +the forward-vs-backward all-gather format asymmetry (fwd gathers rowwise, bwd +columnwise) -- which same-format recipes cannot surface. + +Two comparison kinds with different tolerances: + * distributed-vs-single-node (``_test_*``): inherently loose -- the sharded + side quantizes per-shard and reduces across ranks, so it is never bitwise. + ``_get_tolerances`` matches upstream ``run_numerics.py`` per format. + * hybrid-vs-vanilla (``_test_linear_vs_vanilla``): same topology, so bitwise + (``rtol=0, atol=0``) for forward (all configs) and backward (non-SP). """ import argparse @@ -117,31 +114,75 @@ def _hybrid_mxfp8_qfactory(role): rowwise_quantizer=_make_mxfp8_quantizer(), columnwise_quantizer=_make_mxfp8_quantizer(), ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_mxfp8_quantizer(fp8_dtype=tex.DType.kFloat8E5M2) + # MXFP8 uses E4M3 for every pass (its canonical Format.E4M3) return _make_mxfp8_quantizer() -def _make_nvfp4_quantizer(): - """Default NVFP4Quantizer: no RHT, no stochastic rounding, no 2D - scaling — matches upstream ``run_numerics.py::nvfp4_vanilla()`` which - strips the recipe to bare 1D block scaling for distributed TP - fairness. Same-format hybrid NVFP4 has no E5M2 variant (NVFP4 is a - single format family — E2M1 only), so grad roles reuse the same - NVFP4 quantizer.""" +def _make_nvfp4_bare(): + """Bare NVFP4Quantizer (1D, no RHT/SR/2D), used by the cross-format recipe + to avoid cross-operand RHT-consistency concerns in the mixed MXFP8/NVFP4 + GEMMs.""" return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) +def _make_nvfp4_quantizer(role=None): + """Role-based NVFP4Quantizer mirroring ``nvfp4_quantizer_factory`` / + :class:`NVFP4BlockScaling`, but with 2D quantization disabled. + + Per role: weight = no RHT/SR, input = RHT, grad = RHT + SR. + + ``with_2d_quantization`` is forced off everywhere: the 2D quantize-transpose + kernel has no columnwise-only path, so a hybrid columnwise sub-quantizer + cannot use it. + TODO(negvet): enable 2D for the rowwise direction once + https://github.com/NVIDIA/TransformerEngine/pull/3027 lands. + """ + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + is_weight = is_linear and role.tensor_type == "weight" + is_grad = is_linear and role.tensor_type in ("grad_output", "grad_input") + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + with_rht=not is_weight, + with_post_rht_amax=not is_weight, + with_2d_quantization=False, # TODO(negvet): enable via PR #3027 + stochastic_rounding=is_grad, + with_random_sign_mask=True, + ) + + def _hybrid_nvfp4_qfactory(role): is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight", "output"): + # Same per-role config for both directions (RHT/SR are per role). return HybridQuantizer( - rowwise_quantizer=_make_nvfp4_quantizer(), - columnwise_quantizer=_make_nvfp4_quantizer(), + rowwise_quantizer=_make_nvfp4_quantizer(role), + columnwise_quantizer=_make_nvfp4_quantizer(role), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_nvfp4_quantizer() - return _make_nvfp4_quantizer() + return _make_nvfp4_quantizer(role) + return _make_nvfp4_quantizer(role) + + +def _hybrid_mxfp8_nvfp4_qfactory(role): + """Cross-format: MXFP8 forward (rowwise) + NVFP4 backward (columnwise). + + fprop TN: weight.row(MXFP8) x input.row(MXFP8) -> MXFP8 x MXFP8 + dgrad NN: weight.col(NVFP4) x grad_output.row(NVFP4) -> NVFP4 x NVFP4 + wgrad NT: input.col(NVFP4) x grad_output.col(NVFP4) -> NVFP4 x NVFP4 + + So weight/input = Hybrid(row=MXFP8, col=NVFP4), grad_output = plain NVFP4. + The forward all-gather consumes the MXFP8 rowwise sub-storage and the + backward all-gather the NVFP4 columnwise one -- the fwd-vs-bwd format + asymmetry that same-format recipes cannot surface. + """ + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _make_nvfp4_bare() + # input / weight / output / unknown-None: MXFP8 rowwise + NVFP4 columnwise. + return HybridQuantizer( + rowwise_quantizer=_make_mxfp8_quantizer(), + columnwise_quantizer=_make_nvfp4_bare(), + ) def hybrid_recipe(): @@ -153,33 +194,37 @@ def hybrid_recipe(): return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_qfactory) if QUANTIZATION == "hybrid_nvfp4": return te_recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) + if QUANTIZATION == "hybrid_mxfp8_nvfp4": + return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_nvfp4_qfactory) raise ValueError(f"Unknown hybrid QUANTIZATION={QUANTIZATION!r}") # ── Tolerances ─────────────────────────────────────────────────────── # -# Upstream ``run_numerics.py::_get_tolerances`` uses (0.4, 0.25) for -# fp8_cs (loose because of sequence parallel & amax reduction) and -# (0.125, 0.0625) for other FP8 recipes. Hybrid with same-format -# sub-quantizers should inherit the underlying format's distributed -# behaviour — with slightly looser bounds to absorb the two-pass -# quantization (rowwise and columnwise quantizers run independently, so -# their outputs may differ by ~1 ULP from a single fused-quantize path -# in edge cases). +# These match upstream ``run_numerics.py::_get_tolerances`` exactly, for +# the same TP/SP-distributed-vs-single-node comparison. A same-format +# hybrid inherits the underlying format's distributed behaviour: both the +# distributed and single-node models run the *same* two-pass hybrid recipe, +# so the two-pass quantization cancels in the comparison and the only +# remaining difference is the TP/SP path (per-shard quantization, +# all-gather/reduce-scatter order, and -- for fp8_cs only -- cross-rank +# amax reduction). There is therefore no reason for hybrid to need looser +# bounds than the vanilla format. def _get_tolerances(): if QUANTIZATION == "hybrid_fp8": + # Loose because of sequence parallel & amax reduction (fp8_cs). return {"rtol": 0.4, "atol": 0.25} if QUANTIZATION == "hybrid_mxfp8": - return {"rtol": 0.2, "atol": 0.1} + return {"rtol": 0.125, "atol": 0.0625} if QUANTIZATION == "hybrid_nvfp4": - # Upstream ``run_numerics.py`` uses (0.125, 0.12) for vanilla - # NVFP4 with an open TODO to investigate why the tolerance is so - # large. Hybrid NVFP4 runs the same block-scaled kernel in each - # direction independently; bump atol modestly to absorb the - # two-pass asymmetry without hiding a real regression. - return {"rtol": 0.2, "atol": 0.15} + # Upstream ``run_numerics.py`` uses (0.125, 0.12) for vanilla NVFP4 + # (with an open TODO to investigate why the tolerance is so large). + return {"rtol": 0.125, "atol": 0.12} + if QUANTIZATION == "hybrid_mxfp8_nvfp4": + # Backward GEMMs run in NVFP4 -> inherit the (looser) NVFP4 bounds. + return {"rtol": 0.125, "atol": 0.12} raise ValueError(f"No tolerances for QUANTIZATION={QUANTIZATION!r}") @@ -294,10 +339,10 @@ def _loss_backward(out_single, out_dist): # ── Test 1: te.Linear TP (column + row) × SP (on/off) ──────────────── -def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): +def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16, amax_stress=False): dist_print( f"linear: parallel_mode={parallel_mode} sequence_parallel={sequence_parallel}" - f" dtype={params_dtype}" + f" dtype={params_dtype} amax_stress={amax_stress}" ) torch.manual_seed(12345) @@ -326,6 +371,11 @@ def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): # SP column: input is sharded along batch/sequence dim 0. inp_single = torch.empty((WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if amax_stress and WORLD_RANK == WORLD_SIZE - 1: + # Large element on one rank: its local amax diverges from the + # global one. Hybrid gathers the SP activation whole before + # quantizing, so the output must still match the single-node ref. + inp_dist[-1, -1] = 1.0e3 inp_single = _gather(inp_dist, dim=0).detach() else: inp_dist = inp_single.clone() @@ -341,7 +391,11 @@ def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): out_dist = _gather(out_dist, dim=gather_dim) _loss_backward(out_single, out_dist) - _check_outputs(out_single, out_dist, label=f"linear[{parallel_mode},sp={sequence_parallel}]") + _check_outputs( + out_single, + out_dist, + label=f"linear[{parallel_mode},sp={sequence_parallel},amax_stress={amax_stress}]", + ) # Gradient check is only well-defined in these configurations (the # others need cross-rank synchronization that the test doesn't @@ -355,6 +409,146 @@ def test_linear(): for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: _test_linear(parallel_mode, sequence_parallel) + # Amax corner-case (current scaling only): a large element on one rank makes + # its local amax diverge from the global one. Hybrid gathers SP activations in + # high precision before quantizing, so the SP output must still match + # single-node -- guards against a future regression to quantize-then-gather + # without cross-rank amax reduction. + if QUANTIZATION == "hybrid_fp8": + _test_linear("column", True, amax_stress=True) + + +# ── Test 1b: te.Linear hybrid-vs-vanilla bitwise operand equivalence ─ + + +def vanilla_recipe(): + """Built-in single-format recipe matching the same-format hybrid recipe + for the bitwise ``_test_linear_vs_vanilla`` check: FP8 current scaling and + MXFP8 use their defaults (E4M3 fwd / E5M2 bwd, and E4M3 everywhere); NVFP4 + uses the full recipe with 2D disabled to match the role-based 1D + ``_make_nvfp4_quantizer``.""" + if QUANTIZATION == "hybrid_fp8": + return te_recipe.Float8CurrentScaling() + if QUANTIZATION == "hybrid_mxfp8": + return te_recipe.MXFP8BlockScaling() + if QUANTIZATION == "hybrid_nvfp4": + return te_recipe.NVFP4BlockScaling(disable_2d_quantization=True) + raise ValueError(f"No vanilla recipe for QUANTIZATION={QUANTIZATION!r}") + + +def _backward_not_bitwise_comparable(): + """Whether the recipe's backward can't be compared bitwise to vanilla. + + True only for NVFP4's full recipe, which combines RHT with stochastic + rounding. That pair triggers NVFP4's separate columnwise RNG state + (``need_separate_columnwise_rng`` in the cast backend), and the hybrid + (two-pass) vs vanilla (fused) executions then consume that columnwise + random stream differently. Verified by isolation: neither RHT nor SR alone + diverges -- only the combination, and only on the columnwise gradient + (wgrad); the rowwise gradient (dgrad) stays bitwise. + """ + return QUANTIZATION == "hybrid_nvfp4" + + +def _check_bitwise(actual, expected, label): + """Assert bitwise equality (rtol=0, atol=0), all-reduced across ranks.""" + failed = torch.tensor([0], dtype=torch.uint8, device="cuda") + try: + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + except AssertionError as exc: + dist_print(f"{label}: {exc}", src=WORLD_RANK, error=True) + failed[0] = 1 + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: not bitwise-identical on at least one rank" + + +def _test_linear_vs_vanilla(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): + """Same-format hybrid must match its built-in vanilla recipe **bitwise** + through the *same* TP/SP-distributed ``te.Linear`` (forward in all configs; + backward in the non-SP, non-SR configs). + + Unlike ``_test_linear`` (distributed vs single-node, inherently loose), + this compares hybrid vs vanilla at the same topology, so any non-bitwise + difference is a real hybrid-plumbing bug. Complements the FSDP2 parity test + in ``run_fsdp2_fused_adam.py`` by locking the TP/SP comm path. + """ + dist_print( + f"linear_vs_vanilla: parallel_mode={parallel_mode} sequence_parallel={sequence_parallel}" + ) + + def run(recipe): + # Fresh model per recipe (re-seeded for identical weights): TE caches a + # quantized weight workspace on the module, so reusing one model would + # let the first recipe's cached weight contaminate the second. + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + model = te.Linear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode=parallel_mode, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + torch.manual_seed(34567) + torch.cuda.manual_seed(34567) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + if parallel_mode == "row": + split = HIDDEN_SIZE // WORLD_SIZE + inp = inp[:, WORLD_RANK * split : (WORLD_RANK + 1) * split].clone() + inp.requires_grad_() + + with te.autocast(enabled=True, recipe=recipe): + out = model(inp) + # Fixed, recipe-independent target so both backward graphs match. + torch.manual_seed(54321) + torch.cuda.manual_seed(54321) + target = torch.randn_like(out) + LOSS_FN(out, target).backward() + weight_grads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return out.detach().clone(), inp.grad.detach().clone(), weight_grads + + out_h, dinp_h, wgrads_h = run(hybrid_recipe()) + out_v, dinp_v, wgrads_v = run(vanilla_recipe()) + + tag = f"linear_vs_vanilla[{parallel_mode},sp={sequence_parallel}]" + + # Forward is bitwise-identical in every config (the fprop operand- + # equivalence check: hybrid weight.rowwise/input.rowwise == vanilla). + _check_bitwise(out_h, out_v, f"{tag} forward") + + # Backward is bitwise only without SP and without stochastic rounding; + # both are within training tolerance and covered by the loose + # distributed-vs-single-node check: + # * Under SP, hybrid has no native quantized all-gather, so it routes + # through the BF16 dequant/requant fallback while vanilla gathers native + # per-shard bytes. For per-tensor-scaled formats (FP8 current, NVFP4) the + # requantized scale then differs; MXFP8 (per-block only) is immune. + # TODO(negvet): extend to SP once native hybrid AG lands (tracked in + # HybridQuantizer.supports_only_rowwise_all_gather). + # * NVFP4's full recipe combines RHT + stochastic rounding, which triggers + # a separate columnwise RNG (need_separate_columnwise_rng); the hybrid + # two-pass and vanilla fused paths then consume that columnwise random + # stream differently, so the columnwise gradient (wgrad) rounds + # differently. Neither RHT nor SR alone diverges -- only the pair. + if not sequence_parallel and not _backward_not_bitwise_comparable(): + _check_bitwise(dinp_h, dinp_v, f"{tag} dgrad") + assert len(wgrads_h) == len(wgrads_v), f"{tag}: weight-grad count mismatch" + for i, (gh, gv) in enumerate(zip(wgrads_h, wgrads_v)): + _check_bitwise(gh, gv, f"{tag} wgrad[{i}]") + + +def test_linear_vs_vanilla(): + # Cross-format hybrid has no single built-in vanilla recipe to compare + # against bitwise; it is covered by the distributed-vs-single-node checks. + if QUANTIZATION == "hybrid_mxfp8_nvfp4": + dist_print("linear_vs_vanilla: skipped for cross-format hybrid (no vanilla equivalent)") + return + for parallel_mode in ["column", "row"]: + for sequence_parallel in [False, True]: + _test_linear_vs_vanilla(parallel_mode, sequence_parallel) # ── Test 2: te.LayerNormLinear column + SP ────────────────────────── @@ -404,7 +598,62 @@ def test_layernorm_linear(): _test_layernorm_linear(sequence_parallel) -# ── Test 3: te.TransformerLayer + TP + SP ─────────────────────────── +# ── Test 3: te.LayerNormMLP + TP + SP ─────────────────────────────── + + +def _test_layernorm_mlp(sequence_parallel, params_dtype=torch.bfloat16): + """``te.LayerNormMLP`` with ``set_parallel_mode=True`` and optional SP: + column-parallel FC1 → activation → row-parallel FC2. Isolates the FC1 + hybrid unfused-norm path and the row-parallel FC2 + SP reduce-scatter, + otherwise only exercised transitively inside ``_test_transformer_layer``. + """ + dist_print(f"layernorm_mlp: parallel_mode=set sequence_parallel={sequence_parallel}") + + torch.manual_seed(45678) + torch.cuda.manual_seed(45678) + + model_single = te.LayerNormMLP(HIDDEN_SIZE, FFN_HIDDEN_SIZE, params_dtype=params_dtype).cuda() + model_dist = te.LayerNormMLP( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + + _copy_params(model_dist, model_single) + + if sequence_parallel: + inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_single = _gather(inp_dist, dim=0).detach() + else: + inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp_dist = inp_single.clone() + + out_single, out_dist = _apply_models(model_single, model_dist, inp_single, inp_dist) + + # Row-parallel FC2 output is in the full hidden space; with SP it is + # reduce-scattered along the token dim 0, so gather it back. + if sequence_parallel: + out_dist = _gather(out_dist, dim=0) + + _loss_backward(out_single, out_dist) + _check_outputs(out_single, out_dist, label=f"layernorm_mlp[sp={sequence_parallel}]") + + # Without SP, grads align with the single-node ref (SP needs cross-rank + # grad sync the test doesn't do -- matches run_numerics.py). + if not sequence_parallel: + _check_gradients(model_dist, model_single) + + +def test_layernorm_mlp(): + for sequence_parallel in [False, True]: + _test_layernorm_mlp(sequence_parallel) + + +# ── Test 4: te.TransformerLayer + TP + SP ─────────────────────────── def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): @@ -460,6 +709,11 @@ def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): _loss_backward(out_single, out_dist) _check_outputs(out_single, out_dist, label=f"transformer_layer[sp={sequence_parallel}]") + # Without SP, verify the integration path at the gradient level too (SP + # needs cross-rank grad sync the test doesn't do -- matches run_numerics.py). + if not sequence_parallel: + _check_gradients(model_dist, model_single) + def test_transformer_layer(): for sequence_parallel in [False, True]: @@ -496,13 +750,20 @@ def main(argv=None): "--quantization", type=str, required=True, - choices=["hybrid_fp8", "hybrid_mxfp8", "hybrid_nvfp4"], + choices=["hybrid_fp8", "hybrid_mxfp8", "hybrid_nvfp4", "hybrid_mxfp8_nvfp4"], ) parser.add_argument( "--test", type=str, default="all", - choices=["all", "linear", "layernorm_linear", "transformer_layer"], + choices=[ + "all", + "linear", + "linear_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", + ], help="Run only the named test (speeds up iterative debugging)", ) args = parser.parse_args(argv) @@ -510,7 +771,9 @@ def main(argv=None): test_map = { "linear": test_linear, + "linear_vs_vanilla": test_linear_vs_vanilla, "layernorm_linear": test_layernorm_linear, + "layernorm_mlp": test_layernorm_mlp, "transformer_layer": test_transformer_layer, } if args.test == "all": diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py index 87e21255eb2..a2ebbe6480c 100644 --- a/tests/pytorch/distributed/test_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -73,6 +73,15 @@ def test_hybrid_fp8_linear(): _run_test("hybrid_fp8", "linear") +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_hybrid_fp8_linear_vs_vanilla(): + """Bitwise operand equivalence: same-format hybrid FP8 must match the + built-in ``Float8CurrentScaling`` recipe through the same TP ``te.Linear`` + (forward in all configs; backward in the non-SP configs). Locks the TP + comm path that the FSDP2 parity test does not exercise.""" + _run_test("hybrid_fp8", "linear_vs_vanilla") + + @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_layernorm_linear(): """Column-parallel ``te.LayerNormLinear`` with and without SP. @@ -82,6 +91,15 @@ def test_hybrid_fp8_layernorm_linear(): _run_test("hybrid_fp8", "layernorm_linear") +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_hybrid_fp8_layernorm_mlp(): + """Standalone ``te.LayerNormMLP`` (column FC1 / row FC2) with and + without SP under hybrid FP8. Isolates the MLP block's unfused-norm + and row-parallel reduce-scatter paths, and checks gradients in the + no-SP case.""" + _run_test("hybrid_fp8", "layernorm_mlp") + + @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_transformer_layer(): """Full ``te.TransformerLayer`` with ``set_parallel_mode=True`` and @@ -107,11 +125,21 @@ def test_hybrid_mxfp8_linear(): _run_test("hybrid_mxfp8", "linear") +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8_linear_vs_vanilla(): + _run_test("hybrid_mxfp8", "linear_vs_vanilla") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") def test_hybrid_mxfp8_layernorm_linear(): _run_test("hybrid_mxfp8", "layernorm_linear") +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8_layernorm_mlp(): + _run_test("hybrid_mxfp8", "layernorm_mlp") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") def test_hybrid_mxfp8_transformer_layer(): _run_test("hybrid_mxfp8", "transformer_layer") @@ -143,11 +171,64 @@ def test_hybrid_nvfp4_linear(): _run_test("hybrid_nvfp4", "linear") +@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") +def test_hybrid_nvfp4_linear_vs_vanilla(): + _run_test("hybrid_nvfp4", "linear_vs_vanilla") + + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") def test_hybrid_nvfp4_layernorm_linear(): _run_test("hybrid_nvfp4", "layernorm_linear") +@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") +def test_hybrid_nvfp4_layernorm_mlp(): + _run_test("hybrid_nvfp4", "layernorm_mlp") + + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") def test_hybrid_nvfp4_transformer_layer(): _run_test("hybrid_nvfp4", "transformer_layer") + + +# ────────────────────────────────────────────────────────────────────── +# Cross-format hybrid: MXFP8 forward (rowwise) + NVFP4 backward (columnwise) +# ────────────────────────────────────────────────────────────────────── +# +# Forward and backward all-gather *different* formats (MXFP8 rowwise vs NVFP4 +# columnwise) -- the asymmetry same-format recipes can't surface. Only the +# distributed-vs-single-node checks run (no single vanilla recipe to match a +# cross-format hybrid bitwise). Needs both MXFP8 and NVFP4 hardware support. + +_cross_format_available = mxfp8_available and nvfp4_available +_reason_for_no_cross_format = ( + reason_for_no_mxfp8 if not mxfp8_available else reason_for_no_nvfp4 +) + + +@pytest.mark.skipif( + not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" +) +def test_hybrid_mxfp8_nvfp4_linear(): + _run_test("hybrid_mxfp8_nvfp4", "linear") + + +@pytest.mark.skipif( + not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" +) +def test_hybrid_mxfp8_nvfp4_layernorm_linear(): + _run_test("hybrid_mxfp8_nvfp4", "layernorm_linear") + + +@pytest.mark.skipif( + not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" +) +def test_hybrid_mxfp8_nvfp4_layernorm_mlp(): + _run_test("hybrid_mxfp8_nvfp4", "layernorm_mlp") + + +@pytest.mark.skipif( + not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" +) +def test_hybrid_mxfp8_nvfp4_transformer_layer(): + _run_test("hybrid_mxfp8_nvfp4", "transformer_layer") diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index bfca2dfaa4b..ccbc1041df8 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -3536,7 +3536,10 @@ def test_training_loop_loss_decreases(self): optimizer.step() - assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" def test_training_loop_params_remain_quantized(self): """Params should remain HybridQuantizedTensors after training.""" @@ -3680,7 +3683,10 @@ def test_mixed_format_training_loop(self): loss.backward() optimizer.step() - assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" + # Strictly monotonic decrease + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Loss not strictly decreasing each step: {losses}" for name, p in model.named_parameters(): if "bias" not in name: assert isinstance(p, HybridQuantizedTensor), f"{name} is {type(p).__name__}" diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 24055b35e4a..6d947d2d050 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1653,8 +1653,11 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: raise RuntimeError("Weight quantizer has not been initialized") quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) quantizer.internal = False + # HybridQuantizer is included so its current-scaling / NVFP4 + # sub-quantizers get the same cross-shard amax reduction as the + # vanilla path (no-op for block-scaled sub-quantizers like MXFP8). if is_dtensor and isinstance( - quantizer, (Float8CurrentScalingQuantizer, NVFP4Quantizer) + quantizer, (Float8CurrentScalingQuantizer, NVFP4Quantizer, HybridQuantizer) ): device_mesh = dtensor_param.device_mesh amax_reduction_group = ( diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 297efcf7ead..4ac362aa78c 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1506,6 +1506,15 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: self._customize_quantizers_float8_current_scaling(fwd, recipe) elif recipe.nvfp4(): self._customize_quantizers_nvfp4(fwd, recipe) + # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP + # activations are gathered in high precision and re-quantized whole, so + # every rank already sees the same global amax. + # TODO(negvet): once native quantized all-gather lands (see + # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path + # quantizes per-shard, needing a hybrid branch here that mirrors the + # current-scaling / NVFP4 SP reduction above: + # elif recipe.custom(): + # ... # enable SP amax reduction on the hybrid input/grad quantizer def get_quantizer_roles( self, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 3111d63bf45..f0451351118 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2106,6 +2106,15 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: self._customize_quantizers_float8_current_scaling(fwd, recipe) elif recipe.nvfp4(): self._customize_quantizers_nvfp4(fwd, recipe) + # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP + # activations are gathered in high precision and re-quantized whole, so + # every rank already sees the same global amax. + # TODO(negvet): once native quantized all-gather lands (see + # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path + # quantizes per-shard, needing a hybrid branch here that mirrors the + # current-scaling / NVFP4 SP reduction above: + # elif recipe.custom(): + # ... # enable SP amax reduction on the hybrid input/grad quantizer def get_quantizer_roles( self, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index dcbb9eaf930..a7969132f40 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1694,6 +1694,15 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: self._customize_quantizers_float8_current_scaling(fwd, recipe) elif recipe.nvfp4(): self._customize_quantizers_nvfp4(fwd, recipe) + # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP + # activations are gathered in high precision and re-quantized whole, so + # every rank already sees the same global amax. + # TODO(negvet): once native quantized all-gather lands (see + # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path + # quantizes per-shard, needing a hybrid branch here that mirrors the + # current-scaling / NVFP4 SP reduction above: + # elif recipe.custom(): + # ... # enable SP amax reduction on the hybrid input/grad quantizer def reset_parameters(self, defer_init=False): super().reset_parameters(defer_init=defer_init) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 5e8a1f87832..f4dbb9c9eca 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -48,6 +48,36 @@ def __init__( self.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) self.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + @property + def with_amax_reduction(self) -> bool: + """Whether either sub-quantizer has cross-rank amax reduction enabled.""" + return getattr(self.rowwise_quantizer, "with_amax_reduction", False) or getattr( + self.columnwise_quantizer, "with_amax_reduction", False + ) + + @with_amax_reduction.setter + def with_amax_reduction(self, value: bool) -> None: + # Set on the HybridQuantizer by module / FSDP2 code, but read by the C++ + # kernel off the sub-quantizer that runs -- hence forwarded, not stored here. + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + if hasattr(sub, "with_amax_reduction"): + sub.with_amax_reduction = value + + @property + def amax_reduction_group(self): + """Amax-reduction group of the sub-quantizers, or ``None`` if unset.""" + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + group = getattr(sub, "amax_reduction_group", None) + if group is not None: + return group + return None + + @amax_reduction_group.setter + def amax_reduction_group(self, value) -> None: + for sub in (self.rowwise_quantizer, self.columnwise_quantizer): + if hasattr(sub, "amax_reduction_group"): + sub.amax_reduction_group = value + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: # Gate each sub-quantizer call on the parent usage flag. Sub-quantizers # are pinned to one direction in ``__init__``; the parent flag decides @@ -441,6 +471,14 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m tuple; ``fsdp_post_all_gather`` would slice the gathered flat buffer using those offsets. """ + # Mirror ``Float8Tensor.fsdp_pre_all_gather``: enable cross-shard amax + # reduction so the post-optimizer re-quantization of the sharded weight + # keeps one shared scale across shards (no-op for sub-quantizers without + # amax reduction, e.g. MXFP8). + if mesh is not None and self._quantizer is not None: + self._quantizer.amax_reduction_group = mesh.get_group() + self._quantizer.with_amax_reduction = True + # Quick, targeted error for sub-storages whose FSDP2 support isn't # implemented yet (e.g. NVFP4). Without this, users hit # NotImplementedError from deep inside fsdp_extract_buffers with a From 5892a749f8ddf491d3709fb9a06f8fed4a9fba3c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 3 Jun 2026 13:03:33 +0000 Subject: [PATCH 17/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 29 +++++-------------- .../fsdp2_tests/run_fsdp2_model.py | 12 ++++---- .../pytorch/distributed/test_hybrid_tp_sp.py | 4 +-- 3 files changed, 15 insertions(+), 30 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 591bd643867..3b131fded3c 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1388,11 +1388,7 @@ def run_training(build_fn, recipe_for_autocast): base_first, rtol=0.0, atol=0.0, - msg=lambda m: ( - f"[{hybrid_recipe_name} vs {base_recipe_name}] first forward output not" - f" bitwise-identical (a same-format hybrid must match its vanilla recipe" - f" before any optimizer step): {m}" - ), + msg=lambda m: f"[{hybrid_recipe_name} vs {base_recipe_name}] first forward output not bitwise-identical (a same-format hybrid must match its vanilla recipe before any optimizer step): {m}", ) # (2) Every per-step loss: bitwise-identical across the whole optimizer loop. @@ -1402,10 +1398,7 @@ def run_training(build_fn, recipe_for_autocast): b_loss, rtol=0.0, atol=0.0, - msg=lambda m, s=step: ( - f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} loss not" - f" bitwise-identical to the vanilla recipe: {m}" - ), + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} loss not bitwise-identical to the vanilla recipe: {m}", ) # (3) Backward: every weight-gradient shard at every step bitwise-identical @@ -1423,10 +1416,7 @@ def run_training(build_fn, recipe_for_autocast): b_grad, rtol=0.0, atol=0.0, - msg=lambda m, s=step, i=i: ( - f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i}" - f" gradient not bitwise-identical to the vanilla recipe: {m}" - ), + msg=lambda m, s=step, i=i: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i} gradient not bitwise-identical to the vanilla recipe: {m}", ) @@ -1471,10 +1461,7 @@ def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): gathered[0], rtol=0.0, atol=0.0, - msg=lambda m, n=name, r=r: ( - f"{n}: rank {r} rowwise _scale_inv differs from rank 0 -- cross-shard " - f"amax reduction was not applied to the hybrid current-scaling weight: {m}" - ), + msg=lambda m, n=name, r=r: f"{n}: rank {r} rowwise _scale_inv differs from rank 0 -- cross-shard amax reduction was not applied to the hybrid current-scaling weight: {m}", ) checked += 1 assert checked > 0, "no hybrid current-scaling weights found to check" @@ -1572,13 +1559,13 @@ def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): assert per_rank_out % 32 == 0, ( f"Test setup error: per_rank_out={per_rank_out} (= out_features / world_size, " f"world_size={world_size}) must be a multiple of the MXFP8 block size (32) so the " - f"sharded weight's data stays block-aligned. Pick a per_rank_out divisible by 32." + "sharded weight's data stays block-aligned. Pick a per_rank_out divisible by 32." ) assert per_rank_out % 128 != 0, ( f"Test setup error: per_rank_out={per_rank_out} must NOT be a multiple of 128, or the " - f"rowwise scale-inv needs no alignment padding and this test stops exercising the MXFP8 " - f"unpad-before-gather / pad-after-gather path it exists to cover. Pick a per_rank_out " - f"divisible by 32 but not 128 (e.g. 96)." + "rowwise scale-inv needs no alignment padding and this test stops exercising the MXFP8 " + "unpad-before-gather / pad-after-gather path it exists to cover. Pick a per_rank_out " + "divisible by 32 but not 128 (e.g. 96)." ) for recipe_name in ("HybridMXFP8", "HybridMixed_MXFP8_FP8"): diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 6b3c2c2d2f0..85d339540c2 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -534,9 +534,9 @@ def _hybrid_param_count(): # Params must stay HybridQuantizedTensor after the optimizer step -- guards a # silent dequantize-to-bf16 through the FSDP2 / optimizer path. - assert _hybrid_param_count() == hybrid_count, ( - "HybridQuantizedTensor params lost their quantized type after optimizer.step()" - ) + assert ( + _hybrid_param_count() == hybrid_count + ), "HybridQuantizedTensor params lost their quantized type after optimizer.step()" # FSDP2 quantized all-gather must match a manual fp32 dequant-then-allgather. _check_hybrid_fsdp2_allgather(model) @@ -619,9 +619,9 @@ def _hybrid_param_count(): # Params must survive the split/as_strided/slice reshard dispatch ops with # their hybrid quantized type intact. - assert _hybrid_param_count() == hybrid_count, ( - "HybridQuantizedTensor params lost their quantized type after optimizer.step()" - ) + assert ( + _hybrid_param_count() == hybrid_count + ), "HybridQuantizedTensor params lost their quantized type after optimizer.step()" if __name__ == "__main__": diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py index a2ebbe6480c..711e76063eb 100644 --- a/tests/pytorch/distributed/test_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -201,9 +201,7 @@ def test_hybrid_nvfp4_transformer_layer(): # cross-format hybrid bitwise). Needs both MXFP8 and NVFP4 hardware support. _cross_format_available = mxfp8_available and nvfp4_available -_reason_for_no_cross_format = ( - reason_for_no_mxfp8 if not mxfp8_available else reason_for_no_nvfp4 -) +_reason_for_no_cross_format = reason_for_no_mxfp8 if not mxfp8_available else reason_for_no_nvfp4 @pytest.mark.skipif( From ec7b84cb01e2e04ee7a044098839720a3087c202 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 5 Jun 2026 09:22:51 +0000 Subject: [PATCH 18/97] Enable FSDP2 hybrid protocol for Float8Block tensor Signed-off-by: Evgeny --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 36 ----- .../fsdp2_tests/run_fsdp2_mem_leak.py | 12 -- .../fsdp2_tests/run_fsdp2_model.py | 12 -- .../pytorch/tensor/hybrid_tensor.py | 13 +- .../float8_blockwise_tensor_storage.py | 125 +++++++++++++++++- .../tensor/storage/hybrid_tensor_storage.py | 21 ++- 6 files changed, 156 insertions(+), 63 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 3b131fded3c..d90dad041f8 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1091,13 +1091,6 @@ def test_fused_adam_hybrid_master_weights(hybrid_recipe_name): - Optimizer states are FP32 - Loss decreases over training steps """ - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1) in reset_sharded_param. " - "Same root cause as vanilla Float8BlockScaling + quantized_model_init." - ) - from transformer_engine.pytorch import HybridQuantizedTensor from fsdp2_utils import get_hybrid_recipe_from_string @@ -1163,12 +1156,6 @@ def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name): all-gather hooks are schedule-invariant across both FSDP2 passes, and regression-guards the future P1.1 buffer-split bandwidth optimization. """ - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1) in reset_sharded_param." - ) - from fsdp2_utils import get_hybrid_recipe_from_string hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) @@ -1224,12 +1211,6 @@ def test_fused_adam_hybrid_bf16_vs_hybrid_parity(hybrid_recipe_name): This is a sanity check that hybrid quantized training converges similarly to BF16 training, not a bitwise-exact comparison. """ - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1)." - ) - from fsdp2_utils import get_hybrid_recipe_from_string hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) @@ -1312,11 +1293,6 @@ def test_fused_adam_hybrid_vs_base_recipe_parity(hybrid_recipe_name): fix. Uses a bare ``te.Linear`` stack (see ``_build_linear_parity_stack``) to isolate GEMM-operand quantization. """ - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1) in reset_sharded_param." - ) if hybrid_recipe_name not in _HYBRID_TO_BASE_RECIPE: pytest.skip( f"{hybrid_recipe_name} is cross-format; no single-format vanilla " @@ -1478,12 +1454,6 @@ def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): bytes is therefore bitwise-identical to concatenating the dequantized shards — the tolerance is effectively ``assert_equal``. """ - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1)." - ) - from transformer_engine.pytorch import HybridQuantizedTensor from fsdp2_utils import get_hybrid_recipe_from_string @@ -1624,12 +1594,6 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): """ import torch.distributed.checkpoint as dcp - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1)." - ) - if hybrid_recipe_name == "HybridFP8CurrentScaling": # TODO: preserve hybrid current-scaling primary-weight scales across DCP # by implementing __tensor_flatten__/__tensor_unflatten__ on the quantized diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py index 5aa09ae6c46..e943d48d4ab 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py @@ -486,12 +486,6 @@ def test_hybrid_no_excess_forward_memory(hybrid_recipe_name): Same methodology as test_fp8_temp_accumulation_across_layers but for hybrid quantized tensors. """ - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1)." - ) - from fsdp2_utils import get_hybrid_recipe_from_string hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) @@ -555,12 +549,6 @@ def test_hybrid_no_excess_forward_memory(hybrid_recipe_name): def test_hybrid_transpose_cache_after_backward(hybrid_recipe_name): """Detect transpose caches from hybrid sub-storages persisting after backward.""" - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1)." - ) - from fsdp2_utils import get_hybrid_recipe_from_string hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 85d339540c2..4624c995f1f 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -461,12 +461,6 @@ def test_distributed_hybrid(hybrid_recipe_name): - params keep their hybrid quantized type across optimizer.step(), - FSDP2's quantized all-gather matches a manual fp32 dequant-then-allgather. """ - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1)." - ) - from fsdp2_utils import get_hybrid_recipe_from_string hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) @@ -549,12 +543,6 @@ def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): from FSDP2. This exercises the forward-reshard-backward-reshard cycle where split/as_strided/slice dispatch ops fire every iteration. """ - if hybrid_recipe_name == "HybridFloat8BlockScaling": - pytest.xfail( - "HybridFloat8BlockScaling: Float8BlockwiseQTensor sub-storage loses " - "quantized type through FSDP2 view(-1)." - ) - from fsdp2_utils import get_hybrid_recipe_from_string hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index f4dbb9c9eca..bb3adaedb84 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -496,7 +496,6 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m "Hybrid FSDP2 all-gather is not supported for a " f"{type(sub).__name__} {role} sub-storage: it does not " "implement fsdp_buffer_fields. " - "See hybrid_quantization_fsdp.md section 9 (Gap 5) — " "NVFP4 sub-storages need packed-FP4 dim-0 alignment, " "columnwise dequantization and RHT-cache handling before " "they can be gathered. Use a supported sub-quantizer " @@ -677,6 +676,18 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.view.default: tensor = args[0] shape = args[1] + # Identity view fast-path (FSDP2 reset_sharded_param issues a view + # to the param's own shape). The columnwise sub-storage's own shape + # is transposed relative to the hybrid for some formats (e.g. a 2D + # block-scaled Float8BlockwiseQTensor has shape (N, M) for an + # (M, N) weight). Forwarding the hybrid's row-major shape to it + # would be a spurious last-2-dims change, which 2D block scaling + # cannot represent and so dequantizes the sub-storage to a plain + # tensor -- breaking the FSDP2 sub-storage protocol later. Preserve + # the sub-storages as-is, mirroring the as_strided / slice no-op + # fast paths below. + if list(shape) == list(tensor.shape): + return HybridQuantizedTensor.make_like(tensor) row_view = None col_view = None if tensor._rowwise_storage is not None: diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index ca3913762f9..7161771dab0 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -16,7 +16,7 @@ from ...constants import TE_DType_To_Torch -from ...utils import _empty_tensor +from ...utils import _empty_tensor, round_up_to_nearest_multiple class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): @@ -425,3 +425,126 @@ def get_usages(self) -> Dict[str, bool]: "rowwise": self._rowwise_data is not None, "columnwise": self._columnwise_data is not None, } + + # ── FSDP2 sub-storage buffer protocol ──────────────────────────── + # + # Float8Block stores columnwise data N-major (transposed) for the GEMM, so + # it cannot be dim-0 all-gathered directly. Each direction is made + # self-contained: the columnwise direction fp8-transposes its own data to + # M-major for the gather and back on assign, using only its own buffers (no + # dependency on a rowwise sibling, which in a hybrid tensor may be a + # different format). Block-scale GEMM alignment padding (round-up-to-4) is + # stripped before the gather and re-applied after. Only 2D block scaling is + # supported -- the 1D scale layout has M in dim1, incompatible with FSDP2's + # dim-0 all-gather. + + _FSDP_BLOCK_LEN = 128 + + def _fsdp_logical_mn(self) -> Tuple[int, int]: + """Flattened ``(M, N)`` of this sub-storage's logical shape.""" + shape = self.size() + last_dim = shape[-1] if len(shape) > 0 else 1 + leading = 1 + for dim in shape[:-1]: + leading *= dim + return leading, last_dim + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Fields gathered by FSDP2 for Float8 block scaling (2D scaling only).""" + if not self._is_2D_scaled: + raise NotImplementedError( + "FSDP2 for Float8BlockwiseQTensor requires 2D block scaling " + "(block_scaling_dim=2). 1D block scaling is not supported because " + "its scale layout has M in dim1, which is incompatible with FSDP2 " + "dim-0 all-gather." + ) + fields = [] + if self._rowwise_data is not None: + fields.extend(("_rowwise_data", "_rowwise_scale_inv")) + if self._columnwise_data is not None: + fields.extend(("_columnwise_data", "_columnwise_scale_inv")) + return tuple(fields) + + def fsdp_buffer_fields( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract M-major, alignment-stripped buffers for dim-0 all-gather. + + Rowwise data is already M-major; columnwise data is N-major and is + fp8-transposed to M-major here (and transposed back in + :meth:`fsdp_assign_gathered`). The block-scale round-up-to-4 alignment + padding is stripped so dim-0 concatenation across shards is well-defined. + """ + names = self.fsdp_buffer_fields() + block_len = self._FSDP_BLOCK_LEN + m, n = self._fsdp_logical_mn() + m_tiles = (m + block_len - 1) // block_len + last_tiles = (n + block_len - 1) // block_len + + if self._rowwise_data is not None: + # Rowwise scale is (m_tiles, round_up(last_tiles, 4)); m_tiles sits in + # dim-0 (sharded/gathered) unpadded, the round-up padding is on dim-1 + # (not sharded). Strip dim-1 to the compact tile count. + scale = self._rowwise_scale_inv + if scale is not None and scale.size(1) > last_tiles: + scale = scale[:, :last_tiles].contiguous() + buffers = (self._rowwise_data, scale) + direction = "rowwise" + else: + # Columnwise data is N-major (N, M); transpose to M-major (M, N). + col_data = self._columnwise_data + if not col_data.is_contiguous(): + col_data = col_data.contiguous() + data_m = tex.fp8_transpose(col_data, self._fp8_dtype, out=None) + # Columnwise scale is (last_tiles, round_up(m_tiles, 4)); transpose to + # (round_up(m_tiles, 4), last_tiles) and strip dim-0 to m_tiles so the + # gathered (dim-0) axis is the M-tiles, matching the rowwise layout. + scale = self._columnwise_scale_inv.transpose(0, 1).contiguous() + if scale.size(0) > m_tiles: + scale = scale[:m_tiles].contiguous() + buffers = (data_m, scale) + direction = "columnwise" + + return buffers, {"direction": direction, "field_names": names} + + def fsdp_assign_gathered( + self, + gathered: Tuple[Optional[torch.Tensor], ...], + meta: Dict[str, Any], + ) -> None: + """Write gathered buffers back, re-applying transpose + scale padding. + + Inverse of :meth:`fsdp_extract_buffers`: rowwise re-pads the scale's + last-dim alignment; columnwise transposes the M-major gathered data back + to N-major and re-pads/transposes the scale to the GEMM scale layout + produced by ``get_scale_shape(..., columnwise=True)``. + """ + block_len = self._FSDP_BLOCK_LEN + direction = meta["direction"] + data, scale = gathered + + if direction == "rowwise": + last_dim = data.size(-1) + last_tiles = (last_dim + block_len - 1) // block_len + if scale is not None: + pad = round_up_to_nearest_multiple(last_tiles, 4) - last_tiles + if pad > 0: + scale = torch.nn.functional.pad(scale, (0, pad)) + self._rowwise_data = data + self._rowwise_scale_inv = scale + return + + # Columnwise: gathered data is M-major (M_full, N); transpose to N-major. + data_m = data if data.is_contiguous() else data.contiguous() + self._columnwise_data = tex.fp8_transpose(data_m, self._fp8_dtype, out=None) + m_full = 1 + for dim in data.shape[:-1]: + m_full *= dim + m_tiles_full = (m_full + block_len - 1) // block_len + # Gathered scale is compact (m_tiles_full, last_tiles); transpose to + # (last_tiles, m_tiles_full) and re-pad the M-tile dim to multiple of 4. + scale_t = scale.transpose(0, 1).contiguous() + pad = round_up_to_nearest_multiple(m_tiles_full, 4) - m_tiles_full + if pad > 0: + scale_t = torch.nn.functional.pad(scale_t, (0, pad)) + self._columnwise_scale_inv = scale_t.contiguous() diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index d9f58734509..9d152982c72 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -5,6 +5,7 @@ """Mixin class holding data specific for HybridQuantizedTensor""" from __future__ import annotations +from collections.abc import Iterable from typing import Any, Dict, Optional, Tuple import torch @@ -149,7 +150,25 @@ def device(self): raise RuntimeError("HybridQuantizedTensorStorage has no data") def view(self, *shape): - """View delegates to each sub-storage. Used by FSDP2 reset_sharded_param.""" + """View delegates to each sub-storage. Used by FSDP2 reset_sharded_param. + + Identity views are handled without forwarding a reshape to the + sub-storages: the columnwise sub-storage's own shape is transposed + relative to the hybrid for some formats (e.g. a 2D block-scaled + Float8BlockwiseQTensor has shape ``(N, M)`` for an ``(M, N)`` weight), + so forwarding the hybrid's row-major shape would be a spurious + last-2-dims change that dequantizes it to a plain tensor. + """ + flat_shape = shape[0] if len(shape) == 1 and isinstance(shape[0], Iterable) else shape + if list(flat_shape) == list(self.size()): + return HybridQuantizedTensorStorage( + rowwise_storage=self._rowwise_storage, + columnwise_storage=self._columnwise_storage, + rowwise_quantizer=self._rowwise_quantizer, + columnwise_quantizer=self._columnwise_quantizer, + quantizer=self._quantizer, + fake_dtype=self._dtype, + ) row_view = self._rowwise_storage.view(*shape) if self._rowwise_storage is not None else None col_view = ( self._columnwise_storage.view(*shape) if self._columnwise_storage is not None else None From b99277a3150f829a6848db06f554d8861eb391c1 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 9 Jun 2026 15:30:03 +0000 Subject: [PATCH 19/97] Enable Identity (no-op) quantization Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/fsdp2_utils.py | 17 +- .../fsdp2_tests/run_fsdp2_fused_adam.py | 112 +++ tests/pytorch/distributed/run_hybrid_tp_sp.py | 48 +- .../pytorch/distributed/test_hybrid_tp_sp.py | 13 + tests/pytorch/test_identity_quantizer.py | 935 ++++++++++++++++++ transformer_engine/pytorch/__init__.py | 3 + .../pytorch/cpp_extensions/gemm.py | 23 +- .../quantization_factory_examples.py | 36 + transformer_engine/pytorch/module/base.py | 4 +- transformer_engine/pytorch/quantization.py | 11 +- transformer_engine/pytorch/tensor/__init__.py | 7 + .../pytorch/tensor/identity_tensor.py | 322 ++++++ .../tensor/storage/identity_tensor_storage.py | 152 +++ transformer_engine/pytorch/tensor/utils.py | 108 +- 14 files changed, 1755 insertions(+), 36 deletions(-) create mode 100644 tests/pytorch/test_identity_quantizer.py create mode 100644 transformer_engine/pytorch/tensor/identity_tensor.py create mode 100644 transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 7b71cb04c39..f53c237e406 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -5,7 +5,7 @@ """Shared utility functions for FSDP2 distributed tests.""" import transformer_engine.common.recipe -from transformer_engine.pytorch import HybridQuantizer, QuantizedTensor +from transformer_engine.pytorch import HybridQuantizer, IdentityQuantizer, QuantizedTensor from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( current_scaling_quantizer_factory, float8_block_scaling_quantizer_factory, @@ -61,6 +61,19 @@ def _hybrid_mixed_mxfp8_fp8_qfactory(role): return current_scaling_quantizer_factory(role) +def _hybrid_fp8_current_identity_qfactory(role): + """FP8 current forward + high-precision backward via Identity.""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=current_scaling_quantizer_factory(role), + columnwise_quantizer=IdentityQuantizer(), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return current_scaling_quantizer_factory(role) + + # The qfactories above are registered here as module-level functions (not # lambdas or closures) on purpose: DCP serializes ``CustomRecipe`` via # ``pickle``, and closure-based qfactories (or inner functions capturing state) @@ -71,6 +84,7 @@ def _hybrid_mixed_mxfp8_fp8_qfactory(role): "HybridMXFP8": _hybrid_mxfp8_qfactory, "HybridFloat8BlockScaling": _hybrid_float8_block_qfactory, "HybridMixed_MXFP8_FP8": _hybrid_mixed_mxfp8_fp8_qfactory, + "HybridFP8CurrentScalingIdentity": _hybrid_fp8_current_identity_qfactory, } @@ -86,6 +100,7 @@ def get_hybrid_recipe_from_string(recipe): "HybridMXFP8" — MXFP8 for both directions "HybridFloat8BlockScaling" — Float8 block scaling for both directions "HybridMixed_MXFP8_FP8" — MXFP8 rowwise + FP8 current columnwise + "HybridFP8CurrentScalingIdentity" — FP8 current forward + Identity backward """ if recipe not in _HYBRID_QFACTORIES: raise ValueError( diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index d90dad041f8..358ac8757e9 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1443,6 +1443,113 @@ def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): assert checked > 0, "no hybrid current-scaling weights found to check" +def test_fused_adam_hybrid_identity_fp8_master_weights(): + """FSDP2 + FusedAdam with Hybrid(FP8 current rowwise, Identity columnwise). + + Covers the Identity sub-storage in hybrid FSDP2 all-gather while the FP8 + current rowwise direction validates cross-shard amax reduction via scale + uniformity. + """ + from transformer_engine.pytorch import HybridQuantizedTensor + from transformer_engine.pytorch.tensor.storage.identity_tensor_storage import ( + IdentityTensorStorage, + ) + from fsdp2_utils import get_hybrid_recipe_from_string + + world_size, device = _get_dist_info() + if world_size < 2: + pytest.skip("needs >=2 ranks to validate cross-shard amax reduction") + + hybrid_recipe = get_hybrid_recipe_from_string("HybridFP8CurrentScalingIdentity") + model = _build_linear_parity_stack(hybrid_recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + identity_steps = 2 + for step in range(identity_steps): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + if step < identity_steps - 1: + optimizer.step() + + assert all( + losses[i + 1] < losses[i] for i in range(len(losses) - 1) + ), f"Hybrid Identity/FP8 loss not strictly decreasing: {losses}" + + checked_identity = 0 + checked_scale = 0 + for name, param in model.named_parameters(): + if not ( + isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ): + continue + local = param._local_tensor + assert isinstance(local._columnwise_storage, IdentityTensorStorage) + + scale_inv = getattr(local._rowwise_storage, "_scale_inv", None) + if scale_inv is not None: + local_scale = scale_inv.detach().reshape(-1).clone() + gathered_scales = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered_scales, local_scale) + for r in range(1, world_size): + torch.testing.assert_close( + gathered_scales[r], + gathered_scales[0], + rtol=0.0, + atol=0.0, + msg=lambda m, n=name, r=r: f"{n}: rank {r} rowwise _scale_inv differs from rank 0 for Hybrid(FP8Current, Identity): {m}", + ) + checked_scale += 1 + + local_identity = local._columnwise_storage.dequantize().contiguous() + gathered_identity = [torch.zeros_like(local_identity) for _ in range(world_size)] + dist.all_gather(gathered_identity, local_identity) + manual_full = torch.cat(gathered_identity, dim=0) + + sharded_tensors, metadata = local.fsdp_pre_all_gather( + mesh=None, + orig_size=local.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + all_gather_outputs = [] + for shard in sharded_tensors: + gathered = [torch.zeros_like(shard) for _ in range(world_size)] + dist.all_gather(gathered, shard) + all_gather_outputs.append(torch.cat(gathered, dim=0)) + fsdp_full, _ = local.fsdp_post_all_gather( + tuple(all_gather_outputs), metadata, local.dtype, out=None + ) + assert isinstance(fsdp_full, HybridQuantizedTensor) + full_identity = fsdp_full._columnwise_storage.dequantize() + torch.testing.assert_close( + manual_full.float(), + full_identity[: manual_full.shape[0]].float(), + rtol=0.0, + atol=0.0, + msg=lambda m, n=name: f"{n}: Identity columnwise all-gather mismatch: {m}", + ) + checked_identity += 1 + + assert checked_identity > 0, "no Hybrid(FP8Current, Identity) params found" + assert checked_scale > 0, "no FP8 current rowwise scales found to check" + + def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): """Validate that FSDP2 all-gather + post-gather reconstruction produces correct results by comparing ``unshard(param).dequantize()`` with a manual @@ -1711,6 +1818,9 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): "fused_adam_hybrid_scale_uniform_across_shards": ( test_fused_adam_hybrid_scale_uniform_across_shards ), + "fused_adam_hybrid_identity_fp8_master_weights": ( + test_fused_adam_hybrid_identity_fp8_master_weights + ), "fused_adam_hybrid_allgather_correctness": test_fused_adam_hybrid_allgather_correctness, "fused_adam_hybrid_mxfp8_awkward_shard_shape": ( test_fused_adam_hybrid_mxfp8_awkward_shard_shape @@ -1720,6 +1830,7 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): # Hybrid tests that are NOT parametrized by recipe (they sweep internally). _HYBRID_NON_PARAMETRIZED_TESTS = { + "fused_adam_hybrid_identity_fp8_master_weights", "fused_adam_hybrid_mxfp8_awkward_shard_shape", } @@ -1742,6 +1853,7 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): "HybridMXFP8", "HybridFloat8BlockScaling", "HybridMixed_MXFP8_FP8", + "HybridFP8CurrentScalingIdentity", ], ) args = parser.parse_args() diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index 12203ec39c2..1537c949c5c 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -49,6 +49,7 @@ from transformer_engine.pytorch import ( Float8CurrentScalingQuantizer, HybridQuantizer, + IdentityQuantizer, MXFP8Quantizer, NVFP4Quantizer, ) @@ -118,6 +119,30 @@ def _hybrid_mxfp8_qfactory(role): return _make_mxfp8_quantizer() +def _hybrid_fp8_identity_qfactory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_current_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return _make_fp8_current_quantizer() + + +def _hybrid_mxfp8_identity_qfactory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=_make_mxfp8_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return _make_mxfp8_quantizer() + + def _make_nvfp4_bare(): """Bare NVFP4Quantizer (1D, no RHT/SR/2D), used by the cross-format recipe to avoid cross-operand RHT-consistency concerns in the mixed MXFP8/NVFP4 @@ -192,6 +217,10 @@ def hybrid_recipe(): return te_recipe.CustomRecipe(qfactory=_hybrid_fp8_qfactory) if QUANTIZATION == "hybrid_mxfp8": return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_qfactory) + if QUANTIZATION == "hybrid_fp8_identity": + return te_recipe.CustomRecipe(qfactory=_hybrid_fp8_identity_qfactory) + if QUANTIZATION == "hybrid_mxfp8_identity": + return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_identity_qfactory) if QUANTIZATION == "hybrid_nvfp4": return te_recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) if QUANTIZATION == "hybrid_mxfp8_nvfp4": @@ -213,10 +242,10 @@ def hybrid_recipe(): def _get_tolerances(): - if QUANTIZATION == "hybrid_fp8": + if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): # Loose because of sequence parallel & amax reduction (fp8_cs). return {"rtol": 0.4, "atol": 0.25} - if QUANTIZATION == "hybrid_mxfp8": + if QUANTIZATION in ("hybrid_mxfp8", "hybrid_mxfp8_identity"): return {"rtol": 0.125, "atol": 0.0625} if QUANTIZATION == "hybrid_nvfp4": # Upstream ``run_numerics.py`` uses (0.125, 0.12) for vanilla NVFP4 @@ -414,7 +443,7 @@ def test_linear(): # high precision before quantizing, so the SP output must still match # single-node -- guards against a future regression to quantize-then-gather # without cross-rank amax reduction. - if QUANTIZATION == "hybrid_fp8": + if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): _test_linear("column", True, amax_stress=True) @@ -543,8 +572,8 @@ def run(recipe): def test_linear_vs_vanilla(): # Cross-format hybrid has no single built-in vanilla recipe to compare # against bitwise; it is covered by the distributed-vs-single-node checks. - if QUANTIZATION == "hybrid_mxfp8_nvfp4": - dist_print("linear_vs_vanilla: skipped for cross-format hybrid (no vanilla equivalent)") + if QUANTIZATION in ("hybrid_mxfp8_nvfp4", "hybrid_fp8_identity", "hybrid_mxfp8_identity"): + dist_print("linear_vs_vanilla: skipped for hybrid without a vanilla equivalent") return for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: @@ -750,7 +779,14 @@ def main(argv=None): "--quantization", type=str, required=True, - choices=["hybrid_fp8", "hybrid_mxfp8", "hybrid_nvfp4", "hybrid_mxfp8_nvfp4"], + choices=[ + "hybrid_fp8", + "hybrid_mxfp8", + "hybrid_fp8_identity", + "hybrid_mxfp8_identity", + "hybrid_nvfp4", + "hybrid_mxfp8_nvfp4", + ], ) parser.add_argument( "--test", diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py index 711e76063eb..8dc2995761d 100644 --- a/tests/pytorch/distributed/test_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -109,6 +109,13 @@ def test_hybrid_fp8_transformer_layer(): _run_test("hybrid_fp8", "transformer_layer") +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +def test_hybrid_fp8_identity_linear(): + """Linear-only TP/SP coverage for FP8-current forward plus Identity backward. + Includes the amax-stress branch inside ``run_hybrid_tp_sp.py``.""" + _run_test("hybrid_fp8_identity", "linear") + + # ────────────────────────────────────────────────────────────────────── # Hybrid MXFP8 (rowwise + columnwise same format) # ────────────────────────────────────────────────────────────────────── @@ -145,6 +152,12 @@ def test_hybrid_mxfp8_transformer_layer(): _run_test("hybrid_mxfp8", "transformer_layer") +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +def test_hybrid_mxfp8_identity_linear(): + """Linear-only TP/SP coverage for MXFP8 forward plus Identity backward.""" + _run_test("hybrid_mxfp8_identity", "linear") + + # ────────────────────────────────────────────────────────────────────── # Hybrid NVFP4 (rowwise + columnwise same format, 1D block scaling) # ────────────────────────────────────────────────────────────────────── diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py new file mode 100644 index 00000000000..cf083dd0d65 --- /dev/null +++ b/tests/pytorch/test_identity_quantizer.py @@ -0,0 +1,935 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for IdentityQuantizer (high-precision passthrough) and its use as a +per-direction component of HybridQuantizer to express mixed forward/backward +precision via the CustomRecipe + qfactory machinery. Scoped to te.Linear, +single GPU. +""" + +import io + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + HybridQuantizer, + HybridQuantizedTensor, + IdentityQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) +from transformer_engine.pytorch.tensor.identity_tensor import IdentityTensor +from transformer_engine.pytorch.tensor.storage.identity_tensor_storage import ( + IdentityTensorStorage, +) + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = ( + te.is_fp8_block_scaling_available(return_reason=True) +) + +# ── Module-level qfactories (picklable / autocast-friendly) ────────── + + +def identity_all_factory(role): # pylint: disable=unused-argument + """Whole layer in high precision: Identity for every slot.""" + return IdentityQuantizer() + + +def _fp8_cs(fp8_dtype=tex.DType.kFloat8E4M3): + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + +def _mxfp8(fp8_dtype=tex.DType.kFloat8E4M3): + return MXFP8Quantizer(fp8_dtype=fp8_dtype) + + +def _float8_blockwise(fp8_dtype=tex.DType.kFloat8E4M3): + return Float8BlockQuantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) + + +def _nvfp4(): + return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + + +_HYBRID_IDENTITY_FORMATS = [ + pytest.param( + "fp8_current", + marks=pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}"), + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + ), + pytest.param( + "float8_blockwise", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8Blockwise: {reason_for_no_fp8_block_scaling}", + ), + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), + ), +] + +_HYBRID_IDENTITY_RECOMPUTE_FORMATS = [ + pytest.param( + "fp8_current", + marks=pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}"), + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", + ), + ), +] + + +def _format_quantizer(format_name): + if format_name == "fp8_current": + return _fp8_cs(tex.DType.kFloat8E4M3) + if format_name == "mxfp8": + return _mxfp8(tex.DType.kFloat8E4M3) + if format_name == "float8_blockwise": + return _float8_blockwise(tex.DType.kFloat8E4M3) + if format_name == "nvfp4": + return _nvfp4() + raise ValueError(format_name) + + +def _hybrid_quantized_fwd_identity_bwd_factory(format_name): + def qfactory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + + return qfactory + + +def fwd_hp_bwd_fp8_factory(role): + """High-precision forward, FP8 backward (per-direction via hybrid).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _fp8_cs(tex.DType.kFloat8E5M2) + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + ) + + +def fwd_fp8_bwd_hp_factory(role): + """FP8 forward, high-precision backward (per-direction via hybrid).""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ) + + +def hybrid_all_identity_factory(role): + """All directions high precision, expressed through the hybrid container. + + weight / input / output -> Hybrid(Identity, Identity); grad -> Identity. + Exercises the HybridQuantizedTensor path with Identity sub-storages in both + directions (distinct from the non-hybrid whole-layer-HP path). + """ + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + + +def fp8_fwd_factory(role): + """Plain FP8 current scaling for every slot (E4M3 fwd, E5M2 grad). + + Used with ``backward_override="high_precision"`` as the reference that the + per-direction Identity machinery (``fwd_fp8_bwd_hp_factory``) must reproduce + bitwise: same FP8 forward, high-precision backward. + """ + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _fp8_cs(tex.DType.kFloat8E5M2) + return _fp8_cs(tex.DType.kFloat8E4M3) + + +def _offload_roundtrip(tensor): + from transformer_engine.pytorch.cpu_offload import OffloadableLayerState + + stream = torch.cuda.Stream() + state = OffloadableLayerState(offload_stream=stream) + tid = state.push_tensor(tensor) + state.start_offload() + state.release_activation_forward_gpu_memory() + state.start_reload() + reloaded = state.pop_tensor(tid) + torch.cuda.synchronize() + try: + return reloaded + finally: + state.release_all_memory() + + +# ── Unit tests ─────────────────────────────────────────────────────── + + +class TestIdentityQuantizerUnit: + """IdentityQuantizer / IdentityTensorStorage basic behavior.""" + + def test_quantize_returns_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + out = IdentityQuantizer()(x) + assert isinstance(out, IdentityTensor) + + def test_internal_returns_storage(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + out = q(x) + assert isinstance(out, IdentityTensorStorage) + assert not isinstance(out, IdentityTensor) + + + + def test_dequantize_bitwise_identical(self): + x = torch.randn(4, 32, device="cuda", dtype=torch.bfloat16) + out = IdentityQuantizer()(x) + assert torch.equal(out.dequantize(), x) + + def test_dtype_cast(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.float32) + out = IdentityQuantizer(dtype=torch.bfloat16)(x) + assert out.dequantize().dtype == torch.bfloat16 + + def test_update_usage_is_noop(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + st = q(x) + st.update_usage(rowwise_usage=False, columnwise_usage=True) + assert torch.equal(st.dequantize(), x) + assert st.get_usages() == {"rowwise": True, "columnwise": True} + + def test_save_restore_roundtrip(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + q.internal = True + st = q(x) + tensors, _ = st.prepare_for_saving() + assert st._hp_data is None + leftover = st.restore_from_saved(tensors) + assert leftover == [] + assert torch.equal(st.dequantize(), x) + + def test_update_quantized_inplace(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + q = IdentityQuantizer() + st = q.make_empty((4, 8), dtype=torch.bfloat16, device="cuda") + q.update_quantized(x, st) + assert torch.equal(st.dequantize(), x) + + def test_tensor_ops_preserve_identity_and_values(self): + x = torch.arange(24, device="cuda", dtype=torch.bfloat16).reshape(6, 4) + t = IdentityQuantizer()(x) + + view = t.view(3, 8) + assert isinstance(view, IdentityTensor) + torch.testing.assert_close(view.dequantize(), x.view(3, 8), rtol=0.0, atol=0.0) + + pieces = torch.split(t, 2, dim=0) + assert all(isinstance(piece, IdentityTensor) for piece in pieces) + for piece, ref in zip(pieces, torch.split(x, 2, dim=0)): + torch.testing.assert_close(piece.dequantize(), ref, rtol=0.0, atol=0.0) + + sliced = t[1:5:2] + assert isinstance(sliced, IdentityTensor) + torch.testing.assert_close(sliced.dequantize(), x[1:5:2], rtol=0.0, atol=0.0) + + strided = torch.as_strided(t, (3, 4), (4, 1), 4) + assert isinstance(strided, IdentityTensor) + torch.testing.assert_close(strided.dequantize(), x[1:4], rtol=0.0, atol=0.0) + + cloned = torch.clone(t) + assert isinstance(cloned, IdentityTensor) + torch.testing.assert_close(cloned.dequantize(), x, rtol=0.0, atol=0.0) + + zeros = t.new_zeros((2, 3)) + assert isinstance(zeros, IdentityTensor) + torch.testing.assert_close( + zeros.dequantize(), torch.zeros((2, 3), device="cuda", dtype=x.dtype), rtol=0.0, atol=0.0 + ) + + dst = IdentityQuantizer().make_empty(x.shape, dtype=x.dtype, device="cuda") + dst.copy_(t) + torch.testing.assert_close(dst.dequantize(), x, rtol=0.0, atol=0.0) + + def test_fsdp_pre_post_all_gather_roundtrip(self): + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + sharded_tensors, metadata = t.fsdp_pre_all_gather( + mesh=None, orig_size=t.shape, contiguous_orig_stride=None, module=None, mp_policy=None + ) + gathered, outputs = t.fsdp_post_all_gather(sharded_tensors, metadata, t.dtype, out=None) + assert isinstance(gathered, IdentityTensor) + assert outputs is sharded_tensors + torch.testing.assert_close(gathered.dequantize(), x, rtol=0.0, atol=0.0) + + reuse, _ = t.fsdp_post_all_gather(sharded_tensors, metadata, t.dtype, out=gathered) + assert reuse is gathered + torch.testing.assert_close(reuse.dequantize(), x, rtol=0.0, atol=0.0) + + + def test_cpu_offload_roundtrip_identity_exact(self): + x = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + + reloaded = _offload_roundtrip(t) + + assert isinstance(reloaded, IdentityTensor) + torch.testing.assert_close(reloaded.dequantize(), x, rtol=0.0, atol=0.0) + + def test_replace_raw_data_preserves_identity_values(self): + from transformer_engine.pytorch.tensor.utils import replace_raw_data + + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + new_raw = torch.empty_like(x) + replace_raw_data(t, new_raw) + assert t._hp_data is new_raw + torch.testing.assert_close(t.dequantize(), x, rtol=0.0, atol=0.0) + + def test_quantize_master_weights_identity_exact_nonzero_offset(self): + from transformer_engine.pytorch.tensor.utils import ( + post_all_gather_processing, + quantize_master_weights, + ) + + group = _ensure_single_rank_dp_group() + q = IdentityQuantizer() + weight = q.make_empty((4, 8), dtype=torch.bfloat16, device="cuda") + original = torch.randn_like(weight.dequantize()) + q.update_quantized(original, weight) + + master_full = torch.randn(4, 8, device="cuda", dtype=torch.float32) + start_offset = master_full.numel() // 2 + master_shard = master_full.reshape(-1)[start_offset:].contiguous() + + quantize_master_weights([weight], [master_shard], [start_offset], group=group) + post_all_gather_processing([weight]) + + expected = original.clone() + expected.reshape(-1)[start_offset:] = master_shard.to(torch.bfloat16) + torch.testing.assert_close(weight.dequantize(), expected, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_quantize_master_weights_hybrid_identity_fp8_current(self): + from transformer_engine.pytorch.tensor.utils import ( + post_all_gather_processing, + quantize_master_weights, + ) + + group = _ensure_single_rank_dp_group() + recipe = CustomRecipe(qfactory=fwd_hp_bwd_fp8_factory) + torch.manual_seed(123) + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(32, 32, bias=False, params_dtype=torch.bfloat16).cuda() + weight = model.weight + assert isinstance(weight, HybridQuantizedTensor) + assert isinstance(weight._rowwise_storage, IdentityTensorStorage) + + master = torch.randn_like(weight.dequantize(dtype=torch.float32)).reshape(-1).contiguous() + quantize_master_weights([weight], [master], [0], group=group) + post_all_gather_processing([weight]) + + expected = master.to(torch.bfloat16) + row_deq = weight._rowwise_storage.dequantize().reshape(-1) + col_deq = weight._columnwise_storage.dequantize(dtype=torch.float32).reshape(-1) + torch.testing.assert_close(row_deq, expected, rtol=0.0, atol=0.0) + torch.testing.assert_close(col_deq, master, rtol=0.125, atol=0.1) + + +# ── te.Linear integration ──────────────────────────────────────────── + + +def _make_linears(in_f, out_f, seed=1234, dtype=torch.bfloat16): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + ref = te.Linear(in_f, out_f, params_dtype=dtype).cuda() + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + test = te.Linear(in_f, out_f, params_dtype=dtype).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + return ref, test + + +def _rel_l2_error(actual, reference): + """Relative L2-norm error ``||actual - reference|| / ||reference||``. + + The right metric for comparing a quantized result to a high-precision + reference: element-wise ``rtol`` is meaningless here because reference grads + contain near-zero entries (relative error on ~1e-12 values explodes), while + the aggregate norm error reflects the true quantization noise. + """ + a = actual.float() + b = reference.float() + return (a - b).norm().item() / (b.norm().item() + 1e-12) + + +def _ensure_single_rank_dp_group(): + import pathlib + import tempfile + + if not torch.distributed.is_initialized(): + torch.cuda.set_device(0) + with tempfile.NamedTemporaryFile(delete=False) as f: + rendezvous_file = pathlib.Path(f.name) + torch.distributed.init_process_group( + backend="nccl", + init_method=rendezvous_file.resolve().as_uri(), + rank=0, + world_size=1, + ) + return torch.distributed.GroupMember.WORLD + + +def _fwd_bwd(model, x, recipe=None): + x = x.clone().detach().requires_grad_(True) + if recipe is not None: + with te.autocast(enabled=True, recipe=recipe): + y = model(x) + else: + y = model(x) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +def _fwd_bwd_checkpoint(model, x, recipe, use_reentrant): + x = x.clone().detach().requires_grad_(True) + with te.autocast(enabled=True, recipe=recipe): + if use_reentrant is None: + y = model(x) + else: + y = te.checkpoint(model, x, use_reentrant=use_reentrant) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +_IDENTITY_MODULE_NAMES = ( + "Linear", + "LayerNormLinear", + "LayerNormMLP", + "GroupedLinear", + "TransformerLayer", +) + + +def _make_identity_module(module_name, seed=1234, dtype=torch.bfloat16): + hidden_size = 64 + ffn_hidden_size = 128 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + if module_name == "Linear": + return te.Linear(hidden_size, hidden_size, params_dtype=dtype).cuda() + if module_name == "LayerNormLinear": + return te.LayerNormLinear(hidden_size, hidden_size, params_dtype=dtype).cuda() + if module_name == "LayerNormMLP": + return te.LayerNormMLP(hidden_size, ffn_hidden_size, params_dtype=dtype).cuda() + if module_name == "GroupedLinear": + return te.GroupedLinear(2, hidden_size, hidden_size, params_dtype=dtype).cuda() + if module_name == "TransformerLayer": + return te.TransformerLayer( + hidden_size, + ffn_hidden_size, + 4, + hidden_dropout=0.0, + attention_dropout=0.0, + params_dtype=dtype, + ).cuda() + raise ValueError(module_name) + + +def _make_identity_module_pair(module_name, seed=1234, dtype=torch.bfloat16): + ref = _make_identity_module(module_name, seed=seed, dtype=dtype) + test = _make_identity_module(module_name, seed=seed + 1, dtype=dtype) + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + return ref, test + + +def _identity_module_input(module_name): + torch.manual_seed(7) + if module_name == "TransformerLayer": + return torch.randn(4, 2, 64, device="cuda", dtype=torch.bfloat16) + return torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + + +def _identity_module_forward(module_name, module, x): + if module_name == "GroupedLinear": + m_splits = torch.tensor([8, 8], device="cuda", dtype=torch.int32) + return module(x, m_splits=m_splits) + return module(x) + + +def _fwd_bwd_module(module_name, model, x, recipe=None): + x = x.clone().detach().requires_grad_(True) + if recipe is not None: + with te.autocast(enabled=True, recipe=recipe): + y = _identity_module_forward(module_name, model, x) + else: + y = _identity_module_forward(module_name, model, x) + torch.manual_seed(99) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +class TestIdentityTEModuleCoverage: + """All-Identity recipes should route every TE module through HP-compatible paths.""" + + @pytest.mark.parametrize("module_name", _IDENTITY_MODULE_NAMES) + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(identity_all_factory, id="plain_identity"), + pytest.param(hybrid_all_identity_factory, id="hybrid_identity"), + ], + ) + def test_identity_recipe_matches_bf16_bitwise(self, module_name, qfactory): + ref, test = _make_identity_module_pair(module_name, seed=7300) + x = _identity_module_input(module_name) + recipe = CustomRecipe(qfactory=qfactory) + + y_ref, dx_ref, wg_ref = _fwd_bwd_module(module_name, ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd_module(module_name, test, x, recipe=recipe) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + + +class TestIdentityHybridFormatProtocols: + SHAPE = (256, 256) + OFFLOAD_SHAPE = (1024, 1024) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_FORMATS) + def test_save_restore_keeps_identity_direction_exact(self, format_name): + torch.manual_seed(401) + x = torch.randn(*self.SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + hybrid = q.quantize(x) + expected_row = hybrid._rowwise_storage.dequantize().clone() + expected_col = x.clone() + + tensors, obj = hybrid.prepare_for_saving() + leftover = obj.restore_from_saved(tensors) + + assert leftover == [] + assert isinstance(hybrid._columnwise_storage, IdentityTensorStorage) + torch.testing.assert_close( + hybrid._columnwise_storage.dequantize(), expected_col, rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + hybrid._rowwise_storage.dequantize(), expected_row, rtol=0.0, atol=0.0 + ) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_FORMATS) + def test_cpu_offload_keeps_identity_direction_exact(self, format_name): + torch.manual_seed(402) + x = torch.randn(*self.OFFLOAD_SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + hybrid = q.quantize(x) + expected_row = hybrid._rowwise_storage.dequantize().clone() + + reloaded = _offload_roundtrip(hybrid) + + assert isinstance(reloaded, HybridQuantizedTensor) + assert isinstance(reloaded._columnwise_storage, IdentityTensorStorage) + torch.testing.assert_close( + reloaded._columnwise_storage.dequantize(), x, rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + reloaded._rowwise_storage.dequantize(), expected_row, rtol=0.0, atol=0.0 + ) + + @pytest.mark.parametrize("format_name", ["mxfp8", "float8_blockwise", "nvfp4"]) + def test_quantize_master_weights_per_block_hybrid_identity_rejected(self, format_name): + if format_name == "mxfp8" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + if format_name == "float8_blockwise" and not fp8_block_scaling_available: + pytest.skip(f"Float8Blockwise: {reason_for_no_fp8_block_scaling}") + if format_name == "nvfp4" and not (fp8_available and nvfp4_available): + pytest.skip(f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}") + + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + x = torch.randn(*self.SHAPE, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_format_quantizer(format_name), + columnwise_quantizer=IdentityQuantizer(), + ) + weight = q.quantize(x) + master = torch.randn_like(x, dtype=torch.float32).reshape(-1).contiguous() + + with pytest.raises(NotImplementedError, match="HybridQuantizer"): + quantize_master_weights([weight], [master], [0], group=group) + + @pytest.mark.parametrize("format_name", _HYBRID_IDENTITY_RECOMPUTE_FORMATS) + @pytest.mark.parametrize("use_reentrant", [True, False]) + def test_activation_recompute_matches_no_checkpoint(self, format_name, use_reentrant): + recipe = CustomRecipe(qfactory=_hybrid_quantized_fwd_identity_bwd_factory(format_name)) + ref, test = _make_linears(128, 128, seed=440) + torch.manual_seed(441) + x = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16) + + y_ref, dx_ref, wg_ref = _fwd_bwd_checkpoint(ref, x, recipe, use_reentrant=None) + y_test, dx_test, wg_test = _fwd_bwd_checkpoint(test, x, recipe, use_reentrant=use_reentrant) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_test, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_test) == len(wg_ref) + for g_test, g_ref in zip(wg_test, wg_ref): + torch.testing.assert_close(g_test, g_ref, rtol=0.0, atol=0.0) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +class TestIdentityLinear: + """End-to-end te.Linear with Identity-based recipes.""" + + IN_F = 128 + OUT_F = 128 + BATCH = 64 + + def _input(self): + torch.manual_seed(7) + return torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + + def test_whole_layer_hp_matches_bf16_bitwise(self): + """Identity for every slot => all GEMMs high precision => bitwise-equal + to a plain BF16 te.Linear (no autocast).""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=identity_all_factory)) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_fwd_hp_bwd_fp8_forward_bitwise(self): + """High-precision forward must be bitwise-equal to BF16 forward; the + backward runs in FP8 (finite, close to BF16 within a loose tolerance).""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_h, dx_h, wg_h = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_hp_bwd_fp8_factory)) + + # Forward is high precision -> bitwise equal. + torch.testing.assert_close(y_h, y_ref, rtol=0.0, atol=0.0) + # Backward is FP8 (E4M3 weight col, E5M2 grad) -> relative L2 error vs the + # BF16 reference reflects pure FP8 quant noise. Measured: dgrad ~5.7e-2, + # weight-grad ~5.8e-2 (E4M3 ~6% step). Bound 7e-2 keeps a small margin. + assert torch.isfinite(dx_h).all() + assert _rel_l2_error(dx_h, dx_ref) < 7e-2 + for g, g_ref in zip(wg_h, wg_ref): + assert torch.isfinite(g).all() + if g.dim() == 1: + # Bias grad = sum(dY) is computed in high precision (dY is bitwise + # identical since the forward is bitwise), so it must match exactly. + torch.testing.assert_close(g, g_ref, rtol=0.0, atol=0.0) + else: + assert _rel_l2_error(g, g_ref) < 7e-2 + + def test_fwd_fp8_bwd_hp_runs_and_backward_high_precision(self): + """FP8 forward + high-precision backward. Forward differs from BF16 + (quantized), backward GEMMs run in high precision.""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, _ = _fwd_bwd(ref, x, recipe=None) + y_q, dx_q, wg_q = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_hp_factory)) + + # Forward is FP8 (E4M3) -> relative L2 error vs BF16 is the quant noise. + # Measured ~3.7e-2; bound 5e-2. + assert torch.isfinite(y_q).all() + assert _rel_l2_error(y_q, y_ref) < 5e-2 + # Backward GEMMs run in high precision. dgrad differs from the BF16 + # reference only because the FP8 forward perturbs dY; measured ~1.2e-2, + # bound 3e-2 (the bitwise HP-backward guarantee is locked by the + # backward_override equivalence test below). + assert torch.isfinite(dx_q).all() + assert _rel_l2_error(dx_q, dx_ref) < 3e-2 + for g in wg_q: + assert torch.isfinite(g).all() + + def test_hybrid_all_identity_matches_bf16_bitwise(self): + """All-Identity through the *hybrid* container must be bitwise-equal to a + plain BF16 te.Linear. Complements the non-hybrid whole-layer-HP test: this + exercises HybridQuantizedTensor with Identity sub-storages in both + directions and the per-operand unwrap of every GEMM.""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd( + test, x, recipe=CustomRecipe(qfactory=hybrid_all_identity_factory) + ) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_ref) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_identity_reproduces_backward_override_high_precision_bitwise(self): + """The per-direction Identity machinery must reproduce + ``backward_override="high_precision"`` **bitwise**. + + Both runs quantize the forward to the same FP8 (current scaling) and run + the backward in high precision against the original operands. The Identity + path expresses this per-tensor (weight/input = Hybrid(row=FP8, col=Identity), + grad = Identity); the reference uses the global ``backward_override`` knob. + Identical forward FP8 + identical original HP backward operands => bitwise. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_bo, dx_bo, wg_bo = _fwd_bwd( + ref, + x, + recipe=CustomRecipe(qfactory=fp8_fwd_factory, backward_override="high_precision"), + ) + y_id, dx_id, wg_id = _fwd_bwd( + test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_hp_factory) + ) + + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + + def test_identity_matches_bf16_multistep_training_bitwise(self): + """Multi-step SGD: an all-Identity recipe must track a plain BF16 + te.Linear bitwise across optimizer steps (no drift from workspace caching + or any hidden state).""" + ref, test = _make_linears(self.IN_F, self.OUT_F) + opt_ref = torch.optim.SGD(ref.parameters(), lr=0.1) + opt_test = torch.optim.SGD(test.parameters(), lr=0.1) + recipe = CustomRecipe(qfactory=identity_all_factory) + + for step in range(4): + torch.manual_seed(1000 + step) + x = torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(2000 + step) + target = torch.randn(self.BATCH, self.OUT_F, device="cuda", dtype=torch.bfloat16) + + opt_ref.zero_grad() + y_ref = ref(x) + loss_ref = torch.nn.functional.mse_loss(y_ref, target) + loss_ref.backward() + opt_ref.step() + + opt_test.zero_grad() + with te.autocast(enabled=True, recipe=recipe): + y_test = test(x) + loss_test = torch.nn.functional.mse_loss(y_test, target) + loss_test.backward() + opt_test.step() + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + torch.testing.assert_close(p_test, p_ref, rtol=0.0, atol=0.0, msg=f"step {step}") + + def test_quantized_model_init_identity_matches_bf16_bitwise(self): + """Persistent Identity params from quantized_model_init should match BF16 exactly.""" + torch.manual_seed(314) + ref = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(2718) + with te.quantized_model_init(enabled=True, recipe=recipe): + test = te.Linear( + self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16 + ).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + assert isinstance(p_test, IdentityTensor) + p_test.copy_(p_ref) + + x = self._input() + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=recipe) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_quantized_model_init_identity_training_loss_decreases_bitwise(self): + """All-Identity quantized params train like BF16 and loss decreases.""" + torch.manual_seed(777) + ref = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(888) + with te.quantized_model_init(enabled=True, recipe=recipe): + test = te.Linear( + self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16 + ).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + assert isinstance(p_test, IdentityTensor) + p_test.copy_(p_ref) + + torch.manual_seed(909) + x = torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + target = torch.zeros(self.BATCH, self.OUT_F, device="cuda", dtype=torch.bfloat16) + opt_ref = torch.optim.SGD(ref.parameters(), lr=0.1) + opt_test = torch.optim.SGD(test.parameters(), lr=0.1) + losses_ref = [] + losses_test = [] + + for _ in range(5): + opt_ref.zero_grad() + y_ref = ref(x) + loss_ref = torch.nn.functional.mse_loss(y_ref, target) + loss_ref.backward() + opt_ref.step() + losses_ref.append(loss_ref.detach().clone()) + + opt_test.zero_grad() + with te.autocast(enabled=True, recipe=recipe): + y_test = test(x) + loss_test = torch.nn.functional.mse_loss(y_test, target) + loss_test.backward() + opt_test.step() + losses_test.append(loss_test.detach().clone()) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + torch.testing.assert_close( + p_test.dequantize(), p_ref, rtol=0.0, atol=0.0 + ) + + assert all( + losses_ref[i + 1].item() < losses_ref[i].item() + for i in range(len(losses_ref) - 1) + ), f"BF16 loss did not strictly decrease: {[x.item() for x in losses_ref]}" + for loss_test, loss_ref in zip(losses_test, losses_ref): + torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize("use_reentrant", [True, False]) + def test_identity_activation_recompute_matches_bf16_bitwise(self, use_reentrant): + """All-Identity recompute should be exactly the BF16 no-checkpoint path.""" + ref, test = _make_linears(self.IN_F, self.OUT_F, seed=4242) + recipe = CustomRecipe(qfactory=identity_all_factory) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, wg_id = _fwd_bwd_checkpoint(test, x, recipe, use_reentrant=use_reentrant) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + for g_id, g_ref in zip(wg_id, wg_ref): + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + + def test_quantized_model_init_identity_state_dict_save_load_exact(self): + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(5151) + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + + torch.manual_seed(5152) + x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + out_before = model(x) + + buffer = io.BytesIO() + torch.save(model.state_dict(), buffer) + buffer.seek(0) + + with te.quantized_model_init(enabled=True, recipe=recipe): + model2 = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + model2.load_state_dict(torch.load(buffer)) + + with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): + out_after = model2(x) + + assert isinstance(model2.weight, IdentityTensor) + torch.testing.assert_close(out_after, out_before, rtol=0.0, atol=0.0) + + def test_load_bf16_state_dict_into_identity_model_exact(self): + recipe = CustomRecipe(qfactory=identity_all_factory) + torch.manual_seed(6161) + ref = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + with te.quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() + + model.load_state_dict(ref.state_dict()) + assert isinstance(model.weight, IdentityTensor) + torch.testing.assert_close(model.weight.dequantize(), ref.weight, rtol=0.0, atol=0.0) + + torch.manual_seed(6162) + x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + out_ref = ref(x) + with te.autocast(enabled=True, recipe=recipe): + out_id = model(x) + + torch.testing.assert_close(out_id, out_ref, rtol=0.0, atol=0.0) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index c4d349b506e..9a541a8fd40 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -91,7 +91,10 @@ from transformer_engine.pytorch.tensor import NVFP4Tensor from transformer_engine.pytorch.tensor import HybridQuantizer from transformer_engine.pytorch.tensor import HybridQuantizedTensorStorage +from transformer_engine.pytorch.tensor import IdentityQuantizer +from transformer_engine.pytorch.tensor import IdentityTensorStorage from transformer_engine.pytorch.tensor import HybridQuantizedTensor +from transformer_engine.pytorch.tensor import IdentityTensor from transformer_engine.pytorch.tensor.float8_tensor import ( _make_float8_tensor_in_reduce_ex, ) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index f19f796bb11..154929c87fc 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -18,6 +18,7 @@ from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage +from ..tensor.storage.identity_tensor_storage import IdentityTensorStorage from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizer @@ -133,6 +134,20 @@ def _unwrap_hybrid_B(tensor, layout): return tensor.columnwise_sub_storage +def _materialize_high_precision(tensor): + """Replace an :class:`IdentityTensorStorage` operand with its plain tensor. + + Identity (high-precision passthrough) operands carry an unquantized tensor; + materializing it here routes the matmul through the standard high-precision + GEMM path. Non-identity operands pass through unchanged. Called after the + hybrid unwrap, so a high-precision *direction* of a hybrid tensor is handled + too. + """ + if isinstance(tensor, IdentityTensorStorage): + return tensor.dequantize() + return tensor + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -159,8 +174,8 @@ def general_gemm( transa = layout[0] == "T" transb = layout[1] == "T" - A = _unwrap_hybrid_A(A, layout) - B = _unwrap_hybrid_B(B, layout) + A = _materialize_high_precision(_unwrap_hybrid_A(A, layout)) + B = _materialize_high_precision(_unwrap_hybrid_B(B, layout)) alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) @@ -329,8 +344,8 @@ def general_grouped_gemm( """ num_gemms = len(A) - A = [_unwrap_hybrid_A(a, layout) for a in A] - B = [_unwrap_hybrid_B(b, layout) for b in B] + A = [_materialize_high_precision(_unwrap_hybrid_A(a, layout)) for a in A] + B = [_materialize_high_precision(_unwrap_hybrid_B(b, layout)) for b in B] transa = layout[0] == "T" transb = layout[1] == "T" diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py index d660e5a53bc..9276d077e60 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py @@ -269,3 +269,39 @@ def nvfp4_linear_mxfp8_dpa_factory( return _make_mxfp8_quantizer() return _make_nvfp4_quantizer(role) + + +def high_precision_factory( + role: Optional[QuantizerRole], # pylint: disable=unused-argument +): + """Quantizer factory: run all GEMMs in high precision. + + Dispatch logic: + * every role -> ``IdentityQuantizer`` (no quantization) + """ + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + return IdentityQuantizer() + + +def fwd_high_precision_bwd_mxfp8_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: high-precision forward, MXFP8 backward. + + Dispatch logic: + * ``grad_output`` / ``grad_input`` -> MXFP8 (E4M3, block-32) + * everything else -> ``Hybrid(rowwise=IdentityQuantizer, columnwise=MXFP8)`` + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return _make_mxfp8_quantizer() + + # fprop consumes rowwise high precision; dgrad / wgrad consume columnwise MXFP8. + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=_make_mxfp8_quantizer(), + ) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 6d947d2d050..4408a2435cf 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -47,6 +47,7 @@ from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage @@ -1584,9 +1585,10 @@ def grad_output_preprocess( ): grad_bias = grad_output.dequantize().view(-1, grad_output.shape[-1]).sum(dim=0) else: - if isinstance(quantizer, (Float8BlockQuantizer, HybridQuantizer)): + if isinstance(quantizer, (Float8BlockQuantizer, HybridQuantizer, IdentityQuantizer)): # Float8BlockQuantizer: unfused until cast_transpose + dgrad is ready. # HybridQuantizer: tex.bgrad_quantize doesn't recognize hybrid quantizers. + # IdentityQuantizer: high-precision passthrough; bgrad computed in HP. grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0) else: grad_bias, grad_output = tex.bgrad_quantize(grad_output, quantizer) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index e503b4b560e..0d972b465ea 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1880,18 +1880,17 @@ def make_quantizers(self) -> list: ) roles = [QuantizerRole() for _ in range(self.num_quantizers)] - # qfactory must return a Quantizer or QuantizerRequest for every slot. - # None is not a valid return value — it would silently disable quantization - # for that tensor, risking hard-to-detect performance regressions. - # TODO(negvet): Introduce an explicit IdentityQuantizer for intentional no-op - # quantization. Until then, None is rejected. + # qfactory returns one quantizer-like object per slot; use + # ``IdentityQuantizer`` for intentional high-precision passthrough. raw = [qfactory(roles[i]) for i in range(self.num_quantizers)] for i, q in enumerate(raw): if q is None: raise ValueError( f"CustomRecipe qfactory returned None for slot {i} " f"(role={roles[i]}). Every slot must return a Quantizer " - "instance or a QuantizerRequest." + "instance or a QuantizerRequest. For an intentional no-op " + "(high-precision / unquantized) slot, return an " + "IdentityQuantizer instead of None." ) # -- Delayed scaling sub-state -- diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index 6098649182e..c3355b6c621 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -26,6 +26,8 @@ from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer from .grouped_tensor import GroupedTensor from .hybrid_tensor import HybridQuantizedTensor, HybridQuantizer +from .identity_tensor import IdentityTensor, IdentityQuantizer +from .storage.identity_tensor_storage import IdentityTensorStorage from .utils import cast_master_weights_to_fp8, replace_raw_data __all__ = [ @@ -36,6 +38,7 @@ "Float8BlockQuantizer", "NVFP4Quantizer", "HybridQuantizer", + "IdentityQuantizer", "QuantizedTensorStorage", "Float8TensorStorage", "MXFP8TensorStorage", @@ -43,8 +46,10 @@ "NVFP4TensorStorage", "GroupedTensorStorage", "HybridQuantizedTensorStorage", + "IdentityTensorStorage", "QuantizedTensor", "Float8Tensor", + "IdentityTensor", "MXFP8Tensor", "Float8BlockwiseQTensor", "NVFP4Tensor", @@ -104,5 +109,7 @@ def get_all_tensor_types(): GroupedTensorStorage, HybridQuantizedTensor, HybridQuantizedTensorStorage, + IdentityTensor, + IdentityTensorStorage, ] return all_tensor_types diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py new file mode 100644 index 00000000000..a36c1b6ce31 --- /dev/null +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -0,0 +1,322 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""High-precision passthrough quantizer and tensor. + +``IdentityQuantizer`` is a no-op "quantizer": it keeps the input tensor in its +original high precision instead of casting to a low-precision format. It exists +so the ``CustomRecipe`` + ``qfactory`` machinery can express *unquantized* +tensors and, composed inside a :class:`HybridQuantizer`, *unquantized +directions* (e.g. high-precision forward + quantized backward, or vice versa) +without scattering ``None``/``isinstance`` special-cases across the modules. +""" + +from __future__ import annotations +from typing import Any, Iterable, Optional, Tuple + +import torch +from torch.ops import aten + +from .storage.identity_tensor_storage import IdentityTensorStorage +from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer + + +class IdentityQuantizer(Quantizer): + """Quantizer that performs no quantization (high-precision passthrough). + + Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`) + wrapping the original high-precision tensor. ``general_gemm`` materializes + it back to a plain tensor, so any GEMM consuming it runs in high precision. + + Parameters + ---------- + dtype : torch.dtype, optional + If set, the held tensor is cast to this dtype on quantize. ``None`` + (default) keeps the input's dtype. + rowwise, columnwise : bool + Usage flags (kept for interface compatibility; the single + high-precision buffer serves both directions). + """ + + def __init__( + self, + *, + dtype: Optional[torch.dtype] = None, + rowwise: bool = True, + columnwise: bool = True, + ) -> None: + super().__init__(rowwise=rowwise, columnwise=columnwise) + self.dtype = dtype + + def copy(self) -> "IdentityQuantizer": + """Create shallow copy.""" + quantizer = IdentityQuantizer( + dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def _maybe_cast(self, tensor: torch.Tensor) -> torch.Tensor: + # Detach so the held buffer is plain "data" with no autograd graph edge, + # mirroring the real quantizers (whose quantize kernels emit fresh, + # non-differentiable tensors). Autograd connectivity for the *quantize* + # op is provided separately by ``_QuantizeFunc`` in ``Quantizer.quantize``; + # the surrounding TE module Function computes dgrad/wgrad manually. Without + # the detach the produced tensor aliases a grad-requiring input (e.g. the + # weight workspace returned across the module Function boundary), which + # creates a spurious empty grad edge. + out = tensor.detach() + if self.dtype is not None and out.dtype != self.dtype: + return out.to(self.dtype) + return out + + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensorStorage: + data = self._maybe_cast(tensor) + if self.internal: + return IdentityTensorStorage( + hp_data=data, + fake_dtype=data.dtype, + quantizer=self, + ) + # requires_grad=False: this is the quantized "data" tensor. Autograd + # connectivity is provided by ``_QuantizeFunc`` in ``Quantizer.quantize`` + # (mirrors the real quantizers, which return non-differentiable data). + return IdentityTensor( + data.shape, + data.dtype, + hp_data=data, + quantizer=self, + requires_grad=False, + device=data.device, + ) + + def make_empty( + self, + shape: Iterable[int], + *, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + requires_grad: bool = False, + pin_memory: bool = False, + ) -> "IdentityTensor": + if device is None: + device = torch.device("cuda") + device = torch.device(device) + data = torch.empty(tuple(shape), dtype=dtype, device=device, pin_memory=pin_memory) + return IdentityTensor( + data.shape, + dtype, + hp_data=data, + quantizer=self, + requires_grad=requires_grad, + device=device, + ) + + def update_quantized( + self, + src: torch.Tensor, + dst: QuantizedTensorStorage, + *, + noop_flag: Optional[torch.Tensor] = None, # pylint: disable=unused-argument + ) -> QuantizedTensorStorage: + if not isinstance(dst, IdentityTensorStorage): + raise ValueError( + "IdentityQuantizer can only update IdentityTensorStorage, got" + f" {type(dst).__name__}" + ) + data = self._maybe_cast(src) + if ( + dst._hp_data is not None + and dst._hp_data.shape == data.shape + and dst._hp_data.dtype == data.dtype + and dst._hp_data.device == data.device + ): + dst._hp_data.copy_(data) + else: + dst._hp_data = data.detach() + return dst + + def calibrate(self, tensor: torch.Tensor) -> None: + # No state to calibrate. + return + + def _get_compatible_recipe(self): + # Only reachable via CustomRecipe (qfactory returns IdentityQuantizer). + from transformer_engine.common.recipe import CustomRecipe # avoid circular import + + return CustomRecipe + + +class IdentityTensor(IdentityTensorStorage, QuantizedTensor): + """High-precision passthrough tensor produced by :class:`IdentityQuantizer`. + + Presents as a standard tensor of its nominal dtype; internally it just + holds the original high-precision data (no quantization). + """ + + def __repr__(self, *, tensor_contents=None): + return f"IdentityTensor(dtype={self.dtype}, data={self._hp_data})" + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + return IdentityTensorStorage.dequantize(self, dtype=dtype) + + def view(self, *shape) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + flat_shape = shape[0] if len(shape) == 1 and not isinstance(shape[0], int) else shape + return self._wrap_data_view(self._hp_data.view(*flat_shape)) + + def detach(self) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + return IdentityTensor.make_like(self) + + def clone(self) -> "IdentityTensor": + # pylint: disable=missing-function-docstring + data = self._hp_data.detach().clone() if self._hp_data is not None else None + return IdentityTensor( + self.shape, + self.dtype, + hp_data=data, + quantizer=self._quantizer, + requires_grad=self.requires_grad, + device=self.device, + ) + + @classmethod + def _make_in_reduce_ex( + cls, + hp_data: torch.Tensor, + quantizer: Optional[Quantizer], + dtype: torch.dtype, + shape: torch.Size, + ) -> "IdentityTensor": + """Build IdentityTensor, for use in ``__reduce_ex__``.""" + return IdentityTensor( + shape=shape, + dtype=dtype, + hp_data=hp_data, + quantizer=quantizer, + requires_grad=False, + device=hp_data.device if hp_data is not None else None, + ) + + def __reduce_ex__(self, protocol: int) -> tuple: + """Custom pickling that preserves the high-precision payload.""" + return ( + IdentityTensor._make_in_reduce_ex, + (self._hp_data, self._quantizer, self.dtype, self.shape), + ) + + def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): + """Extract the high-precision buffer for FSDP2 all-gather.""" + return (self._hp_data,), (self._quantizer,) + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional["IdentityTensor"] = None, + ): + """Rebuild IdentityTensor from the gathered high-precision buffer.""" + (data,) = all_gather_outputs + (quantizer,) = metadata + if out is not None: + out._hp_data = data + else: + out = IdentityTensor( + shape=data.shape, + dtype=param_dtype, + hp_data=data, + quantizer=quantizer, + requires_grad=False, + device=data.device, + ) + return out, all_gather_outputs + + def _wrap_data_view(self, data: torch.Tensor) -> "IdentityTensor": + return IdentityTensor( + shape=data.shape, + dtype=self.dtype, + hp_data=data, + quantizer=self._quantizer, + requires_grad=self.requires_grad, + device=data.device, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if kwargs is None: + kwargs = {} + + if func == aten.detach.default: + return args[0].detach() + + if func == aten.clone.default: + return args[0].clone() + + if func == aten.view.default: + tensor = args[0] + shape = args[1] + if list(shape) == list(tensor.shape): + return IdentityTensor.make_like(tensor) + return tensor._wrap_data_view(tensor._hp_data.view(*shape)) + + if func == aten.split.Tensor: + tensor = args[0] + split_size = args[1] + dim = kwargs.get("dim", args[2] if len(args) > 2 else 0) + return [ + tensor._wrap_data_view(piece) + for piece in torch.split(tensor._hp_data, split_size, dim=dim) + ] + + if func == aten.as_strided.default: + tensor = args[0] + shape = args[1] + strides = args[2] + storage_offset = kwargs.get("storage_offset", args[3] if len(args) > 3 else 0) + if ( + tuple(shape) == tuple(tensor.shape) + and tuple(strides) == tuple(tensor.stride()) + and storage_offset == tensor.storage_offset() + ): + return IdentityTensor.make_like(tensor) + return tensor._wrap_data_view( + torch.as_strided(tensor._hp_data, shape, strides, storage_offset) + ) + + if func == aten.slice.Tensor: + tensor = args[0] + dim = args[1] + start = args[2] + end = args[3] + step = args[4] if len(args) > 4 else 1 + if start == 0 and end == tensor.size(dim) and step == 1: + return IdentityTensor.make_like(tensor) + return tensor._wrap_data_view(aten.slice.Tensor(tensor._hp_data, dim, start, end, step)) + + if func == aten.copy_.default: + dst, src = args[0], args[1] + if isinstance(dst, IdentityTensor) and isinstance(src, IdentityTensor): + dst._hp_data.copy_(src._hp_data, *args[2:], **kwargs) + return dst + + if func == aten.new_zeros.default: + tensor = args[0] + new_shape = args[1] + if tensor._quantizer is not None: + out = tensor._quantizer.make_empty( + new_shape, + dtype=kwargs.get("dtype") or tensor.dtype, + device=kwargs.get("device") or tensor.device, + pin_memory=bool(kwargs.get("pin_memory", False)), + ) + out._hp_data.zero_() + return out + + return super().__torch_dispatch__(func, types, args, kwargs) diff --git a/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py new file mode 100644 index 00000000000..4cb7fc092a1 --- /dev/null +++ b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py @@ -0,0 +1,152 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Mixin class holding data for IdentityTensor (high-precision passthrough).""" + +from __future__ import annotations +from typing import Any, Dict, Optional, Tuple + +import torch + +from ...quantized_tensor import QuantizedTensorStorage, Quantizer +from ...utils import _empty_tensor + + +class IdentityTensorStorage(QuantizedTensorStorage): + """Passthrough storage that holds a high-precision (unquantized) tensor. + + Produced by :class:`IdentityQuantizer`. It implements the + ``QuantizedTensorStorage`` interface so it can flow through the same + module / GEMM / save-for-backward / FSDP machinery as the real quantized + storages, but it performs no quantization: it simply carries the original + high-precision tensor. ``general_gemm`` materializes it back to that plain + tensor (so the matmul runs in high precision). + + The data is direction-agnostic -- the same tensor serves both the rowwise + and columnwise directions (the GEMM transposes via its layout flags), so a + single buffer is stored. This is what lets a ``HybridQuantizer`` mix one + quantized direction with one high-precision direction. + """ + + _hp_data: Optional[torch.Tensor] + _quantizer: Optional[Quantizer] + + def __new__( + cls, + *args, + hp_data: Optional[torch.Tensor], + fake_dtype: Optional[torch.dtype] = None, + quantizer: Optional[Quantizer] = None, + **kwargs, + ): + if cls is IdentityTensorStorage: + instance = object.__new__(cls) + if fake_dtype is not None: + instance._dtype = fake_dtype + elif hp_data is not None: + instance._dtype = hp_data.dtype + else: + instance._dtype = torch.float32 + else: + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) + instance._hp_data = hp_data + instance._quantizer = quantizer.copy() if quantizer is not None else None + return instance + + def clear(self): + """Deallocate the held tensor's memory.""" + if self._hp_data is not None: + self._hp_data.data = _empty_tensor() + + def copy_from_storage(self, src: QuantizedTensorStorage) -> None: + """Copy data from another IdentityTensorStorage.""" + if not isinstance(src, IdentityTensorStorage): + raise TypeError("copy_from_storage expects IdentityTensorStorage") + if self._hp_data is not None and src._hp_data is not None: + self._hp_data.copy_(src._hp_data) + + def get_metadata(self) -> Dict[str, Any]: + """Get this tensor's metadata.""" + return { + "hp_data": self._hp_data, + "quantizer": self._quantizer, + "fake_dtype": self._dtype, + } + + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], "IdentityTensorStorage"]: + """Prepare the tensor base for saving for backward.""" + tensors = [self._hp_data] + self._hp_data = None + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + """Restore the held tensor from the saved tensors list.""" + self._hp_data = tensors[0] + return tensors[1:] + + def get_data_tensors(self, rowwise_data: bool = True, columnwise_data: bool = True): + """Get this tensor's data. The single HP buffer serves both directions.""" + if rowwise_data and columnwise_data: + return self._hp_data, None + if rowwise_data: + return self._hp_data + if columnwise_data: + return self._hp_data + raise ValueError("No data to get, both rowwise_data and columnwise_data are False") + + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + """Return the held high-precision tensor (no-op dequantization).""" + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data to dequantize") + if dtype is not None and self._hp_data.dtype != dtype: + return self._hp_data.to(dtype) + return self._hp_data + + def update_usage( + self, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ): + """No-op: the single high-precision buffer serves both directions.""" + # High-precision data is not direction-specific, so there is nothing + # to drop or synthesize. Honor the request only insofar as keeping the + # buffer (a request to drop both would leave no data, which is invalid). + + def get_usages(self) -> Dict[str, bool]: + """Get the usage of the tensor.""" + has_data = self._hp_data is not None + return {"rowwise": has_data, "columnwise": has_data} + + def size(self, *args, **kwargs): + # pylint: disable=missing-function-docstring + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + return self._hp_data.size(*args, **kwargs) + + @property + def device(self): + """Return the device of the held tensor.""" + if self._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data!") + return self._hp_data.device + + def view(self, *shape): + # pylint: disable=missing-function-docstring + flat_shape = shape[0] if len(shape) == 1 and not isinstance(shape[0], int) else shape + return IdentityTensorStorage( + hp_data=self._hp_data.view(*flat_shape) if self._hp_data is not None else None, + fake_dtype=self._dtype, + quantizer=self._quantizer, + ) + + def fsdp_buffer_fields(self) -> Tuple[str, ...]: + """Field gathered by FSDP2 for the high-precision passthrough.""" + return ("_hp_data",) + + def __repr__(self): + return f"IdentityTensorStorage(dtype={self._dtype}, data={self._hp_data})" diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 6a1cd57c7ab..5c595175a85 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -20,6 +20,8 @@ from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer from .hybrid_tensor import HybridQuantizedTensor, HybridQuantizer +from .identity_tensor import IdentityQuantizer +from .storage.identity_tensor_storage import IdentityTensorStorage from ..optimizers.multi_tensor_apply import multi_tensor_applier from ..utils import is_non_tn_fp8_gemm_supported from ..constants import NVFP4_BLOCK_SCALING_SIZE @@ -66,6 +68,18 @@ def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): new_raw_data.detach().copy_(old_rowwise) tensor._rowwise_data = new_raw_data del old_rowwise + elif isinstance(tensor, IdentityTensorStorage): + old_raw_data = tensor._hp_data + if old_raw_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + if old_raw_data.dtype != new_raw_data.dtype: + raise ValueError( + "The data types of raw data don't match: " + f"old dtype={old_raw_data.dtype}, new dtype={new_raw_data.dtype}" + ) + new_raw_data.detach().copy_(old_raw_data) + tensor._hp_data = new_raw_data + del old_raw_data elif isinstance(tensor, MXFP8Tensor): raise NotImplementedError("replace_raw_data for MXFP8Tensor is not supported yet") elif isinstance(tensor, HybridQuantizedTensor): @@ -124,6 +138,7 @@ def quantize_master_weights( blockwise_scaling_params = [] mxfp8_scaling_params = [] nvfp4_params = [] + identity_params = [] if fsdp_shard_model_weights is None: use_fsdp_shard_model_weights = False @@ -178,6 +193,10 @@ def quantize_master_weights( mxfp8_scaling_params.append( (model_weight, master_weight, start_offset, fsdp_shard_model_weight) ) + elif isinstance(quantizer, IdentityQuantizer): + identity_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) elif isinstance(quantizer, HybridQuantizer): _route_hybrid_to_buckets( model_weight, @@ -186,6 +205,7 @@ def quantize_master_weights( fsdp_shard_model_weight, delayed_scaling_params=delayed_scaling_params, current_scaling_params=current_scaling_params, + identity_params=identity_params, ) else: raise ValueError(f"quantize_master_weights for {type(quantizer)} is not supported yet") @@ -201,6 +221,8 @@ def quantize_master_weights( _cast_master_weights_to_fp8_mxfp8_scaling(mxfp8_scaling_params, *extra_args) if len(nvfp4_params) > 0: _cast_master_weights_to_nvfp4_2d(nvfp4_params, *extra_args) + if len(identity_params) > 0: + _cast_master_weights_to_identity(identity_params, *extra_args) def cast_master_weights_to_fp8( @@ -816,6 +838,53 @@ def _cast_master_weights_to_nvfp4_2d( ) +def _identity_storage_data(tensor): + if not isinstance(tensor, IdentityTensorStorage): + raise TypeError(f"Expected IdentityTensorStorage, got {type(tensor).__name__}") + if tensor._hp_data is None: + raise RuntimeError("IdentityTensorStorage has no data") + return tensor._hp_data + + +def _cast_master_weights_to_identity( + params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False +): + del group, manual_post_all_gather_processing + + for model_weight, master_weight, start_offset, model_weight_fragment in params: + if master_weight is None: + continue + if start_offset is None: + raise ValueError("start_offset must not be None when master_weight is provided") + if start_offset < 0: + raise ValueError(f"start_offset must be non-negative, got {start_offset}") + end_offset = start_offset + master_weight.numel() + if end_offset > model_weight.numel(): + raise ValueError( + f"end_offset ({end_offset}) exceeds model_weight numel ({model_weight.numel()}), " + f"start_offset={start_offset}, master_weight numel={master_weight.numel()}" + ) + + if use_fsdp_shard_model_weights: + target = model_weight_fragment + if target is None: + raise RuntimeError("FSDP shard model weight is required for Identity writeback") + if isinstance(target, IdentityTensorStorage): + target_flat = _identity_storage_data(target).reshape(-1) + else: + target_flat = target.reshape(-1) + target_slice = target_flat[: master_weight.numel()] + else: + target_slice = _identity_storage_data(model_weight).reshape(-1)[start_offset:end_offset] + + if target_slice.numel() != master_weight.numel(): + raise ValueError( + f"Identity target slice has {target_slice.numel()} elements, " + f"but master_weight has {master_weight.numel()}" + ) + target_slice.copy_(master_weight.reshape(-1)) + + def _cast_master_weights_to_fp8_mxfp8_scaling( params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False ): # pylint: disable=unused-argument @@ -963,9 +1032,10 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # bucket matching its own sub-quantizer type. Row and col make their own decisions and # can mix any pair of currently-supported sub-quantizers. # -# Supported (per-tensor Float8 sub-quantizers, in any per-direction combination): +# Supported (per-tensor Float8 or Identity sub-quantizers, any direction): # - Float8Quantizer (delayed scaling) # - Float8CurrentScalingQuantizer (current scaling) +# - IdentityQuantizer (high-precision passthrough) # # Per-tensor Float8 works because `_cast_master_weights_to_fp8_{delayed,current}_scaling` # accept any Float8Tensor (single direction is fine — each entry is one Float8Tensor @@ -974,11 +1044,9 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # independent entries (into the same bucket for same-format, or into different # buckets for cross-format Float8 — e.g. delayed row + current col). # -# Single-direction hybrid (only one sub-storage populated, e.g. after -# `update_usage(columnwise=False)`) routes the present direction only — the -# per-direction loop skips dropped sub-storages. Both-None hybrids raise ValueError. -# Per-block sub-quantizers still hit their per-direction TODO regardless of single -# vs both direction. +# Identity routes to an exact copy bucket. Single-direction hybrid (only one +# sub-storage populated) routes the present direction only. Both-None hybrids +# raise ValueError. Per-block sub-quantizers still hit their per-direction TODO. # # Not supported (raise NotImplementedError per-direction + TODO): # @@ -1032,6 +1100,7 @@ def _route_hybrid_to_buckets( *, delayed_scaling_params, current_scaling_params, + identity_params, ): """Decompose a `HybridQuantizedTensor` into per-direction entries and route each into the appropriate per-format bucket used by `quantize_master_weights`. @@ -1064,13 +1133,22 @@ def _route_hybrid_to_buckets( ): if sub_storage is None: continue - entry = (sub_storage, master_weight, start_offset, fsdp_shard_model_weight) + shard_fragment = fsdp_shard_model_weight + if shard_fragment is not None and isinstance(shard_fragment, HybridQuantizedTensor): + shard_fragment = ( + shard_fragment._rowwise_storage + if direction == "rowwise" + else shard_fragment._columnwise_storage + ) + entry = (sub_storage, master_weight, start_offset, shard_fragment) if isinstance(sub_q, Float8Quantizer): # Delayed scaling: the per-format helper iterates entries # independently and does a per-DP amax all-reduce across the bucket. delayed_scaling_params.append(entry) elif isinstance(sub_q, Float8CurrentScalingQuantizer): current_scaling_params.append(entry) + elif isinstance(sub_q, IdentityQuantizer): + identity_params.append(entry) elif isinstance(sub_q, MXFP8Quantizer): # TODO(hybrid-mxfp8-distopt): the distopt cast kernels are # bidirectional, so a single-direction hybrid sub-storage cannot be @@ -1119,12 +1197,8 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten For NVFP4 tensors, uses batched multi-tensor processing to reduce CPU overhead. - For `HybridQuantizedTensor`, recurses per-direction so that each - sub-storage's native post-processing runs (e.g. Float8 Hopper transpose-cache - pre-creation). Per-block sub-quantizers are rejected at - `quantize_master_weights` time, so by the time we reach here each present - sub-storage is a `Float8Tensor` and the recursive call hits the native - Float8 branch above. + For `HybridQuantizedTensor`, recurses per-direction so each present + sub-storage runs its native post-processing. Identity sub-storages are no-op. """ if not isinstance(model_weights, list): model_weights = [model_weights] @@ -1147,13 +1221,11 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten elif isinstance(model_weight, MXFP8Tensor): # MXFP8 scaling: no need to do anything. pass + elif isinstance(model_weight, IdentityTensorStorage): + pass elif isinstance(model_weight, HybridQuantizedTensor): - # Per-direction post-processing: each Float8 sub-storage routes - # through the recursive call (None / other-type sub-storages are - # silently skipped by the isinstance filter — they would have been - # rejected upstream in `quantize_master_weights`). for sub in (model_weight._rowwise_storage, model_weight._columnwise_storage): - if isinstance(sub, Float8Tensor): + if sub is not None: post_all_gather_processing(sub) elif isinstance(model_weight, QuantizedTensor): raise ValueError(f"post_processing for {type(model_weight)} is not supported") From 8cc3332249addd5f0c32bc17d1ff568d965128af Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 9 Jun 2026 16:12:28 +0000 Subject: [PATCH 20/97] Bug fixing Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/fsdp2_utils.py | 7 + .../fsdp2_tests/run_fsdp2_model.py | 73 +++++ tests/pytorch/distributed/run_hybrid_tp_sp.py | 132 ++++++++- .../pytorch/distributed/test_hybrid_tp_sp.py | 28 ++ tests/pytorch/test_hybrid_quantization.py | 95 +++++++ tests/pytorch/test_identity_quantizer.py | 266 +++++++++++++++++- transformer_engine/pytorch/__init__.py | 12 + .../pytorch/module/grouped_linear.py | 157 ++++++++++- .../pytorch/module/layernorm_linear.py | 3 + .../pytorch/module/layernorm_mlp.py | 8 +- .../pytorch/tensor/float8_tensor.py | 5 +- .../pytorch/tensor/hybrid_tensor.py | 68 +++-- .../pytorch/tensor/identity_tensor.py | 41 +-- .../pytorch/tensor/mxfp8_tensor.py | 25 +- .../float8_blockwise_tensor_storage.py | 2 +- .../tensor/storage/hybrid_tensor_storage.py | 5 + transformer_engine/pytorch/tensor/utils.py | 8 +- 17 files changed, 855 insertions(+), 80 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index f53c237e406..257760da37b 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -74,6 +74,11 @@ def _hybrid_fp8_current_identity_qfactory(role): return current_scaling_quantizer_factory(role) +def _identity_qfactory(role): # pylint: disable=unused-argument + """High-precision passthrough for every quantizer slot.""" + return IdentityQuantizer() + + # The qfactories above are registered here as module-level functions (not # lambdas or closures) on purpose: DCP serializes ``CustomRecipe`` via # ``pickle``, and closure-based qfactories (or inner functions capturing state) @@ -85,6 +90,7 @@ def _hybrid_fp8_current_identity_qfactory(role): "HybridFloat8BlockScaling": _hybrid_float8_block_qfactory, "HybridMixed_MXFP8_FP8": _hybrid_mixed_mxfp8_fp8_qfactory, "HybridFP8CurrentScalingIdentity": _hybrid_fp8_current_identity_qfactory, + "Identity": _identity_qfactory, } @@ -101,6 +107,7 @@ def get_hybrid_recipe_from_string(recipe): "HybridFloat8BlockScaling" — Float8 block scaling for both directions "HybridMixed_MXFP8_FP8" — MXFP8 rowwise + FP8 current columnwise "HybridFP8CurrentScalingIdentity" — FP8 current forward + Identity backward + "Identity" — high-precision passthrough for every slot """ if recipe not in _HYBRID_QFACTORIES: raise ValueError( diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 4624c995f1f..2c97b04d12f 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -536,6 +536,79 @@ def _hybrid_param_count(): _check_hybrid_fsdp2_allgather(model) +def test_distributed_hybrid_identity_all(): + """FSDP2 training/all-gather with an all-Identity CustomRecipe. + + This is the high-precision passthrough baseline for the hybrid/identity + tensor plumbing: quantized_model_init should produce IdentityTensor local + shards, optimizer steps should preserve that type, and FSDP2 all-gather + should reconstruct the same high-precision values as a manual gather. + """ + from transformer_engine.pytorch.tensor.identity_tensor import IdentityTensor + from fsdp2_utils import get_hybrid_recipe_from_string + + identity_recipe = get_hybrid_recipe_from_string("Identity") + world_size = int(os.environ.get("WORLD_SIZE", "1")) + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + device="meta", + ) + with te.quantized_model_init(enabled=True, recipe=identity_recipe): + model = torch.nn.Sequential( + *[te.TransformerLayer(512, 2048, 8, **kwargs) for _ in range(2)] + ) + + custom_attrs = save_custom_attrs(model) + mesh = get_device_mesh(world_size, [world_size]) + model = shard_model_with_fsdp2(model, mesh) + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + + def _identity_param_count(): + return sum( + 1 + for p in model.parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, IdentityTensor) + ) + + identity_count = _identity_param_count() + assert identity_count > 0, "No IdentityTensor local tensors after FSDP2 sharding" + + optimizer = optim.Adam(model.parameters(), lr=1e-3) + input_data = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + + losses = [] + for iteration in range(3): + optimizer.zero_grad() + with te.autocast(enabled=True, recipe=identity_recipe): + output = model(input_data) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + loss_val = loss.item() + assert math.isfinite(loss_val), f"Non-finite Identity loss at iter {iteration}: {loss_val}" + losses.append(loss_val) + dist_print(f"Identity iteration {iteration} completed with loss {loss_val}") + + assert losses[-1] < losses[0], f"Identity loss did not decrease: {losses}" + assert ( + _identity_param_count() == identity_count + ), "IdentityTensor params lost their quantized type after optimizer.step()" + + _check_hybrid_fsdp2_allgather(model) + + def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): """FSDP2 training with hybrid params and reshard_after_forward=True. diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index 1537c949c5c..48986ced752 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -143,6 +143,10 @@ def _hybrid_mxfp8_identity_qfactory(role): return _make_mxfp8_quantizer() +def _identity_qfactory(role): # pylint: disable=unused-argument + return IdentityQuantizer() + + def _make_nvfp4_bare(): """Bare NVFP4Quantizer (1D, no RHT/SR/2D), used by the cross-format recipe to avoid cross-operand RHT-consistency concerns in the mixed MXFP8/NVFP4 @@ -221,6 +225,8 @@ def hybrid_recipe(): return te_recipe.CustomRecipe(qfactory=_hybrid_fp8_identity_qfactory) if QUANTIZATION == "hybrid_mxfp8_identity": return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_identity_qfactory) + if QUANTIZATION == "identity": + return te_recipe.CustomRecipe(qfactory=_identity_qfactory) if QUANTIZATION == "hybrid_nvfp4": return te_recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) if QUANTIZATION == "hybrid_mxfp8_nvfp4": @@ -242,6 +248,10 @@ def hybrid_recipe(): def _get_tolerances(): + if QUANTIZATION == "identity": + # Same tolerance as upstream distributed BF16 numerics: TP row + # reductions can accumulate in a different order from the single-node ref. + return {"rtol": 1.6e-2, "atol": 1.0e-5} if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): # Loose because of sequence parallel & amax reduction (fp8_cs). return {"rtol": 0.4, "atol": 0.25} @@ -572,7 +582,12 @@ def run(recipe): def test_linear_vs_vanilla(): # Cross-format hybrid has no single built-in vanilla recipe to compare # against bitwise; it is covered by the distributed-vs-single-node checks. - if QUANTIZATION in ("hybrid_mxfp8_nvfp4", "hybrid_fp8_identity", "hybrid_mxfp8_identity"): + if QUANTIZATION in ( + "identity", + "hybrid_mxfp8_nvfp4", + "hybrid_fp8_identity", + "hybrid_mxfp8_identity", + ): dist_print("linear_vs_vanilla: skipped for hybrid without a vanilla equivalent") return for parallel_mode in ["column", "row"]: @@ -580,6 +595,116 @@ def test_linear_vs_vanilla(): _test_linear_vs_vanilla(parallel_mode, sequence_parallel) +def _same_format_parity_supported(): + return QUANTIZATION in ("hybrid_fp8", "hybrid_mxfp8") + + +def _check_same_topology_parity(out_h, dinp_h, model_h, out_v, dinp_v, model_v, tag, *, check_grads): + # Larger modules use different fused/unfused norm paths between hybrid and + # vanilla, so numerical parity is the meaningful contract here. Linear keeps + # the stricter bitwise check above. + _check_outputs(out_v, out_h, label=f"{tag} forward") + if check_grads: + _check_outputs(dinp_v, dinp_h, label=f"{tag} dgrad") + _check_gradients(model_h, model_v) + + +def _test_layernorm_linear_vs_vanilla(sequence_parallel, params_dtype=torch.bfloat16): + if not _same_format_parity_supported(): + dist_print("layernorm_linear_vs_vanilla: skipped for recipe without vanilla equivalent") + return + dist_print(f"layernorm_linear_vs_vanilla: sequence_parallel={sequence_parallel}") + + def run(recipe_obj): + torch.manual_seed(23456) + torch.cuda.manual_seed(23456) + model = te.LayerNormLinear( + HIDDEN_SIZE, + HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + parallel_mode="column", + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + torch.manual_seed(45670) + torch.cuda.manual_seed(45670) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp.requires_grad_() + with te.autocast(enabled=True, recipe=recipe_obj): + out = model(inp) + torch.manual_seed(45671) + torch.cuda.manual_seed(45671) + LOSS_FN(out, torch.randn_like(out)).backward() + return model, out.detach().clone(), inp.grad.detach().clone() + + model_h, out_h, dinp_h = run(hybrid_recipe()) + model_v, out_v, dinp_v = run(vanilla_recipe()) + _check_same_topology_parity( + out_h, + dinp_h, + model_h, + out_v, + dinp_v, + model_v, + f"layernorm_linear_vs_vanilla[sp={sequence_parallel}]", + check_grads=not sequence_parallel, + ) + + +def test_layernorm_linear_vs_vanilla(): + for sequence_parallel in [False, True]: + _test_layernorm_linear_vs_vanilla(sequence_parallel) + + +def _test_layernorm_mlp_vs_vanilla(sequence_parallel, params_dtype=torch.bfloat16): + if not _same_format_parity_supported(): + dist_print("layernorm_mlp_vs_vanilla: skipped for recipe without vanilla equivalent") + return + dist_print(f"layernorm_mlp_vs_vanilla: sequence_parallel={sequence_parallel}") + + def run(recipe_obj): + torch.manual_seed(45678) + torch.cuda.manual_seed(45678) + model = te.LayerNormMLP( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + tp_size=WORLD_SIZE, + tp_group=NCCL_WORLD, + set_parallel_mode=True, + sequence_parallel=sequence_parallel, + params_dtype=params_dtype, + ).cuda() + torch.manual_seed(56780) + torch.cuda.manual_seed(56780) + inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) + inp.requires_grad_() + with te.autocast(enabled=True, recipe=recipe_obj): + out = model(inp) + torch.manual_seed(56781) + torch.cuda.manual_seed(56781) + LOSS_FN(out, torch.randn_like(out)).backward() + return model, out.detach().clone(), inp.grad.detach().clone() + + model_h, out_h, dinp_h = run(hybrid_recipe()) + model_v, out_v, dinp_v = run(vanilla_recipe()) + _check_same_topology_parity( + out_h, + dinp_h, + model_h, + out_v, + dinp_v, + model_v, + f"layernorm_mlp_vs_vanilla[sp={sequence_parallel}]", + check_grads=not sequence_parallel, + ) + + +def test_layernorm_mlp_vs_vanilla(): + for sequence_parallel in [False, True]: + _test_layernorm_mlp_vs_vanilla(sequence_parallel) + + # ── Test 2: te.LayerNormLinear column + SP ────────────────────────── @@ -784,6 +909,7 @@ def main(argv=None): "hybrid_mxfp8", "hybrid_fp8_identity", "hybrid_mxfp8_identity", + "identity", "hybrid_nvfp4", "hybrid_mxfp8_nvfp4", ], @@ -796,6 +922,8 @@ def main(argv=None): "all", "linear", "linear_vs_vanilla", + "layernorm_linear_vs_vanilla", + "layernorm_mlp_vs_vanilla", "layernorm_linear", "layernorm_mlp", "transformer_layer", @@ -808,6 +936,8 @@ def main(argv=None): test_map = { "linear": test_linear, "linear_vs_vanilla": test_linear_vs_vanilla, + "layernorm_linear_vs_vanilla": test_layernorm_linear_vs_vanilla, + "layernorm_mlp_vs_vanilla": test_layernorm_mlp_vs_vanilla, "layernorm_linear": test_layernorm_linear, "layernorm_mlp": test_layernorm_mlp, "transformer_layer": test_transformer_layer, diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py index 8dc2995761d..23ae42d8325 100644 --- a/tests/pytorch/distributed/test_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -82,6 +82,21 @@ def test_hybrid_fp8_linear_vs_vanilla(): _run_test("hybrid_fp8", "linear_vs_vanilla") +def test_hybrid_fp8_layernorm_linear_vs_vanilla(): + """Same-topology numerical parity against vanilla FP8 for LayerNormLinear. + + This extends the Linear bitwise operand check to the unfused-norm hybrid + module path; exact bitwise parity is not required because vanilla may use + fused quantized norm while hybrid routes through high-precision norm. + """ + _run_test("hybrid_fp8", "layernorm_linear_vs_vanilla") + + +def test_hybrid_fp8_layernorm_mlp_vs_vanilla(): + """Same-topology numerical parity against vanilla FP8 for LayerNormMLP.""" + _run_test("hybrid_fp8", "layernorm_mlp_vs_vanilla") + + @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_layernorm_linear(): """Column-parallel ``te.LayerNormLinear`` with and without SP. @@ -116,6 +131,11 @@ def test_hybrid_fp8_identity_linear(): _run_test("hybrid_fp8_identity", "linear") +def test_identity_all_modules(): + """All-Identity TP/SP end-to-end coverage for every supported TE module.""" + _run_test("identity", "all") + + # ────────────────────────────────────────────────────────────────────── # Hybrid MXFP8 (rowwise + columnwise same format) # ────────────────────────────────────────────────────────────────────── @@ -137,6 +157,14 @@ def test_hybrid_mxfp8_linear_vs_vanilla(): _run_test("hybrid_mxfp8", "linear_vs_vanilla") +def test_hybrid_mxfp8_layernorm_linear_vs_vanilla(): + _run_test("hybrid_mxfp8", "layernorm_linear_vs_vanilla") + + +def test_hybrid_mxfp8_layernorm_mlp_vs_vanilla(): + _run_test("hybrid_mxfp8", "layernorm_mlp_vs_vanilla") + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") def test_hybrid_mxfp8_layernorm_linear(): _run_test("hybrid_mxfp8", "layernorm_linear") diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index ccbc1041df8..95e5f555bd4 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -101,6 +101,15 @@ def test_creation(self): assert isinstance(hq.rowwise_quantizer, Float8CurrentScalingQuantizer) assert isinstance(hq.columnwise_quantizer, NVFP4Quantizer) + def test_rejects_same_sub_quantizer_instance_for_both_directions(self): + quantizer = _make_fp8_quantizer() + + with pytest.raises(ValueError, match="requires distinct rowwise and columnwise"): + HybridQuantizer(rowwise_quantizer=quantizer, columnwise_quantizer=quantizer) + + assert quantizer.rowwise_usage is True + assert quantizer.columnwise_usage is True + def test_compatible_recipe_is_custom_recipe(self): hq = _make_hybrid_quantizer_fp8_row_fp4_col() assert hq._get_compatible_recipe() is recipe.CustomRecipe @@ -4234,6 +4243,82 @@ def _get_dispatch_hybrid_param(config_name): raise ValueError(f"Unknown config: {config_name}") +@requires_fp8 +class TestFloat8TransposeOnlySplit: + """Regression coverage for columnwise-only Float8 split metadata. + + A columnwise-only per-tensor Float8 sub-storage may have ``_data=None`` and + store its bytes in ``_transpose`` with physical shape ``[K, M]``. Splitting + that tensor must still produce pieces whose wrapper shape is the logical + row-major shape ``[M_i, K]``; otherwise HybridQuantizedTensor uses the + transposed shape when rowwise storage is absent. + """ + + @staticmethod + def _make_transpose_only_float8_tensor(shape=(12, 16)): + m, k = shape + data_transpose = torch.empty((k, m), dtype=torch.uint8, device="cuda") + return Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=data_transpose, + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + requires_grad=False, + device="cuda", + ) + + @pytest.mark.parametrize( + "split_size,dim,expected_shapes,expected_transpose_shapes", + [ + (5, 0, [(5, 16), (5, 16), (2, 16)], [(16, 5), (16, 5), (16, 2)]), + (6, 1, [(12, 6), (12, 6), (12, 4)], [(6, 12), (6, 12), (4, 12)]), + ], + ) + def test_float8_split_uses_logical_shape_for_transpose_only_storage( + self, split_size, dim, expected_shapes, expected_transpose_shapes + ): + tensor = self._make_transpose_only_float8_tensor() + + pieces = torch.split(tensor, split_size, dim=dim) + + assert [tuple(piece.shape) for piece in pieces] == expected_shapes + assert [ + tuple(piece._transpose.shape) for piece in pieces + ] == expected_transpose_shapes + assert all(piece._data is None for piece in pieces) + assert all(piece._transpose_invalid is False for piece in pieces) + + def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self): + columnwise = self._make_transpose_only_float8_tensor() + hybrid = HybridQuantizedTensor( + shape=columnwise.shape, + dtype=columnwise.dtype, + rowwise_storage=None, + columnwise_storage=columnwise, + rowwise_quantizer=None, + columnwise_quantizer=None, + quantizer=None, + device="cuda", + ) + + pieces = torch.split(hybrid, 5, dim=0) + + assert [tuple(piece.shape) for piece in pieces] == [ + (5, 16), + (5, 16), + (2, 16), + ] + assert all(piece.rowwise_sub_storage is None for piece in pieces) + assert [ + tuple(piece.columnwise_sub_storage.shape) for piece in pieces + ] == [(5, 16), (5, 16), (2, 16)] + assert [ + tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces + ] == [(16, 5), (16, 5), (16, 2)] + + @requires_fp8 class TestHybridTorchDispatchFSDP2Ops: """Test aten ops that FSDP2 relies on to preserve the HybridQuantizedTensor type. @@ -4399,6 +4484,8 @@ def _make_fsdp_protocol_param(config_name): r = _hybrid_custom_recipe(_fp8_row_factory, _fp8_col_factory, _fp8_grad_factory) elif config_name == "mxfp8_fp8": r = _hybrid_custom_recipe(_mxfp8_factory, _fp8_col_factory, _fp8_grad_factory) + elif config_name == "block_fp8": + r = recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) else: raise ValueError(f"Unknown config: {config_name}") with quantized_model_init(enabled=True, recipe=r): @@ -4409,6 +4496,8 @@ def _make_fsdp_protocol_param(config_name): _fsdp_protocol_configs = [pytest.param("fp8_fp8", id="same-format")] if mxfp8_available: _fsdp_protocol_configs.append(pytest.param("mxfp8_fp8", id="mixed-mxfp8-fp8")) +if fp8_block_scaling_available: + _fsdp_protocol_configs.append(pytest.param("block_fp8", id="same-format-block-fp8")) @requires_fp8 @@ -4742,6 +4831,12 @@ def test_scale_refresh_across_iterations(self): "weight; the scale-refresh invariant is not being exercised" ) + @pytest.mark.xfail( + reason=( + "Hybrid FSDP2 does not support NVFP4 sub-storages yet; NVFP4 uses " + "dedicated tensor hooks and does not implement the hybrid fsdp_buffer_fields protocol." + ) + ) def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): """Hybrid FSDP2 with an NVFP4 sub-storage must raise a clear error. diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index cf083dd0d65..3f5a6a6909f 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -37,6 +37,7 @@ te.is_fp8_block_scaling_available(return_reason=True) ) + # ── Module-level qfactories (picklable / autocast-friendly) ────────── @@ -217,7 +218,160 @@ def test_internal_returns_storage(self): assert isinstance(out, IdentityTensorStorage) assert not isinstance(out, IdentityTensor) + def test_grouped_split_all_identity_uses_plain_tensor_views(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _split_quantize_with_identity_fallback, + ) + + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + m_splits = [3, 5] + quantizers = [IdentityQuantizer(), IdentityQuantizer()] + + out = _split_quantize_with_identity_fallback( + x, m_splits, quantizers, activation_dtype=torch.bfloat16 + ) + + assert all(isinstance(t, torch.Tensor) for t in out) + assert not any(isinstance(t, IdentityTensorStorage) for t in out) + for actual, expected in zip(out, torch.split(x, m_splits)): + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + cast_quantizers = [ + IdentityQuantizer(dtype=torch.float32), + IdentityQuantizer(dtype=torch.float32), + ] + cast_out = _split_quantize_with_identity_fallback( + x, m_splits, cast_quantizers, activation_dtype=torch.bfloat16 + ) + assert all(isinstance(t, IdentityTensorStorage) for t in cast_out) + assert all(t.dequantize().dtype == torch.float32 for t in cast_out) + + def test_grouped_split_rejects_mixed_identity_and_quantized_operands(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _split_quantize_with_identity_fallback, + ) + + x = torch.empty(64, 64, dtype=torch.bfloat16) + m_splits = [32, 32] + cases = [ + [IdentityQuantizer(), _mxfp8(tex.DType.kFloat8E4M3)], + [ + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ), + HybridQuantizer( + rowwise_quantizer=_mxfp8(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ), + ], + ] + + for quantizers in cases: + with pytest.raises(ValueError, match="mixes Identity-backed and non-Identity-backed"): + _split_quantize_with_identity_fallback( + x, + m_splits, + quantizers, + activation_dtype=torch.bfloat16, + ) + + def test_hybrid_split_forwards_disable_bulk_allocation_to_both_directions( + self, monkeypatch + ): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + from transformer_engine.pytorch.module.grouped_linear import _hybrid_split_quantize + + calls = [] + + def fake_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocation=False): + calls.append(disable_bulk_allocation) + return [ + quantizer(tensor_part) + for tensor_part, quantizer in zip(torch.split(tensor, m_splits), quantizers) + ] + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", fake_split_quantize) + x = torch.randn(8, 16, dtype=torch.bfloat16) + m_splits = [3, 5] + quantizers = [ + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + for _ in m_splits + ] + + out = _hybrid_split_quantize( + x, + m_splits, + quantizers, + disable_bulk_allocation=True, + ) + + assert calls == [True, True] + assert len(out) == len(m_splits) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_grouped_linear_cpu_offload_disables_bulk_allocation_for_hybrid_input( + self, monkeypatch + ): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + class StopAfterFlagCapture(RuntimeError): + pass + + def qfactory(role): + if role is not None and role.module_type == "grouped_linear": + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + ) + return _fp8_cs(tex.DType.kFloat8E4M3) + + calls = [] + + def fake_hybrid_split_quantize( + tensor, m_splits, quantizers, *, disable_bulk_allocation=False + ): + del tensor, m_splits, quantizers + calls.append(disable_bulk_allocation) + raise StopAfterFlagCapture("captured hybrid split kwargs") + + monkeypatch.setattr(grouped_linear, "is_cpu_offload_enabled", lambda: True) + monkeypatch.setattr(grouped_linear, "_hybrid_split_quantize", fake_hybrid_split_quantize) + + model = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + + with pytest.raises(StopAfterFlagCapture): + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + model(x, m_splits=m_splits) + + assert calls == [True] + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_grouped_linear_rejects_mixed_identity_weight_quantizers(self): + weight_count = 0 + + def qfactory(role): + nonlocal weight_count + if role is not None and role.module_type == "grouped_linear": + if role.tensor_type == "weight": + weight_count += 1 + if weight_count == 1: + return IdentityQuantizer() + return _mxfp8(tex.DType.kFloat8E4M3) + return _mxfp8(tex.DType.kFloat8E4M3) + model = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + + with pytest.raises(ValueError, match="mixes Identity-backed and non-Identity-backed"): + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + model(x, m_splits=m_splits) def test_dequantize_bitwise_identical(self): x = torch.randn(4, 32, device="cuda", dtype=torch.bfloat16) @@ -306,6 +460,57 @@ def test_fsdp_pre_post_all_gather_roundtrip(self): assert reuse is gathered torch.testing.assert_close(reuse.dequantize(), x, rtol=0.0, atol=0.0) + def test_torch_weights_only_load_preserves_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + t = IdentityQuantizer()(x) + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, IdentityTensor) + assert isinstance(loaded._quantizer, IdentityQuantizer) + torch.testing.assert_close(loaded.dequantize(), x, rtol=0.0, atol=0.0) + + def test_torch_weights_only_load_preserves_hybrid_identity_tensor(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, HybridQuantizedTensor) + assert isinstance(loaded._quantizer, HybridQuantizer) + assert isinstance(loaded._rowwise_storage, IdentityTensor) + assert isinstance(loaded._columnwise_storage, IdentityTensor) + torch.testing.assert_close(loaded.dequantize(), x, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_torch_weights_only_load_preserves_hybrid_mxfp8_identity_tensor(self): + x = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=_mxfp8(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + expected = t.dequantize() + buffer = io.BytesIO() + torch.save(t, buffer) + buffer.seek(0) + + loaded = torch.load(buffer, weights_only=True) + + assert isinstance(loaded, HybridQuantizedTensor) + assert isinstance(loaded._quantizer, HybridQuantizer) + assert isinstance(loaded._columnwise_storage, IdentityTensor) + torch.testing.assert_close(loaded.dequantize(), expected, rtol=0.0, atol=0.0) def test_cpu_offload_roundtrip_identity_exact(self): x = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) @@ -544,12 +749,65 @@ def test_identity_recipe_matches_bf16_bitwise(self, module_name, qfactory): y_ref, dx_ref, wg_ref = _fwd_bwd_module(module_name, ref, x, recipe=None) y_id, dx_id, wg_id = _fwd_bwd_module(module_name, test, x, recipe=recipe) - torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) - torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + # Linear / LayerNormLinear / GroupedLinear route through the same HP + # math with Identity and should stay bitwise exact. Composite modules + # can select different fused/unfused BF16 kernel paths after prior FP8 + # tests have warmed TE/CUDA state, so require tight BF16 numerical + # parity instead of order-dependent bitwise identity. + kwargs = ( + {"rtol": 0.0, "atol": 0.0} + if module_name in ("Linear", "LayerNormLinear", "GroupedLinear") + else {"rtol": 2.0e-2, "atol": 8.0e-3} + ) + torch.testing.assert_close(y_id, y_ref, **kwargs) + torch.testing.assert_close(dx_id, dx_ref, **kwargs) assert len(wg_id) == len(wg_ref) for g_id, g_ref in zip(wg_id, wg_ref): - torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(g_id, g_ref, **kwargs) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_grouped_linear_mxfp8_forward_identity_backward_matches_override(self): + def mxfp8_all_factory(role): # pylint: disable=unused-argument + return _mxfp8(tex.DType.kFloat8E4M3) + + def run(model, x, recipe): + x = x.detach().clone().requires_grad_(True) + m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) + with te.autocast(enabled=True, recipe=recipe): + y = model(x, m_splits=m_splits) + torch.manual_seed(9001) + target = torch.randn_like(y) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + wgrads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None] + return y.detach().clone(), x.grad.detach().clone(), wgrads + + torch.manual_seed(8300) + ref = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + torch.manual_seed(8301) + test = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() + with torch.no_grad(): + for p_test, p_ref in zip(test.parameters(), ref.parameters()): + p_test.copy_(p_ref) + + torch.manual_seed(8302) + x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + y_bo, dx_bo, wg_bo = run( + ref, + x, + CustomRecipe(qfactory=mxfp8_all_factory, backward_override="high_precision"), + ) + y_id, dx_id, wg_id = run( + test, + x, + CustomRecipe(qfactory=_hybrid_quantized_fwd_identity_bwd_factory("mxfp8")), + ) + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) class TestIdentityHybridFormatProtocols: @@ -906,7 +1164,7 @@ def test_quantized_model_init_identity_state_dict_save_load_exact(self): with te.quantized_model_init(enabled=True, recipe=recipe): model2 = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16).cuda() - model2.load_state_dict(torch.load(buffer)) + model2.load_state_dict(torch.load(buffer, weights_only=True)) with torch.no_grad(), te.autocast(enabled=True, recipe=recipe): out_after = model2(x) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 9a541a8fd40..8b62cd2d0a1 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -107,6 +107,12 @@ from transformer_engine.pytorch.tensor.float8_blockwise_tensor import ( _make_float8_blockwise_tensor_in_reduce_ex, ) +from transformer_engine.pytorch.tensor.hybrid_tensor import ( + _make_hybrid_quantized_tensor_in_reduce_ex, +) +from transformer_engine.pytorch.tensor.identity_tensor import ( + _make_identity_tensor_in_reduce_ex, +) try: torch._dynamo.config.error_on_nested_jit_trace = False @@ -131,6 +137,8 @@ MXFP8TensorStorage, NVFP4TensorStorage, Float8BlockwiseQTensorStorage, + HybridQuantizedTensorStorage, + IdentityTensorStorage, # Quantizer types embedded in metadata Quantizer, Float8Quantizer, @@ -138,6 +146,8 @@ MXFP8Quantizer, NVFP4Quantizer, Float8BlockQuantizer, + HybridQuantizer, + IdentityQuantizer, # pybind11 enum used as Quantizer.dtype tex.DType, # __reduce_ex__ reconstructors (module-level functions). @@ -145,6 +155,8 @@ _make_mxfp8_tensor_in_reduce_ex, _make_nvfp4_tensor_in_reduce_ex, _make_float8_blockwise_tensor_in_reduce_ex, + _make_hybrid_quantized_tensor_in_reduce_ex, + _make_identity_tensor_in_reduce_ex, ] ) except (ImportError, AttributeError): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 26d6f01e660..475c7e25758 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -54,6 +54,7 @@ from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.identity_tensor import IdentityQuantizer from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, @@ -65,6 +66,91 @@ from ...debug.pytorch.debug_state import TEDebugState +def _uses_identity_quantizer(quantizer): + """Whether a quantizer, including a hybrid sub-quantizer, is Identity-backed.""" + if quantizer is None: + return False + if isinstance(quantizer, IdentityQuantizer): + return True + if isinstance(quantizer, HybridQuantizer): + return _uses_identity_quantizer( + quantizer.rowwise_quantizer + ) or _uses_identity_quantizer(quantizer.columnwise_quantizer) + return False + + +def _has_identity_quantizer_list(quantizers): + """Whether any quantizer in a grouped list uses Identity.""" + return any(_uses_identity_quantizer(q) for q in quantizers) + + +def _identity_quantizer_signature(quantizer): + """Identity usage per GEMM direction: (rowwise, columnwise).""" + if isinstance(quantizer, HybridQuantizer): + return ( + _uses_identity_quantizer(quantizer.rowwise_quantizer), + _uses_identity_quantizer(quantizer.columnwise_quantizer), + ) + identity = isinstance(quantizer, IdentityQuantizer) + return (identity, identity) + + +def _check_uniform_identity_quantizer_list(quantizers): + """Reject grouped lists that mix Identity-backed and quantized directions.""" + signatures = [_identity_quantizer_signature(q) for q in quantizers] + if not any(rowwise or columnwise for rowwise, columnwise in signatures): + return + if all(signature == signatures[0] for signature in signatures): + return + raise ValueError( + "GroupedLinear quantizer list mixes Identity-backed and non-Identity-backed" + f" directions: {signatures}. This combination is not supported because" + " grouped GEMM requires a uniform scaling mode for every tensor in each" + " operand list. Make the CustomRecipe `qfactory` return Identity" + " consistently for every expert in the grouped operand, or return no" + " Identity quantizers for that operand." + ) + + +def _is_plain_identity_passthrough_list(quantizers, activation_dtype): + """Whether every quantizer is a plain Identity passthrough.""" + return bool(quantizers) and all( + isinstance(q, IdentityQuantizer) and (q.dtype is None or q.dtype == activation_dtype) + for q in quantizers + ) + + +def _split_quantize_with_identity_fallback( + tensor, + m_splits, + quantizers, + activation_dtype, + *, + disable_bulk_allocation=False, +): + """Split+quantize, avoiding native C++ split kernels for Identity quantizers.""" + # No Identity anywhere: use the native grouped split+quantize kernel. + if not _has_identity_quantizer_list(quantizers): + return tex.split_quantize( + tensor, + m_splits, + quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + + _check_uniform_identity_quantizer_list(quantizers) + tensor = cast_if_needed(tensor, activation_dtype) + # Plain all-Identity passthrough: match the native high-precision BF16 split path. + if _is_plain_identity_passthrough_list(quantizers, activation_dtype): + return torch.split(tensor, m_splits) + + # Uniform Identity-backed wrappers still need per-split Python quantizer calls. + return [ + quantizer(tensor_part) if quantizer is not None else tensor_part + for tensor_part, quantizer in zip(torch.split(tensor, m_splits), quantizers) + ] + + def _is_hybrid_quantizer_list(quantizers): """Classify a GroupedLinear quantizer list as hybrid-uniform or plain-uniform. @@ -110,7 +196,7 @@ def _is_hybrid_quantizer_list(quantizers): ) -def _hybrid_split_quantize(tensor, m_splits, quantizers): +def _hybrid_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocation=False): """Grouped split+quantize for an **all-hybrid** quantizer list. Precondition: every ``q`` in ``quantizers`` is a ``HybridQuantizer``. @@ -135,8 +221,18 @@ def _hybrid_split_quantize(tensor, m_splits, quantizers): row_quantizers = [q.rowwise_quantizer for q in quantizers] col_quantizers = [q.columnwise_quantizer for q in quantizers] - row_results = tex.split_quantize(tensor, m_splits, row_quantizers) - col_results = tex.split_quantize(tensor, m_splits, col_quantizers) + row_results = tex.split_quantize( + tensor, + m_splits, + row_quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + col_results = tex.split_quantize( + tensor, + m_splits, + col_quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) return [ HybridStorage( @@ -565,6 +661,11 @@ def forward( for output_quantizer in output_quantizers: output_quantizer.set_usage(rowwise=True, columnwise=False) + if fp8 and not debug: + _check_uniform_identity_quantizer_list(input_quantizers) + _check_uniform_identity_quantizer_list(weight_quantizers) + _check_uniform_identity_quantizer_list(grad_output_quantizers) + # Initialize input tensors in_features = weights[0].size(-1) if inp.size(-1) != in_features: @@ -614,19 +715,26 @@ def forward( inp_view = inp.reshape(-1, in_features) inputmats: list - hybrid = _is_hybrid_quantizer_list(input_quantizers) + identity = _has_identity_quantizer_list(input_quantizers) + hybrid = False if identity else _is_hybrid_quantizer_list(input_quantizers) if fp8 and not debug and not hybrid: # Disable bulk allocation when CPU offloading is active: offloading skips small # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. - inputmats = tex.split_quantize( + inputmats = _split_quantize_with_identity_fallback( inp_view, m_splits, input_quantizers, + activation_dtype, disable_bulk_allocation=cpu_offloading, ) elif fp8 and hybrid: - inputmats = _hybrid_split_quantize(inp_view, m_splits, input_quantizers) + inputmats = _hybrid_split_quantize( + inp_view, + m_splits, + input_quantizers, + disable_bulk_allocation=cpu_offloading, + ) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, input_quantizers, m_splits, activation_dtype @@ -1028,12 +1136,19 @@ def backward( grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) grad_output = [None] * ctx.num_gemms grad_biases = [None] * ctx.num_gemms - grad_output_hybrid = _is_hybrid_quantizer_list(ctx.grad_output_quantizers) + grad_output_identity = _has_identity_quantizer_list(ctx.grad_output_quantizers) + grad_output_hybrid = ( + False + if grad_output_identity + else _is_hybrid_quantizer_list(ctx.grad_output_quantizers) + ) if ctx.fp8 and not ctx.debug and not grad_output_hybrid: if ctx.use_bias: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) recipe = ctx.fp8_recipe - if recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8(): + if not grad_output_identity and ( + recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8() + ): # Fused bias grad + quantize kernel for i in range(ctx.num_gemms): grad_biases[i], grad_output[i] = tex.bgrad_quantize( @@ -1044,17 +1159,19 @@ def backward( # Unfused bias grad and multi-tensor quantize for i in range(ctx.num_gemms): grad_biases[i] = grad_output_mats[i].sum(dim=0) - grad_output = tex.split_quantize( + grad_output = _split_quantize_with_identity_fallback( grad_output_view, ctx.m_splits, ctx.grad_output_quantizers, + ctx.activation_dtype, ) else: # Multi-tensor quantize - grad_output = tex.split_quantize( + grad_output = _split_quantize_with_identity_fallback( grad_output_view, ctx.m_splits, ctx.grad_output_quantizers, + ctx.activation_dtype, ) elif ctx.fp8 and grad_output_hybrid: if ctx.use_bias: @@ -1065,6 +1182,7 @@ def backward( grad_output_view, ctx.m_splits, ctx.grad_output_quantizers, + disable_bulk_allocation=ctx.cpu_offloading, ) elif ctx.debug: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) @@ -1174,12 +1292,25 @@ def backward( else: input_quantizer.set_usage(rowwise=False, columnwise=True) inputmats: list - input_hybrid = _is_hybrid_quantizer_list(ctx.input_quantizers) + input_identity = _has_identity_quantizer_list(ctx.input_quantizers) + input_hybrid = ( + False + if input_identity + else _is_hybrid_quantizer_list(ctx.input_quantizers) + ) if ctx.fp8 and not ctx.debug and not input_hybrid: - inputmats = tex.split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) + inputmats = _split_quantize_with_identity_fallback( + inp_view, + ctx.m_splits, + ctx.input_quantizers, + ctx.activation_dtype, + ) elif ctx.fp8 and input_hybrid: inputmats = _hybrid_split_quantize( - inp_view, ctx.m_splits, ctx.input_quantizers + inp_view, + ctx.m_splits, + ctx.input_quantizers, + disable_bulk_allocation=ctx.cpu_offloading, ) elif ctx.debug: inputmats = DebugQuantizer.multi_tensor_quantize( diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 4ac362aa78c..558d10edabc 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -67,6 +67,7 @@ from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer from ..cpu_offload import ( is_cpu_offload_enabled, start_offload, @@ -219,6 +220,7 @@ def forward( # or if a gather of ln_out must be in high precision. custom = is_custom(input_quantizer) hybrid = isinstance(input_quantizer, HybridQuantizer) + identity = isinstance(input_quantizer, IdentityQuantizer) with_quantized_norm = ( fp8 and not debug @@ -227,6 +229,7 @@ def forward( and backward_override is None and not custom # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() and not hybrid + and not identity ) # Apply normalization diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index f0451351118..0be7a857d97 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -72,6 +72,7 @@ from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer from ._common import apply_normalization, WeightGradStore from ..cpu_offload import ( is_cpu_offload_enabled, @@ -407,6 +408,7 @@ def _forward( custom = is_custom(fc1_input_quantizer) hybrid = isinstance(fc1_input_quantizer, HybridQuantizer) + identity = isinstance(fc1_input_quantizer, IdentityQuantizer) with_quantized_norm = ( fp8 and not debug @@ -414,6 +416,7 @@ def _forward( and not return_layernorm_output_gathered and not custom and not hybrid + and not identity ) # Apply normalization @@ -1428,7 +1431,10 @@ def fc2_wgrad_gemm( if ctx.fp8: # TODO float8 blockwise current scaling (as well as custom quantizers) has no bgrad fusion for now if ( - isinstance(ctx.fc1_grad_output_quantizer, Float8BlockQuantizer) + isinstance( + ctx.fc1_grad_output_quantizer, + (Float8BlockQuantizer, IdentityQuantizer), + ) or ctx.fp8_recipe.custom() ): fc1_bias_grad = dact.view(-1, dact.shape[-1]).sum(dim=0) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index f9af524e0f2..5a040c17c7b 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -674,7 +674,10 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): shape=( split_tensor.shape if split_tensor is not None - else split_transpose_tensor.shape + else ( + *split_transpose_tensor.shape[1:], + split_transpose_tensor.shape[0], + ) ), ) for split_tensor, split_transpose_tensor in zip(func_out, t_func_out) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index bb3adaedb84..072cd8782b2 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -29,6 +29,19 @@ class HybridQuantizer(Quantizer): columnwise_quantizer : Quantizer Quantizer for the columnwise direction (e.g. NVFP4Quantizer). + Notes + ----- + ``HybridQuantizer`` pins each sub-quantizer to its designated direction by + mutating its usage flags, so it takes ownership of the supplied quantizer + instances. The rowwise and columnwise quantizers must be distinct objects. + If both directions need shared state, construct two quantizer instances that + reference the same external state object. + + Reusing a sub-quantizer instance across multiple ``HybridQuantizer`` objects + is unsupported by contract. Catching that robustly would require copying or + ownership tracking, both of which are more intrusive, so only the direct + rowwise/columnwise aliasing case is enforced. + """ rowwise_quantizer: Quantizer @@ -41,6 +54,12 @@ def __init__( columnwise_quantizer: Quantizer, ) -> None: super().__init__(rowwise=True, columnwise=True) + if rowwise_quantizer is columnwise_quantizer: + raise ValueError( + "HybridQuantizer requires distinct rowwise and columnwise quantizer" + " instances. If both directions need shared state, construct two" + " quantizer objects that reference the same shared state." + ) self.rowwise_quantizer = rowwise_quantizer self.columnwise_quantizer = columnwise_quantizer @@ -396,28 +415,6 @@ def detach(self) -> HybridQuantizedTensor: def get_metadata(self) -> Dict[str, Any]: return HybridQuantizedTensorStorage.get_metadata(self) - @classmethod - def _make_in_reduce_ex( - cls, - rowwise_storage: Optional[QuantizedTensorStorage], - columnwise_storage: Optional[QuantizedTensorStorage], - rowwise_quantizer: Optional[Quantizer], - columnwise_quantizer: Optional[Quantizer], - quantizer: Optional[Quantizer], - dtype: torch.dtype, - shape: torch.Size, - ) -> HybridQuantizedTensor: - """Build HybridQuantizedTensor, for use in ``__reduce_ex__``.""" - return HybridQuantizedTensor( - shape=shape, - dtype=dtype, - rowwise_storage=rowwise_storage, - columnwise_storage=columnwise_storage, - rowwise_quantizer=rowwise_quantizer, - columnwise_quantizer=columnwise_quantizer, - quantizer=quantizer, - ) - def __reduce_ex__(self, protocol: int) -> tuple: """Custom pickling. @@ -436,7 +433,7 @@ def __reduce_ex__(self, protocol: int) -> tuple: themselves. """ return ( - HybridQuantizedTensor._make_in_reduce_ex, + _make_hybrid_quantized_tensor_in_reduce_ex, ( self._rowwise_storage, self._columnwise_storage, @@ -450,7 +447,9 @@ def __reduce_ex__(self, protocol: int) -> tuple: # ── FSDP2 protocol ────────────────────────────────────────────── - def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): + def fsdp_pre_all_gather( # pylint: disable=unused-argument + self, mesh, orig_size, contiguous_orig_stride, module, mp_policy + ): """Extract plain tensor buffers from both sub-storages for FSDP2 all-gather. Always send both directions. This gives a stable buffer count/shape @@ -827,3 +826,24 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ) return super().__torch_dispatch__(func, types, args, kwargs) + + +def _make_hybrid_quantized_tensor_in_reduce_ex( + rowwise_storage: Optional[QuantizedTensorStorage], + columnwise_storage: Optional[QuantizedTensorStorage], + rowwise_quantizer: Optional[Quantizer], + columnwise_quantizer: Optional[Quantizer], + quantizer: Optional[Quantizer], + dtype: torch.dtype, + shape: torch.Size, +) -> HybridQuantizedTensor: + """Reconstruct a ``HybridQuantizedTensor`` from its ``__reduce_ex__`` payload.""" + return HybridQuantizedTensor( + shape=shape, + dtype=dtype, + rowwise_storage=rowwise_storage, + columnwise_storage=columnwise_storage, + rowwise_quantizer=rowwise_quantizer, + columnwise_quantizer=columnwise_quantizer, + quantizer=quantizer, + ) diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index a36c1b6ce31..004fcfb1a26 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -185,32 +185,16 @@ def clone(self) -> "IdentityTensor": device=self.device, ) - @classmethod - def _make_in_reduce_ex( - cls, - hp_data: torch.Tensor, - quantizer: Optional[Quantizer], - dtype: torch.dtype, - shape: torch.Size, - ) -> "IdentityTensor": - """Build IdentityTensor, for use in ``__reduce_ex__``.""" - return IdentityTensor( - shape=shape, - dtype=dtype, - hp_data=hp_data, - quantizer=quantizer, - requires_grad=False, - device=hp_data.device if hp_data is not None else None, - ) - def __reduce_ex__(self, protocol: int) -> tuple: """Custom pickling that preserves the high-precision payload.""" return ( - IdentityTensor._make_in_reduce_ex, + _make_identity_tensor_in_reduce_ex, (self._hp_data, self._quantizer, self.dtype, self.shape), ) - def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): + def fsdp_pre_all_gather( # pylint: disable=unused-argument + self, mesh, orig_size, contiguous_orig_stride, module, mp_policy + ): """Extract the high-precision buffer for FSDP2 all-gather.""" return (self._hp_data,), (self._quantizer,) @@ -320,3 +304,20 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): return out return super().__torch_dispatch__(func, types, args, kwargs) + + +def _make_identity_tensor_in_reduce_ex( + hp_data: torch.Tensor, + quantizer: Optional[Quantizer], + dtype: torch.dtype, + shape: torch.Size, +) -> IdentityTensor: + """Reconstruct an ``IdentityTensor`` from its ``__reduce_ex__`` payload.""" + return IdentityTensor( + shape=shape, + dtype=dtype, + hp_data=hp_data, + quantizer=quantizer, + requires_grad=False, + device=hp_data.device if hp_data is not None else None, + ) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 1b356b5aac2..6f5adf31777 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -411,16 +411,9 @@ def _split_data(data): row_data_splits = _split_data(tensor._rowwise_data) col_data_splits = _split_data(tensor._columnwise_data) - scale_invs = [tensor._rowwise_scale_inv, tensor._columnwise_scale_inv] - split_sizes_for_scale = [split_size, split_size // MXFP8_BLOCK_SCALING_SIZE] - padding_multiples = [128, 4] - scale_splits = [] - for scale_inv, scale_split_size, pad_multiple in zip( - scale_invs, split_sizes_for_scale, padding_multiples - ): + def _split_scale_inv(scale_inv, scale_split_size, pad_multiple): if scale_inv is None: - scale_splits.append(None) - continue + return None scale_inv_out = list( scale_inv.__torch_dispatch__( func, @@ -436,8 +429,18 @@ def _split_data(data): scale_inv_out[idx] = torch.nn.functional.pad( split_scale_inv_out, (0, 0, 0, pad_dim0) ) - scale_splits.append(scale_inv_out) - row_scale_splits, col_scale_splits = scale_splits + return scale_inv_out + + row_scale_splits = _split_scale_inv( + tensor._rowwise_scale_inv, + split_size, + 128, + ) + col_scale_splits = _split_scale_inv( + tensor._columnwise_scale_inv, + split_size // MXFP8_BLOCK_SCALING_SIZE, + 4, + ) ref_splits = row_data_splits if row_data_splits is not None else col_data_splits num_splits = len(ref_splits) diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 7161771dab0..a9810c29e7c 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -465,7 +465,7 @@ def fsdp_buffer_fields(self) -> Tuple[str, ...]: fields.extend(("_columnwise_data", "_columnwise_scale_inv")) return tuple(fields) - def fsdp_buffer_fields( + def fsdp_extract_buffers( self, ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: """Extract M-major, alignment-stripped buffers for dim-0 all-gather. diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index 9d152982c72..03af00184d7 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -115,6 +115,7 @@ def restore_from_saved( return tensors def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + """Dequantize using the first available sub-storage.""" if dtype is None: dtype = self._dtype if self._rowwise_storage is not None: @@ -124,6 +125,7 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: raise RuntimeError("HybridQuantizedTensorStorage has no data to dequantize") def get_data_tensors(self): + """Return raw data tensors from both available sub-storages.""" row_tensors = () col_tensors = () if self._rowwise_storage is not None: @@ -135,6 +137,7 @@ def get_data_tensors(self): return row_tensors + col_tensors def size(self, *args, **kwargs): + """Return the logical size from the first available sub-storage.""" if self._rowwise_storage is not None: return self._rowwise_storage.size(*args, **kwargs) if self._columnwise_storage is not None: @@ -143,6 +146,7 @@ def size(self, *args, **kwargs): @property def device(self): + """Return the device from the first available sub-storage.""" if self._rowwise_storage is not None: return self._rowwise_storage.device if self._columnwise_storage is not None: @@ -183,6 +187,7 @@ def view(self, *shape): ) def get_metadata(self) -> Dict[str, Any]: + """Return constructor metadata for make_like and serialization paths.""" return { "rowwise_storage": self._rowwise_storage, "columnwise_storage": self._columnwise_storage, diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 5c595175a85..c24137d56c3 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1172,10 +1172,10 @@ def _route_hybrid_to_buckets( "block above _route_hybrid_to_buckets for details." ) elif isinstance(sub_q, Float8BlockQuantizer): - # TODO(hybrid-fp8-blockwise): same shape as the NVFP4 secondary - # blocker (and only that one — no kernel-level construction - # blocker for Block FP8). Python-side post-AG fix. See top-of-file - # TODO block for details. + # Pending hybrid-fp8-blockwise work: same shape as the NVFP4 + # secondary blocker (and only that one — no kernel-level construction + # blocker for Block FP8). Python-side post-AG fix. See the + # _route_hybrid_to_buckets design note above for details. raise NotImplementedError( "quantize_master_weights for HybridQuantizer with Float8BlockQuantizer " f"{direction} sub-quantizer is not supported yet. See the TODO " From 9b444d0149e6142140eb3c90345d56f7260da9fd Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 10 Jun 2026 16:48:46 +0000 Subject: [PATCH 21/97] More fixes Signed-off-by: Evgeny --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_hybrid_quantization.py | 49 +++++ tests/pytorch/test_identity_quantizer.py | 169 ++++++++++++++++++ .../attention/dot_product_attention/utils.py | 17 +- .../pytorch/module/grouped_linear.py | 38 ++-- .../pytorch/tensor/hybrid_tensor.py | 164 +++++++++++++++++ .../pytorch/tensor/identity_tensor.py | 11 ++ 7 files changed, 431 insertions(+), 18 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a23ca4647b3..2b502413ab8 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -50,6 +50,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_e python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 95e5f555bd4..463de1c4360 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -2456,6 +2456,55 @@ def test_hybrid_split_quantize_rejects_plain_element(self): assert "HybridQuantizer" in msg assert "Float8CurrentScalingQuantizer" in msg + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.parametrize( + ("usage", "expected"), + [ + pytest.param((True, False), {"rowwise": True, "columnwise": False}, id="rowwise"), + pytest.param((False, True), {"rowwise": False, "columnwise": True}, id="columnwise"), + pytest.param((True, True), {"rowwise": True, "columnwise": True}, id="both"), + ], + ) + def test_hybrid_split_quantize_respects_parent_usage_flags(self, usage, expected): + from transformer_engine.pytorch.module.grouped_linear import ( + _hybrid_split_quantize, + ) + + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=usage[0], columnwise=usage[1]) + + out = _hybrid_split_quantize(tensor, [16, 16], quantizers) + + assert [storage.get_usages() for storage in out] == [expected, expected] + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_hybrid_split_quantize_rejects_mixed_parent_usage_flags(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _hybrid_split_quantize, + ) + + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + for _ in range(2) + ] + quantizers[0].set_usage(rowwise=True, columnwise=False) + quantizers[1].set_usage(rowwise=True, columnwise=True) + + with pytest.raises(ValueError, match="mixed parent usage flags"): + _hybrid_split_quantize(tensor, [16, 16], quantizers) + # =========================================================================== # Quantized Parameters (quantized_model_init) tests for hybrid quantization diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 3f5a6a6909f..65df088269b 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -373,6 +373,175 @@ def qfactory(role): with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): model(x, m_splits=m_splits) + def test_identity_contiguous_preserves_wrapper_and_values(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16).t() + t = IdentityQuantizer()(x) + + out = t.contiguous() + + assert isinstance(out, IdentityTensor) + assert out.is_contiguous() + torch.testing.assert_close(out.dequantize(), x.contiguous(), rtol=0.0, atol=0.0) + + def test_hybrid_identity_contiguous_preserves_wrapper_and_values(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + + out = t.contiguous() + + assert out is t + assert isinstance(out, HybridQuantizedTensor) + assert isinstance(out.rowwise_sub_storage, IdentityTensor) + assert isinstance(out.columnwise_sub_storage, IdentityTensor) + torch.testing.assert_close(out.dequantize(), x, rtol=0.0, atol=0.0) + + def test_hybrid_identity_cpu_preserves_nested_storage_types(self): + x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + t = q(x) + + out = t.cpu() + + assert isinstance(out, HybridQuantizedTensor) + assert out.device.type == "cpu" + assert isinstance(out.rowwise_sub_storage, IdentityTensor) + assert isinstance(out.columnwise_sub_storage, IdentityTensor) + assert out.rowwise_sub_storage.device.type == "cpu" + assert out.columnwise_sub_storage.device.type == "cpu" + torch.testing.assert_close(out.dequantize(), x.cpu(), rtol=0.0, atol=0.0) + assert len(out.get_data_tensors()) == 4 + out.copy_(torch.ones_like(x, device="cpu")) + torch.testing.assert_close( + out.dequantize(), torch.ones_like(x, device="cpu"), rtol=0.0, atol=0.0 + ) + + def test_hybrid_quantizer_copy_preserves_parent_flags(self): + q = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + q.set_usage(rowwise=True, columnwise=False) + q.internal = True + q.optimize_for_gemm = True + + out = q.copy() + + assert isinstance(out, HybridQuantizer) + assert out is not q + assert out.rowwise_quantizer is not q.rowwise_quantizer + assert out.columnwise_quantizer is not q.columnwise_quantizer + assert out.rowwise_usage is True + assert out.columnwise_usage is False + assert out.internal is True + assert out.optimize_for_gemm is True + assert out.rowwise_quantizer.rowwise_usage is True + assert out.rowwise_quantizer.columnwise_usage is False + assert out.columnwise_quantizer.rowwise_usage is False + assert out.columnwise_quantizer.columnwise_usage is True + + def test_te_ops_basic_linear_accepts_hybrid_identity_quantized_weight(self): + import transformer_engine.pytorch.ops as te_ops + + def qfactory(role): # pylint: disable=unused-argument + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + + custom_recipe = CustomRecipe(qfactory=qfactory) + with te.quantized_model_init(enabled=True, recipe=custom_recipe): + op = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) + + x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with te.autocast(enabled=True, recipe=custom_recipe): + y = op(x) + y.sum().backward() + + assert isinstance(op.weight, HybridQuantizedTensor) + assert x.grad is not None + + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(lambda role: IdentityQuantizer(), id="identity"), + pytest.param( + lambda role: HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ), + id="hybrid_identity", + ), + ], + ) + def test_te_ops_quantize_then_gelu_accepts_identity_backed_tensors(self, qfactory): + import transformer_engine.pytorch.ops as te_ops + + model = te_ops.Sequential(te_ops.Quantize(forward=True), te_ops.GELU()) + x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + y = model(x) + + assert isinstance(y, torch.Tensor) + assert y.shape == x.shape + + def test_hybrid_fsdp_rejects_storage_only_sub_storages(self): + row_quantizer = IdentityQuantizer() + col_quantizer = IdentityQuantizer() + row_quantizer.internal = True + col_quantizer.internal = True + q = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + ) + t = q(torch.randn(8, 16, device="cuda", dtype=torch.bfloat16)) + + with pytest.raises(NotImplementedError, match="storage-only rowwise sub-storage"): + t.fsdp_pre_all_gather( + mesh=None, + orig_size=t.shape, + contiguous_orig_stride=t.stride(), + module=None, + mp_policy=None, + ) + + def test_hybrid_quantizer_rejects_nested_quantizer_requests(self): + from transformer_engine.pytorch.quantization import DelayedScalingRequest + + with pytest.raises(TypeError, match="does not support nested QuantizerRequest"): + HybridQuantizer( + rowwise_quantizer=DelayedScalingRequest(), + columnwise_quantizer=IdentityQuantizer(), + ) + + def test_fp8_dpa_rejects_identity_quantizer_with_type_error(self): + from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + META_DO, + META_DP, + META_DQKV, + META_O, + META_QKV, + META_S, + ) + + n_fwd = max(META_QKV, META_S, META_O) + 1 + n_bwd = max(META_DO, META_DP, META_DQKV) + 1 + quantizers = { + "scaling_fwd": [IdentityQuantizer() for _ in range(n_fwd)], + "scaling_bwd": [IdentityQuantizer() for _ in range(n_bwd)], + } + + with pytest.raises(TypeError, match="FP8 attention requires FP8-compatible quantizers"): + dpa_utils.get_attention_quantizers(True, quantizers) + def test_dequantize_bitwise_identical(self): x = torch.randn(4, 32, device="cuda", dtype=torch.bfloat16) out = IdentityQuantizer()(x) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 989b65f190c..fd3df99ea04 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2435,14 +2435,15 @@ def get_attention_quantizers(fp8, quantizers): ]: if _q is None and _name in _allow_none: continue - assert isinstance(_q, _fp8_types), ( - "FP8 attention requires FP8-compatible quantizers for all DPA tensor slots, " - f"but {_name} quantizer is {type(_q).__name__}. " - "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an " - "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or " - "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles " - "(boundary slots like O output and dQKV grad-input)." - ) + if not isinstance(_q, _fp8_types): + raise TypeError( + "FP8 attention requires FP8-compatible quantizers for all DPA tensor slots, " + f"but {_name} quantizer is {type(_q).__name__}. " + "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an " + "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or " + "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles " + "(boundary slots like O output and dQKV grad-input)." + ) return QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 475c7e25758..f184d7b6668 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -218,20 +218,38 @@ def _hybrid_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocat f" Got types: {[type(q).__name__ for q in quantizers]}" ) + usage_signatures = [(q.rowwise_usage, q.columnwise_usage) for q in quantizers] + if not all(signature == usage_signatures[0] for signature in usage_signatures): + raise ValueError( + "GroupedLinear HybridQuantizer list has mixed parent usage flags " + f"{usage_signatures}. This is not supported by the grouped " + "split-quantize path; all experts for a grouped operand must " + "request the same rowwise/columnwise directions." + ) + rowwise_enabled, columnwise_enabled = usage_signatures[0] + row_quantizers = [q.rowwise_quantizer for q in quantizers] col_quantizers = [q.columnwise_quantizer for q in quantizers] - row_results = tex.split_quantize( - tensor, - m_splits, - row_quantizers, - disable_bulk_allocation=disable_bulk_allocation, + row_results = ( + tex.split_quantize( + tensor, + m_splits, + row_quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + if rowwise_enabled + else [None] * len(quantizers) ) - col_results = tex.split_quantize( - tensor, - m_splits, - col_quantizers, - disable_bulk_allocation=disable_bulk_allocation, + col_results = ( + tex.split_quantize( + tensor, + m_splits, + col_quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + if columnwise_enabled + else [None] * len(quantizers) ) return [ diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 072cd8782b2..2004a80b567 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -54,6 +54,28 @@ def __init__( columnwise_quantizer: Quantizer, ) -> None: super().__init__(rowwise=True, columnwise=True) + from transformer_engine.pytorch.quantization import QuantizerRequest # local import + + for role, quantizer in ( + ("rowwise", rowwise_quantizer), + ("columnwise", columnwise_quantizer), + ): + if isinstance(quantizer, QuantizerRequest): + raise TypeError( + "HybridQuantizer does not support nested QuantizerRequest " + f"objects yet; got {type(quantizer).__name__} for the {role} " + "direction. Delayed scaling in CustomRecipe is currently " + "supported only when the qfactory returns DelayedScalingRequest " + "as a top-level slot. Resolving delayed-scaling requests inside " + "HybridQuantizer is future work; pass a concrete Quantizer " + "instance instead." + ) + if not isinstance(quantizer, Quantizer): + raise TypeError( + "HybridQuantizer requires concrete Quantizer instances for " + f"both directions, but the {role} argument is " + f"{type(quantizer).__name__}." + ) if rowwise_quantizer is columnwise_quantizer: raise ValueError( "HybridQuantizer requires distinct rowwise and columnwise quantizer" @@ -67,6 +89,20 @@ def __init__( self.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) self.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + def copy(self) -> "HybridQuantizer": + """Create a shallow copy, preserving parent and sub-quantizer state.""" + quantizer = HybridQuantizer( + rowwise_quantizer=self.rowwise_quantizer.copy(), + columnwise_quantizer=self.columnwise_quantizer.copy(), + ) + quantizer.set_usage( + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + @property def with_amax_reduction(self) -> bool: """Whether either sub-quantizer has cross-rank amax reduction enabled.""" @@ -415,6 +451,94 @@ def detach(self) -> HybridQuantizedTensor: def get_metadata(self) -> Dict[str, Any]: return HybridQuantizedTensorStorage.get_metadata(self) + @staticmethod + def _move_metadata_value( + value: Any, + *, + target_device: torch.device, + non_blocking: bool, + pin_memory: bool, + ) -> Any: + if isinstance(value, torch.Tensor): + value = value.to(device=target_device, non_blocking=non_blocking) + if pin_memory and target_device.type == "cpu": + value = value.pin_memory() + return value + + @classmethod + def _move_sub_storage( + cls, + sub_storage: Optional[QuantizedTensorStorage], + *, + target_device: torch.device, + non_blocking: bool, + pin_memory: bool, + ) -> Optional[QuantizedTensorStorage]: + if sub_storage is None: + return None + metadata = { + key: cls._move_metadata_value( + value, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + for key, value in sub_storage.get_metadata().items() + } + if isinstance(sub_storage, QuantizedTensor): + metadata.update( + { + "shape": sub_storage.shape, + "dtype": sub_storage.dtype, + "requires_grad": sub_storage.requires_grad, + "device": target_device, + } + ) + return type(sub_storage)(**metadata) + + def contiguous( + self, + memory_format: torch.memory_format = torch.contiguous_format, + ) -> "HybridQuantizedTensor": + """Return a HybridQuantizedTensor with contiguous sub-storages.""" + + def _contiguous_sub( + role: str, + sub_storage: Optional[QuantizedTensorStorage], + ) -> Optional[QuantizedTensorStorage]: + if sub_storage is None: + return None + if not isinstance(sub_storage, torch.Tensor): + raise ValueError( + "HybridQuantizedTensor.contiguous does not support storage-only " + f"{role} sub-storage {type(sub_storage).__name__}. This path is " + "only supported for tensor sub-storages." + ) + try: + return sub_storage.contiguous(memory_format=memory_format) + except (NotImplementedError, ValueError) as err: + raise ValueError( + "HybridQuantizedTensor.contiguous could not make the " + f"{role} sub-storage {type(sub_storage).__name__} contiguous " + f"with memory_format={memory_format}." + ) from err + + row = _contiguous_sub("rowwise", self._rowwise_storage) + col = _contiguous_sub("columnwise", self._columnwise_storage) + if row is self._rowwise_storage and col is self._columnwise_storage: + return self + return HybridQuantizedTensor( + shape=self.shape, + dtype=self.dtype, + rowwise_storage=row, + columnwise_storage=col, + rowwise_quantizer=self._rowwise_quantizer, + columnwise_quantizer=self._columnwise_quantizer, + quantizer=self._quantizer, + requires_grad=self.requires_grad, + device=self.device, + ) + def __reduce_ex__(self, protocol: int) -> tuple: """Custom pickling. @@ -488,6 +612,14 @@ def fsdp_pre_all_gather( # pylint: disable=unused-argument ): if sub is None: continue + if not isinstance(sub, QuantizedTensor): + raise NotImplementedError( + "Hybrid FSDP2 all-gather does not support storage-only " + f"{role} sub-storage {type(sub).__name__}. This usually means " + "a HybridQuantizer sub-quantizer had internal=True; use " + "tensor sub-storages for Hybrid FSDP2 or disable Hybrid FSDP2 " + "for this parameter." + ) try: sub.fsdp_buffer_fields() except NotImplementedError as err: @@ -671,6 +803,38 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.detach.default: return args[0].detach() + if func == aten._to_copy.default: + tensor = args[0] + kw = dict(kwargs) if kwargs else {} + dtype = kw.get("dtype", None) + if dtype is None or dtype == tensor.dtype: + target_device = torch.device(kw.get("device", tensor.device) or tensor.device) + pin_memory = bool(kw.get("pin_memory", False)) + non_blocking = bool(kw.get("non_blocking", False)) + row = cls._move_sub_storage( + tensor._rowwise_storage, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + col = cls._move_sub_storage( + tensor._columnwise_storage, + target_device=target_device, + non_blocking=non_blocking, + pin_memory=pin_memory, + ) + return HybridQuantizedTensor( + shape=tensor.shape, + dtype=tensor.dtype, + rowwise_storage=row, + columnwise_storage=col, + rowwise_quantizer=tensor._rowwise_quantizer, + columnwise_quantizer=tensor._columnwise_quantizer, + quantizer=tensor._quantizer, + requires_grad=tensor.requires_grad, + device=target_device, + ) + # ── FSDP2: view ────────────────────────────────────────────── if func == aten.view.default: tensor = args[0] diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index 004fcfb1a26..92ef1bdb051 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -185,6 +185,17 @@ def clone(self) -> "IdentityTensor": device=self.device, ) + def contiguous( + self, + memory_format: torch.memory_format = torch.contiguous_format, + ) -> "IdentityTensor": + """Return an IdentityTensor with contiguous high-precision storage.""" + if self._hp_data is not None and self._hp_data.is_contiguous( + memory_format=memory_format + ): + return self + return self._wrap_data_view(self._hp_data.contiguous(memory_format=memory_format)) + def __reduce_ex__(self, protocol: int) -> tuple: """Custom pickling that preserves the high-precision payload.""" return ( From 9f20d3591d379a1ad7bc6b329856810a4ec7128c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 16:55:06 +0000 Subject: [PATCH 22/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/run_hybrid_tp_sp.py | 4 ++- tests/pytorch/test_hybrid_quantization.py | 20 ++++++----- tests/pytorch/test_identity_quantizer.py | 36 +++++++------------ transformer_engine/pytorch/module/base.py | 4 ++- .../pytorch/module/grouped_linear.py | 10 +++--- .../pytorch/tensor/identity_tensor.py | 7 ++-- 6 files changed, 36 insertions(+), 45 deletions(-) diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index 48986ced752..b38028a36b8 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -599,7 +599,9 @@ def _same_format_parity_supported(): return QUANTIZATION in ("hybrid_fp8", "hybrid_mxfp8") -def _check_same_topology_parity(out_h, dinp_h, model_h, out_v, dinp_v, model_v, tag, *, check_grads): +def _check_same_topology_parity( + out_h, dinp_h, model_h, out_v, dinp_v, model_v, tag, *, check_grads +): # Larger modules use different fused/unfused norm paths between hybrid and # vanilla, so numerical parity is the meaningful contract here. Linear keeps # the stricter bitwise check above. diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 463de1c4360..bcc53339625 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -4333,9 +4333,7 @@ def test_float8_split_uses_logical_shape_for_transpose_only_storage( pieces = torch.split(tensor, split_size, dim=dim) assert [tuple(piece.shape) for piece in pieces] == expected_shapes - assert [ - tuple(piece._transpose.shape) for piece in pieces - ] == expected_transpose_shapes + assert [tuple(piece._transpose.shape) for piece in pieces] == expected_transpose_shapes assert all(piece._data is None for piece in pieces) assert all(piece._transpose_invalid is False for piece in pieces) @@ -4360,12 +4358,16 @@ def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self) (2, 16), ] assert all(piece.rowwise_sub_storage is None for piece in pieces) - assert [ - tuple(piece.columnwise_sub_storage.shape) for piece in pieces - ] == [(5, 16), (5, 16), (2, 16)] - assert [ - tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces - ] == [(16, 5), (16, 5), (16, 2)] + assert [tuple(piece.columnwise_sub_storage.shape) for piece in pieces] == [ + (5, 16), + (5, 16), + (2, 16), + ] + assert [tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces] == [ + (16, 5), + (16, 5), + (16, 2), + ] @requires_fp8 diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 65df088269b..b21041cbc73 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -33,8 +33,8 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) -fp8_block_scaling_available, reason_for_no_fp8_block_scaling = ( - te.is_fp8_block_scaling_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True ) @@ -276,9 +276,7 @@ def test_grouped_split_rejects_mixed_identity_and_quantized_operands(self): activation_dtype=torch.bfloat16, ) - def test_hybrid_split_forwards_disable_bulk_allocation_to_both_directions( - self, monkeypatch - ): + def test_hybrid_split_forwards_disable_bulk_allocation_to_both_directions(self, monkeypatch): import transformer_engine.pytorch.module.grouped_linear as grouped_linear from transformer_engine.pytorch.module.grouped_linear import _hybrid_split_quantize @@ -607,7 +605,10 @@ def test_tensor_ops_preserve_identity_and_values(self): zeros = t.new_zeros((2, 3)) assert isinstance(zeros, IdentityTensor) torch.testing.assert_close( - zeros.dequantize(), torch.zeros((2, 3), device="cuda", dtype=x.dtype), rtol=0.0, atol=0.0 + zeros.dequantize(), + torch.zeros((2, 3), device="cuda", dtype=x.dtype), + rtol=0.0, + atol=0.0, ) dst = IdentityQuantizer().make_empty(x.shape, dtype=x.dtype, device="cuda") @@ -1022,9 +1023,7 @@ def test_cpu_offload_keeps_identity_direction_exact(self, format_name): assert isinstance(reloaded, HybridQuantizedTensor) assert isinstance(reloaded._columnwise_storage, IdentityTensorStorage) - torch.testing.assert_close( - reloaded._columnwise_storage.dequantize(), x, rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(reloaded._columnwise_storage.dequantize(), x, rtol=0.0, atol=0.0) torch.testing.assert_close( reloaded._rowwise_storage.dequantize(), expected_row, rtol=0.0, atol=0.0 ) @@ -1181,9 +1180,7 @@ def test_identity_reproduces_backward_override_high_precision_bitwise(self): x, recipe=CustomRecipe(qfactory=fp8_fwd_factory, backward_override="high_precision"), ) - y_id, dx_id, wg_id = _fwd_bwd( - test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_hp_factory) - ) + y_id, dx_id, wg_id = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_hp_factory)) torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) @@ -1231,9 +1228,7 @@ def test_quantized_model_init_identity_matches_bf16_bitwise(self): recipe = CustomRecipe(qfactory=identity_all_factory) torch.manual_seed(2718) with te.quantized_model_init(enabled=True, recipe=recipe): - test = te.Linear( - self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16 - ).cuda() + test = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() with torch.no_grad(): for p_test, p_ref in zip(test.parameters(), ref.parameters()): assert isinstance(p_test, IdentityTensor) @@ -1255,9 +1250,7 @@ def test_quantized_model_init_identity_training_loss_decreases_bitwise(self): recipe = CustomRecipe(qfactory=identity_all_factory) torch.manual_seed(888) with te.quantized_model_init(enabled=True, recipe=recipe): - test = te.Linear( - self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16 - ).cuda() + test = te.Linear(self.IN_F, self.OUT_F, bias=False, params_dtype=torch.bfloat16).cuda() with torch.no_grad(): for p_test, p_ref in zip(test.parameters(), ref.parameters()): assert isinstance(p_test, IdentityTensor) @@ -1290,13 +1283,10 @@ def test_quantized_model_init_identity_training_loss_decreases_bitwise(self): torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) for p_test, p_ref in zip(test.parameters(), ref.parameters()): - torch.testing.assert_close( - p_test.dequantize(), p_ref, rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(p_test.dequantize(), p_ref, rtol=0.0, atol=0.0) assert all( - losses_ref[i + 1].item() < losses_ref[i].item() - for i in range(len(losses_ref) - 1) + losses_ref[i + 1].item() < losses_ref[i].item() for i in range(len(losses_ref) - 1) ), f"BF16 loss did not strictly decrease: {[x.item() for x in losses_ref]}" for loss_test, loss_ref in zip(losses_test, losses_ref): torch.testing.assert_close(loss_test, loss_ref, rtol=0.0, atol=0.0) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index f5f9ec0e0e1..58e3b09b3f6 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1660,7 +1660,9 @@ def grad_output_preprocess( ): grad_bias = grad_output.dequantize().view(-1, grad_output.shape[-1]).sum(dim=0) else: - if isinstance(quantizer, (Float8BlockQuantizer, HybridQuantizer, IdentityQuantizer)): + if isinstance( + quantizer, (Float8BlockQuantizer, HybridQuantizer, IdentityQuantizer) + ): # Float8BlockQuantizer: unfused until cast_transpose + dgrad is ready. # HybridQuantizer: tex.bgrad_quantize doesn't recognize hybrid quantizers. # IdentityQuantizer: high-precision passthrough; bgrad computed in HP. diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 4ab22663dbf..9ad013f97b3 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -76,9 +76,9 @@ def _uses_identity_quantizer(quantizer): if isinstance(quantizer, IdentityQuantizer): return True if isinstance(quantizer, HybridQuantizer): - return _uses_identity_quantizer( - quantizer.rowwise_quantizer - ) or _uses_identity_quantizer(quantizer.columnwise_quantizer) + return _uses_identity_quantizer(quantizer.rowwise_quantizer) or _uses_identity_quantizer( + quantizer.columnwise_quantizer + ) return False @@ -1312,9 +1312,7 @@ def backward( inputmats: list input_identity = _has_identity_quantizer_list(ctx.input_quantizers) input_hybrid = ( - False - if input_identity - else _is_hybrid_quantizer_list(ctx.input_quantizers) + False if input_identity else _is_hybrid_quantizer_list(ctx.input_quantizers) ) if ctx.fp8 and not ctx.debug and not input_hybrid: inputmats = _split_quantize_with_identity_fallback( diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index 92ef1bdb051..cb9ea2e0b0c 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -125,8 +125,7 @@ def update_quantized( ) -> QuantizedTensorStorage: if not isinstance(dst, IdentityTensorStorage): raise ValueError( - "IdentityQuantizer can only update IdentityTensorStorage, got" - f" {type(dst).__name__}" + f"IdentityQuantizer can only update IdentityTensorStorage, got {type(dst).__name__}" ) data = self._maybe_cast(src) if ( @@ -190,9 +189,7 @@ def contiguous( memory_format: torch.memory_format = torch.contiguous_format, ) -> "IdentityTensor": """Return an IdentityTensor with contiguous high-precision storage.""" - if self._hp_data is not None and self._hp_data.is_contiguous( - memory_format=memory_format - ): + if self._hp_data is not None and self._hp_data.is_contiguous(memory_format=memory_format): return self return self._wrap_data_view(self._hp_data.contiguous(memory_format=memory_format)) From 04965f57dfe77fb0bf7cf8d011dd36fd9095523f Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 11 Jun 2026 15:14:23 +0000 Subject: [PATCH 23/97] Minor fixes Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/conftest.py | 19 ++++++- .../fsdp2_tests/run_fsdp2_fused_adam.py | 15 ----- tests/pytorch/test_cpu_offloading.py | 16 +++--- tests/pytorch/test_hybrid_quantization.py | 56 +++++++++++++++++++ .../pytorch/cpp_extensions/gemm.py | 19 +++++++ transformer_engine/pytorch/module/base.py | 4 ++ .../pytorch/tensor/hybrid_tensor.py | 28 ++++++---- 7 files changed, 120 insertions(+), 37 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index 94bcc43f060..621972df33e 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -52,6 +52,13 @@ def _check_nvfp4_support(): ("HybridMixed_MXFP8_FP8", fp8.check_mxfp8_support), ] +_HYBRID_FLOAT8_BLOCK_FSDP2_XFAIL_REASON = ( + "HybridFloat8BlockScaling + FSDP2 is not supported when dim-0 shards split " + "128-row Float8Block scale tiles. Each shard stores independently rounded " + "scale rows, but the gathered tensor expects scale rows rounded from the " + "global row count, so naively all-gathered scale buffers have the wrong shape." +) + def _parametrize_recipes(): params = [] @@ -67,9 +74,15 @@ def _parametrize_hybrid_recipes(): params = [] for name, check_fn in _HYBRID_RECIPE_CONFIGS: supported, reason = check_fn() - params.append( - pytest.param(name, id=name, marks=pytest.mark.skipif(not supported, reason=reason)) - ) + marks = [pytest.mark.skipif(not supported, reason=reason)] + if name == "HybridFloat8BlockScaling": + marks.append( + pytest.mark.xfail( + raises=RuntimeError, + reason=_HYBRID_FLOAT8_BLOCK_FSDP2_XFAIL_REASON, + ) + ) + params.append(pytest.param(name, id=name, marks=marks)) return params diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 358ac8757e9..fae36e81a53 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1701,21 +1701,6 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): """ import torch.distributed.checkpoint as dcp - if hybrid_recipe_name == "HybridFP8CurrentScaling": - # TODO: preserve hybrid current-scaling primary-weight scales across DCP - # by implementing __tensor_flatten__/__tensor_unflatten__ on the quantized - # tensor stack (HybridQuantizedTensor + its Float8Tensor sub-storages) so - # DCP serializes fp8 data + fp32 _scale_inv as explicit tensor leaves - # instead of round-tripping through a dequantized bf16 weight. - pytest.xfail( - "HybridFP8CurrentScaling: hybrid current-scaling primary-weight " - "_scale_inv is not preserved across DCP. DCP stores each weight as a " - "dequantized bf16 leaf (no scale leaf) and re-quantizes on load; this " - "is idempotent for a single per-tensor Float8Tensor (vanilla " - "round-trips bitwise) but not for the hybrid's two sub-storages, so " - "the loaded output diverges ~5e-2. torch.save/load and the FP32 " - "master weight are unaffected. See the TODO above for the fix." - ) from fsdp2_utils import get_hybrid_recipe_from_string diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 3cd098fbdc2..b58c20b7b67 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -40,13 +40,13 @@ def _hybrid_fp8_mxfp8_qfactory(role): if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( rowwise_quantizer=te.Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" + te.DType.kFloat8E4M3, device="cuda" ), - columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2) - return te.Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + return te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E5M2) + return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") def _hybrid_mxfp8_nvfp4_qfactory(role): @@ -59,12 +59,12 @@ def _hybrid_mxfp8_nvfp4_qfactory(role): is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( - rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - columnwise_quantizer=te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3), + columnwise_quantizer=te.NVFP4Quantizer(fp4_dtype=te.DType.kFloat4E2M1), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) - return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + return te.NVFP4Quantizer(fp4_dtype=te.DType.kFloat4E2M1) + return te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) def nvfp4_row_scaled(): diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index bcc53339625..2b4826597d3 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -425,6 +425,62 @@ def test_clear_is_idempotent(self, input_tensor): assert n == 0 +@requires_fp8 +class TestHybridTensorShapeOps: + """Shape ops that preserve supported Hybrid sub-storages.""" + + def test_fp8_current_non_noop_slice_and_narrow_preserve_hybrid(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 0, 32, 1) + narrowed = tensor.narrow(0, 16, 32) + + assert isinstance(sliced, HybridQuantizedTensor) + assert isinstance(narrowed, HybridQuantizedTensor) + torch.testing.assert_close(sliced.dequantize(), dequantized[:32], rtol=0, atol=0) + torch.testing.assert_close(narrowed.dequantize(), dequantized[16:48], rtol=0, atol=0) + + def test_full_span_step_slice_is_not_treated_as_noop(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 0, tensor.size(0), 2) + + assert isinstance(sliced, HybridQuantizedTensor) + torch.testing.assert_close(sliced.dequantize(), dequantized[::2], rtol=0, atol=0) + + def test_same_shape_as_strided_with_offset_is_not_treated_as_noop(self): + torch.manual_seed(42) + x = torch.randn(64, 32, dtype=torch.bfloat16, device="cuda") + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + tensor = quantizer.quantize(x) + dequantized = tensor.dequantize() + + base_view = torch.ops.aten.slice.Tensor(tensor, 0, 0, 63, 1) + shifted = torch.ops.aten.as_strided.default( + base_view, base_view.shape, base_view.stride(), x.stride(0) + ) + + assert isinstance(shifted, HybridQuantizedTensor) + torch.testing.assert_close(shifted.dequantize(), dequantized[1:], rtol=0, atol=0) + + @requires_fp8_and_nvfp4 class TestHybridDetachIsolation: """``HybridQuantizedTensor.detach()`` must produce a hybrid whose diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 79b9b9ffd88..7ce9bffaed5 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -17,6 +17,8 @@ from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom +from ..tensor.hybrid_tensor import HybridQuantizer +from ..tensor.identity_tensor import IdentityQuantizer from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage from ..tensor.storage.identity_tensor_storage import IdentityTensorStorage from ..custom_recipes.gemm import custom_gemm @@ -148,6 +150,21 @@ def _materialize_high_precision(tensor): return tensor +def _reject_unsupported_output_quantizer(quantization_params): + """Reject output quantizers that the native C++ GEMM path cannot convert.""" + if isinstance(quantization_params, (HybridQuantizer, IdentityQuantizer)): + quantizer_name = type(quantization_params).__name__ + # TODO(negvet): Lower HybridQuantizer output to its native rowwise + # sub-quantizer, and IdentityQuantizer output to an unquantized/no-op + # path, once the returned tensor contract is defined for these boundary + # roles. + raise NotImplementedError( + f"{quantizer_name} is not supported as a native GEMM output quantizer. " + "Return a TE-native quantizer for output/grad_input roles or disable " + "quantized GEMM output for this boundary." + ) + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -222,6 +239,8 @@ def general_gemm( A = A.get_tensor(not transa) B = B.get_tensor(transb) + _reject_unsupported_output_quantizer(quantization_params) + # Use bfloat16 as default bias_dtype bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 58e3b09b3f6..5b9c420fe71 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1044,6 +1044,8 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # Return early if recipe state matches recipe if self.fp8_meta_tensors_initialized: recipe_state = self.fp8_meta[fp8_meta_tensor_key] + # Follow-up: Match built-in recipes by full config, not just RecipeState type, so + # same-class mid-training changes rebuild quantizers/workspaces correctly. if recipe.delayed() and isinstance(recipe_state, DelayedScalingRecipeState): self.adjust_amax_history_length(recipe.amax_history_len, fwd=fwd) return @@ -1060,6 +1062,8 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return if recipe.custom() and isinstance(recipe_state, CustomRecipeState): + # Follow-up: Compare CustomRecipe/qfactory config here. qfactory changes made + # mid-training currently do not take effect because stale quantizers are reused. return # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 2004a80b567..c8809e514c0 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -917,26 +917,32 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): tensor = args[0] shape = args[1] strides = args[2] + storage_offset = kwargs.get("storage_offset", args[3] if len(args) > 3 else None) + if storage_offset is None: + storage_offset = tensor.storage_offset() if ( - len(shape) == len(strides) == 2 - and tuple(strides) == (shape[-1], 1) - and tuple(shape) == tuple(tensor.size()) + tuple(shape) == tuple(tensor.size()) + and tuple(strides) == tuple(tensor.stride()) + and storage_offset == tensor.storage_offset() ): return HybridQuantizedTensor.make_like(tensor) - return cls._delegate_reshape_op( - func, tensor, args, kwargs - ) or super().__torch_dispatch__(func, types, args, kwargs) + out = cls._delegate_reshape_op(func, tensor, args, kwargs) + if out is not None: + return out + return super().__torch_dispatch__(func, types, args, kwargs) if func == aten.slice.Tensor: tensor = args[0] dim = args[1] start = args[2] - length = args[3] - if start == 0 and length == tensor.size(dim): + end = args[3] + step = args[4] if len(args) > 4 else 1 + if start == 0 and end == tensor.size(dim) and step == 1: return HybridQuantizedTensor.make_like(tensor) - return cls._delegate_reshape_op( - func, tensor, args, kwargs - ) or super().__torch_dispatch__(func, types, args, kwargs) + out = cls._delegate_reshape_op(func, tensor, args, kwargs) + if out is not None: + return out + return super().__torch_dispatch__(func, types, args, kwargs) # ── FSDP2: copy_ ───────────────────────────────────────────── # Fast path for hybrid-to-hybrid (FSDP2 fills buffer allocated via From 4a79a077cb6ddbeccee32e4081ff42ee1e711c25 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Jun 2026 15:15:49 +0000 Subject: [PATCH 24/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py | 1 - tests/pytorch/test_cpu_offloading.py | 4 +--- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index fae36e81a53..4fa4d232029 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1701,7 +1701,6 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): """ import torch.distributed.checkpoint as dcp - from fsdp2_utils import get_hybrid_recipe_from_string hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index b58c20b7b67..0f6ba43fbaf 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -39,9 +39,7 @@ def _hybrid_fp8_mxfp8_qfactory(role): is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( - rowwise_quantizer=te.Float8CurrentScalingQuantizer( - te.DType.kFloat8E4M3, device="cuda" - ), + rowwise_quantizer=te.Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda"), columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): From 58d6ad04c738dbac042563f496ea2d48c94646ab Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 12 Jun 2026 13:15:32 +0000 Subject: [PATCH 25/97] Handle transpose-only Float8 in hybrid distopt Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 111 +++++++++++++++++ .../pytorch/tensor/float8_tensor.py | 6 + transformer_engine/pytorch/tensor/utils.py | 113 +++++++++++++++++- 3 files changed, 228 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 2b4826597d3..43f3dd2ae77 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -3019,6 +3019,117 @@ class TestHybridQuantizeMasterWeights: * Both sub-storages dropped (caller bug: nothing left to cast) """ + @staticmethod + def _make_transpose_only_float8_weight(shape, quantizer, *, fill_value=173): + """Build a Hopper/L40-style columnwise-only Float8Tensor. + + Blackwell keeps ``_data`` populated for columnwise-only Float8, so these + tests synthesize the older architecture layout directly: logical + ``[M, K]`` shape with the only live FP8 bytes in ``_transpose[K, M]``. + """ + rows, cols = shape + return Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=torch.full( + (cols, rows), fill_value, dtype=torch.uint8, device="cuda" + ), + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + + @staticmethod + def _scatter_expected_logical_bytes(initial_transpose, fp8_bytes, logical_shape, start_offset): + rows, cols = logical_shape + expected = initial_transpose.clone() + expected_2d = expected.reshape(cols, rows) + remaining = fp8_bytes.numel() + logical_offset = start_offset + src_offset = 0 + while remaining > 0: + row = logical_offset // cols + col = logical_offset % cols + n = min(remaining, cols - col) + expected_2d[col : col + n, row].copy_(fp8_bytes[src_offset : src_offset + n]) + logical_offset += n + src_offset += n + remaining -= n + return expected + + @staticmethod + def _reference_fp8_bytes(master, scale, dtype=torch.bfloat16): + quantizer = Float8Quantizer( + scale=scale, + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + raw = torch.empty((1, master.numel()), dtype=torch.uint8, device="cuda") + temp = quantizer.create_tensor_from_data(raw, dtype) + quantizer.update_quantized(master.reshape(1, -1), temp) + return temp._data.reshape(-1) + + def test_fp8_current_transpose_only_nonzero_offset(self): + """Current-scaling distopt update handles Hopper-style columnwise storage. + + Regression for ``model_weight.reshape(-1)`` reaching + ``Float8Tensor._ReshapeFunc.forward`` and dereferencing ``_data=None``. + The nonzero offset spans multiple logical rows, so the test also checks + row-major shard bytes are scattered into transposed storage correctly. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + start_offset = 5 + master = torch.linspace(-2.0, 2.0, steps=17, dtype=torch.float32, device="cuda") + quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, device="cuda", rowwise=False, columnwise=True + ) + weight = self._make_transpose_only_float8_weight(shape, quantizer) + initial = weight._transpose.clone() + + quantize_master_weights([weight], [master], [start_offset], group=group) + + scale = torch.reciprocal(weight._scale_inv.detach().clone()) + fp8_bytes = self._reference_fp8_bytes(master.to(weight.dtype), scale, weight.dtype) + expected = self._scatter_expected_logical_bytes(initial, fp8_bytes, shape, start_offset) + assert weight._data is None + assert weight._transpose_invalid is False + assert torch.equal(weight._transpose, expected) + + def test_fp8_delayed_transpose_only_nonzero_offset(self): + """Delayed-scaling distopt update handles Hopper-style columnwise storage. + + Regression for the direct ``model_weight._data.view(-1)`` path in the + delayed-scaling helper. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + start_offset = 6 + master = torch.linspace(-3.0, 1.0, steps=15, dtype=torch.float32, device="cuda") + quantizer = _make_delayed_quantizer(tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=False, columnwise=True) + weight = self._make_transpose_only_float8_weight(shape, quantizer) + initial = weight._transpose.clone() + + quantize_master_weights([weight], [master], [start_offset], group=group) + + fp8_bytes = self._reference_fp8_bytes( + master.to(weight.dtype), weight._get_quantizer().scale, weight.dtype + ) + expected = self._scatter_expected_logical_bytes(initial, fp8_bytes, shape, start_offset) + assert weight._data is None + assert weight._transpose_invalid is False + assert torch.equal(weight._transpose, expected) + # ---------- Positive tests (same-format) ---------- def test_fp8_current_same_format_full_master(self): diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 3e06a52830c..9c16e5b9a66 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -523,6 +523,12 @@ def _reset_caches(self) -> None: Set transpose cache as invalid. Should be called after any in-place operation. """ + if self._data is None and self._transpose is not None: + # Columnwise-only Float8 tensors on Hopper / L40 store their only + # live FP8 payload in _transpose. Treat it as primary storage, not + # as a derived cache that can be invalidated. + self._transpose_invalid = False + return self._transpose_invalid = True def remove_caches(self) -> None: diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index e29e2701141..199417df443 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -98,6 +98,107 @@ def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): raise ValueError(f"replace_raw_data for {type(tensor)} is not supported yet") +def _is_float8_transpose_only(tensor: QuantizedTensor) -> bool: + """Whether a Float8 tensor stores its live payload only in _transpose.""" + return ( + isinstance(tensor, Float8Tensor) + and tensor._data is None + and tensor._transpose is not None + and not tensor._transpose_invalid + ) + + +def _validate_flat_fragment(model_weight: QuantizedTensor, master_weight: torch.Tensor, start_offset): + """Validate a flat logical shard and return its exclusive end offset.""" + if start_offset is None: + raise ValueError("start_offset must not be None when master_weight is provided") + if start_offset < 0: + raise ValueError(f"start_offset must be non-negative, got {start_offset}") + end_offset = start_offset + master_weight.numel() + if end_offset > model_weight.numel(): + raise ValueError( + f"end_offset ({end_offset}) exceeds model_weight numel ({model_weight.numel()}), " + f"start_offset={start_offset}, master_weight numel={master_weight.numel()}" + ) + return end_offset + + +def _cast_master_weight_to_rowwise_fp8_bytes( + master_weight: torch.Tensor, + model_weight: Float8Tensor, + quantizer: Float8Quantizer, +) -> torch.Tensor: + """Cast a flat master shard to row-major FP8 bytes using ``quantizer`` scale state.""" + rowwise_quantizer = Float8Quantizer( + scale=quantizer.scale, + amax=quantizer.amax, + fp8_dtype=quantizer.dtype, + rowwise=True, + columnwise=False, + ) + raw = torch.empty((1, master_weight.numel()), dtype=torch.uint8, device=model_weight.device) + temp = rowwise_quantizer.create_tensor_from_data(raw, model_weight.dtype) + rowwise_quantizer.update_quantized(master_weight.reshape(1, -1), temp) + if temp._data is None: + raise RuntimeError("Expected rowwise Float8 temporary to populate _data") + return temp._data.reshape(-1) + + +def _update_transpose_only_float8_flat_fragment( + model_weight: QuantizedTensor, + master_weight: torch.Tensor, + start_offset, + quantizer: Float8Quantizer, +) -> bool: + """Update a logical flat shard in a transpose-only Float8 tensor. + + Hopper / L40 columnwise-only Float8 sub-storages keep their live FP8 + bytes in ``_transpose`` with physical shape ``[K, rows]`` for a logical + ``[rows, K]`` tensor. A row-major logical shard is not contiguous in that + storage, so cast the shard once and scatter the resulting FP8 bytes by + logical row into the transpose buffer. + """ + if not _is_float8_transpose_only(model_weight): + return False + + _validate_flat_fragment(model_weight, master_weight, start_offset) + numel = master_weight.numel() + if numel == 0: + return True + + shape = tuple(model_weight.shape) + if len(shape) == 0: + raise ValueError("Float8 scalar transpose-only flat update is not supported") + logical_cols = int(shape[-1]) + if logical_cols <= 0 or model_weight.numel() % logical_cols != 0: + raise ValueError(f"Invalid Float8 logical shape for transpose-only update: {shape}") + logical_rows = model_weight.numel() // logical_cols + + transpose = model_weight._transpose + if transpose.numel() != model_weight.numel(): + raise ValueError( + "Float8 transpose-only storage has unexpected numel: " + f"transpose={transpose.numel()}, logical={model_weight.numel()}" + ) + transpose_2d = transpose.reshape(logical_cols, logical_rows) + fp8_bytes = _cast_master_weight_to_rowwise_fp8_bytes(master_weight, model_weight, quantizer) + + remaining = numel + logical_offset = start_offset + src_offset = 0 + while remaining > 0: + row = logical_offset // logical_cols + col = logical_offset % logical_cols + n = min(remaining, logical_cols - col) + transpose_2d[col : col + n, row].copy_(fp8_bytes[src_offset : src_offset + n]) + logical_offset += n + src_offset += n + remaining -= n + + model_weight._transpose_invalid = False + return True + + def quantize_master_weights( model_weights, master_weights, @@ -301,6 +402,10 @@ def _cast_master_weights_to_fp8_delayed_scaling( # master_weight may be smaller than model_weight because it could be distributed across # multiple ranks. So we need to create a dummy weight using the raw data from model_weight. if not use_fsdp_shard_model_weights: + if _update_transpose_only_float8_flat_fragment( + model_weight, master_weight, start_offset, quantizer + ): + continue shard_model_weight_raw = model_weight._data.view(-1)[start_offset:end_offset] shard_model_weight_fp8 = quantizer.create_tensor_from_data( shard_model_weight_raw.view(1, -1), @@ -446,13 +551,17 @@ def _cast_master_weights_to_fp8_current_scaling( # Cast master weight to FP8 end_offset = start_offset + master_weight.numel() - if not use_fsdp_shard_model_weights: - model_weight_fragment = model_weight.reshape(-1)[start_offset:end_offset] quantizer = Float8Quantizer( scale=scale, amax=torch.Tensor(), fp8_dtype=model_weight._fp8_dtype, ) + if not use_fsdp_shard_model_weights: + if _update_transpose_only_float8_flat_fragment( + model_weight, master_weight, start_offset, quantizer + ): + continue + model_weight_fragment = model_weight.reshape(-1)[start_offset:end_offset] if use_fsdp_shard_model_weights and not isinstance(model_weight_fragment, Float8Tensor): # NOTE: The fsdp shard model weight may be a unit8 tensor instead of # a float8 tensor. We should handle this situation properly. From 4b9d1f8285ee11538a5607c72edf38c472bf7ceb Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 12 Jun 2026 14:53:01 +0000 Subject: [PATCH 26/97] Refactor hybrid tensors to own only parent quantizer Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 33 +++++++-- .../pytorch/module/grouped_linear.py | 6 +- .../pytorch/tensor/hybrid_tensor.py | 74 +++++-------------- .../tensor/storage/hybrid_tensor_storage.py | 30 ++++---- transformer_engine/pytorch/tensor/utils.py | 4 +- 5 files changed, 63 insertions(+), 84 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 43f3dd2ae77..17fc2ebe67f 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -101,6 +101,22 @@ def test_creation(self): assert isinstance(hq.rowwise_quantizer, Float8CurrentScalingQuantizer) assert isinstance(hq.columnwise_quantizer, NVFP4Quantizer) + def test_hybrid_storage_and_tensor_require_parent_quantizer(self): + with pytest.raises(TypeError, match="requires a parent HybridQuantizer"): + HybridQuantizedTensorStorage( + rowwise_storage=None, + columnwise_storage=None, + quantizer=None, + ) + with pytest.raises(TypeError, match="requires a parent HybridQuantizer"): + HybridQuantizedTensor( + shape=(1, 1), + dtype=torch.bfloat16, + rowwise_storage=None, + columnwise_storage=None, + quantizer=None, + ) + def test_rejects_same_sub_quantizer_instance_for_both_directions(self): quantizer = _make_fp8_quantizer() @@ -3241,8 +3257,8 @@ def test_fp8_delayed_row_current_col_full_master(self): assert weight._rowwise_storage is not None assert weight._columnwise_storage is not None - assert isinstance(weight._rowwise_quantizer, Float8Quantizer) - assert isinstance(weight._columnwise_quantizer, Float8CurrentScalingQuantizer) + assert isinstance(weight._quantizer.rowwise_quantizer, Float8Quantizer) + assert isinstance(weight._quantizer.columnwise_quantizer, Float8CurrentScalingQuantizer) dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) @@ -3270,8 +3286,8 @@ def test_fp8_current_row_delayed_col_full_master(self): assert weight._rowwise_storage is not None assert weight._columnwise_storage is not None - assert isinstance(weight._rowwise_quantizer, Float8CurrentScalingQuantizer) - assert isinstance(weight._columnwise_quantizer, Float8Quantizer) + assert isinstance(weight._quantizer.rowwise_quantizer, Float8CurrentScalingQuantizer) + assert isinstance(weight._quantizer.columnwise_quantizer, Float8Quantizer) dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) @@ -4506,14 +4522,17 @@ def test_float8_split_uses_logical_shape_for_transpose_only_storage( def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self): columnwise = self._make_transpose_only_float8_tensor() + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + quantizer.set_usage(rowwise=False, columnwise=True) hybrid = HybridQuantizedTensor( shape=columnwise.shape, dtype=columnwise.dtype, rowwise_storage=None, columnwise_storage=columnwise, - rowwise_quantizer=None, - columnwise_quantizer=None, - quantizer=None, + quantizer=quantizer, device="cuda", ) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 9ad013f97b3..0b46985e355 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -259,16 +259,12 @@ def _hybrid_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocat HybridStorage( rowwise_storage=row, columnwise_storage=col, - rowwise_quantizer=rq, - columnwise_quantizer=cq, quantizer=q, fake_dtype=tensor.dtype, ) - for row, col, rq, cq, q in zip( + for row, col, q in zip( row_results, col_results, - row_quantizers, - col_quantizers, quantizers, ) ] diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index c8809e514c0..61fa694e4c4 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -146,8 +146,6 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: return HybridQuantizedTensorStorage( rowwise_storage=rowwise_result, columnwise_storage=columnwise_result, - rowwise_quantizer=self.rowwise_quantizer, - columnwise_quantizer=self.columnwise_quantizer, quantizer=self, fake_dtype=tensor.dtype, ) @@ -157,8 +155,6 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: dtype=tensor.dtype, rowwise_storage=rowwise_result, columnwise_storage=columnwise_result, - rowwise_quantizer=self.rowwise_quantizer, - columnwise_quantizer=self.columnwise_quantizer, quantizer=self, ) @@ -197,8 +193,6 @@ def make_empty( device=device, rowwise_storage=rowwise_empty, columnwise_storage=columnwise_empty, - rowwise_quantizer=self.rowwise_quantizer, - columnwise_quantizer=self.columnwise_quantizer, quantizer=self, ) @@ -350,12 +344,8 @@ class HybridQuantizedTensor(HybridQuantizedTensorStorage, QuantizedTensor): Sub-storage for rowwise quantized data. columnwise_storage : QuantizedTensorStorage Sub-storage for columnwise quantized data. - rowwise_quantizer : Quantizer, optional - Quantizer used for the rowwise sub-storage. - columnwise_quantizer : Quantizer, optional - Quantizer used for the columnwise sub-storage. - quantizer : HybridQuantizer, optional - Parent hybrid quantizer. + quantizer : HybridQuantizer + Parent hybrid quantizer that owns the rowwise and columnwise sub-quantizers. requires_grad : bool, default = False Whether to compute gradients for this tensor. @@ -366,9 +356,7 @@ def __new__( *args, rowwise_storage: Optional[QuantizedTensorStorage], columnwise_storage: Optional[QuantizedTensorStorage], - rowwise_quantizer: Optional[Quantizer] = None, - columnwise_quantizer: Optional[Quantizer] = None, - quantizer: Optional[Quantizer] = None, + quantizer: HybridQuantizer, **kwargs, ): instance = super().__new__( @@ -376,8 +364,6 @@ def __new__( *args, rowwise_storage=rowwise_storage, columnwise_storage=columnwise_storage, - rowwise_quantizer=rowwise_quantizer, - columnwise_quantizer=columnwise_quantizer, quantizer=quantizer, **kwargs, ) @@ -443,8 +429,6 @@ def detach(self) -> HybridQuantizedTensor: dtype=self.dtype, rowwise_storage=row, columnwise_storage=col, - rowwise_quantizer=self._rowwise_quantizer, - columnwise_quantizer=self._columnwise_quantizer, quantizer=self._quantizer, ) @@ -532,8 +516,6 @@ def _contiguous_sub( dtype=self.dtype, rowwise_storage=row, columnwise_storage=col, - rowwise_quantizer=self._rowwise_quantizer, - columnwise_quantizer=self._columnwise_quantizer, quantizer=self._quantizer, requires_grad=self.requires_grad, device=self.device, @@ -561,8 +543,6 @@ def __reduce_ex__(self, protocol: int) -> tuple: ( self._rowwise_storage, self._columnwise_storage, - self._rowwise_quantizer, - self._columnwise_quantizer, self._quantizer, self.dtype, self.shape, @@ -598,7 +578,7 @@ def fsdp_pre_all_gather( # pylint: disable=unused-argument # reduction so the post-optimizer re-quantization of the sharded weight # keeps one shared scale across shards (no-op for sub-quantizers without # amax reduction, e.g. MXFP8). - if mesh is not None and self._quantizer is not None: + if mesh is not None: self._quantizer.amax_reduction_group = mesh.get_group() self._quantizer.with_amax_reduction = True @@ -651,8 +631,6 @@ def fsdp_pre_all_gather( # pylint: disable=unused-argument col_meta, self._rowwise_storage, # original sharded sub-storage (for make_like on iter-1) self._columnwise_storage, - self._rowwise_quantizer, - self._columnwise_quantizer, self._quantizer, ) return sharded_tensors, metadata @@ -680,11 +658,12 @@ def fsdp_post_all_gather( col_meta, orig_row_sub, orig_col_sub, - row_quantizer, - col_quantizer, hybrid_quantizer, ) = metadata + row_quantizer = hybrid_quantizer.rowwise_quantizer + col_quantizer = hybrid_quantizer.columnwise_quantizer + row_gathered = all_gather_outputs[:n_row_buffers] col_gathered = all_gather_outputs[n_row_buffers:] @@ -714,10 +693,10 @@ def _sync_usage(sub_storage, sub_quantizer): # Iteration 2+: in-place field update on existing sub-storages if out._rowwise_storage is not None and row_meta is not None: out._rowwise_storage.fsdp_assign_gathered(row_gathered, row_meta) - _sync_usage(out._rowwise_storage, out._rowwise_quantizer) + _sync_usage(out._rowwise_storage, out._quantizer.rowwise_quantizer) if out._columnwise_storage is not None and col_meta is not None: out._columnwise_storage.fsdp_assign_gathered(col_gathered, col_meta) - _sync_usage(out._columnwise_storage, out._columnwise_quantizer) + _sync_usage(out._columnwise_storage, out._quantizer.columnwise_quantizer) else: # First iteration: clone the original sharded sub-storages via make_like, # then write gathered (full-size) buffers via each sub-storage's own @@ -748,8 +727,6 @@ def _sync_usage(sub_storage, sub_quantizer): dtype=param_dtype, rowwise_storage=row_sub, columnwise_storage=col_sub, - rowwise_quantizer=row_quantizer, - columnwise_quantizer=col_quantizer, quantizer=hybrid_quantizer, ) @@ -790,8 +767,6 @@ def _delegate(sub): dtype=tensor.dtype, rowwise_storage=row_out, columnwise_storage=col_out, - rowwise_quantizer=tensor._rowwise_quantizer, - columnwise_quantizer=tensor._columnwise_quantizer, quantizer=tensor._quantizer, ) @@ -828,8 +803,6 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): dtype=tensor.dtype, rowwise_storage=row, columnwise_storage=col, - rowwise_quantizer=tensor._rowwise_quantizer, - columnwise_quantizer=tensor._columnwise_quantizer, quantizer=tensor._quantizer, requires_grad=tensor.requires_grad, device=target_device, @@ -862,8 +835,6 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): dtype=tensor.dtype, rowwise_storage=row_view, columnwise_storage=col_view, - rowwise_quantizer=tensor._rowwise_quantizer, - columnwise_quantizer=tensor._columnwise_quantizer, quantizer=tensor._quantizer, ) @@ -897,8 +868,6 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): dtype=tensor.dtype, rowwise_storage=row_pieces[i] if row_pieces is not None else None, columnwise_storage=col_pieces[i] if col_pieces is not None else None, - rowwise_quantizer=tensor._rowwise_quantizer, - columnwise_quantizer=tensor._columnwise_quantizer, quantizer=tensor._quantizer, ) for i in range(num_pieces) @@ -962,15 +931,14 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.new_zeros.default: tensor = args[0] new_shape = args[1] - if tensor._quantizer is not None: - # FSDP2 allocates new_zeros buffers as all-gather destinations - # that are immediately overwritten by copy_. Use make_empty - # (uninitialized storage with the right container shape/fields). - return tensor._quantizer.make_empty( - new_shape, - dtype=tensor.dtype, - device=tensor.device, - ) + # FSDP2 allocates new_zeros buffers as all-gather destinations + # that are immediately overwritten by copy_. Use make_empty + # (uninitialized storage with the right container shape/fields). + return tensor._quantizer.make_empty( + new_shape, + dtype=tensor.dtype, + device=tensor.device, + ) # ── FSDP2: clone ───────────────────────────────────────────── if func == aten.clone.default: @@ -990,8 +958,6 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): dtype=tensor.dtype, rowwise_storage=row_clone, columnwise_storage=col_clone, - rowwise_quantizer=tensor._rowwise_quantizer, - columnwise_quantizer=tensor._columnwise_quantizer, quantizer=tensor._quantizer, ) @@ -1001,9 +967,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): def _make_hybrid_quantized_tensor_in_reduce_ex( rowwise_storage: Optional[QuantizedTensorStorage], columnwise_storage: Optional[QuantizedTensorStorage], - rowwise_quantizer: Optional[Quantizer], - columnwise_quantizer: Optional[Quantizer], - quantizer: Optional[Quantizer], + quantizer: HybridQuantizer, dtype: torch.dtype, shape: torch.Size, ) -> HybridQuantizedTensor: @@ -1013,7 +977,5 @@ def _make_hybrid_quantized_tensor_in_reduce_ex( dtype=dtype, rowwise_storage=rowwise_storage, columnwise_storage=columnwise_storage, - rowwise_quantizer=rowwise_quantizer, - columnwise_quantizer=columnwise_quantizer, quantizer=quantizer, ) diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index 03af00184d7..9e94b654fcb 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -24,21 +24,31 @@ class HybridQuantizedTensorStorage(QuantizedTensorStorage): _rowwise_storage: Optional[QuantizedTensorStorage] _columnwise_storage: Optional[QuantizedTensorStorage] - _rowwise_quantizer: Optional[Quantizer] - _columnwise_quantizer: Optional[Quantizer] - _quantizer: Optional[Quantizer] + _quantizer: Quantizer def __new__( cls, *args, rowwise_storage: Optional[QuantizedTensorStorage], columnwise_storage: Optional[QuantizedTensorStorage], - rowwise_quantizer: Optional[Quantizer] = None, - columnwise_quantizer: Optional[Quantizer] = None, - quantizer: Optional[Quantizer] = None, + quantizer: Quantizer, fake_dtype: Optional[torch.dtype] = None, **kwargs, ): + if quantizer is None or not isinstance(quantizer, Quantizer): + raise TypeError( + "HybridQuantizedTensorStorage requires a parent HybridQuantizer; " + f"got {type(quantizer).__name__}." + ) + if not hasattr(quantizer, "rowwise_quantizer") or not hasattr( + quantizer, "columnwise_quantizer" + ): + raise TypeError( + "HybridQuantizedTensorStorage requires a parent HybridQuantizer " + "with rowwise_quantizer and columnwise_quantizer attributes; " + f"got {type(quantizer).__name__}." + ) + if cls is HybridQuantizedTensorStorage: instance = object.__new__(cls) instance._dtype = fake_dtype if fake_dtype is not None else torch.float32 @@ -47,8 +57,6 @@ def __new__( instance._rowwise_storage = rowwise_storage instance._columnwise_storage = columnwise_storage - instance._rowwise_quantizer = rowwise_quantizer - instance._columnwise_quantizer = columnwise_quantizer instance._quantizer = quantizer return instance @@ -168,8 +176,6 @@ def view(self, *shape): return HybridQuantizedTensorStorage( rowwise_storage=self._rowwise_storage, columnwise_storage=self._columnwise_storage, - rowwise_quantizer=self._rowwise_quantizer, - columnwise_quantizer=self._columnwise_quantizer, quantizer=self._quantizer, fake_dtype=self._dtype, ) @@ -180,8 +186,6 @@ def view(self, *shape): return HybridQuantizedTensorStorage( rowwise_storage=row_view, columnwise_storage=col_view, - rowwise_quantizer=self._rowwise_quantizer, - columnwise_quantizer=self._columnwise_quantizer, quantizer=self._quantizer, fake_dtype=self._dtype, ) @@ -191,8 +195,6 @@ def get_metadata(self) -> Dict[str, Any]: return { "rowwise_storage": self._rowwise_storage, "columnwise_storage": self._columnwise_storage, - "rowwise_quantizer": self._rowwise_quantizer, - "columnwise_quantizer": self._columnwise_quantizer, "quantizer": self._quantizer, "fake_dtype": self._dtype, } diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 199417df443..ca54c0fe09f 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1223,8 +1223,8 @@ def _route_hybrid_to_buckets( """ row_sub = model_weight._rowwise_storage col_sub = model_weight._columnwise_storage - sub_q_row = model_weight._rowwise_quantizer - sub_q_col = model_weight._columnwise_quantizer + sub_q_row = model_weight._quantizer.rowwise_quantizer + sub_q_col = model_weight._quantizer.columnwise_quantizer if row_sub is None and col_sub is None: raise ValueError( From 03fd6c0cb7893a621a364ab29de55a852a4611c8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 15 Jun 2026 10:11:57 +0000 Subject: [PATCH 27/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 4 +--- transformer_engine/pytorch/tensor/utils.py | 4 +++- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 17fc2ebe67f..e4e940f5e04 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -3048,9 +3048,7 @@ def _make_transpose_only_float8_weight(shape, quantizer, *, fill_value=173): shape=shape, dtype=torch.bfloat16, data=None, - data_transpose=torch.full( - (cols, rows), fill_value, dtype=torch.uint8, device="cuda" - ), + data_transpose=torch.full((cols, rows), fill_value, dtype=torch.uint8, device="cuda"), fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), fp8_dtype=tex.DType.kFloat8E4M3, quantizer=quantizer, diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index ca54c0fe09f..1d9adbe3403 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -108,7 +108,9 @@ def _is_float8_transpose_only(tensor: QuantizedTensor) -> bool: ) -def _validate_flat_fragment(model_weight: QuantizedTensor, master_weight: torch.Tensor, start_offset): +def _validate_flat_fragment( + model_weight: QuantizedTensor, master_weight: torch.Tensor, start_offset +): """Validate a flat logical shard and return its exclusive end offset.""" if start_offset is None: raise ValueError("start_offset must not be None when master_weight is provided") From e4210688cb92d2a21fda38734ab2f3fa660b7dfc Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 16 Jun 2026 12:28:33 +0000 Subject: [PATCH 28/97] Enable full hybrid nvfp4 recipe (col only nvfp4) + minor fixes Signed-off-by: Evgeny --- tests/pytorch/distributed/run_hybrid_tp_sp.py | 239 +++--------------- .../pytorch/distributed/test_hybrid_tp_sp.py | 101 ++------ tests/pytorch/test_cpu_offloading.py | 9 +- tests/pytorch/test_cpu_offloading_v1.py | 9 +- tests/pytorch/test_hybrid_quantization.py | 150 +++++++---- .../quantization_factory_examples.py | 112 ++++---- .../quantization_recipes_base.py | 2 + .../pytorch/tensor/float8_tensor.py | 72 ++++-- .../tensor/storage/float8_tensor_storage.py | 25 +- .../tensor/storage/hybrid_tensor_storage.py | 2 + 10 files changed, 299 insertions(+), 422 deletions(-) diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index b38028a36b8..4697b4d80b3 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -4,34 +4,7 @@ # # See LICENSE for license information. -""" -Distributed TP/SP numerics tests for hybrid quantization. - -Launched via torchrun from ``test_hybrid_tp_sp.py``. Each test compares a -tensor-parallel (optionally sequence-parallel) TE module running a hybrid -``CustomRecipe`` against a single-node reference module running the same -recipe. Weights are synchronized via ``_copy_params`` (shared with -``run_numerics.py``), so any drift between the two paths is a hybrid- -specific TP/SP issue rather than an initialization artifact. - -Test surface: ``te.Linear`` (column/row x SP on/off, plus a bitwise -hybrid-vs-vanilla operand-equivalence check), ``te.LayerNormLinear``, -``te.LayerNormMLP``, and ``te.TransformerLayer`` (all with SP on/off). The -non-attention tests also compare per-parameter gradients in the no-SP configs, -where grads align directly with the single-node reference. - -Recipes: same-format (FP8-current, MXFP8, NVFP4) for a clean signal and the -bitwise check, plus a cross-format one (MXFP8 fwd / NVFP4 bwd) that exercises -the forward-vs-backward all-gather format asymmetry (fwd gathers rowwise, bwd -columnwise) -- which same-format recipes cannot surface. - -Two comparison kinds with different tolerances: - * distributed-vs-single-node (``_test_*``): inherently loose -- the sharded - side quantizes per-shard and reduces across ranks, so it is never bitwise. - ``_get_tolerances`` matches upstream ``run_numerics.py`` per format. - * hybrid-vs-vanilla (``_test_linear_vs_vanilla``): same topology, so bitwise - (``rtol=0, atol=0``) for forward (all configs) and backward (non-SP). -""" +"""Distributed TP/SP coverage for hybrid quantization.""" import argparse import datetime @@ -46,16 +19,17 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.common import recipe as te_recipe +from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + nvfp4_quantizer_factory, +) from transformer_engine.pytorch import ( Float8CurrentScalingQuantizer, HybridQuantizer, IdentityQuantizer, MXFP8Quantizer, - NVFP4Quantizer, ) -# Reuse helpers from run_numerics.py (sibling import — same pattern as -# run_numerics.py's own `from run_layer_with_overlap import _compare_tensors`). +# Sibling helper shared with distributed numerics tests. TEST_ROOT = Path(__file__).parent.resolve() sys.path.insert(0, str(TEST_ROOT)) from run_layer_with_overlap import _compare_tensors # noqa: E402 @@ -79,11 +53,7 @@ # ── Hybrid recipe factories ────────────────────────────────────────── # -# Both rowwise and columnwise sub-quantizers use the same format so the -# observed distributed numerics only reflect TP/SP interactions and not -# cross-format composition noise. For comparison against vanilla built-in -# recipes that have well-understood TP/SP tolerances, see upstream -# ``run_numerics.py``. +# Factories stay small; these tests target TP/SP plumbing. def _make_fp8_current_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): @@ -95,8 +65,7 @@ def _make_mxfp8_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): def _hybrid_fp8_qfactory(role): - """FP8 current scaling in both directions for fwd roles; E5M2 for - grad roles (standard Hybrid:HYBRID format pairing).""" + """FP8 current scaling; backward operands use E5M2.""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( @@ -115,7 +84,6 @@ def _hybrid_mxfp8_qfactory(role): rowwise_quantizer=_make_mxfp8_quantizer(), columnwise_quantizer=_make_mxfp8_quantizer(), ) - # MXFP8 uses E4M3 for every pass (its canonical Format.E4M3) return _make_mxfp8_quantizer() @@ -147,76 +115,32 @@ def _identity_qfactory(role): # pylint: disable=unused-argument return IdentityQuantizer() -def _make_nvfp4_bare(): - """Bare NVFP4Quantizer (1D, no RHT/SR/2D), used by the cross-format recipe - to avoid cross-operand RHT-consistency concerns in the mixed MXFP8/NVFP4 - GEMMs.""" - return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) - - -def _make_nvfp4_quantizer(role=None): - """Role-based NVFP4Quantizer mirroring ``nvfp4_quantizer_factory`` / - :class:`NVFP4BlockScaling`, but with 2D quantization disabled. - - Per role: weight = no RHT/SR, input = RHT, grad = RHT + SR. - - ``with_2d_quantization`` is forced off everywhere: the 2D quantize-transpose - kernel has no columnwise-only path, so a hybrid columnwise sub-quantizer - cannot use it. - TODO(negvet): enable 2D for the rowwise direction once - https://github.com/NVIDIA/TransformerEngine/pull/3027 lands. - """ - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - is_weight = is_linear and role.tensor_type == "weight" - is_grad = is_linear and role.tensor_type in ("grad_output", "grad_input") - return NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, - with_rht=not is_weight, - with_post_rht_amax=not is_weight, - with_2d_quantization=False, # TODO(negvet): enable via PR #3027 - stochastic_rounding=is_grad, - with_random_sign_mask=True, - ) - - def _hybrid_nvfp4_qfactory(role): is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight", "output"): - # Same per-role config for both directions (RHT/SR are per role). return HybridQuantizer( - rowwise_quantizer=_make_nvfp4_quantizer(role), - columnwise_quantizer=_make_nvfp4_quantizer(role), + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_nvfp4_quantizer(role) - return _make_nvfp4_quantizer(role) + return nvfp4_quantizer_factory(role) + return nvfp4_quantizer_factory(role) def _hybrid_mxfp8_nvfp4_qfactory(role): - """Cross-format: MXFP8 forward (rowwise) + NVFP4 backward (columnwise). - - fprop TN: weight.row(MXFP8) x input.row(MXFP8) -> MXFP8 x MXFP8 - dgrad NN: weight.col(NVFP4) x grad_output.row(NVFP4) -> NVFP4 x NVFP4 - wgrad NT: input.col(NVFP4) x grad_output.col(NVFP4) -> NVFP4 x NVFP4 - - So weight/input = Hybrid(row=MXFP8, col=NVFP4), grad_output = plain NVFP4. - The forward all-gather consumes the MXFP8 rowwise sub-storage and the - backward all-gather the NVFP4 columnwise one -- the fwd-vs-bwd format - asymmetry that same-format recipes cannot surface. - """ + """Cross-format recipe: MXFP8 rowwise, NVFP4 columnwise.""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_nvfp4_bare() - # input / weight / output / unknown-None: MXFP8 rowwise + NVFP4 columnwise. + return nvfp4_quantizer_factory(role) + # Forward/boundary roles keep both formats available. return HybridQuantizer( rowwise_quantizer=_make_mxfp8_quantizer(), - columnwise_quantizer=_make_nvfp4_bare(), + columnwise_quantizer=nvfp4_quantizer_factory(role), ) def hybrid_recipe(): - """Fresh CustomRecipe instance per call (mirrors - ``run_numerics.quantization_recipe`` lifetime contract).""" + """Return a fresh CustomRecipe for the selected test recipe.""" if QUANTIZATION == "hybrid_fp8": return te_recipe.CustomRecipe(qfactory=_hybrid_fp8_qfactory) if QUANTIZATION == "hybrid_mxfp8": @@ -236,21 +160,13 @@ def hybrid_recipe(): # ── Tolerances ─────────────────────────────────────────────────────── # -# These match upstream ``run_numerics.py::_get_tolerances`` exactly, for -# the same TP/SP-distributed-vs-single-node comparison. A same-format -# hybrid inherits the underlying format's distributed behaviour: both the -# distributed and single-node models run the *same* two-pass hybrid recipe, -# so the two-pass quantization cancels in the comparison and the only -# remaining difference is the TP/SP path (per-shard quantization, -# all-gather/reduce-scatter order, and -- for fp8_cs only -- cross-rank -# amax reduction). There is therefore no reason for hybrid to need looser -# bounds than the vanilla format. +# Mostly upstream ``run_numerics.py`` tolerances. NVFP4 needs a slightly +# larger atol for the measured column+SP Linear output. def _get_tolerances(): if QUANTIZATION == "identity": - # Same tolerance as upstream distributed BF16 numerics: TP row - # reductions can accumulate in a different order from the single-node ref. + # BF16 TP reductions accumulate in a different order. return {"rtol": 1.6e-2, "atol": 1.0e-5} if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): # Loose because of sequence parallel & amax reduction (fp8_cs). @@ -258,9 +174,8 @@ def _get_tolerances(): if QUANTIZATION in ("hybrid_mxfp8", "hybrid_mxfp8_identity"): return {"rtol": 0.125, "atol": 0.0625} if QUANTIZATION == "hybrid_nvfp4": - # Upstream ``run_numerics.py`` uses (0.125, 0.12) for vanilla NVFP4 - # (with an open TODO to investigate why the tolerance is so large). - return {"rtol": 0.125, "atol": 0.12} + # Measured column+SP Linear output max abs is ~0.144. + return {"rtol": 0.125, "atol": 0.15} if QUANTIZATION == "hybrid_mxfp8_nvfp4": # Backward GEMMs run in NVFP4 -> inherit the (looser) NVFP4 bounds. return {"rtol": 0.125, "atol": 0.12} @@ -278,12 +193,7 @@ def dist_print(msg, src=None, error=False): def _gather(tensor, dim=0): - """All-gather with gradient scaling, matching - ``run_numerics.py::_gather``. Required because - ``torch.distributed.nn.functional.all_gather`` multiplies gradients - by WORLD_SIZE on the backward pass — so gradients in the - ``output_distributed`` backward would be WORLD_SIZE× too large - compared to ``output_single_node``.""" + """All-gather with ``run_numerics.py`` gradient scaling.""" class HalfGradient(torch.autograd.Function): @staticmethod @@ -300,10 +210,7 @@ def backward(ctx, grad_output): def _copy_params(model_distributed, model_single): - """Shard the single-node parameters into the TP-split distributed - model. Same algorithm as ``run_numerics.py::_copy_params``: for each - dim where shapes differ between the two params, slice the single- - node param along that dim using ``WORLD_RANK``.""" + """Copy single-node parameters into the local TP shard.""" for dp, sp in zip(model_distributed.parameters(), model_single.parameters()): with torch.no_grad(): to_copy = sp @@ -356,10 +263,7 @@ def _check_gradients(model_dist, model_single): def _apply_models(model_single, model_dist, inp_single, inp_dist, **kwargs): - """Run both models under te.autocast with a fresh hybrid recipe each - time. Both models see the same recipe instance-shape (CustomRecipe - with the same qfactory), but get independently-constructed - quantizers — matching how real training would instantiate them.""" + """Run both models with fresh CustomRecipe instances.""" inp_single.requires_grad_() inp_dist.requires_grad_() with te.autocast(enabled=True, recipe=hybrid_recipe()): @@ -400,7 +304,7 @@ def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16, _copy_params(model_dist, model_single) - # Prepare inputs matching run_numerics._test_linear's conventions. + # Match run_numerics._test_linear input layouts. inp_single = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) if parallel_mode == "row": split = HIDDEN_SIZE // WORLD_SIZE @@ -411,9 +315,7 @@ def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16, inp_single = torch.empty((WORLD_SIZE * BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) inp_dist = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) if amax_stress and WORLD_RANK == WORLD_SIZE - 1: - # Large element on one rank: its local amax diverges from the - # global one. Hybrid gathers the SP activation whole before - # quantizing, so the output must still match the single-node ref. + # One-rank outlier before SP gather. inp_dist[-1, -1] = 1.0e3 inp_single = _gather(inp_dist, dim=0).detach() else: @@ -436,10 +338,7 @@ def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16, label=f"linear[{parallel_mode},sp={sequence_parallel},amax_stress={amax_stress}]", ) - # Gradient check is only well-defined in these configurations (the - # others need cross-rank synchronization that the test doesn't - # perform — see run_numerics.py::_test_linear line 725 for the - # matching gate). + # Other configs need extra grad sync, matching run_numerics.py. if parallel_mode == "column" or not sequence_parallel: _check_gradients(model_dist, model_single) @@ -448,11 +347,7 @@ def test_linear(): for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: _test_linear(parallel_mode, sequence_parallel) - # Amax corner-case (current scaling only): a large element on one rank makes - # its local amax diverge from the global one. Hybrid gathers SP activations in - # high precision before quantizing, so the SP output must still match - # single-node -- guards against a future regression to quantize-then-gather - # without cross-rank amax reduction. + # Current-scaling amax stress: one rank owns an outlier before SP gather. if QUANTIZATION in ("hybrid_fp8", "hybrid_fp8_identity"): _test_linear("column", True, amax_stress=True) @@ -461,31 +356,18 @@ def test_linear(): def vanilla_recipe(): - """Built-in single-format recipe matching the same-format hybrid recipe - for the bitwise ``_test_linear_vs_vanilla`` check: FP8 current scaling and - MXFP8 use their defaults (E4M3 fwd / E5M2 bwd, and E4M3 everywhere); NVFP4 - uses the full recipe with 2D disabled to match the role-based 1D - ``_make_nvfp4_quantizer``.""" + """Built-in recipe for same-format hybrid-vs-vanilla checks.""" if QUANTIZATION == "hybrid_fp8": return te_recipe.Float8CurrentScaling() if QUANTIZATION == "hybrid_mxfp8": return te_recipe.MXFP8BlockScaling() if QUANTIZATION == "hybrid_nvfp4": - return te_recipe.NVFP4BlockScaling(disable_2d_quantization=True) + return te_recipe.NVFP4BlockScaling() raise ValueError(f"No vanilla recipe for QUANTIZATION={QUANTIZATION!r}") def _backward_not_bitwise_comparable(): - """Whether the recipe's backward can't be compared bitwise to vanilla. - - True only for NVFP4's full recipe, which combines RHT with stochastic - rounding. That pair triggers NVFP4's separate columnwise RNG state - (``need_separate_columnwise_rng`` in the cast backend), and the hybrid - (two-pass) vs vanilla (fused) executions then consume that columnwise - random stream differently. Verified by isolation: neither RHT nor SR alone - diverges -- only the combination, and only on the columnwise gradient - (wgrad); the rowwise gradient (dgrad) stays bitwise. - """ + """NVFP4 backward consumes SR RNG differently in hybrid vs vanilla.""" return QUANTIZATION == "hybrid_nvfp4" @@ -502,15 +384,7 @@ def _check_bitwise(actual, expected, label): def _test_linear_vs_vanilla(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16): - """Same-format hybrid must match its built-in vanilla recipe **bitwise** - through the *same* TP/SP-distributed ``te.Linear`` (forward in all configs; - backward in the non-SP, non-SR configs). - - Unlike ``_test_linear`` (distributed vs single-node, inherently loose), - this compares hybrid vs vanilla at the same topology, so any non-bitwise - difference is a real hybrid-plumbing bug. Complements the FSDP2 parity test - in ``run_fsdp2_fused_adam.py`` by locking the TP/SP comm path. - """ + """Same-topology Linear check: forward bitwise, backward where comparable.""" dist_print( f"linear_vs_vanilla: parallel_mode={parallel_mode} sequence_parallel={sequence_parallel}" ) @@ -554,24 +428,11 @@ def run(recipe): tag = f"linear_vs_vanilla[{parallel_mode},sp={sequence_parallel}]" - # Forward is bitwise-identical in every config (the fprop operand- - # equivalence check: hybrid weight.rowwise/input.rowwise == vanilla). + # Same-topology fprop operands should match bitwise. _check_bitwise(out_h, out_v, f"{tag} forward") - # Backward is bitwise only without SP and without stochastic rounding; - # both are within training tolerance and covered by the loose - # distributed-vs-single-node check: - # * Under SP, hybrid has no native quantized all-gather, so it routes - # through the BF16 dequant/requant fallback while vanilla gathers native - # per-shard bytes. For per-tensor-scaled formats (FP8 current, NVFP4) the - # requantized scale then differs; MXFP8 (per-block only) is immune. - # TODO(negvet): extend to SP once native hybrid AG lands (tracked in - # HybridQuantizer.supports_only_rowwise_all_gather). - # * NVFP4's full recipe combines RHT + stochastic rounding, which triggers - # a separate columnwise RNG (need_separate_columnwise_rng); the hybrid - # two-pass and vanilla fused paths then consume that columnwise random - # stream differently, so the columnwise gradient (wgrad) rounds - # differently. Neither RHT nor SR alone diverges -- only the pair. + # Skip backward bitwise when SP changes the gather path or NVFP4 consumes + # columnwise SR RNG differently. Loose TP/SP checks cover those cases. if not sequence_parallel and not _backward_not_bitwise_comparable(): _check_bitwise(dinp_h, dinp_v, f"{tag} dgrad") assert len(wgrads_h) == len(wgrads_v), f"{tag}: weight-grad count mismatch" @@ -580,8 +441,7 @@ def run(recipe): def test_linear_vs_vanilla(): - # Cross-format hybrid has no single built-in vanilla recipe to compare - # against bitwise; it is covered by the distributed-vs-single-node checks. + # These recipes have no same-format vanilla bitwise target. if QUANTIZATION in ( "identity", "hybrid_mxfp8_nvfp4", @@ -602,9 +462,7 @@ def _same_format_parity_supported(): def _check_same_topology_parity( out_h, dinp_h, model_h, out_v, dinp_v, model_v, tag, *, check_grads ): - # Larger modules use different fused/unfused norm paths between hybrid and - # vanilla, so numerical parity is the meaningful contract here. Linear keeps - # the stricter bitwise check above. + # Larger modules may use different fused/unfused norm paths. _check_outputs(out_v, out_h, label=f"{tag} forward") if check_grads: _check_outputs(dinp_v, dinp_h, label=f"{tag} dgrad") @@ -711,10 +569,7 @@ def test_layernorm_mlp_vs_vanilla(): def _test_layernorm_linear(sequence_parallel, params_dtype=torch.bfloat16): - """Column-parallel LayerNormLinear. Exercises the SP all-gather path - that runs BEFORE quantization for hybrid (since - ``with_quantized_norm=False`` for HybridQuantizer — see - ``layernorm_linear.py:220``).""" + """Column-parallel LayerNormLinear with optional SP.""" dist_print(f"layernorm_linear: parallel_mode=column sequence_parallel={sequence_parallel}") torch.manual_seed(23456) @@ -758,11 +613,7 @@ def test_layernorm_linear(): def _test_layernorm_mlp(sequence_parallel, params_dtype=torch.bfloat16): - """``te.LayerNormMLP`` with ``set_parallel_mode=True`` and optional SP: - column-parallel FC1 → activation → row-parallel FC2. Isolates the FC1 - hybrid unfused-norm path and the row-parallel FC2 + SP reduce-scatter, - otherwise only exercised transitively inside ``_test_transformer_layer``. - """ + """LayerNormMLP with set_parallel_mode=True and optional SP.""" dist_print(f"layernorm_mlp: parallel_mode=set sequence_parallel={sequence_parallel}") torch.manual_seed(45678) @@ -798,8 +649,7 @@ def _test_layernorm_mlp(sequence_parallel, params_dtype=torch.bfloat16): _loss_backward(out_single, out_dist) _check_outputs(out_single, out_dist, label=f"layernorm_mlp[sp={sequence_parallel}]") - # Without SP, grads align with the single-node ref (SP needs cross-rank - # grad sync the test doesn't do -- matches run_numerics.py). + # SP grads need extra sync, matching run_numerics.py. if not sequence_parallel: _check_gradients(model_dist, model_single) @@ -813,11 +663,7 @@ def test_layernorm_mlp(): def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): - """Integration test: full TransformerLayer with TP and optional SP. - Hits LayerNormLinear(QKV), DPA, and LayerNormMLP all with hybrid - quantizers. If any of the unfused/hybrid code paths break something - visible to the backward graph, this catches it with a concrete - forward-output mismatch.""" + """TransformerLayer integration with TP and optional SP.""" dist_print(f"transformer_layer: parallel_mode=set sequence_parallel={sequence_parallel}") torch.manual_seed(34567) @@ -865,8 +711,7 @@ def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): _loss_backward(out_single, out_dist) _check_outputs(out_single, out_dist, label=f"transformer_layer[sp={sequence_parallel}]") - # Without SP, verify the integration path at the gradient level too (SP - # needs cross-rank grad sync the test doesn't do -- matches run_numerics.py). + # SP grads need extra sync, matching run_numerics.py. if not sequence_parallel: _check_gradients(model_dist, model_single) diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py index 23ae42d8325..8cc78d793a1 100644 --- a/tests/pytorch/distributed/test_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -2,19 +2,7 @@ # # See LICENSE for license information. -"""Pytest driver for hybrid quantization TP/SP distributed tests. - -Launches ``run_hybrid_tp_sp.py`` via ``torchrun --nproc_per_node=N`` -and asserts a zero exit code. Rank-level numerical checks are performed -inside ``run_hybrid_tp_sp.py`` and propagated via ``dist.all_reduce`` -with ``ReduceOp.MAX`` on a failure flag, so a failure on any rank -fails the whole subprocess (and thus the pytest assertion). - -Mirrors the ``test_numerics.py`` ↔ ``run_numerics.py`` split pattern but -scoped to hybrid recipes only. Isolated from the main ``run_numerics.py`` -harness so that adding hybrid-specific cases here doesn't perturb the -larger vanilla-recipe test matrix. -""" +"""Pytest launcher for hybrid TP/SP distributed tests.""" import os import subprocess @@ -51,83 +39,54 @@ def _run_test(quantization: str, test: str = "all"): # ────────────────────────────────────────────────────────────────────── -# Hybrid FP8 current scaling (rowwise + columnwise same format) +# Hybrid FP8 current scaling # ────────────────────────────────────────────────────────────────────── -# -# FP8 current scaling is stateless per-tensor (amax computed from the -# live tensor) and therefore uses amax reduction across TP ranks when -# sequence parallelism is on. This is the tightest integration of -# hybrid with the distributed codepath: each rank computes an -# independent amax, reduces it across TP group, then both sub- -# quantizers (rowwise + columnwise) use the same reduced scale. If -# ``HybridQuantizer`` mis-plumbs the amax reduction through its two -# inner quantizers, we'd see numerical drift vs single-node. +# Exercises TP amax reduction and SP gather paths. @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_linear(): - """TP column + row × SP on/off for ``te.Linear`` under hybrid FP8 - current-scaling. Fine-grained: this runs first (cheapest, most - likely to surface TP-path hybrid bugs) so a failure here tells us - to stop before the more-expensive TransformerLayer run.""" + """Linear TP/SP coverage for hybrid FP8 current scaling.""" _run_test("hybrid_fp8", "linear") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_linear_vs_vanilla(): - """Bitwise operand equivalence: same-format hybrid FP8 must match the - built-in ``Float8CurrentScaling`` recipe through the same TP ``te.Linear`` - (forward in all configs; backward in the non-SP configs). Locks the TP - comm path that the FSDP2 parity test does not exercise.""" + """Same-topology Linear parity against Float8CurrentScaling.""" _run_test("hybrid_fp8", "linear_vs_vanilla") def test_hybrid_fp8_layernorm_linear_vs_vanilla(): - """Same-topology numerical parity against vanilla FP8 for LayerNormLinear. - - This extends the Linear bitwise operand check to the unfused-norm hybrid - module path; exact bitwise parity is not required because vanilla may use - fused quantized norm while hybrid routes through high-precision norm. - """ + """Same-topology LayerNormLinear parity against vanilla FP8.""" _run_test("hybrid_fp8", "layernorm_linear_vs_vanilla") def test_hybrid_fp8_layernorm_mlp_vs_vanilla(): - """Same-topology numerical parity against vanilla FP8 for LayerNormMLP.""" + """Same-topology LayerNormMLP parity against vanilla FP8.""" _run_test("hybrid_fp8", "layernorm_mlp_vs_vanilla") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_layernorm_linear(): - """Column-parallel ``te.LayerNormLinear`` with and without SP. - Probes the all-gather-before-quantize path that - ``layernorm_linear.py`` disables the fused norm for when - ``isinstance(input_quantizer, HybridQuantizer)``.""" + """LayerNormLinear TP/SP coverage for hybrid FP8.""" _run_test("hybrid_fp8", "layernorm_linear") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_layernorm_mlp(): - """Standalone ``te.LayerNormMLP`` (column FC1 / row FC2) with and - without SP under hybrid FP8. Isolates the MLP block's unfused-norm - and row-parallel reduce-scatter paths, and checks gradients in the - no-SP case.""" + """LayerNormMLP TP/SP coverage for hybrid FP8.""" _run_test("hybrid_fp8", "layernorm_mlp") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_transformer_layer(): - """Full ``te.TransformerLayer`` with ``set_parallel_mode=True`` and - SP on/off. Integration check hitting LayerNormLinear(QKV) → DPA → - LayerNormMLP → row-parallel output projection all under hybrid - FP8.""" + """TransformerLayer TP/SP coverage for hybrid FP8.""" _run_test("hybrid_fp8", "transformer_layer") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_identity_linear(): - """Linear-only TP/SP coverage for FP8-current forward plus Identity backward. - Includes the amax-stress branch inside ``run_hybrid_tp_sp.py``.""" + """Linear TP/SP coverage for FP8 forward plus Identity backward.""" _run_test("hybrid_fp8_identity", "linear") @@ -137,14 +96,9 @@ def test_identity_all_modules(): # ────────────────────────────────────────────────────────────────────── -# Hybrid MXFP8 (rowwise + columnwise same format) +# Hybrid MXFP8 # ────────────────────────────────────────────────────────────────────── -# -# MXFP8 is per-block (32-element microblocks), stateless, no amax -# reduction. Simpler distributed behaviour than FP8 current scaling, -# but exercises the ``[128, 4]`` / ``[4, 128]`` scale alignment padding -# through TP shards (each rank sees its own dim-0 slice which may not -# be a multiple of 128). +# Covers per-block scale layout through TP shards. @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") @@ -182,29 +136,14 @@ def test_hybrid_mxfp8_transformer_layer(): @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") def test_hybrid_mxfp8_identity_linear(): - """Linear-only TP/SP coverage for MXFP8 forward plus Identity backward.""" + """Linear TP/SP coverage for MXFP8 forward plus Identity backward.""" _run_test("hybrid_mxfp8_identity", "linear") # ────────────────────────────────────────────────────────────────────── -# Hybrid NVFP4 (rowwise + columnwise same format, 1D block scaling) +# Hybrid NVFP4 # ────────────────────────────────────────────────────────────────────── -# -# NVFP4 is the Rubin-era target recipe: 4-bit data (E2M1) with FP8 block -# scales on 16-element microblocks. The default ``NVFP4Quantizer()`` is -# 1D block scaling only — no RHT, no stochastic rounding, no 2D block -# scaling — matching upstream ``run_numerics.py::nvfp4_vanilla()``. -# Those more-sophisticated knobs are orthogonal to hybrid composition -# and can be layered in separately once baseline distributed NVFP4 -# hybrid is stable. -# -# Unlike FP8 current scaling, NVFP4 does not reduce amax across TP ranks -# (block-level scales are computed per-microblock locally), so SP -# amax-reduction issues don't apply. The tight interaction to watch is -# the packed FP4 dim-0 alignment in the TP shard — each rank sees a -# weight slice that may not be naturally aligned to the NVFP4 block -# boundary, and hybrid quantizes twice (rowwise + columnwise) on that -# shard. +# Same-format NVFP4 coverage with base role-wise settings. @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") @@ -233,13 +172,9 @@ def test_hybrid_nvfp4_transformer_layer(): # ────────────────────────────────────────────────────────────────────── -# Cross-format hybrid: MXFP8 forward (rowwise) + NVFP4 backward (columnwise) +# Cross-format hybrid: MXFP8 rowwise + NVFP4 columnwise # ────────────────────────────────────────────────────────────────────── -# -# Forward and backward all-gather *different* formats (MXFP8 rowwise vs NVFP4 -# columnwise) -- the asymmetry same-format recipes can't surface. Only the -# distributed-vs-single-node checks run (no single vanilla recipe to match a -# cross-format hybrid bitwise). Needs both MXFP8 and NVFP4 hardware support. +# No single vanilla recipe exists for bitwise comparison. _cross_format_available = mxfp8_available and nvfp4_available _reason_for_no_cross_format = reason_for_no_mxfp8 if not mxfp8_available else reason_for_no_nvfp4 diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 0f6ba43fbaf..e6170145e96 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -19,6 +19,9 @@ from transformer_engine.pytorch.fp8 import FP8GlobalStateManager import transformer_engine.pytorch as te from transformer_engine.common import recipe +from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + nvfp4_quantizer_factory, +) from utils import ModelConfig, recipe_id, skip_unsupported_backward_override # Check supported quantization schemes @@ -51,17 +54,17 @@ def _hybrid_mxfp8_nvfp4_qfactory(role): """Hybrid CustomRecipe factory: MXFP8 rowwise + NVFP4 columnwise. Mirrors ``mxfp8_fwd_nvfp4_bwd_quantizer_factory`` from - ``custom_recipes/quantization_nvfp4.py``. grad_output uses plain NVFP4 + ``custom_recipes/quantization_factory_examples.py``. grad_output uses plain NVFP4 (both directions) so wgrad's columnwise operand matches. """ is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3), - columnwise_quantizer=te.NVFP4Quantizer(fp4_dtype=te.DType.kFloat4E2M1), + columnwise_quantizer=nvfp4_quantizer_factory(role), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return te.NVFP4Quantizer(fp4_dtype=te.DType.kFloat4E2M1) + return nvfp4_quantizer_factory(role) return te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index a5e5df8d200..d6a4e02a875 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -13,6 +13,9 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.common import recipe +from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + nvfp4_quantizer_factory, +) from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported from utils import ModelConfig, get_available_attention_backends @@ -48,17 +51,17 @@ def _hybrid_mxfp8_nvfp4_qfactory(role): """Hybrid CustomRecipe factory: MXFP8 rowwise + NVFP4 columnwise. Mirrors the ``mxfp8_fwd_nvfp4_bwd_quantizer_factory`` headline recipe - from ``custom_recipes/quantization_nvfp4.py``. grad_output uses plain + from ``custom_recipes/quantization_factory_examples.py``. grad_output uses plain NVFP4 (both directions) so wgrad's columnwise operand matches. """ is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight", "output"): return te.HybridQuantizer( rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - columnwise_quantizer=te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + columnwise_quantizer=nvfp4_quantizer_factory(role), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return te.NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + return nvfp4_quantizer_factory(role) return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index e4e940f5e04..7aad5670bf5 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -11,6 +11,9 @@ import transformer_engine_torch as tex from transformer_engine.common import recipe +from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + nvfp4_quantizer_factory, +) from transformer_engine.pytorch import ( autocast, quantized_model_init, @@ -1151,18 +1154,13 @@ def hybrid_block_fp8_factory(role): ) + @pytest.mark.skipif( not (fp8_available and nvfp4_available), reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", ) class TestHybridGemmBitwiseIdenticalNVFP4: - """Hybrid quantizer with NVFP4 in both directions must produce - bitwise-identical results to the vanilla NVFP4BlockScaling recipe. - - RHT, stochastic rounding, and 2D quantization are disabled so the - test is fully deterministic and two independent quantizer instances - produce the same output. - """ + """Same-format hybrid NVFP4 must match vanilla with seeded SR.""" def test_linear_fwd_bwd_matches_vanilla_nvfp4(self): torch.manual_seed(202) @@ -1177,11 +1175,9 @@ def test_linear_fwd_bwd_matches_vanilla_nvfp4(self): inp_ref = base_inp.clone().detach().requires_grad_(True) inp_hybrid = base_inp.clone().detach().requires_grad_(True) - ref_recipe = recipe.NVFP4BlockScaling( - disable_rht=True, - disable_stochastic_rounding=True, - disable_2d_quantization=True, - ) + ref_recipe = recipe.NVFP4BlockScaling() + torch.manual_seed(1202) + torch.cuda.manual_seed_all(1202) with autocast(enabled=True, recipe=ref_recipe): out_ref = model_ref(inp_ref) out_ref.float().sum().backward() @@ -1190,15 +1186,17 @@ def hybrid_nvfp4_factory(role): if ( role is not None and role.module_type in ("linear", "grouped_linear") - and role.tensor_type in ("grad_output", "grad_input") + and role.tensor_type == "grad_output" ): - return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + return nvfp4_quantizer_factory(role) return HybridQuantizer( - rowwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), - columnwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), ) hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_factory) + torch.manual_seed(1202) + torch.cuda.manual_seed_all(1202) with autocast(enabled=True, recipe=hybrid_recipe): out_hybrid = model_hybrid(inp_hybrid) out_hybrid.float().sum().backward() @@ -1220,9 +1218,7 @@ def hybrid_nvfp4_factory(role): ) def test_linear_fwd_bwd_all_roles_hybrid(self): - """All roles (including grad_output) use HybridQuantizer with NVFP4 both - directions. Validates per-operand unwrap produces bitwise-identical results - when grad_output is hybrid.""" + """Exercise grad_output as a HybridQuantizer too.""" torch.manual_seed(203) in_features, out_features, batch = 128, 128, 32 @@ -1235,22 +1231,22 @@ def test_linear_fwd_bwd_all_roles_hybrid(self): inp_ref = base_inp.clone().detach().requires_grad_(True) inp_hybrid = base_inp.clone().detach().requires_grad_(True) - ref_recipe = recipe.NVFP4BlockScaling( - disable_rht=True, - disable_stochastic_rounding=True, - disable_2d_quantization=True, - ) + ref_recipe = recipe.NVFP4BlockScaling() + torch.manual_seed(1203) + torch.cuda.manual_seed_all(1203) with autocast(enabled=True, recipe=ref_recipe): out_ref = model_ref(inp_ref) out_ref.float().sum().backward() def hybrid_nvfp4_all_roles_factory(role): return HybridQuantizer( - rowwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), - columnwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), ) hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_all_roles_factory) + torch.manual_seed(1203) + torch.cuda.manual_seed_all(1203) with autocast(enabled=True, recipe=hybrid_recipe): out_hybrid = model_hybrid(inp_hybrid) out_hybrid.float().sum().backward() @@ -3359,21 +3355,19 @@ def test_nvfp4_rowwise_raises(self): documented in the TODO block above ``_route_hybrid_to_buckets`` in tensor/utils.py. - NOTE: this test exercises the rejection at the ``quantize_master_weights`` - entrypoint, but ``quantized_model_init`` with 2D NVFP4 hybrid already - fails earlier (single-direction 2D NVFP4 quantize is rejected by the - kernel). Use ``with_2d_quantization=False`` so the param can be - constructed and the rejection surfaces at the cast site we care about. + NOTE: after PR #3027, single-direction 2D NVFP4 construction works, + so this test now reaches the intended ``quantize_master_weights`` + rejection while using the base weight scaling mode. """ from transformer_engine.pytorch.tensor.utils import quantize_master_weights group = _ensure_single_rank_dp_group() hybrid_recipe = _hybrid_custom_recipe( row_factory=lambda: NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True ), col_factory=lambda: NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True ), grad_factory=lambda: NVFP4Quantizer( fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False @@ -3850,12 +3844,20 @@ class TestHybridMixedFormatQuantizedParams: """Quantized params with genuinely different formats per direction.""" def _build_mixed_model(self, in_features=256, out_features=256): - """MXFP8 rowwise (for fprop TN) + NVFP4 columnwise (for wgrad NT).""" - hybrid_recipe = _hybrid_custom_recipe( - row_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - col_factory=lambda: NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), - grad_factory=lambda: NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), - ) + """MXFP8 rowwise (fprop) + role-aware NVFP4 columnwise (bwd).""" + + def qfactory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + if is_linear and role.tensor_type == "grad_output": + return nvfp4_quantizer_factory(role) + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + + hybrid_recipe = recipe.CustomRecipe(qfactory=qfactory) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() return model, hybrid_recipe @@ -4008,13 +4010,13 @@ def _hybrid_block_fp8_qfactory(role): def _hybrid_nvfp4_qfactory(role): - """Hybrid NVFP4 (E2M1 both dirs).""" + """Hybrid NVFP4 (E2M1 both dirs, base role behavior).""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1) + if is_linear and role.tensor_type == "grad_output": + return nvfp4_quantizer_factory(role) return HybridQuantizer( - rowwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), - columnwise_quantizer=NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1), + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), ) @@ -4085,6 +4087,8 @@ def _test_equivalence(self): x = torch.randn(4, 32, self.hidden_size, dtype=torch.bfloat16, device="cuda") target = torch.randn_like(x) + torch.manual_seed(199) + torch.cuda.manual_seed_all(199) losses_ref, masters_ref = self._run_training_loop( model_ref, self._vanilla_recipe(), @@ -4092,6 +4096,8 @@ def _test_equivalence(self): target, self.num_steps, ) + torch.manual_seed(199) + torch.cuda.manual_seed_all(199) losses_hyb, masters_hyb = self._run_training_loop( model_hyb, self._hybrid_recipe(), @@ -4200,17 +4206,10 @@ def test_equivalence(self): reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", ) class TestQuantizedParamsEquivalenceNVFP4(_QuantizedParamsEquivalenceBase): - """Vanilla NVFP4BlockScaling vs hybrid NVFP4 (same format both dirs). - - RHT, stochastic rounding, and 2D quantization disabled for determinism. - """ + """Vanilla NVFP4BlockScaling vs same-format hybrid NVFP4.""" def _vanilla_recipe(self): - return recipe.NVFP4BlockScaling( - disable_rht=True, - disable_stochastic_rounding=True, - disable_2d_quantization=True, - ) + return recipe.NVFP4BlockScaling() def _hybrid_recipe(self): return recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) @@ -4518,6 +4517,51 @@ def test_float8_split_uses_logical_shape_for_transpose_only_storage( assert all(piece._data is None for piece in pieces) assert all(piece._transpose_invalid is False for piece in pieces) + def test_float8_view_preserves_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + viewed = tensor.view(12, 16) + + assert tuple(viewed.shape) == (12, 16) + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_aten_view_preserves_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + viewed = torch.ops.aten.view.default(tensor, [12, 16]) + + assert tuple(viewed.shape) == (12, 16) + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_storage_view_preserves_transpose_only_storage(self): + data_transpose = torch.empty((16, 12), dtype=torch.uint8, device="cuda") + storage = Float8TensorStorage( + data=None, + data_transpose=data_transpose, + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + fake_dtype=torch.bfloat16, + ) + + viewed = storage.view(torch.Size((12, 16))) + + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_shape_changing_view_raises_for_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + with pytest.raises(NotImplementedError, match="columnwise-only data"): + tensor.view(6, 32) + def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self): columnwise = self._make_transpose_only_float8_tensor() quantizer = HybridQuantizer( diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py index 956fed26734..2347ca8de94 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py @@ -13,11 +13,19 @@ from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + mxfp8_fwd_nvfp4_bwd_quantizer_factory, nvfp4_linear_mxfp8_grouped_linear_factory, nvfp4_linear_fp8_dpa_factory, nvfp4_linear_mxfp8_dpa_factory, + high_precision_factory, + fwd_high_precision_bwd_mxfp8_factory, ) + # Hybrid per-direction recipe: MXFP8 for fprop, NVFP4 for dgrad/wgrad + recipe = CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_quantizer_factory) + with autocast(recipe=recipe): + output = model(input) + # Mixed module types: NVFP4 for Linear, MXFP8 for GroupedLinear recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_grouped_linear_factory) with autocast(recipe=recipe): @@ -32,6 +40,11 @@ recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True) with autocast(recipe=recipe): output = model(input) + + # High precision forward, MXFP8 backward + recipe = CustomRecipe(qfactory=fwd_high_precision_bwd_mxfp8_factory) + with autocast(recipe=recipe): + output = model(input) """ from __future__ import annotations @@ -40,6 +53,35 @@ from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType +from .quantization_recipes_base import mxfp8_quantizer_factory, nvfp4_quantizer_factory + + +def mxfp8_fwd_nvfp4_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: MXFP8 forward, NVFP4 backward. + + Per-GEMM format consumption: + * fprop: ``weight.row(MXFP8) x input.row(MXFP8)`` + * dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` + * wgrad: ``input.col(NVFP4) x grad_output.col(NVFP4)`` + + The NVFP4 side mirrors :func:`nvfp4_quantizer_factory` role semantics: + weights use 2D scaling, activations/grads use the base 1D RHT/SR settings. + ``HybridQuantizer`` then pins each sub-quantizer to the direction that is + actually consumed. + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight", "output"): + return HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return nvfp4_quantizer_factory(role) + return mxfp8_quantizer_factory(role) def nvfp4_linear_mxfp8_grouped_linear_factory( @@ -59,63 +101,9 @@ def nvfp4_linear_mxfp8_grouped_linear_factory( is_grouped_linear = role is not None and role.module_type == "grouped_linear" if is_grouped_linear: - return _make_mxfp8_quantizer() - - return _make_nvfp4_quantizer(role) - - -def _make_mxfp8_quantizer(): - """Return an MXFP8 quantizer with default settings (E4M3, block-32, E8M0 scales).""" - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - return MXFP8Quantizer( - fp8_dtype=DType.kFloat8E4M3, - ) - + return mxfp8_quantizer_factory(role) -def _make_nvfp4_quantizer(role: Optional[QuantizerRole]): - """Return an NVFP4 quantizer configured per tensor role. - - Mirrors :class:`NVFP4BlockScaling` recipe defaults. - """ - from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - - is_linear = role is not None and role.module_type == "linear" - is_weight = is_linear and role.tensor_type == "weight" - is_grad = is_linear and role.tensor_type == "grad_output" - - if is_weight: - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - with_rht=False, - with_post_rht_amax=False, - with_2d_quantization=True, - stochastic_rounding=False, - with_random_sign_mask=True, - ) - - if is_grad: - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, - with_rht=True, - with_post_rht_amax=True, - with_2d_quantization=False, - stochastic_rounding=True, - with_random_sign_mask=True, - ) - - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, - with_rht=True, - with_post_rht_amax=True, - with_2d_quantization=False, - stochastic_rounding=False, - with_random_sign_mask=True, - ) + return nvfp4_quantizer_factory(role) def nvfp4_linear_fp8_dpa_factory( @@ -196,7 +184,7 @@ def nvfp4_linear_fp8_dpa_factory( device="cuda", ) - return _make_nvfp4_quantizer(role) + return nvfp4_quantizer_factory(role) def nvfp4_linear_mxfp8_dpa_factory( @@ -253,7 +241,7 @@ def nvfp4_linear_mxfp8_dpa_factory( """ is_dpa = role is not None and role.module_type == "dpa" if is_dpa: - return _make_mxfp8_quantizer() + return mxfp8_quantizer_factory(role) # DPA boundary slots (O output / dQKV grad-input): emitted by DPA with # empty `module_type` and a `name` like ".dpa_output". The fused @@ -265,9 +253,9 @@ def nvfp4_linear_mxfp8_dpa_factory( and ("dpa_output" in role.name or "dpa_grad_input" in role.name) ) if is_dpa_boundary: - return _make_mxfp8_quantizer() + return mxfp8_quantizer_factory(role) - return _make_nvfp4_quantizer(role) + return nvfp4_quantizer_factory(role) def high_precision_factory( @@ -297,10 +285,10 @@ def fwd_high_precision_bwd_mxfp8_factory( is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_mxfp8_quantizer() + return mxfp8_quantizer_factory(role) # fprop consumes rowwise high precision; dgrad / wgrad consume columnwise MXFP8. return HybridQuantizer( rowwise_quantizer=IdentityQuantizer(), - columnwise_quantizer=_make_mxfp8_quantizer(), + columnwise_quantizer=mxfp8_quantizer_factory(role), ) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py b/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py index 45febaf413d..bebb394269b 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py @@ -146,6 +146,8 @@ def nvfp4_quantizer_factory( if is_weight: return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, + rowwise=True, + columnwise=True, with_rht=False, with_post_rht_amax=False, with_2d_quantization=True, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 9c16e5b9a66..0531a0ec71e 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -37,6 +37,19 @@ } +def _canonical_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: + """Resolve PyTorch view shape syntax, including ``-1`` inference.""" + return torch.empty(tuple(input_shape), device="meta").view(*shape).shape + + +def _columnwise_shape_for(rowwise_shape: Iterable[int]) -> torch.Size: + """Physical columnwise FP8 shape for a logical rowwise shape.""" + shape = torch.Size(rowwise_shape) + if len(shape) == 0: + return shape + return torch.Size((shape[-1], *shape[:-1])) + + class Float8Quantizer(Quantizer): """Builder class for FP8 tensors with per-tensor delayed scaling @@ -573,24 +586,35 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.view.default: tensor = args[0] data = tensor._data - out_data = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - out_shape = out_data.size() + out_data = None + if data is not None: + out_data = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) + out_shape = out_data.size() + else: + out_shape = _canonical_view_shape(tensor.shape, args[1:]) + out_transpose = None if tensor._transpose_invalid else tensor._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if ( - out_transpose_shape[0] != out_shape[-1] - or out_transpose_shape[1:] != out_shape[:-1] - ): + view_shape_for_transpose = _columnwise_shape_for(out_shape) + if out_transpose.shape != view_shape_for_transpose: + if data is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data is only supported " + "when the requested shape preserves the columnwise layout" + ) out_transpose = None else: - view_shape_for_transpose = [out_shape[-1]] + list(out_shape[:-1]) out_transpose = out_transpose.view(*view_shape_for_transpose) + if data is None and out_transpose is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data requires a valid " + "columnwise buffer" + ) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, @@ -1068,16 +1092,28 @@ def forward( ctx.shape = tensor.shape if shape is None: return tensor.detach() - out_data = tensor._data.view(*shape) - out_shape = out_data.size() + out_data = None + if tensor._data is not None: + out_data = tensor._data.view(*shape) + out_shape = out_data.size() + else: + out_shape = _canonical_view_shape(tensor.shape, shape) out_transpose = None if tensor._transpose_invalid else tensor._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if out_transpose_shape[0] != out_shape[-1] or out_transpose_shape[1:] != out_shape[:-1]: + view_shape_for_transpose = _columnwise_shape_for(out_shape) + if out_transpose.shape != view_shape_for_transpose: + if tensor._data is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data is only supported " + "when the requested shape preserves the columnwise layout" + ) out_transpose = None else: - view_shape_for_transpose = [shape[-1]] + list(shape[:-1]) out_transpose = out_transpose.view(*view_shape_for_transpose) + if tensor._data is None and out_transpose is None: + raise NotImplementedError( + "Float8Tensor view with columnwise-only data requires a valid columnwise buffer" + ) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 8a030c6db8b..b5578a75391 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -192,12 +192,31 @@ def device(self): def view(self, shape: torch.Size): # pylint: disable=missing-function-docstring - out_data = self._data.view(shape) + out_data = self._data.view(shape) if self._data is not None else None + if out_data is not None: + out_shape = out_data.size() + else: + out_shape = torch.empty(tuple(self.size()), device="meta").view(shape).shape out_transpose = None if self._transpose_invalid else self._transpose if out_transpose is not None: - out_transpose_shape = out_transpose.size() - if out_transpose_shape[0] != shape[-1] or out_transpose_shape[1:] != shape[:-1]: + if len(out_shape) == 0: + view_shape_for_transpose = out_shape + else: + view_shape_for_transpose = torch.Size((out_shape[-1], *out_shape[:-1])) + if out_transpose.shape != view_shape_for_transpose: + if self._data is None: + raise NotImplementedError( + "Float8TensorStorage view with columnwise-only data is only " + "supported when the requested shape preserves the columnwise layout" + ) out_transpose = None + else: + out_transpose = out_transpose.view(*view_shape_for_transpose) + if self._data is None and out_transpose is None: + raise NotImplementedError( + "Float8TensorStorage view with columnwise-only data requires a valid " + "columnwise buffer" + ) return Float8TensorStorage( data=out_data, diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index 9e94b654fcb..e376591cdbb 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -126,6 +126,8 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """Dequantize using the first available sub-storage.""" if dtype is None: dtype = self._dtype + # TODO: choose the best native sub-storage for dequantization, preferring # pylint: disable=fixme + # identity/high-precision and FP8-ish formats over FP4 when both exist. if self._rowwise_storage is not None: return self._rowwise_storage.dequantize(dtype=dtype) if self._columnwise_storage is not None: From e3dbd1589f9fdbe37df8a1600b564b5a4be6fde5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 16 Jun 2026 12:34:36 +0000 Subject: [PATCH 29/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 1 - transformer_engine/pytorch/tensor/float8_tensor.py | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 7aad5670bf5..8b7d9f2667b 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -1154,7 +1154,6 @@ def hybrid_block_fp8_factory(role): ) - @pytest.mark.skipif( not (fp8_available and nvfp4_available), reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 0531a0ec71e..f3be2484a28 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -612,8 +612,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): out_transpose = out_transpose.view(*view_shape_for_transpose) if data is None and out_transpose is None: raise NotImplementedError( - "Float8Tensor view with columnwise-only data requires a valid " - "columnwise buffer" + "Float8Tensor view with columnwise-only data requires a valid columnwise buffer" ) return Float8Tensor( shape=out_shape, From 9d0e7c70644879f9efbc9d07ce197f36966660ad Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 17 Jun 2026 10:33:11 +0000 Subject: [PATCH 30/97] quantization_factory_base/zoo Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/fsdp2_utils.py | 4 +- tests/pytorch/distributed/run_hybrid_tp_sp.py | 2 +- .../pytorch/distributed/run_numerics_exact.py | 2 +- tests/pytorch/test_cpu_offloading.py | 4 +- tests/pytorch/test_cpu_offloading_v1.py | 4 +- tests/pytorch/test_custom_recipe.py | 41 +++++++++++- tests/pytorch/test_hybrid_quantization.py | 2 +- ...s_base.py => quantization_factory_base.py} | 20 +++++- ...xamples.py => quantization_factory_zoo.py} | 65 +++++++------------ 9 files changed, 91 insertions(+), 53 deletions(-) rename transformer_engine/pytorch/custom_recipes/{quantization_recipes_base.py => quantization_factory_base.py} (89%) rename transformer_engine/pytorch/custom_recipes/{quantization_factory_examples.py => quantization_factory_zoo.py} (86%) diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 257760da37b..39b7b0ad088 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -6,7 +6,7 @@ import transformer_engine.common.recipe from transformer_engine.pytorch import HybridQuantizer, IdentityQuantizer, QuantizedTensor -from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( current_scaling_quantizer_factory, float8_block_scaling_quantizer_factory, mxfp8_quantizer_factory, @@ -98,7 +98,7 @@ def get_hybrid_recipe_from_string(recipe): """Build a CustomRecipe wrapping a module-level (picklable) hybrid qfactory. Each hybrid qfactory composes one or two role-aware base factories from - ``quantization_recipes_base`` per direction; per-role behavior is delegated + ``quantization_factory_base`` per direction; per-role behavior is delegated to the base factory and the hybrid layer only decides the direction pairing. Supported values: diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index 4697b4d80b3..dad5e619278 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -19,7 +19,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.common import recipe as te_recipe -from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, ) from transformer_engine.pytorch import ( diff --git a/tests/pytorch/distributed/run_numerics_exact.py b/tests/pytorch/distributed/run_numerics_exact.py index 15ae2dae63b..6af992dac28 100644 --- a/tests/pytorch/distributed/run_numerics_exact.py +++ b/tests/pytorch/distributed/run_numerics_exact.py @@ -58,7 +58,7 @@ def get_nvfp4_quantizer_factory(): Mirrors the canonical "branch on what we care about, default fall-through" pattern from - ``transformer_engine.pytorch.custom_recipes.quantization_recipes_base``; + ``transformer_engine.pytorch.custom_recipes.quantization_factory_base``; every slot gets a real :class:`NVFP4QuantizerRef` (``CustomRecipeState`` rejects ``None`` returns). diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index e6170145e96..697ee06b81b 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -19,7 +19,7 @@ from transformer_engine.pytorch.fp8 import FP8GlobalStateManager import transformer_engine.pytorch as te from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, ) from utils import ModelConfig, recipe_id, skip_unsupported_backward_override @@ -54,7 +54,7 @@ def _hybrid_mxfp8_nvfp4_qfactory(role): """Hybrid CustomRecipe factory: MXFP8 rowwise + NVFP4 columnwise. Mirrors ``mxfp8_fwd_nvfp4_bwd_quantizer_factory`` from - ``custom_recipes/quantization_factory_examples.py``. grad_output uses plain NVFP4 + ``custom_recipes/quantization_factory_zoo.py``. grad_output uses plain NVFP4 (both directions) so wgrad's columnwise operand matches. """ is_linear = role is not None and role.module_type in ("linear", "grouped_linear") diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index d6a4e02a875..bf39cfa3e77 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -13,7 +13,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, ) from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends @@ -51,7 +51,7 @@ def _hybrid_mxfp8_nvfp4_qfactory(role): """Hybrid CustomRecipe factory: MXFP8 rowwise + NVFP4 columnwise. Mirrors the ``mxfp8_fwd_nvfp4_bwd_quantizer_factory`` headline recipe - from ``custom_recipes/quantization_factory_examples.py``. grad_output uses plain + from ``custom_recipes/quantization_factory_zoo.py``. grad_output uses plain NVFP4 (both directions) so wgrad's columnwise operand matches. """ is_linear = role is not None and role.module_type in ("linear", "grouped_linear") diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 3e6fdb816b1..f2a6bebae34 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -19,12 +19,13 @@ ) from transformer_engine.pytorch.quantization import QuantizerRole import transformer_engine.pytorch.ops as te_ops -from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( current_scaling_quantizer_factory, mxfp8_quantizer_factory, float8_block_scaling_quantizer_factory, nvfp4_quantizer_factory, delayed_scaling_quantizer_factory, + high_precision_factory, ) from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import ( nvfp4_ref_rht_2d_quantizer_factory, @@ -477,6 +478,40 @@ def test_factory_matches_nvfp4(): _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) +def test_factory_matches_high_precision(): + """high_precision_factory (all IdentityQuantizer) should produce bit-identical + results to running with no quantization (plain BF16, no autocast). + + Unlike the other ``test_factory_matches_*`` cases, the reference here is not a + built-in recipe but the unquantized path: IdentityQuantizer leaves tensors + untouched, so the GEMMs run in high precision exactly as they would with + autocast disabled. + + Note: the custom side still runs inside ``autocast(enabled=True, ...)``, which + asserts FP8 availability on entry, so this test skips on non-FP8 hardware even + though no tensor is actually quantized. + """ + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + model_ref, model_cus, inp_ref, inp_cus = _make_pair() + + # Reference: no autocast -> high-precision GEMMs, no quantization. + with autocast(enabled=False): + out_ref = model_ref(inp_ref) + out_ref.float().sum().backward() + out_ref = out_ref.clone() + grad_ref = inp_ref.grad.clone() + pgrads_ref = {n: p.grad.clone() for n, p in model_ref.named_parameters() if p.grad is not None} + + out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( + model_cus, inp_cus, recipe.CustomRecipe(qfactory=high_precision_factory) + ) + + _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) + + def test_custom_recipe_quantization_targets(): """Validate fine-grained per-module quantization targeting via QuantizerRole. @@ -1090,7 +1125,7 @@ def test_custom_recipe_dpa_fp8(): Float8Quantizer, Float8CurrentScalingQuantizer, ) - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) @@ -1217,7 +1252,7 @@ def test_custom_recipe_dpa_mxfp8(): from transformer_engine.pytorch.quantization import CustomRecipeState from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( nvfp4_linear_mxfp8_dpa_factory, ) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 8b7d9f2667b..5b83f6b0b71 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -11,7 +11,7 @@ import transformer_engine_torch as tex from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, ) from transformer_engine.pytorch import ( diff --git a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py similarity index 89% rename from transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py rename to transformer_engine/pytorch/custom_recipes/quantization_factory_base.py index bebb394269b..0a2b8f99334 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py @@ -14,7 +14,7 @@ from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, ) @@ -32,6 +32,24 @@ from ..constants import DType +def high_precision_factory( + role: Optional[QuantizerRole], # pylint: disable=unused-argument +) -> "IdentityQuantizer": + """Factory that runs all GEMMs in high precision (no quantization). + + Returns an :class:`IdentityQuantizer` for every slot, so no tensor is + quantized. This is the simplest base factory and a good starting point to + branch from: keep most roles in high precision and selectively override the + ones you want to quantize. + + Dispatch logic: + * every role -> ``IdentityQuantizer`` (no quantization) + """ + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + return IdentityQuantizer() + + def delayed_scaling_quantizer_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "DelayedScalingRequest": diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py similarity index 86% rename from transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py rename to transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index 2347ca8de94..c6c06131a4b 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -3,48 +3,46 @@ # See LICENSE for license information. """ -Quantizer factory examples. +Quantizer factory zoo. -Demonstrates how to use the ``CustomRecipe`` + ``qfactory`` interface to apply -*different* quantization recipes to different module/tensor types/instances within the same model. +A collection of composed/mixed-recipe factories built on top of the single-format +building blocks. They demonstrate how to use the ``CustomRecipe`` + ``qfactory`` +interface to apply *different* quantization recipes to different +module/tensor types/instances within the same model. + +.. warning:: + + Use these with caution. These are **not** official, supported recipes + provided by Transformer Engine -- they are illustrative examples meant to + inspire your own experiments, not drop-in production defaults. While most + of the factories here are grounded in some evidence or rationale (see the + per-factory docstrings), they have not been broadly validated for accuracy, + convergence, or performance across models and hardware. Treat them as + starting points: benchmark and verify on your own workload before relying on + any of them. Usage:: from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( mxfp8_fwd_nvfp4_bwd_quantizer_factory, - nvfp4_linear_mxfp8_grouped_linear_factory, - nvfp4_linear_fp8_dpa_factory, nvfp4_linear_mxfp8_dpa_factory, - high_precision_factory, - fwd_high_precision_bwd_mxfp8_factory, ) - # Hybrid per-direction recipe: MXFP8 for fprop, NVFP4 for dgrad/wgrad + # Linear-only recipe (no attention quantization): the qfactory is the only knob. recipe = CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_quantizer_factory) with autocast(recipe=recipe): output = model(input) - # Mixed module types: NVFP4 for Linear, MXFP8 for GroupedLinear - recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_grouped_linear_factory) - with autocast(recipe=recipe): - output = model(input) - - # NVFP4 for Linear, FP8 current-scaling + delayed-scaling for DPA - recipe = CustomRecipe(qfactory=nvfp4_linear_fp8_dpa_factory, fp8_dpa=True) - with autocast(recipe=recipe): - output = model(input) - - # NVFP4 for Linear, MXFP8 for DPA + # Recipe that also quantizes DotProductAttention: set ``fp8_dpa=True`` so the + # attention GEMMs request quantizers from the factory (DPA roles) too. recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True) with autocast(recipe=recipe): output = model(input) - # High precision forward, MXFP8 backward - recipe = CustomRecipe(qfactory=fwd_high_precision_bwd_mxfp8_factory) - with autocast(recipe=recipe): - output = model(input) + # The other factories in this module follow the same two patterns; see their + # docstrings for the exact per-role dispatch. """ from __future__ import annotations @@ -53,7 +51,7 @@ from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType -from .quantization_recipes_base import mxfp8_quantizer_factory, nvfp4_quantizer_factory +from .quantization_factory_base import mxfp8_quantizer_factory, nvfp4_quantizer_factory def mxfp8_fwd_nvfp4_bwd_quantizer_factory( @@ -145,7 +143,7 @@ def nvfp4_linear_fp8_dpa_factory( from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) @@ -228,7 +226,7 @@ def nvfp4_linear_mxfp8_dpa_factory( from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( nvfp4_linear_mxfp8_dpa_factory, ) @@ -258,19 +256,6 @@ def nvfp4_linear_mxfp8_dpa_factory( return nvfp4_quantizer_factory(role) -def high_precision_factory( - role: Optional[QuantizerRole], # pylint: disable=unused-argument -): - """Quantizer factory: run all GEMMs in high precision. - - Dispatch logic: - * every role -> ``IdentityQuantizer`` (no quantization) - """ - from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer - - return IdentityQuantizer() - - def fwd_high_precision_bwd_mxfp8_factory( role: Optional[QuantizerRole], ): From cb3ab497a41975d580cf6631f08cd6eb16c92a8a Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 22 Jun 2026 07:23:43 +0000 Subject: [PATCH 31/97] Enable columnwise_source and hybrid recipes Signed-off-by: Evgeny --- tests/pytorch/test_custom_recipe.py | 6 +- tests/pytorch/test_hybrid_quantization.py | 267 +++++++++++++++++ tests/pytorch/test_identity_quantizer.py | 107 +++++++ .../quantization_factory_zoo.py | 278 +++++++++++++++--- .../pytorch/tensor/hybrid_tensor.py | 62 +++- 5 files changed, 673 insertions(+), 47 deletions(-) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index f2a6bebae34..52140b4d9ff 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -1100,7 +1100,7 @@ def test_role_change_does_not_invalidate_when_role_unchanged(): def test_custom_recipe_dpa_fp8(): """DotProductAttention forward+backward with CustomRecipe and role-based mixed quantizers. - Uses the nvfp4_linear_fp8_dpa_factory which dispatches: + Uses the nvfp4_linear_mixed_fp8_dpa_factory which dispatches: * DPA S/dP slots -> DelayedScalingRequest (stateful) * DPA QKV/O/dO/dQKV slots -> Float8CurrentScalingQuantizer * Linear slots -> NVFP4Quantizer @@ -1126,7 +1126,7 @@ def test_custom_recipe_dpa_fp8(): Float8CurrentScalingQuantizer, ) from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - nvfp4_linear_fp8_dpa_factory, + nvfp4_linear_mixed_fp8_dpa_factory, ) torch.manual_seed(42) @@ -1145,7 +1145,7 @@ def test_custom_recipe_dpa_fp8(): out_proj = Linear(H, H, params_dtype=torch.bfloat16, bias=False, name="proj").cuda() custom_recipe = recipe.CustomRecipe( - qfactory=nvfp4_linear_fp8_dpa_factory, + qfactory=nvfp4_linear_mixed_fp8_dpa_factory, fp8_dpa=True, ) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 5b83f6b0b71..2003d3e5f8a 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -29,6 +29,7 @@ NVFP4Quantizer, HybridQuantizer, HybridQuantizedTensor, + IdentityQuantizer, HybridQuantizedTensorStorage, Float8Tensor, Float8TensorStorage, @@ -53,11 +54,21 @@ reason=f"FP8: {reason_for_no_fp8}", ) +requires_nvfp4 = pytest.mark.skipif( + not nvfp4_available, + reason=f"NVFP4: {reason_for_no_nvfp4}", +) + requires_fp8_and_nvfp4 = pytest.mark.skipif( not (fp8_available and nvfp4_available), reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", ) +requires_mxfp8_and_nvfp4 = pytest.mark.skipif( + not (mxfp8_available and nvfp4_available), + reason=f"MXFP8: {reason_for_no_mxfp8}; NVFP4: {reason_for_no_nvfp4}", +) + def _make_fp8_quantizer(*, rowwise=True, columnwise=True): return Float8CurrentScalingQuantizer( @@ -92,6 +103,164 @@ def _make_hybrid_quantizer_fp4_row_fp8_col(): ) +@requires_mxfp8_and_nvfp4 +class TestComposerStyleFactory: + """Composer 2-style row-scaled NVFP4 forward + MXFP8 backward dispatch.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory, + ) + + return nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory(role) + + @pytest.mark.parametrize( + "tensor_type,row_scaled", + [("input", True), ("weight", False)], + ) + def test_grouped_linear_forward_roles_use_nvfp4_rowwise_mxfp8_columnwise( + self, tensor_type, row_scaled + ): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory( + QuantizerRole(module_type="grouped_linear", tensor_type=tensor_type) + ) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "original" + assert isinstance(quantizer.rowwise_quantizer, NVFP4Quantizer) + assert isinstance(quantizer.columnwise_quantizer, MXFP8Quantizer) + assert quantizer.rowwise_quantizer.row_scaled_nvfp4 is row_scaled + assert quantizer.rowwise_quantizer.with_rht is False + assert quantizer.rowwise_quantizer.with_post_rht_amax is False + assert quantizer.rowwise_quantizer.with_2d_quantization is False + assert quantizer.rowwise_quantizer.stochastic_rounding is False + assert quantizer.rowwise_quantizer.rowwise_usage is True + assert quantizer.rowwise_quantizer.columnwise_usage is False + assert quantizer.columnwise_quantizer.rowwise_usage is False + assert quantizer.columnwise_quantizer.columnwise_usage is True + + @pytest.mark.parametrize("tensor_type", ["input", "output", "weight"]) + def test_regular_linear_forward_roles_fall_back_to_mxfp8(self, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory( + QuantizerRole(module_type="linear", tensor_type=tensor_type) + ) + + assert isinstance(quantizer, MXFP8Quantizer) + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["grad_output", "grad_input"]) + def test_backward_roles_use_mxfp8(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole( + module_type=module_type, tensor_type=tensor_type + )) + + assert isinstance(quantizer, MXFP8Quantizer) + + def test_non_linear_roles_fall_back_to_mxfp8(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type="dpa", tensor_type="qkv")) + + assert isinstance(quantizer, MXFP8Quantizer) + + +@requires_nvfp4 +class TestNVFP4WeightDoubleQuantFactory: + """Base NVFP4 recipe with W.T sourced from dequantized 1D W.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_1d_double_quantized_weight_quantizer_factory, + ) + + return nvfp4_1d_double_quantized_weight_quantizer_factory(role) + + @staticmethod + def _assert_plain_1d_nvfp4(quantizer): + assert isinstance(quantizer, NVFP4Quantizer) + assert quantizer.row_scaled_nvfp4 is False + assert quantizer.with_rht is False + assert quantizer.with_post_rht_amax is False + assert quantizer.with_2d_quantization is False + assert quantizer.stochastic_rounding is False + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + def test_weight_roles_use_rowwise_dequantized_source(self, module_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory( + QuantizerRole(module_type=module_type, tensor_type="weight") + ) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + self._assert_plain_1d_nvfp4(quantizer.rowwise_quantizer) + self._assert_plain_1d_nvfp4(quantizer.columnwise_quantizer) + assert quantizer.rowwise_quantizer.rowwise_usage is True + assert quantizer.rowwise_quantizer.columnwise_usage is False + assert quantizer.columnwise_quantizer.rowwise_usage is False + assert quantizer.columnwise_quantizer.columnwise_usage is True + + @staticmethod + def _assert_matches_base_nvfp4_recipe(quantizer, expected): + assert isinstance(quantizer, NVFP4Quantizer) + assert isinstance(expected, NVFP4Quantizer) + assert quantizer.dtype == expected.dtype + assert quantizer.rowwise_usage == expected.rowwise_usage + assert quantizer.columnwise_usage == expected.columnwise_usage + assert quantizer.with_rht == expected.with_rht + assert quantizer.with_post_rht_amax == expected.with_post_rht_amax + assert quantizer.with_2d_quantization == expected.with_2d_quantization + assert quantizer.stochastic_rounding == expected.stochastic_rounding + assert quantizer.row_scaled_nvfp4 == expected.row_scaled_nvfp4 + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["input", "output", "grad_output", "grad_input"]) + def test_non_weight_linear_roles_match_base_nvfp4_recipe(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + nvfp4_quantizer_factory, + ) + + role = QuantizerRole(module_type=module_type, tensor_type=tensor_type) + quantizer = self._factory(role) + expected = nvfp4_quantizer_factory(role) + + self._assert_matches_base_nvfp4_recipe(quantizer, expected) + + def test_non_linear_roles_match_base_nvfp4_recipe(self): + from transformer_engine.pytorch.quantization import QuantizerRole + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + nvfp4_quantizer_factory, + ) + + role = QuantizerRole(module_type="dpa", tensor_type="qkv") + quantizer = self._factory(role) + expected = nvfp4_quantizer_factory(role) + + self._assert_matches_base_nvfp4_recipe(quantizer, expected) + + def test_weight_quantizer_produces_both_nvfp4_storages(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + torch.manual_seed(2026) + src = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + quantizer = self._factory(QuantizerRole(module_type="linear", tensor_type="weight")) + + out = quantizer.quantize(src) + + assert isinstance(out.rowwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) + assert isinstance(out.columnwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) + + @requires_fp8_and_nvfp4 class TestHybridQuantizerConstruction: """Test construction and basic properties of hybrid quantizer.""" @@ -180,6 +349,104 @@ def test_supports_only_rowwise_all_gather_nvfp4_both(self): assert hq.supports_only_rowwise_all_gather() is True +@requires_fp8 +class TestHybridColumnwiseSource: + """Test columnwise source provenance in HybridQuantizer.""" + + @pytest.fixture + def input_tensor(self): + torch.manual_seed(90210) + return torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + + @staticmethod + def _make_quantizer(columnwise_source="original"): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source=columnwise_source, + ) + + def test_default_columnwise_source_original(self, input_tensor): + hq = self._make_quantizer() + out = hq.quantize(input_tensor) + + assert hq.columnwise_source == "original" + torch.testing.assert_close( + out.columnwise_sub_storage.dequantize(), input_tensor, rtol=0.0, atol=0.0 + ) + + def test_invalid_columnwise_source_raises(self): + with pytest.raises(ValueError, match="columnwise_source"): + HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="dequantized", + ) + + def test_copy_preserves_columnwise_source(self): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.set_usage(rowwise=False, columnwise=True) + + copied = hq.copy() + + assert copied.columnwise_source == "rowwise_dequantized" + assert copied.rowwise_usage is False + assert copied.columnwise_usage is True + + def test_rowwise_dequantized_identity_columnwise_matches_rowwise(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + out = hq.quantize(input_tensor) + + rowwise_dq = out.rowwise_sub_storage.dequantize() + columnwise_dq = out.columnwise_sub_storage.dequantize() + torch.testing.assert_close(columnwise_dq, rowwise_dq, rtol=0.0, atol=0.0) + + def test_columnwise_only_uses_transient_rowwise_source(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.set_usage(rowwise=False, columnwise=True) + expected = _make_fp8_quantizer().quantize(input_tensor).dequantize() + + out = hq.quantize(input_tensor) + + assert out.rowwise_sub_storage is None + assert out.columnwise_sub_storage is not None + torch.testing.assert_close( + out.columnwise_sub_storage.dequantize(), expected, rtol=0.0, atol=0.0 + ) + + def test_update_quantized_columnwise_only_uses_transient_rowwise_source(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + dst = hq.quantize(input_tensor) + rowwise_before = dst.rowwise_sub_storage.dequantize().clone() + new_src = torch.randn_like(input_tensor) * 8 + expected = _make_fp8_quantizer().quantize(new_src).dequantize() + + hq.set_usage(rowwise=False, columnwise=True) + hq.update_quantized(new_src, dst) + + torch.testing.assert_close( + dst.rowwise_sub_storage.dequantize(), rowwise_before, rtol=0.0, atol=0.0 + ) + torch.testing.assert_close( + dst.columnwise_sub_storage.dequantize(), expected, rtol=0.0, atol=0.0 + ) + + def test_update_quantized_uses_updated_rowwise_storage(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + dst = hq.quantize(input_tensor) + new_src = torch.randn_like(input_tensor) * 8 + + hq.set_usage(rowwise=True, columnwise=True) + hq.update_quantized(new_src, dst) + + torch.testing.assert_close( + dst.columnwise_sub_storage.dequantize(), + dst.rowwise_sub_storage.dequantize(), + rtol=0.0, + atol=0.0, + ) + + @requires_fp8_and_nvfp4 class TestHybridQuantize: """Test quantization via HybridQuantizer.""" diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index b21041cbc73..ab98e8e1c4f 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -153,6 +153,18 @@ def fwd_fp8_bwd_hp_factory(role): ) +def fwd_fp8_bwd_rowwise_dequantized_hp_factory(role): + """FP8 forward, high-precision backward from dequantized fprop values.""" + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return IdentityQuantizer() + return HybridQuantizer( + rowwise_quantizer=_fp8_cs(tex.DType.kFloat8E4M3), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + + def hybrid_all_identity_factory(role): """All directions high precision, expressed through the hybrid container. @@ -1188,6 +1200,101 @@ def test_identity_reproduces_backward_override_high_precision_bitwise(self): for g_id, g_bo in zip(wg_id, wg_bo): torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + def test_identity_reproduces_backward_override_dequantized_bitwise(self): + """Per-direction Identity must reproduce ``backward_override=dequantized``. + + Both runs quantize forward to the same FP8 values. The reference saves + rowwise fprop payloads and dequantizes them for high-precision backward; + the hybrid path stores Identity columnwise data sourced from the rowwise + dequantized value and uses Identity grad tensors. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_bo, dx_bo, wg_bo = _fwd_bwd( + ref, + x, + recipe=CustomRecipe(qfactory=fp8_fwd_factory, backward_override="dequantized"), + ) + y_id, dx_id, wg_id = _fwd_bwd( + test, x, recipe=CustomRecipe(qfactory=fwd_fp8_bwd_rowwise_dequantized_hp_factory) + ) + + torch.testing.assert_close(y_id, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_id) == len(wg_bo) + for g_id, g_bo in zip(wg_id, wg_bo): + torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") + def test_zoo_mxfp8_dequantized_bwd_matches_backward_override_bitwise(self): + """The zoo MXFP8-dequantized recipe should reproduce + ``backward_override=dequantized`` bitwise for Linear. + """ + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_dequantized_bwd_quantizer_factory, + ) + + def mxfp8_all_factory(role): # pylint: disable=unused-argument + return _mxfp8(tex.DType.kFloat8E4M3) + + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_bo, dx_bo, wg_bo = _fwd_bwd( + ref, + x, + recipe=CustomRecipe(qfactory=mxfp8_all_factory, backward_override="dequantized"), + ) + y_zoo, dx_zoo, wg_zoo = _fwd_bwd( + test, + x, + recipe=CustomRecipe(qfactory=mxfp8_fwd_dequantized_bwd_quantizer_factory), + ) + + torch.testing.assert_close(y_zoo, y_bo, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_zoo, dx_bo, rtol=0.0, atol=0.0) + assert len(wg_zoo) == len(wg_bo) + for g_zoo, g_bo in zip(wg_zoo, wg_bo): + torch.testing.assert_close(g_zoo, g_bo, rtol=0.0, atol=0.0) + + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") + def test_zoo_nvfp4_row_scaled_dequantized_bwd_matches_recipe_bitwise(self): + """The zoo row-scaled NVFP4-dequantized recipe should reproduce + the built-in row-scaled NVFP4 recipe with dequantized backward bitwise. + """ + from transformer_engine.common import recipe as te_recipe + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory, + ) + + ref_recipe = te_recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="dequantized", + ) + ref_recipe.fp4_quant_fwd_inp = te_recipe.QParams() + ref_recipe.fp4_quant_fwd_weight = te_recipe.QParams() + ref_recipe.fp4_quant_bwd_grad = te_recipe.QParams() + + ref, test = _make_linears(self.IN_F, self.OUT_F) + x = self._input() + + y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=ref_recipe) + y_zoo, dx_zoo, wg_zoo = _fwd_bwd( + test, + x, + recipe=CustomRecipe(qfactory=nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory), + ) + + torch.testing.assert_close(y_zoo, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_zoo, dx_ref, rtol=0.0, atol=0.0) + assert len(wg_zoo) == len(wg_ref) + for g_zoo, g_ref in zip(wg_zoo, wg_ref): + torch.testing.assert_close(g_zoo, g_ref, rtol=0.0, atol=0.0) + def test_identity_matches_bf16_multistep_training_bitwise(self): """Multi-step SGD: an all-Identity recipe must track a plain BF16 te.Linear bitwise across optimizer steps (no drift from workspace caching diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index c6c06131a4b..af90f8813fc 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -6,9 +6,18 @@ Quantizer factory zoo. A collection of composed/mixed-recipe factories built on top of the single-format -building blocks. They demonstrate how to use the ``CustomRecipe`` + ``qfactory`` +building blocks. They demonstrate how to use the ``CustomRecipe`` + ``qfactory`` interface to apply *different* quantization recipes to different -module/tensor types/instances within the same model. +module/tensor types/instances within the same model, and how to use +``HybridQuantizer`` when rowwise and columnwise tensor directions should use +different formats or sources. + +Organization: + * Linear / grouped-linear recipes: from conservative to more aggressive quantization. + * RL-oriented recipes: dequantized-backward and MoE-inspired recipes that + favor more precision or explicit source control in backward GEMMs. + * Linear + attention recipes: factories that also cover ``DotProductAttention`` + roles and require ``CustomRecipe(..., fp8_dpa=True)``. .. warning:: @@ -54,6 +63,34 @@ from .quantization_factory_base import mxfp8_quantizer_factory, nvfp4_quantizer_factory +# ----------------------------------------------------------------------------- +# Linear / GroupedLinear Recipes +# ----------------------------------------------------------------------------- + + +def high_precision_fwd_mxfp8_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: high-precision forward, MXFP8 backward. + + Dispatch logic: + * ``grad_output`` -> MXFP8 (E4M3, block-32) + * everything else -> ``Hybrid(rowwise=IdentityQuantizer, columnwise=MXFP8)`` + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "grad_output": + return mxfp8_quantizer_factory(role) + + # fprop consumes rowwise high precision; dgrad / wgrad consume columnwise MXFP8. + return HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=mxfp8_quantizer_factory(role), + ) + + def mxfp8_fwd_nvfp4_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -72,39 +109,221 @@ def mxfp8_fwd_nvfp4_bwd_quantizer_factory( from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): + if is_linear and role.tensor_type in ("input", "weight"): return HybridQuantizer( rowwise_quantizer=mxfp8_quantizer_factory(role), columnwise_quantizer=nvfp4_quantizer_factory(role), ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): + if is_linear and role.tensor_type == "grad_output": return nvfp4_quantizer_factory(role) return mxfp8_quantizer_factory(role) -def nvfp4_linear_mxfp8_grouped_linear_factory( +def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): + """NVFP4 quantizer without RHT, stochastic rounding, or 2D scaling.""" + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + return NVFP4Quantizer( + fp4_dtype=DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=row_scaled_nvfp4, + ) + + +def nvfp4_1d_double_quantized_weight_quantizer_factory( role: Optional[QuantizerRole], ): - """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``GroupedLinear``. + """Quantizer factory: NVFP4 recipe with 1D weight double quantization. Dispatch logic: - * ``role.module_type == "grouped_linear"`` -> MXFP8 (E4M3, block-32) - * everything else (``"linear"`` or unknown) -> NVFP4 (E2M1) + * ``linear`` / ``grouped_linear`` ``weight`` -> + ``Hybrid(rowwise=plain 1D NVFP4, columnwise=plain 1D NVFP4, + columnwise_source="rowwise_dequantized")`` + * everything else -> :func:`nvfp4_quantizer_factory` + + ``W.T`` is quantized from the dequantized forward/rowwise ``W`` value + instead of directly from the original high-precision weight. In + ``HybridQuantizer`` terms, that source choice is expressed with + ``columnwise_source="rowwise_dequantized"``. + + All non-weight roles keep the standard NVFP4 factory behavior, including RHT + for inputs and stochastic rounding for gradients. The weight override uses + plain 1D NVFP4 in both directions: no RHT, stochastic rounding, row-scaled + activations, or 2D weight scaling. + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_source="rowwise_dequantized", + ) + return nvfp4_quantizer_factory(role) + + +# ----------------------------------------------------------------------------- +# RL-Oriented Recipes +# ----------------------------------------------------------------------------- + + +def mxfp8_fwd_dequantized_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: MXFP8 forward, high-precision dequantized backward. - NVFP4 settings follow the built-in ``NVFP4BlockScaling`` defaults: - * Weights: 2D quantization (16x16), no RHT, no stochastic rounding - * Inputs: 1D quantization, RHT enabled, no stochastic rounding - * Grads: 1D quantization, RHT enabled, stochastic rounding enabled + This expresses the linear/grouped-linear equivalent of + ``backward_override="dequantized"`` through per-direction quantizers: + + * ``input`` / ``weight`` -> + ``Hybrid(rowwise=MXFP8, columnwise=Identity, columnwise_source="rowwise_dequantized")`` + * ``grad_output`` -> ``IdentityQuantizer`` + * everything else -> MXFP8 + + Backward GEMMs use high-precision operands dequantized from the MXFP8 + forward payload, avoiding gradient quantization. + + This recipe targets RL-style training use cases and is motivated by + NVIDIA/TransformerEngine#2644, where ``backward_override="dequantized"`` + was introduced: + https://github.com/NVIDIA/TransformerEngine/pull/2644 """ - is_grouped_linear = role is not None and role.module_type == "grouped_linear" + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("input", "weight"): + return HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return IdentityQuantizer() + return mxfp8_quantizer_factory(role) + - if is_grouped_linear: +def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: row-scaled NVFP4 forward, dequantized backward. + + This expresses the linear/grouped-linear equivalent of + ``NVFP4BlockScaling(row_scaled_activation=True, + backward_override="dequantized")`` through per-direction quantizers: + + * ``input`` -> + ``Hybrid(rowwise=row-scaled NVFP4, columnwise=Identity, + columnwise_source="rowwise_dequantized")`` + * ``weight`` -> + ``Hybrid(rowwise=plain NVFP4, columnwise=Identity, + columnwise_source="rowwise_dequantized")`` + * ``grad_output`` -> ``IdentityQuantizer`` + * everything else -> plain NVFP4 + + Row-scaled NVFP4 is fprop-only, so the forward quantizers avoid RHT, + stochastic rounding, and 2D scaling. + + This recipe targets RL-style training use cases and builds on + NVIDIA/TransformerEngine#2931, which introduced row-scaled NVFP4: + https://github.com/NVIDIA/TransformerEngine/pull/2931 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return IdentityQuantizer() + return _plain_nvfp4_quantizer() + + +def nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: row-scaled NVFP4 forward, MXFP8 backward. + + This RL-related recipe is inspired by the Composer 2 MoE grouped-GEMM + recipe described in arXiv:2603.24477. + + Derived from the report: Composer 2 describes row-scaled NVFP4 for the MoE + forward pass and standard MXFP8 for the MoE backward pass. This factory maps + that format split onto ``GroupedLinear`` roles. + + Assumed here: regular non-MoE ``Linear`` layers use the MXFP8 fallback. The + public report does not specify the precision used for non-MoE linears. + + Dispatch logic: + + * ``GroupedLinear`` ``input`` -> + ``Hybrid(rowwise=row-scaled NVFP4, columnwise=MXFP8)`` + * ``GroupedLinear`` ``weight`` -> + ``Hybrid(rowwise=plain NVFP4, columnwise=MXFP8)`` + * regular ``Linear`` -> MXFP8 + * ``grad_output`` -> MXFP8 + * everything else -> MXFP8 + + Row-scaled NVFP4 is fprop-only, so the forward NVFP4 quantizers avoid + RHT, stochastic rounding, and 2D scaling. By default, ``HybridQuantizer`` + leaves ``columnwise_source="original"``, so MXFP8 backward operands are + quantized from the original high-precision tensor. The paper does not + specify how Composer materializes MXFP8 backward operands from saved + training tensors. + + Users who want MXFP8 backward operands to include the forward NVFP4 + quantization source can construct the same ``HybridQuantizer`` with + ``columnwise_source="rowwise_dequantized"``. That makes columnwise MXFP8 + quantization consume the dequantized rowwise NVFP4 value instead of the + original high-precision tensor, which can be useful when the backward path + should reflect the forward quantization error (might be helpful for convergence). + + Composer 2 Technical Report: + https://arxiv.org/abs/2603.24477 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + + is_grouped_linear = role is not None and role.module_type == "grouped_linear" + is_linear = role is not None and role.module_type == "linear" + if is_grouped_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), + columnwise_quantizer=mxfp8_quantizer_factory(role), + ) + if is_grouped_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=mxfp8_quantizer_factory(role), + ) + if is_grouped_linear and role.tensor_type == "grad_output": return mxfp8_quantizer_factory(role) + if is_linear: + return mxfp8_quantizer_factory(role) + return mxfp8_quantizer_factory(role) - return nvfp4_quantizer_factory(role) + +# ----------------------------------------------------------------------------- +# Linear + Attention Recipes +# ----------------------------------------------------------------------------- -def nvfp4_linear_fp8_dpa_factory( +def nvfp4_linear_mixed_fp8_dpa_factory( role: Optional[QuantizerRole], ): """Quantizer factory: NVFP4 for ``Linear``, mixed FP8 for ``DotProductAttention``. @@ -144,11 +363,11 @@ def nvfp4_linear_fp8_dpa_factory( from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - nvfp4_linear_fp8_dpa_factory, + nvfp4_linear_mixed_fp8_dpa_factory, ) recipe = CustomRecipe( - qfactory=nvfp4_linear_fp8_dpa_factory, + qfactory=nvfp4_linear_mixed_fp8_dpa_factory, fp8_dpa=True, ) with autocast(recipe=recipe): @@ -254,26 +473,3 @@ def nvfp4_linear_mxfp8_dpa_factory( return mxfp8_quantizer_factory(role) return nvfp4_quantizer_factory(role) - - -def fwd_high_precision_bwd_mxfp8_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: high-precision forward, MXFP8 backward. - - Dispatch logic: - * ``grad_output`` / ``grad_input`` -> MXFP8 (E4M3, block-32) - * everything else -> ``Hybrid(rowwise=IdentityQuantizer, columnwise=MXFP8)`` - """ - from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer - from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer - - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return mxfp8_quantizer_factory(role) - - # fprop consumes rowwise high precision; dgrad / wgrad consume columnwise MXFP8. - return HybridQuantizer( - rowwise_quantizer=IdentityQuantizer(), - columnwise_quantizer=mxfp8_quantizer_factory(role), - ) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 61fa694e4c4..56e17829553 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -5,7 +5,7 @@ """Tensor class with hybrid quantized data (different formats for rowwise vs columnwise)""" from __future__ import annotations -from typing import Any, Dict, Iterable, Optional, Tuple +from typing import Any, Dict, Iterable, Literal, Optional, Tuple import torch @@ -28,6 +28,11 @@ class HybridQuantizer(Quantizer): Quantizer for the rowwise direction (e.g. MXFP8Quantizer). columnwise_quantizer : Quantizer Quantizer for the columnwise direction (e.g. NVFP4Quantizer). + columnwise_source : {"original", "rowwise_dequantized"}, default = "original" + Source tensor for columnwise quantization. ``"original"`` quantizes + columnwise directly from the input tensor. ``"rowwise_dequantized"`` + quantizes rowwise first, dequantizes the rowwise result, then uses that + value as the columnwise source. Notes ----- @@ -42,16 +47,31 @@ class HybridQuantizer(Quantizer): ownership tracking, both of which are more intrusive, so only the direct rowwise/columnwise aliasing case is enforced. + Example + ------- + MXFP8 forward plus high-precision backward from the rowwise-dequantized + forward value can be expressed as:: + + HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer, + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + """ + _COLUMNWISE_SOURCES = ("original", "rowwise_dequantized") + rowwise_quantizer: Quantizer columnwise_quantizer: Quantizer + columnwise_source: Literal["original", "rowwise_dequantized"] def __init__( self, *, rowwise_quantizer: Quantizer, columnwise_quantizer: Quantizer, + columnwise_source: Literal["original", "rowwise_dequantized"] = "original", ) -> None: super().__init__(rowwise=True, columnwise=True) from transformer_engine.pytorch.quantization import QuantizerRequest # local import @@ -82,18 +102,35 @@ def __init__( " instances. If both directions need shared state, construct two" " quantizer objects that reference the same shared state." ) + if columnwise_source not in self._COLUMNWISE_SOURCES: + raise ValueError( + "HybridQuantizer columnwise_source must be one of " + f"{self._COLUMNWISE_SOURCES}, got {columnwise_source!r}." + ) self.rowwise_quantizer = rowwise_quantizer self.columnwise_quantizer = columnwise_quantizer + self.columnwise_source = columnwise_source # Pin each sub-quantizer to its designated direction self.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) self.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f"rowwise_usage={self.rowwise_usage}, " + f"columnwise_usage={self.columnwise_usage}, " + f"columnwise_source={self.columnwise_source!r}, " + f"internal={self.internal}, " + ")" + ) + def copy(self) -> "HybridQuantizer": """Create a shallow copy, preserving parent and sub-quantizer state.""" quantizer = HybridQuantizer( rowwise_quantizer=self.rowwise_quantizer.copy(), columnwise_quantizer=self.columnwise_quantizer.copy(), + columnwise_source=self.columnwise_source, ) quantizer.set_usage( rowwise=self.rowwise_usage, @@ -133,13 +170,25 @@ def amax_reduction_group(self, value) -> None: if hasattr(sub, "amax_reduction_group"): sub.amax_reduction_group = value + def _columnwise_src_from_rowwise( + self, + tensor: torch.Tensor, + rowwise_result: Optional[Any], + ) -> torch.Tensor: + if rowwise_result is None: + rowwise_result = self.rowwise_quantizer.quantize(tensor) + return rowwise_result.dequantize() + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: # Gate each sub-quantizer call on the parent usage flag. Sub-quantizers # are pinned to one direction in ``__init__``; the parent flag decides # whether to invoke them. rowwise_result = self.rowwise_quantizer.quantize(tensor) if self.rowwise_usage else None + columnwise_src = tensor + if self.columnwise_usage and self.columnwise_source == "rowwise_dequantized": + columnwise_src = self._columnwise_src_from_rowwise(tensor, rowwise_result) columnwise_result = ( - self.columnwise_quantizer.quantize(tensor) if self.columnwise_usage else None + self.columnwise_quantizer.quantize(columnwise_src) if self.columnwise_usage else None ) if self.internal: @@ -213,11 +262,18 @@ def update_quantized( "HybridQuantizer can only update HybridQuantizedTensorStorage, got" f" {type(dst).__name__}" ) + rowwise_result_for_columnwise = None if self.rowwise_usage and dst._rowwise_storage is not None: self.rowwise_quantizer.update_quantized(src, dst._rowwise_storage, noop_flag=noop_flag) + rowwise_result_for_columnwise = dst._rowwise_storage if self.columnwise_usage and dst._columnwise_storage is not None: + columnwise_src = src + if self.columnwise_source == "rowwise_dequantized": + columnwise_src = self._columnwise_src_from_rowwise( + src, rowwise_result_for_columnwise + ) self.columnwise_quantizer.update_quantized( - src, dst._columnwise_storage, noop_flag=noop_flag + columnwise_src, dst._columnwise_storage, noop_flag=noop_flag ) return dst From 10aecf02da3cbea93b5497e7ac73d7387839dcb5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 22 Jun 2026 07:46:20 +0000 Subject: [PATCH 32/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 2003d3e5f8a..c7fff873c8c 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -146,9 +146,7 @@ def test_grouped_linear_forward_roles_use_nvfp4_rowwise_mxfp8_columnwise( def test_regular_linear_forward_roles_fall_back_to_mxfp8(self, tensor_type): from transformer_engine.pytorch.quantization import QuantizerRole - quantizer = self._factory( - QuantizerRole(module_type="linear", tensor_type=tensor_type) - ) + quantizer = self._factory(QuantizerRole(module_type="linear", tensor_type=tensor_type)) assert isinstance(quantizer, MXFP8Quantizer) @@ -157,9 +155,7 @@ def test_regular_linear_forward_roles_fall_back_to_mxfp8(self, tensor_type): def test_backward_roles_use_mxfp8(self, module_type, tensor_type): from transformer_engine.pytorch.quantization import QuantizerRole - quantizer = self._factory(QuantizerRole( - module_type=module_type, tensor_type=tensor_type - )) + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type=tensor_type)) assert isinstance(quantizer, MXFP8Quantizer) @@ -196,9 +192,7 @@ def _assert_plain_1d_nvfp4(quantizer): def test_weight_roles_use_rowwise_dequantized_source(self, module_type): from transformer_engine.pytorch.quantization import QuantizerRole - quantizer = self._factory( - QuantizerRole(module_type=module_type, tensor_type="weight") - ) + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type="weight")) assert isinstance(quantizer, HybridQuantizer) assert quantizer.columnwise_source == "rowwise_dequantized" From cd381ef88dc176653809be0cb1e8b1e82a8d63c5 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 23 Jun 2026 05:39:09 -0700 Subject: [PATCH 33/97] Respect quantizer veto for save_original_inp Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 184 +++++++++++++++++- .../pytorch/module/grouped_linear.py | 18 ++ transformer_engine/pytorch/module/linear.py | 13 ++ .../pytorch/quantized_tensor.py | 10 + .../pytorch/tensor/hybrid_tensor.py | 5 + 5 files changed, 229 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index c7fff873c8c..62355dd6397 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -103,6 +103,29 @@ def _make_hybrid_quantizer_fp4_row_fp8_col(): ) +class _CountingIdentityQuantizer(IdentityQuantizer): + """Identity quantizer that counts quantize calls for regression tests.""" + + def __init__(self, counter, *, dtype=None, rowwise=True, columnwise=True): + super().__init__(dtype=dtype, rowwise=rowwise, columnwise=columnwise) + self.counter = counter + + def copy(self): + quantizer = _CountingIdentityQuantizer( + self.counter, + dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def quantize_impl(self, tensor): + self.counter["calls"] += 1 + return super().quantize_impl(tensor) + + @requires_mxfp8_and_nvfp4 class TestComposerStyleFactory: """Composer 2-style row-scaled NVFP4 forward + MXFP8 backward dispatch.""" @@ -167,6 +190,43 @@ def test_non_linear_roles_fall_back_to_mxfp8(self): assert isinstance(quantizer, MXFP8Quantizer) +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") +class TestMXFP8FwdDequantizedBwdFactory: + """MXFP8 forward + dequantized backward qfactory dispatch.""" + + @staticmethod + def _factory(role): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_dequantized_bwd_quantizer_factory, + ) + + return mxfp8_fwd_dequantized_bwd_quantizer_factory(role) + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + @pytest.mark.parametrize("tensor_type", ["input", "weight"]) + def test_forward_roles_require_saved_forward_quantized_value(self, module_type, tensor_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type=tensor_type)) + + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + assert isinstance(quantizer.rowwise_quantizer, MXFP8Quantizer) + assert isinstance(quantizer.columnwise_quantizer, IdentityQuantizer) + assert quantizer.allows_save_original_input_for_backward() is False + + @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) + def test_grad_output_role_allows_save_original_input(self, module_type): + from transformer_engine.pytorch.quantization import QuantizerRole + + quantizer = self._factory( + QuantizerRole(module_type=module_type, tensor_type="grad_output") + ) + + assert isinstance(quantizer, IdentityQuantizer) + assert quantizer.allows_save_original_input_for_backward() is True + + @requires_nvfp4 class TestNVFP4WeightDoubleQuantFactory: """Base NVFP4 recipe with W.T sourced from dequantized 1D W.""" @@ -254,7 +314,6 @@ def test_weight_quantizer_produces_both_nvfp4_storages(self): assert isinstance(out.rowwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) assert isinstance(out.columnwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) - @requires_fp8_and_nvfp4 class TestHybridQuantizerConstruction: """Test construction and basic properties of hybrid quantizer.""" @@ -377,6 +436,18 @@ def test_invalid_columnwise_source_raises(self): columnwise_source="dequantized", ) + def test_default_quantizer_allows_save_original_input_for_backward(self): + assert IdentityQuantizer().allows_save_original_input_for_backward() is True + + @pytest.mark.parametrize( + "columnwise_source,allowed", + [("original", True), ("rowwise_dequantized", False)], + ) + def test_hybrid_save_original_input_policy(self, columnwise_source, allowed): + hq = self._make_quantizer(columnwise_source=columnwise_source) + + assert hq.allows_save_original_input_for_backward() is allowed + def test_copy_preserves_columnwise_source(self): hq = self._make_quantizer(columnwise_source="rowwise_dequantized") hq.set_usage(rowwise=False, columnwise=True) @@ -441,6 +512,54 @@ def test_update_quantized_uses_updated_rowwise_storage(self, input_tensor): ) +@requires_fp8 +class TestHybridSaveOriginalInputPolicy: + """Module-level save_original_input policy for hybrid qfactory inputs.""" + + @staticmethod + def _counting_identity_hybrid_recipe(counter): + def factory(role): + if ( + role is not None + and role.module_type == "linear" + and role.tensor_type == "input" + ): + return HybridQuantizer( + rowwise_quantizer=_CountingIdentityQuantizer(counter), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + return IdentityQuantizer() + + return recipe.CustomRecipe(qfactory=factory) + + def test_linear_save_original_input_veto_uses_saved_forward_quantized_input(self): + torch.manual_seed(205) + counter = {"calls": 0} + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): + with autocast( + enabled=True, + recipe=self._counting_identity_hybrid_recipe(counter), + ): + out = model(inp) + + calls_after_forward = counter["calls"] + assert calls_after_forward > 0 + + out.float().sum().backward() + + assert counter["calls"] == calls_after_forward + + @requires_fp8_and_nvfp4 class TestHybridQuantize: """Test quantization via HybridQuantizer.""" @@ -1344,6 +1463,69 @@ def hybrid_mxfp8_factory(role): ) + def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwise(self): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_dequantized_bwd_quantizer_factory, + ) + + torch.manual_seed(204) + in_features, out_features, batch = 128, 128, 32 + + model_ref = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + model_qfactory = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + model_qfactory.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_qfactory = base_inp.clone().detach().requires_grad_(True) + + ref_recipe = recipe.MXFP8BlockScaling() + ref_recipe.backward_override = "dequantized" + with autocast(enabled=True, recipe=ref_recipe): + out_ref = model_ref(inp_ref) + + qfactory_recipe = recipe.CustomRecipe( + qfactory=mxfp8_fwd_dequantized_bwd_quantizer_factory + ) + with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): + with autocast(enabled=True, recipe=qfactory_recipe): + out_qfactory = model_qfactory(inp_qfactory) + + assert torch.equal(out_ref, out_qfactory), ( + f"Forward mismatch: max diff = {(out_ref - out_qfactory).abs().max().item()}" + ) + + out_ref.float().sum().backward() + out_qfactory.float().sum().backward() + + assert inp_ref.grad is not None and inp_qfactory.grad is not None + assert torch.equal(inp_ref.grad, inp_qfactory.grad), ( + f"Input grad mismatch: max diff = " + f"{(inp_ref.grad - inp_qfactory.grad).abs().max().item()}" + ) + for name, p_ref in dict(model_ref.named_parameters()).items(): + p_qfactory = dict(model_qfactory.named_parameters())[name] + assert ( + p_ref.grad is not None and p_qfactory.grad is not None + ), f"Missing gradient for param {name!r}" + assert torch.equal(p_ref.grad, p_qfactory.grad), ( + f"Param {name!r} grad mismatch: max diff = " + f"{(p_ref.grad - p_qfactory.grad).abs().max().item()}" + ) + + @pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) class TestHybridGemmBitwiseIdenticalBlockFP8: """Hybrid quantizer with Block FP8 in both directions must produce diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 699e9dd6aca..cfa6c2eb2d6 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -646,6 +646,24 @@ def forward( biases = weights_and_biases[num_gemms:] device = inp.device weight_requires_grad = weights[0].requires_grad + backward_needs_input = is_grad_enabled and weight_requires_grad + if save_original_input and backward_needs_input: + disallowed_input_quantizer = next( + ( + q + for q in input_quantizers + if q is not None and not q.allows_save_original_input_for_backward() + ), + None, + ) + if disallowed_input_quantizer is not None: + warnings.warn( + "Ignoring save_original_input=True because the input quantizer requires " + "the forward quantized activation for backward " + f"({disallowed_input_quantizer}).", + stacklevel=2, + ) + save_original_input = False # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index a5b3842d2b1..2672486476f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -298,6 +298,19 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) backward_needs_input = is_grad_enabled and weight.requires_grad + if ( + save_original_input + and backward_needs_input + and input_quantizer is not None + and not input_quantizer.allows_save_original_input_for_backward() + ): + warnings.warn( + "Ignoring save_original_input=True because the input quantizer requires " + "the forward quantized activation for backward " + f"({input_quantizer}).", + stacklevel=2, + ) + save_original_input = False with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index cc2baf1eb59..9a47efc080d 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -445,6 +445,16 @@ def supports_only_rowwise_all_gather(self) -> bool: """Returns True if the quantizer supports only rowwise all-gather""" return False + def allows_save_original_input_for_backward(self) -> bool: + """Whether ``save_original_input`` preserves backward operand semantics. + + Modules may use ``save_original_input`` as a memory optimization by saving + the high-precision input and re-preparing the backward operand later. Some + quantizers require the exact forward-produced quantized value for backward, + so replacing it with the original tensor would change recipe semantics. + """ + return True + def is_quantizable(self, inp: torch.Tensor) -> bool: # pylint: disable=unused-argument """Whether tensor supports quantized all-gather diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 56e17829553..56d0581cd07 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -350,6 +350,11 @@ def supports_only_rowwise_all_gather(self) -> bool: return True return False + def allows_save_original_input_for_backward(self) -> bool: + # TODO: Add an explicit recompute-from-original policy for deterministic + # quantizers that can trade backward compute for saved activation memory. + return self.columnwise_source == "original" + def _get_compatible_recipe(self): # HybridQuantizer is only reachable via CustomRecipe (the qfactory # returns HybridQuantizer per role). Checking that the autocast recipe From 06dcdd43176e237d52d598155f8730ae1db27cd5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 23 Jun 2026 12:41:01 +0000 Subject: [PATCH 34/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 24 ++++++++--------------- 1 file changed, 8 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 62355dd6397..a9895e05608 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -219,9 +219,7 @@ def test_forward_roles_require_saved_forward_quantized_value(self, module_type, def test_grad_output_role_allows_save_original_input(self, module_type): from transformer_engine.pytorch.quantization import QuantizerRole - quantizer = self._factory( - QuantizerRole(module_type=module_type, tensor_type="grad_output") - ) + quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type="grad_output")) assert isinstance(quantizer, IdentityQuantizer) assert quantizer.allows_save_original_input_for_backward() is True @@ -314,6 +312,7 @@ def test_weight_quantizer_produces_both_nvfp4_storages(self): assert isinstance(out.rowwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) assert isinstance(out.columnwise_sub_storage, (NVFP4TensorStorage, NVFP4Tensor)) + @requires_fp8_and_nvfp4 class TestHybridQuantizerConstruction: """Test construction and basic properties of hybrid quantizer.""" @@ -519,11 +518,7 @@ class TestHybridSaveOriginalInputPolicy: @staticmethod def _counting_identity_hybrid_recipe(counter): def factory(role): - if ( - role is not None - and role.module_type == "linear" - and role.tensor_type == "input" - ): + if role is not None and role.module_type == "linear" and role.tensor_type == "input": return HybridQuantizer( rowwise_quantizer=_CountingIdentityQuantizer(counter), columnwise_quantizer=IdentityQuantizer(), @@ -1462,7 +1457,6 @@ def hybrid_mxfp8_factory(role): f"{(p_ref.grad - p_hyb.grad).abs().max().item()}" ) - def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwise(self): from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( mxfp8_fwd_dequantized_bwd_quantizer_factory, @@ -1496,23 +1490,21 @@ def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwis with autocast(enabled=True, recipe=ref_recipe): out_ref = model_ref(inp_ref) - qfactory_recipe = recipe.CustomRecipe( - qfactory=mxfp8_fwd_dequantized_bwd_quantizer_factory - ) + qfactory_recipe = recipe.CustomRecipe(qfactory=mxfp8_fwd_dequantized_bwd_quantizer_factory) with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): with autocast(enabled=True, recipe=qfactory_recipe): out_qfactory = model_qfactory(inp_qfactory) - assert torch.equal(out_ref, out_qfactory), ( - f"Forward mismatch: max diff = {(out_ref - out_qfactory).abs().max().item()}" - ) + assert torch.equal( + out_ref, out_qfactory + ), f"Forward mismatch: max diff = {(out_ref - out_qfactory).abs().max().item()}" out_ref.float().sum().backward() out_qfactory.float().sum().backward() assert inp_ref.grad is not None and inp_qfactory.grad is not None assert torch.equal(inp_ref.grad, inp_qfactory.grad), ( - f"Input grad mismatch: max diff = " + "Input grad mismatch: max diff = " f"{(inp_ref.grad - inp_qfactory.grad).abs().max().item()}" ) for name, p_ref in dict(model_ref.named_parameters()).items(): From 4d553d4fd0c38593b0d9494f34da22fc813549b6 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 23 Jun 2026 07:33:56 -0700 Subject: [PATCH 35/97] IMprove hybrid factory tests Signed-off-by: Evgeny --- tests/pytorch/test_identity_quantizer.py | 309 +++++++++++++++++------ 1 file changed, 238 insertions(+), 71 deletions(-) diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index ab98e8e1c4f..d58fee3f63a 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -4,8 +4,8 @@ """Tests for IdentityQuantizer (high-precision passthrough) and its use as a per-direction component of HybridQuantizer to express mixed forward/backward -precision via the CustomRecipe + qfactory machinery. Scoped to te.Linear, -single GPU. +precision via the CustomRecipe + qfactory machinery. Scoped to single-GPU +TE GEMM modules. """ import io @@ -1081,6 +1081,242 @@ def test_activation_recompute_matches_no_checkpoint(self, format_name, use_reent torch.testing.assert_close(g_test, g_ref, rtol=0.0, atol=0.0) +_ZOO_DEQUANTIZED_MODULES = [ + pytest.param("Linear", id="linear"), + pytest.param("LayerNormLinear", id="layernorm_linear"), + pytest.param("GroupedLinear", id="grouped_linear"), + pytest.param( + "LayerNormMLP", + marks=pytest.mark.xfail( + reason="LayerNormMLP does not support built-in backward_override=dequantized", + strict=True, + ), + id="layernorm_mlp", + ), + pytest.param("LayerNormLinearLinear", id="layernorm_linear_linear"), +] + +_ZOO_DEQUANTIZED_CASES = [ + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}"), + id="mxfp8", + ), + pytest.param( + "nvfp4_row_scaled", + marks=pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}"), + id="nvfp4_row_scaled", + ), +] + + +def _make_zoo_dequantized_module(module_name, *, save_original_input=False): + hidden_size = 128 + if module_name == "Linear": + return te.Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ).cuda() + if module_name == "LayerNormLinear": + return te.LayerNormLinear(hidden_size, hidden_size, params_dtype=torch.bfloat16).cuda() + if module_name == "GroupedLinear": + return te.GroupedLinear( + 2, + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ).cuda() + if module_name == "LayerNormMLP": + return te.LayerNormMLP( + hidden_size, + 2 * hidden_size, + params_dtype=torch.bfloat16, + ).cuda() + if module_name == "LayerNormLinearLinear": + return torch.nn.Sequential( + te.LayerNormLinear(hidden_size, hidden_size, params_dtype=torch.bfloat16), + te.Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + save_original_input=save_original_input, + ), + ).cuda() + raise ValueError(module_name) + + +def _make_zoo_dequantized_module_pair(module_name): + torch.manual_seed(9400) + torch.cuda.manual_seed(9400) + ref = _make_zoo_dequantized_module(module_name, save_original_input=False) + torch.manual_seed(9401) + torch.cuda.manual_seed(9401) + test = _make_zoo_dequantized_module(module_name, save_original_input=True) + test.load_state_dict(ref.state_dict()) + return ref, test + + +def _zoo_dequantized_input(module_name): + torch.manual_seed(9402) + batch = 128 if module_name == "GroupedLinear" else 64 + return torch.randn(batch, 128, device="cuda", dtype=torch.bfloat16) + + +def _zoo_dequantized_forward(module_name, module, inp): + if module_name == "GroupedLinear": + m_splits = torch.tensor([64, 64], device="cuda", dtype=torch.int32) + return module(inp, m_splits=m_splits) + return module(inp) + + +def _fwd_bwd_zoo_dequantized_module(module_name, module, inp, recipe): + inp = inp.detach().clone().requires_grad_(True) + with te.autocast(enabled=True, recipe=recipe): + out = _zoo_dequantized_forward(module_name, module, inp) + torch.manual_seed(9403) + grad = torch.randn_like(out) + out.backward(grad) + param_grads = { + name: param.grad.detach().clone() + for name, param in module.named_parameters() + if param.grad is not None + } + return out.detach().clone(), inp.grad.detach().clone(), param_grads + + +def _mxfp8_all_qfactory(role): # pylint: disable=unused-argument + return _mxfp8(tex.DType.kFloat8E4M3) + + +def _zoo_dequantized_qfactory(case_name): + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_dequantized_bwd_quantizer_factory, + nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory, + ) + + if case_name == "mxfp8": + return mxfp8_fwd_dequantized_bwd_quantizer_factory + if case_name == "nvfp4_row_scaled": + return nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory + raise ValueError(case_name) + + +def _zoo_dequantized_recipes(case_name): + from transformer_engine.common import recipe as te_recipe + + if case_name == "mxfp8": + return ( + CustomRecipe(qfactory=_mxfp8_all_qfactory, backward_override="dequantized"), + CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)), + ) + if case_name == "nvfp4_row_scaled": + ref_recipe = te_recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="dequantized", + ) + ref_recipe.fp4_quant_fwd_inp = te_recipe.QParams() + ref_recipe.fp4_quant_fwd_weight = te_recipe.QParams() + ref_recipe.fp4_quant_bwd_grad = te_recipe.QParams() + return ( + ref_recipe, + CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)), + ) + raise ValueError(case_name) + + +def _assert_zoo_layernorm_mlp_role_semantics(case_name): + module = _make_zoo_dequantized_module("LayerNormMLP") + qfactory = _zoo_dequantized_qfactory(case_name) + + fwd_roles = module.get_quantizer_roles(fwd=True, num_quantizers=6) + fwd_gemm_roles = [ + role + for role in fwd_roles + if role is not None and role.tensor_type in ("input", "weight") + ] + assert len(fwd_gemm_roles) == 5 + for role in fwd_gemm_roles: + quantizer = qfactory(role) + assert isinstance(quantizer, HybridQuantizer) + assert quantizer.columnwise_source == "rowwise_dequantized" + assert isinstance(quantizer.columnwise_quantizer, IdentityQuantizer) + if case_name == "mxfp8": + assert isinstance(quantizer.rowwise_quantizer, MXFP8Quantizer) + else: + assert isinstance(quantizer.rowwise_quantizer, NVFP4Quantizer) + assert quantizer.rowwise_quantizer.row_scaled_nvfp4 == (role.tensor_type == "input") + + bwd_roles = module.get_quantizer_roles(fwd=False, num_quantizers=4) + grad_output_roles = [ + role for role in bwd_roles if role is not None and role.tensor_type == "grad_output" + ] + assert len(grad_output_roles) == 3 + for role in grad_output_roles: + assert isinstance(qfactory(role), IdentityQuantizer) + + +class TestZooDequantizedBackwardFactoryModuleCoverage: + """Zoo dequantized-backward recipes should match base recipes across TE modules.""" + + @pytest.mark.parametrize("case_name", _ZOO_DEQUANTIZED_CASES) + @pytest.mark.parametrize("module_name", _ZOO_DEQUANTIZED_MODULES) + def test_matches_base_recipe_bitwise(self, case_name, module_name): + ref, test = _make_zoo_dequantized_module_pair(module_name) + inp = _zoo_dequantized_input(module_name) + ref_recipe, qfactory_recipe = _zoo_dequantized_recipes(case_name) + + y_ref, dx_ref, grads_ref = _fwd_bwd_zoo_dequantized_module( + module_name, ref, inp, ref_recipe + ) + if module_name in ("Linear", "GroupedLinear", "LayerNormLinearLinear"): + with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): + y_test, dx_test, grads_test = _fwd_bwd_zoo_dequantized_module( + module_name, test, inp, qfactory_recipe + ) + else: + y_test, dx_test, grads_test = _fwd_bwd_zoo_dequantized_module( + module_name, test, inp, qfactory_recipe + ) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_test, dx_ref, rtol=0.0, atol=0.0) + assert grads_test.keys() == grads_ref.keys() + for name, grad_ref in grads_ref.items(): + torch.testing.assert_close( + grads_test[name], + grad_ref, + rtol=0.0, + atol=0.0, + msg=f"{case_name}/{module_name} grad mismatch for {name}", + ) + + @pytest.mark.parametrize("case_name", _ZOO_DEQUANTIZED_CASES) + def test_layernorm_mlp_runs_and_uses_dequantized_backward_roles(self, case_name): + _assert_zoo_layernorm_mlp_role_semantics(case_name) + + module = _make_zoo_dequantized_module("LayerNormMLP") + inp = _zoo_dequantized_input("LayerNormMLP") + recipe = CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)) + out, dx, grads = _fwd_bwd_zoo_dequantized_module( + "LayerNormMLP", module, inp, recipe + ) + + assert torch.isfinite(out).all() + assert torch.isfinite(dx).all() + expected_grads = { + name for name, param in module.named_parameters() if param.requires_grad + } + assert grads.keys() == expected_grads + for grad in grads.values(): + assert torch.isfinite(grad).all() + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) class TestIdentityLinear: """End-to-end te.Linear with Identity-based recipes.""" @@ -1226,75 +1462,6 @@ def test_identity_reproduces_backward_override_dequantized_bitwise(self): for g_id, g_bo in zip(wg_id, wg_bo): torch.testing.assert_close(g_id, g_bo, rtol=0.0, atol=0.0) - @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") - def test_zoo_mxfp8_dequantized_bwd_matches_backward_override_bitwise(self): - """The zoo MXFP8-dequantized recipe should reproduce - ``backward_override=dequantized`` bitwise for Linear. - """ - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - mxfp8_fwd_dequantized_bwd_quantizer_factory, - ) - - def mxfp8_all_factory(role): # pylint: disable=unused-argument - return _mxfp8(tex.DType.kFloat8E4M3) - - ref, test = _make_linears(self.IN_F, self.OUT_F) - x = self._input() - - y_bo, dx_bo, wg_bo = _fwd_bwd( - ref, - x, - recipe=CustomRecipe(qfactory=mxfp8_all_factory, backward_override="dequantized"), - ) - y_zoo, dx_zoo, wg_zoo = _fwd_bwd( - test, - x, - recipe=CustomRecipe(qfactory=mxfp8_fwd_dequantized_bwd_quantizer_factory), - ) - - torch.testing.assert_close(y_zoo, y_bo, rtol=0.0, atol=0.0) - torch.testing.assert_close(dx_zoo, dx_bo, rtol=0.0, atol=0.0) - assert len(wg_zoo) == len(wg_bo) - for g_zoo, g_bo in zip(wg_zoo, wg_bo): - torch.testing.assert_close(g_zoo, g_bo, rtol=0.0, atol=0.0) - - @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") - def test_zoo_nvfp4_row_scaled_dequantized_bwd_matches_recipe_bitwise(self): - """The zoo row-scaled NVFP4-dequantized recipe should reproduce - the built-in row-scaled NVFP4 recipe with dequantized backward bitwise. - """ - from transformer_engine.common import recipe as te_recipe - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory, - ) - - ref_recipe = te_recipe.NVFP4BlockScaling( - disable_rht=True, - disable_stochastic_rounding=True, - disable_2d_quantization=True, - row_scaled_activation=True, - backward_override="dequantized", - ) - ref_recipe.fp4_quant_fwd_inp = te_recipe.QParams() - ref_recipe.fp4_quant_fwd_weight = te_recipe.QParams() - ref_recipe.fp4_quant_bwd_grad = te_recipe.QParams() - - ref, test = _make_linears(self.IN_F, self.OUT_F) - x = self._input() - - y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=ref_recipe) - y_zoo, dx_zoo, wg_zoo = _fwd_bwd( - test, - x, - recipe=CustomRecipe(qfactory=nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory), - ) - - torch.testing.assert_close(y_zoo, y_ref, rtol=0.0, atol=0.0) - torch.testing.assert_close(dx_zoo, dx_ref, rtol=0.0, atol=0.0) - assert len(wg_zoo) == len(wg_ref) - for g_zoo, g_ref in zip(wg_zoo, wg_ref): - torch.testing.assert_close(g_zoo, g_ref, rtol=0.0, atol=0.0) - def test_identity_matches_bf16_multistep_training_bitwise(self): """Multi-step SGD: an all-Identity recipe must track a plain BF16 te.Linear bitwise across optimizer steps (no drift from workspace caching From bd5d752c182aad51aa5b717e95dd69cd5eada4e6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 23 Jun 2026 14:35:05 +0000 Subject: [PATCH 36/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_identity_quantizer.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index d58fee3f63a..7d96d0720a7 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -1236,9 +1236,7 @@ def _assert_zoo_layernorm_mlp_role_semantics(case_name): fwd_roles = module.get_quantizer_roles(fwd=True, num_quantizers=6) fwd_gemm_roles = [ - role - for role in fwd_roles - if role is not None and role.tensor_type in ("input", "weight") + role for role in fwd_roles if role is not None and role.tensor_type in ("input", "weight") ] assert len(fwd_gemm_roles) == 5 for role in fwd_gemm_roles: @@ -1303,15 +1301,11 @@ def test_layernorm_mlp_runs_and_uses_dequantized_backward_roles(self, case_name) module = _make_zoo_dequantized_module("LayerNormMLP") inp = _zoo_dequantized_input("LayerNormMLP") recipe = CustomRecipe(qfactory=_zoo_dequantized_qfactory(case_name)) - out, dx, grads = _fwd_bwd_zoo_dequantized_module( - "LayerNormMLP", module, inp, recipe - ) + out, dx, grads = _fwd_bwd_zoo_dequantized_module("LayerNormMLP", module, inp, recipe) assert torch.isfinite(out).all() assert torch.isfinite(dx).all() - expected_grads = { - name for name, param in module.named_parameters() if param.requires_grad - } + expected_grads = {name for name, param in module.named_parameters() if param.requires_grad} assert grads.keys() == expected_grads for grad in grads.values(): assert torch.isfinite(grad).all() From 295324577ce2c2c9c7d37d88f0034d1371934a63 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 26 Jun 2026 06:09:34 -0700 Subject: [PATCH 37/97] Fix attention factories and align with env-var setup Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 357 ++++++++++++++++++ .../dot_product_attention.py | 73 ++++ .../pytorch/attention/multi_head_attention.py | 40 +- .../quantization_factory_zoo.py | 77 ++-- 4 files changed, 506 insertions(+), 41 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index a9895e05608..8efa033865e 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -1518,6 +1518,363 @@ def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwis ) +@requires_fp8_and_nvfp4 +class TestAttentionFactoryNativeRecipeParity: + """Linear + DPA qfactories should match native DPA recipe switches bitwise.""" + + batch = 2 + seq_len = 128 + hidden_size = 128 + num_heads = 4 + kv_channels = hidden_size // num_heads + + class _LinearDPALinear(torch.nn.Module): + def __init__(self, hidden_size, num_heads, kv_channels): + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_heads + self.kv_channels = kv_channels + self.qkv_proj = Linear( + hidden_size, + 3 * hidden_size, + params_dtype=torch.bfloat16, + bias=False, + name="qkv", + ).cuda() + self.dpa = te.DotProductAttention( + num_heads, + kv_channels, + attention_dropout=0.0, + qkv_format="bshd", + name="core_attention", + ).cuda() + self.out_proj = Linear( + hidden_size, + hidden_size, + params_dtype=torch.bfloat16, + bias=False, + name="proj", + ).cuda() + + def forward(self, inp): + batch, seq_len, _ = inp.shape + qkv = self.qkv_proj(inp).view( + batch, + seq_len, + 3, + self.num_heads, + self.kv_channels, + ) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + attn_out = self.dpa(q, k, v, qkv_format="bshd").reshape( + batch, + seq_len, + self.hidden_size, + ) + return self.out_proj(attn_out) + + @staticmethod + def _set_native_dpa_recipe(monkeypatch, recipe_name): + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + monkeypatch.setenv("NVTE_DPA_FP8_RECIPE", recipe_name) + monkeypatch.setenv("NVTE_DPA_FP8_FORMAT", "HYBRID") + monkeypatch.setenv("NVTE_DPA_FP8DS_AMAX_ALGO", "most_recent") + monkeypatch.setenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1") + monkeypatch.setenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") + monkeypatch.setattr(dpa_module, "_dpa_fp8_recipe", recipe_name) + monkeypatch.setattr(dpa_module, "_dpa_fp8_format", recipe.Format.HYBRID) + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_amax_algo", "most_recent") + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_amax_histlen", 1) + monkeypatch.setattr(dpa_module, "_dpa_fp8ds_reduce_amax", True) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", recipe_name) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_dpa", False) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_mha", False) + + @staticmethod + def _clear_native_dpa_recipe(monkeypatch): + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + + monkeypatch.delenv("NVTE_DPA_FP8_RECIPE", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8_FORMAT", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_AMAX_ALGO", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_AMAX_HISTLEN", raising=False) + monkeypatch.delenv("NVTE_DPA_FP8DS_REDUCE_AMAX", raising=False) + monkeypatch.setattr(dpa_module, "_dpa_fp8_recipe", "") + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", "") + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_dpa", False) + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe_mha", False) + + @staticmethod + def _assert_equal(actual, expected, label): + assert torch.equal(actual, expected), ( + f"{label} mismatch: max diff = " + f"{(actual.float() - expected.float()).abs().max().item()}" + ) + + def _run_model(self, model, inp, grad, fp8_recipe, seed): + run_inp = inp.clone().detach().requires_grad_(True) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + with autocast(enabled=True, recipe=fp8_recipe): + out = model(run_inp) + out.backward(grad) + local_recipes = [ + type(r).__name__ for r in model.dpa.fp8_meta.get("local_recipes", []) + ] + return ( + out.detach().clone(), + run_inp.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + local_recipes, + ) + + @pytest.mark.parametrize( + "case_name,native_dpa_recipe,qfactory_name", + [ + ( + "mixed_fp8_dpa", + "Float8CurrentScaling", + "nvfp4_linear_mixed_fp8_dpa_factory", + ), + ( + "mxfp8_dpa", + "MXFP8BlockScaling", + "nvfp4_linear_mxfp8_dpa_factory", + ), + ], + ) + def test_linear_dpa_linear_matches_native_env_recipe_bitwise( + self, + monkeypatch, + case_name, + native_dpa_recipe, + qfactory_name, + ): + if case_name == "mxfp8_dpa" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.custom_recipes import quantization_factory_zoo + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + self._set_native_dpa_recipe(monkeypatch, native_dpa_recipe) + + torch.manual_seed(2201) + model_native = self._LinearDPALinear( + self.hidden_size, + self.num_heads, + self.kv_channels, + ) + model_qfactory = self._LinearDPALinear( + self.hidden_size, + self.num_heads, + self.kv_channels, + ) + model_qfactory.load_state_dict(model_native.state_dict()) + + torch.manual_seed(2202) + base_inp = torch.randn( + self.batch, + self.seq_len, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(base_inp) + + native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) + qfactory = getattr(quantization_factory_zoo, qfactory_name) + qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) + + native_out, native_dx, native_grads, native_local_recipes = self._run_model( + model_native, + base_inp, + grad, + native_recipe, + seed=2203, + ) + self._clear_native_dpa_recipe(monkeypatch) + qfactory_out, qfactory_dx, qfactory_grads, qfactory_local_recipes = self._run_model( + model_qfactory, + base_inp, + grad, + qfactory_recipe, + seed=2203, + ) + + expected_local_recipes = ( + ["Float8CurrentScaling", "DelayedScaling"] + if native_dpa_recipe == "Float8CurrentScaling" + else ["MXFP8BlockScaling"] + ) + assert native_local_recipes == expected_local_recipes + assert qfactory_local_recipes == expected_local_recipes + + self._assert_equal(qfactory_out, native_out, f"{case_name} output") + self._assert_equal(qfactory_dx, native_dx, f"{case_name} input grad") + assert qfactory_grads.keys() == native_grads.keys() + for name, native_grad in native_grads.items(): + self._assert_equal( + qfactory_grads[name], + native_grad, + f"{case_name} param grad {name}", + ) + + def _run_mha_model(self, model, inp, grad, fp8_recipe, seed): + run_inp = inp.clone().detach().requires_grad_(True) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + with autocast(enabled=True, recipe=fp8_recipe): + out = model(run_inp, attn_mask_type="no_mask") + if isinstance(out, tuple): + out = out[0] + out.backward(grad) + return ( + out.detach().clone(), + run_inp.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + ) + + @pytest.mark.parametrize( + "case_name,native_dpa_recipe,qfactory_name,expected_flags", + [ + ( + "mixed_fp8_dpa", + "Float8CurrentScaling", + "nvfp4_linear_mixed_fp8_dpa_factory", + (False, False, False), + ), + ( + "mxfp8_dpa", + "MXFP8BlockScaling", + "nvfp4_linear_mxfp8_dpa_factory", + (False, False, False), + ), + ], + ) + def test_multihead_attention_matches_native_env_recipe_bitwise( + self, + monkeypatch, + case_name, + native_dpa_recipe, + qfactory_name, + expected_flags, + ): + if case_name == "mxfp8_dpa" and not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.custom_recipes import quantization_factory_zoo + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + recorded_flags = [] + orig_update_roles = mha_module.MultiheadAttention._update_output_quantizer_roles + + def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output): + recorded_flags.append((qkv_fp8_output, dpa_fp8_output, proj_fp8_grad)) + return orig_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + + monkeypatch.setattr( + mha_module.MultiheadAttention, + "_update_output_quantizer_roles", + _recording_update_roles, + ) + + self._set_native_dpa_recipe(monkeypatch, native_dpa_recipe) + + torch.manual_seed(2301) + model_native = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + model_qfactory = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + model_qfactory.load_state_dict(model_native.state_dict()) + + torch.manual_seed(2302) + base_inp = torch.randn( + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(base_inp) + + native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) + qfactory = getattr(quantization_factory_zoo, qfactory_name) + qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) + + native_out, native_dx, native_grads = self._run_mha_model( + model_native, + base_inp, + grad, + native_recipe, + seed=2303, + ) + native_flags = recorded_flags[-1] + self._clear_native_dpa_recipe(monkeypatch) + qfactory_out, qfactory_dx, qfactory_grads = self._run_mha_model( + model_qfactory, + base_inp, + grad, + qfactory_recipe, + seed=2303, + ) + qfactory_flags = recorded_flags[-1] + + assert native_flags == expected_flags + assert qfactory_flags == expected_flags + self._assert_equal(qfactory_out, native_out, f"{case_name} MHA output") + self._assert_equal(qfactory_dx, native_dx, f"{case_name} MHA input grad") + assert qfactory_grads.keys() == native_grads.keys() + for name, native_grad in native_grads.items(): + self._assert_equal( + qfactory_grads[name], + native_grad, + f"{case_name} MHA param grad {name}", + ) + + @pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) class TestHybridGemmBitwiseIdenticalBlockFP8: """Hybrid quantizer with Block FP8 in both directions must produce diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 03008bb2d75..cbe2a88281b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -86,6 +86,74 @@ "_alibi_bias_require_update": False, } + +def _infer_custom_dpa_local_recipes( + fp8_recipe: Recipe, + fp8_meta: Dict[str, Any], + quantizers: Dict[str, Any], +) -> Optional[List[Recipe]]: + """Infer native-equivalent DPA recipe labels for CustomRecipe control-flow. + + CustomRecipe owns quantizer construction, but DPA backend selection and a + few fused-attention branches still dispatch on recipe predicates. Attach + local recipe labels that match the qfactory DPA quantizer family while + keeping the actual qfactory-created quantizers untouched. + """ + try: + qkv_quantizer = quantizers["scaling_fwd"][dpa_utils.META_QKV] + except (KeyError, IndexError, TypeError): + return None + + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, + ) + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + if isinstance(qkv_quantizer, MXFP8Quantizer): + return [ + MXFP8BlockScaling( + fp8_format=fp8_recipe.fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + ] + + def _delayed_scaling_recipe() -> Optional[DelayedScaling]: + fwd_state = fp8_meta.get("scaling_fwd") + ds_recipe = getattr(fwd_state, "_inner_delayed_scaling_recipe", None) + if ds_recipe is None: + return None + return DelayedScaling( + fp8_format=ds_recipe.fp8_format, + margin=ds_recipe.margin, + amax_history_len=ds_recipe.amax_history_len, + amax_compute_algo=ds_recipe.amax_compute_algo, + scaling_factor_compute_algo=ds_recipe.scaling_factor_compute_algo, + reduce_amax=ds_recipe.reduce_amax, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + + if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): + ds_recipe = _delayed_scaling_recipe() + if ds_recipe is not None: + return [ + Float8CurrentScaling( + fp8_format=fp8_recipe.fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ), + ds_recipe, + ] + + if isinstance(qkv_quantizer, Float8Quantizer): + ds_recipe = _delayed_scaling_recipe() + if ds_recipe is not None: + return [ds_recipe] + + return None + """ This feature is **experimental** and subject to change. @@ -634,6 +702,11 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): super().init_fp8_metadata(num_gemms=num_gemms) + local_recipes = _infer_custom_dpa_local_recipes( + fp8_recipe, self.fp8_meta, self.quantizers + ) + if local_recipes is not None: + self.fp8_meta["local_recipes"] = local_recipes return # switch/append recipe: fp8_recipe stays unchanged, but DPA.fp8_meta["recipe"] may be set to diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 70ae9dfc21a..9fc6f6895be 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -43,6 +43,33 @@ _dpa_fp8_recipe_mha = os.getenv("NVTE_DPA_FP8_RECIPE_MHA", "0") == "1" +def _infer_custom_dpa_scaling(recipe, dpa_name: str) -> Tuple[bool, bool]: + """Infer DPA quantizer family for CustomRecipe boundary decisions. + + MultiheadAttention must decide boundary behavior before + DotProductAttention builds its per-slot quantizers. Built-in recipes expose + the DPA quantizer family through recipe predicates; CustomRecipe exposes it + only through the qfactory role contract, so probe the DPA QKV role. + """ + if not recipe.custom() or not getattr(recipe, "fp8_dpa", False): + return False, False + + try: + qkv_quantizer = recipe.qfactory( + QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) + ) + except Exception: # pragma: no cover - preserve existing fallback behavior + return False, False + + from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return ( + isinstance(qkv_quantizer, Float8CurrentScalingQuantizer), + isinstance(qkv_quantizer, MXFP8Quantizer), + ) + + class MultiheadAttention(torch.nn.Module): r""" Multi-head Attention (MHA), including Query, @@ -869,12 +896,20 @@ def forward( # ====================== fp8 = FP8GlobalStateManager.is_fp8_enabled() + custom_recipe = False if _dpa_fp8_recipe == "": fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() + custom_recipe = fp8_recipe.custom() fp8_dpa = fp8_recipe.fp8_dpa fp8_mha = fp8_recipe.fp8_mha float8_current_scaling = fp8_recipe.float8_current_scaling() mxfp8_scaling = fp8_recipe.mxfp8() + if custom_recipe and fp8_dpa: + custom_float8_cs, custom_mxfp8 = _infer_custom_dpa_scaling( + fp8_recipe, self.core_attention.name or "" + ) + float8_current_scaling = custom_float8_cs + mxfp8_scaling = custom_mxfp8 else: fp8_dpa = _dpa_fp8_recipe_dpa fp8_mha = _dpa_fp8_recipe_mha @@ -896,7 +931,10 @@ def forward( # DPA: produce FP8 output to take advantage of O amax from DPA; Projection Gemm can take FP8 or F16 inputs # 1. FP8DS/FP8CS recipe: produce FP8 output # 2. MXFP8 recipe: produce F16 output; again, due to quantization dimensions mismatch - dpa_fp8_output = fp8 and (fp8_dpa or fp8_mha) and not mxfp8_scaling + # For CustomRecipe, fp8_dpa only controls DPA-internal quantization. + # External MHA boundary tensors become FP8 only when fp8_mha is enabled. + dpa_fp8_output_enabled = fp8_mha if custom_recipe else (fp8_dpa or fp8_mha) + dpa_fp8_output = fp8 and dpa_fp8_output_enabled and not mxfp8_scaling # Projection Gemm: match DPA output except # 1. FP8CS recipe: produce F16 grads; again, due to cuBLAS limitation proj_fp8_grad = dpa_fp8_output and not float8_current_scaling diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index af90f8813fc..5dc5ec8579c 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -12,10 +12,12 @@ ``HybridQuantizer`` when rowwise and columnwise tensor directions should use different formats or sources. +Factories are ordered from conservative to more aggressive quantization. + Organization: - * Linear / grouped-linear recipes: from conservative to more aggressive quantization. - * RL-oriented recipes: dequantized-backward and MoE-inspired recipes that - favor more precision or explicit source control in backward GEMMs. + * Linear / grouped-linear recipes (pre-training). Favor more precision + on the forward pass. + * RL-oriented recipes: Favor more precision in backward GEMMs. * Linear + attention recipes: factories that also cover ``DotProductAttention`` roles and require ``CustomRecipe(..., fp8_dpa=True)``. @@ -64,7 +66,7 @@ # ----------------------------------------------------------------------------- -# Linear / GroupedLinear Recipes +# Linear / GroupedLinear Recipes (pre-training) # ----------------------------------------------------------------------------- @@ -346,9 +348,9 @@ def nvfp4_linear_mixed_fp8_dpa_factory( Dispatch logic: * ``role.module_type == "dpa"`` with ``tensor_type in ("s", "dp")`` - -> FP8 delayed scaling (stateful amax tracking) - * ``role.module_type == "dpa"`` (QKV, dO) - -> FP8 current scaling (E4M3) + -> FP8 delayed scaling (HYBRID, most_recent, history length 1) + * other DPA roles + -> FP8 current scaling (HYBRID: E4M3 fwd, E5M2 bwd) * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in ``role.name``) -> FP8 current scaling placeholder. The fused attention kernel requires FP8-compatible quantizers in all DPA slots, even when the output is @@ -373,33 +375,30 @@ def nvfp4_linear_mixed_fp8_dpa_factory( with autocast(recipe=recipe): output = model(input) """ + from transformer_engine.common.recipe import Format from transformer_engine.pytorch.quantization import DelayedScalingRequest from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer is_dpa = role is not None and role.module_type == "dpa" - is_softmax_or_dp = is_dpa and role.tensor_type in ("s", "dp") - - if is_softmax_or_dp: - return DelayedScalingRequest() + is_dpa_boundary = role is not None and not role.module_type and ( + "dpa_output" in role.name or "dpa_grad_input" in role.name + ) - if is_dpa: - return Float8CurrentScalingQuantizer( - fp8_dtype=DType.kFloat8E4M3, - device="cuda", + # Native NVFP4 + FP8CS attention uses delayed scaling for S/dP. + if is_dpa and role.tensor_type in ("s", "dp"): + return DelayedScalingRequest( + fp8_format=Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + reduce_amax=True, ) - # DPA boundary slots (O output / dQKV grad-input): the fused attention - # kernel only supports FP8 quantizers here, regardless of the linear recipe. - is_dpa_boundary = ( - role is not None - and not role.module_type - and ("dpa_output" in role.name or "dpa_grad_input" in role.name) - ) - if is_dpa_boundary: - return Float8CurrentScalingQuantizer( - fp8_dtype=DType.kFloat8E4M3, - device="cuda", + if is_dpa or is_dpa_boundary: + is_bwd_role = (is_dpa and role.tensor_type in ("do", "dp", "dqkv")) or ( + is_dpa_boundary and "dpa_grad_input" in role.name ) + fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 + return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") return nvfp4_quantizer_factory(role) @@ -430,7 +429,7 @@ def nvfp4_linear_mxfp8_dpa_factory( =========== ============================================================ Dispatch logic: - * ``role.module_type == "dpa"`` -> MXFP8 (E4M3, block-32) + * ``role.module_type == "dpa"`` -> MXFP8 (HYBRID: E4M3 fwd, E5M2 bwd) The MXFP8 fused-attention kernel handles the S/dP slots internally, so any quantizer returned for those roles is later nulled out by ``get_attention_quantizers``. Returning MXFP8 is @@ -456,20 +455,18 @@ def nvfp4_linear_mxfp8_dpa_factory( with autocast(recipe=recipe): output = model(input) """ - is_dpa = role is not None and role.module_type == "dpa" - if is_dpa: - return mxfp8_quantizer_factory(role) + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - # DPA boundary slots (O output / dQKV grad-input): emitted by DPA with - # empty `module_type` and a `name` like ".dpa_output". The fused - # attention kernel requires an FP8-compatible quantizer here even when - # the downstream consumer is unknown. - is_dpa_boundary = ( - role is not None - and not role.module_type - and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = role is not None and not role.module_type and ( + "dpa_output" in role.name or "dpa_grad_input" in role.name ) - if is_dpa_boundary: - return mxfp8_quantizer_factory(role) + + if is_dpa or is_dpa_boundary: + is_bwd_role = (is_dpa and role.tensor_type in ("do", "dp", "dqkv")) or ( + is_dpa_boundary and "dpa_grad_input" in role.name + ) + fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 + return MXFP8Quantizer(fp8_dtype=fp8_dtype) return nvfp4_quantizer_factory(role) From 0043a63cdca1d63f7fb1a397279e77afdb6bb70f Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 29 Jun 2026 02:59:20 -0700 Subject: [PATCH 38/97] Update quantization factory zoo Signed-off-by: Evgeny --- tests/pytorch/test_custom_recipe.py | 8 ++-- tests/pytorch/test_hybrid_quantization.py | 14 +++---- .../quantization_factory_zoo.py | 40 +++++++++---------- 3 files changed, 30 insertions(+), 32 deletions(-) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 52140b4d9ff..79ca5fe6126 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -1098,9 +1098,9 @@ def test_role_change_does_not_invalidate_when_role_unchanged(): def test_custom_recipe_dpa_fp8(): - """DotProductAttention forward+backward with CustomRecipe and role-based mixed quantizers. + """DotProductAttention forward+backward with CustomRecipe and role-based DPA quantizers. - Uses the nvfp4_linear_mixed_fp8_dpa_factory which dispatches: + Uses the nvfp4_linear_fp8_dpa_factory which dispatches: * DPA S/dP slots -> DelayedScalingRequest (stateful) * DPA QKV/O/dO/dQKV slots -> Float8CurrentScalingQuantizer * Linear slots -> NVFP4Quantizer @@ -1126,7 +1126,7 @@ def test_custom_recipe_dpa_fp8(): Float8CurrentScalingQuantizer, ) from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - nvfp4_linear_mixed_fp8_dpa_factory, + nvfp4_linear_fp8_dpa_factory, ) torch.manual_seed(42) @@ -1145,7 +1145,7 @@ def test_custom_recipe_dpa_fp8(): out_proj = Linear(H, H, params_dtype=torch.bfloat16, bias=False, name="proj").cuda() custom_recipe = recipe.CustomRecipe( - qfactory=nvfp4_linear_mixed_fp8_dpa_factory, + qfactory=nvfp4_linear_fp8_dpa_factory, fp8_dpa=True, ) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 8efa033865e..95c7c9f1b76 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -133,10 +133,10 @@ class TestComposerStyleFactory: @staticmethod def _factory(role): from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory, + nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory, ) - return nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory(role) + return nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory(role) @pytest.mark.parametrize( "tensor_type,row_scaled", @@ -152,7 +152,7 @@ def test_grouped_linear_forward_roles_use_nvfp4_rowwise_mxfp8_columnwise( ) assert isinstance(quantizer, HybridQuantizer) - assert quantizer.columnwise_source == "original" + assert quantizer.columnwise_source == "rowwise_dequantized" assert isinstance(quantizer.rowwise_quantizer, NVFP4Quantizer) assert isinstance(quantizer.columnwise_quantizer, MXFP8Quantizer) assert quantizer.rowwise_quantizer.row_scaled_nvfp4 is row_scaled @@ -1644,9 +1644,9 @@ def _run_model(self, model, inp, grad, fp8_recipe, seed): "case_name,native_dpa_recipe,qfactory_name", [ ( - "mixed_fp8_dpa", + "fp8_dpa", "Float8CurrentScaling", - "nvfp4_linear_mixed_fp8_dpa_factory", + "nvfp4_linear_fp8_dpa_factory", ), ( "mxfp8_dpa", @@ -1758,9 +1758,9 @@ def _run_mha_model(self, model, inp, grad, fp8_recipe, seed): "case_name,native_dpa_recipe,qfactory_name,expected_flags", [ ( - "mixed_fp8_dpa", + "fp8_dpa", "Float8CurrentScaling", - "nvfp4_linear_mixed_fp8_dpa_factory", + "nvfp4_linear_fp8_dpa_factory", (False, False, False), ), ( diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index 5dc5ec8579c..3c25f3cd9aa 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -257,10 +257,10 @@ def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( return _plain_nvfp4_quantizer() -def nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory( +def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( role: Optional[QuantizerRole], ): - """Quantizer factory: row-scaled NVFP4 forward, MXFP8 backward. + """Quantizer factory: row-scaled NVFP4 forward, dequantized-source MXFP8 backward. This RL-related recipe is inspired by the Composer 2 MoE grouped-GEMM recipe described in arXiv:2603.24477. @@ -275,26 +275,22 @@ def nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory( Dispatch logic: * ``GroupedLinear`` ``input`` -> - ``Hybrid(rowwise=row-scaled NVFP4, columnwise=MXFP8)`` + ``Hybrid(rowwise=row-scaled NVFP4, columnwise=MXFP8, + columnwise_source="rowwise_dequantized")`` * ``GroupedLinear`` ``weight`` -> - ``Hybrid(rowwise=plain NVFP4, columnwise=MXFP8)`` + ``Hybrid(rowwise=plain NVFP4, columnwise=MXFP8, + columnwise_source="rowwise_dequantized")`` * regular ``Linear`` -> MXFP8 * ``grad_output`` -> MXFP8 * everything else -> MXFP8 Row-scaled NVFP4 is fprop-only, so the forward NVFP4 quantizers avoid - RHT, stochastic rounding, and 2D scaling. By default, ``HybridQuantizer`` - leaves ``columnwise_source="original"``, so MXFP8 backward operands are - quantized from the original high-precision tensor. The paper does not - specify how Composer materializes MXFP8 backward operands from saved - training tensors. - - Users who want MXFP8 backward operands to include the forward NVFP4 - quantization source can construct the same ``HybridQuantizer`` with - ``columnwise_source="rowwise_dequantized"``. That makes columnwise MXFP8 - quantization consume the dequantized rowwise NVFP4 value instead of the - original high-precision tensor, which can be useful when the backward path - should reflect the forward quantization error (might be helpful for convergence). + RHT, stochastic rounding, and 2D scaling. This sample assumes MXFP8 + backward operands are quantized from the dequantized rowwise NVFP4 forward + value, expressed with ``columnwise_source="rowwise_dequantized"``. To + quantize MXFP8 backward operands from the original high-precision tensor + instead, construct the same ``HybridQuantizer`` while omitting + ``columnwise_source="rowwise_dequantized"``. Composer 2 Technical Report: https://arxiv.org/abs/2603.24477 @@ -307,11 +303,13 @@ def nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory( return HybridQuantizer( rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), columnwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_source="rowwise_dequantized", ) if is_grouped_linear and role.tensor_type == "weight": return HybridQuantizer( rowwise_quantizer=_plain_nvfp4_quantizer(), columnwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_source="rowwise_dequantized", ) if is_grouped_linear and role.tensor_type == "grad_output": return mxfp8_quantizer_factory(role) @@ -325,10 +323,10 @@ def nvfp4_row_scaled_fwd_mxfp8_bwd_quantizer_factory( # ----------------------------------------------------------------------------- -def nvfp4_linear_mixed_fp8_dpa_factory( +def nvfp4_linear_fp8_dpa_factory( role: Optional[QuantizerRole], ): - """Quantizer factory: NVFP4 for ``Linear``, mixed FP8 for ``DotProductAttention``. + """Quantizer factory: NVFP4 for ``Linear``, FP8 for ``DotProductAttention``. This factory demonstrates how to use ``CustomRecipe`` with ``fp8_dpa=True`` to combine NVFP4 quantization for linear layers with FP8 attention. @@ -365,11 +363,11 @@ def nvfp4_linear_mixed_fp8_dpa_factory( from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - nvfp4_linear_mixed_fp8_dpa_factory, + nvfp4_linear_fp8_dpa_factory, ) recipe = CustomRecipe( - qfactory=nvfp4_linear_mixed_fp8_dpa_factory, + qfactory=nvfp4_linear_fp8_dpa_factory, fp8_dpa=True, ) with autocast(recipe=recipe): @@ -384,7 +382,7 @@ def nvfp4_linear_mixed_fp8_dpa_factory( "dpa_output" in role.name or "dpa_grad_input" in role.name ) - # Native NVFP4 + FP8CS attention uses delayed scaling for S/dP. + # Native NVFP4 + FP8 attention uses delayed scaling for S/dP. if is_dpa and role.tensor_type in ("s", "dp"): return DelayedScalingRequest( fp8_format=Format.HYBRID, From 8c23fedbfd1ac3c1f22ac10f07a10b326a904d0b Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 30 Jun 2026 13:47:39 +0000 Subject: [PATCH 39/97] Resolve comments Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/conftest.py | 2 +- tests/pytorch/test_hybrid_quantization.py | 60 ++++++++----------- .../pytorch/attention/multi_head_attention.py | 2 +- .../pytorch/cpp_extensions/gemm.py | 14 ++--- transformer_engine/pytorch/module/base.py | 4 +- .../pytorch/module/layernorm_linear.py | 2 +- .../pytorch/module/layernorm_mlp.py | 2 +- transformer_engine/pytorch/module/linear.py | 2 +- .../pytorch/tensor/hybrid_tensor.py | 30 +++++++--- transformer_engine/pytorch/tensor/utils.py | 12 ++-- 10 files changed, 68 insertions(+), 62 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index 621972df33e..bbc5bd246de 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -53,7 +53,7 @@ def _check_nvfp4_support(): ] _HYBRID_FLOAT8_BLOCK_FSDP2_XFAIL_REASON = ( - "HybridFloat8BlockScaling + FSDP2 is not supported when dim-0 shards split " + "Tracked by #3158: HybridFloat8BlockScaling + FSDP2 is not supported when dim-0 shards split " "128-row Float8Block scale tiles. Each shard stores independently rounded " "scale rows, but the gathered tensor expects scale rows rounded from the " "global row count, so naively all-gathered scale buffers have the wrong shape." diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 95c7c9f1b76..083c445952a 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -14,6 +14,9 @@ from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, ) +from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_quantizer_factory, +) from transformer_engine.pytorch import ( autocast, quantized_model_init, @@ -3635,8 +3638,8 @@ def test_quantized_param_survives_multiple_forward_passes(self): # Covers the supported (same-format) cases and the rejected (cross-format, # missing sub-storage, unsupported sub-quantizer) cases. The supported subset # is the first incremental hybrid integration with the distributed-optimizer -# quantized-param all-gather flow. Cross-format support is deferred to a -# follow-up; the tests below pin the NotImplementedError contract so the +# quantized-param all-gather flow. Cross-format support is deferred to +# follow-up #3158; the tests below pin the NotImplementedError contract so the # rejection messaging stays clear as the feature evolves. # --------------------------------------------------------------------------- @@ -3735,7 +3738,7 @@ def _hybrid_recipe_fp8_current_row_delayed_col(): def _hybrid_recipe_mxfp8(): - """Same-format MXFP8 on both directions (rejected today; TODO).""" + """Same-format MXFP8 on both directions (rejected today; TODO #3158).""" return _hybrid_custom_recipe( row_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), @@ -3744,7 +3747,7 @@ def _hybrid_recipe_mxfp8(): def _hybrid_recipe_blockwise(): - """Same-format Float8Blockwise on both directions (rejected today; TODO).""" + """Same-format Float8Blockwise on both directions (rejected today; TODO #3158).""" return _hybrid_custom_recipe( row_factory=lambda: Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True @@ -4083,7 +4086,7 @@ def test_fp8_current_row_delayed_col_full_master(self): # of which direction they appear in. See the per-direction rejection tests # below (``test_mxfp8_*_raises`` covers both rowwise and columnwise rejection # of MXFP8; ``test_nvfp4_*_raises`` and ``test_blockwise_*_raises`` similarly). - # The TODO block above ``_route_hybrid_to_buckets`` in tensor/utils.py + # The TODO #3158 block above ``_route_hybrid_to_buckets`` in tensor/utils.py # documents the upstream constraints (single-direction cast helper / kernel # support) whose unblocker drops per-block format support in for free. @@ -4097,7 +4100,7 @@ def test_mxfp8_rowwise_raises(self): ``model_weight`` has BOTH ``_rowwise_*`` and ``_columnwise_*`` populated (the underlying partial-cast kernel is bidirectional), while a hybrid sub-storage is single-direction by construction. See - ``TODO(hybrid-mxfp8-distopt)`` in tensor/utils.py for the unblocker shape. + ``TODO(#3158, hybrid-mxfp8-distopt)`` in tensor/utils.py for the unblocker shape. """ from transformer_engine.pytorch.tensor.utils import quantize_master_weights @@ -4143,7 +4146,7 @@ def test_nvfp4_rowwise_raises(self): """NVFP4 in the rowwise sub-quantizer is rejected per-direction. The NVFP4 cast path is blocked on a pair of upstream constraints - documented in the TODO block above ``_route_hybrid_to_buckets`` in + documented in the TODO #3158 block above ``_route_hybrid_to_buckets`` in tensor/utils.py. NOTE: after PR #3027, single-direction 2D NVFP4 construction works, @@ -4176,7 +4179,7 @@ def test_nvfp4_rowwise_raises(self): ) def test_blockwise_rowwise_raises(self): """Float8BlockQuantizer in the rowwise sub-quantizer is rejected - per-direction (no e2e factory uses it; TODO marker in tensor/utils.py). + per-direction (no e2e factory uses it; TODO #3158 marker in tensor/utils.py). """ from transformer_engine.pytorch.tensor.utils import quantize_master_weights @@ -5014,35 +5017,22 @@ def test_equivalence(self): # --------------------------------------------------------------------------- -# Module-level qfactories (picklable, required for checkpoint serialization). - - -def _checkpoint_hybrid_fp8_qfactory(role): - """Module-level qfactory (picklable) for checkpoint tests.""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), - columnwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") +# Module-level qfactories give TE-to-TE quantized-param checkpoints a stable +# importable reference for any pickled quantizer/recipe metadata. Portable BF16 +# checkpoint loading should not depend on importing these factories. -@requires_fp8 +@requires_mxfp8_and_nvfp4 class TestHybridCheckpoint: """Test state_dict save/load round-trips for models with hybrid quantized params.""" - def _hybrid_fp8_recipe(self): - return recipe.CustomRecipe(qfactory=_checkpoint_hybrid_fp8_qfactory) + def _hybrid_checkpoint_recipe(self): + return recipe.CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_quantizer_factory) def test_state_dict_save_load_roundtrip(self): """state_dict → save → load → same model should produce identical outputs.""" torch.manual_seed(42) - hybrid_recipe = self._hybrid_fp8_recipe() + hybrid_recipe = self._hybrid_checkpoint_recipe() with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() @@ -5066,7 +5056,7 @@ def test_state_dict_save_load_roundtrip(self): def test_state_dict_contains_weight(self): """state_dict should contain the weight key.""" - hybrid_recipe = self._hybrid_fp8_recipe() + hybrid_recipe = self._hybrid_checkpoint_recipe() with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() @@ -5080,7 +5070,7 @@ def test_load_bf16_state_dict_into_hybrid_model(self): model initialized with quantized_model_init. """ torch.manual_seed(42) - hybrid_recipe = self._hybrid_fp8_recipe() + hybrid_recipe = self._hybrid_checkpoint_recipe() # Create BF16 model and get its state_dict ref_model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() @@ -5107,7 +5097,7 @@ def test_state_dict_torch_save_load(self): import os torch.manual_seed(42) - hybrid_recipe = self._hybrid_fp8_recipe() + hybrid_recipe = self._hybrid_checkpoint_recipe() with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(128, 128, params_dtype=torch.bfloat16).cuda() @@ -5140,7 +5130,7 @@ def test_checkpoint_resume_training(self): import os torch.manual_seed(42) - hybrid_recipe = self._hybrid_fp8_recipe() + hybrid_recipe = self._hybrid_checkpoint_recipe() with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -5903,7 +5893,7 @@ def test_scale_refresh_across_iterations(self): @pytest.mark.xfail( reason=( - "Hybrid FSDP2 does not support NVFP4 sub-storages yet; NVFP4 uses " + "Tracked by #3158: Hybrid FSDP2 does not support NVFP4 sub-storages yet; NVFP4 uses " "dedicated tensor hooks and does not implement the hybrid fsdp_buffer_fields protocol." ) ) @@ -6407,7 +6397,7 @@ def fn(model, inp): # / ``skip_fp8_weight_update`` correctly across the recompute phase, # which is why the ``te.checkpoint`` tests above pass bitwise. # - # Keeping the xfail'd tests here: + # Keeping the xfail'd tests here (tracked by #3158): # 1. pins the boundary — users hitting this failure get a clear # diagnosis and pointer to ``te.checkpoint``; # 2. becomes a regression signal if the underlying cache-vs- @@ -6426,7 +6416,7 @@ def fn(model, inp): " incompatible with TE's weight-workspace cache: cache-miss" " on the first forward saves a different tensor count than" " cache-hit on recompute. Use te.checkpoint instead (tested" - " above)." + " above). Tracked by #3158." ), ) diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 9fc6f6895be..e495c3831ab 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -58,7 +58,7 @@ def _infer_custom_dpa_scaling(recipe, dpa_name: str) -> Tuple[bool, bool]: qkv_quantizer = recipe.qfactory( QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) ) - except Exception: # pragma: no cover - preserve existing fallback behavior + except Exception: # pragma: no cover, pylint: disable=broad-exception-caught return False, False from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 7ce9bffaed5..753a76d1fb1 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -136,11 +136,11 @@ def _unwrap_hybrid_B(tensor, layout): return tensor.columnwise_sub_storage -def _materialize_high_precision(tensor): +def _unwrap_identity_tensor(tensor): """Replace an :class:`IdentityTensorStorage` operand with its plain tensor. Identity (high-precision passthrough) operands carry an unquantized tensor; - materializing it here routes the matmul through the standard high-precision + unwrapping it here routes the matmul through the standard high-precision GEMM path. Non-identity operands pass through unchanged. Called after the hybrid unwrap, so a high-precision *direction* of a hybrid tensor is handled too. @@ -154,7 +154,7 @@ def _reject_unsupported_output_quantizer(quantization_params): """Reject output quantizers that the native C++ GEMM path cannot convert.""" if isinstance(quantization_params, (HybridQuantizer, IdentityQuantizer)): quantizer_name = type(quantization_params).__name__ - # TODO(negvet): Lower HybridQuantizer output to its native rowwise + # TODO(#3158): Lower HybridQuantizer output to its native rowwise # sub-quantizer, and IdentityQuantizer output to an unquantized/no-op # path, once the returned tensor contract is defined for these boundary # roles. @@ -191,8 +191,8 @@ def general_gemm( transa = layout[0] == "T" transb = layout[1] == "T" - A = _materialize_high_precision(_unwrap_hybrid_A(A, layout)) - B = _materialize_high_precision(_unwrap_hybrid_B(B, layout)) + A = _unwrap_identity_tensor(_unwrap_hybrid_A(A, layout)) + B = _unwrap_identity_tensor(_unwrap_hybrid_B(B, layout)) alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) @@ -363,8 +363,8 @@ def general_grouped_gemm( """ num_gemms = len(A) - A = [_materialize_high_precision(_unwrap_hybrid_A(a, layout)) for a in A] - B = [_materialize_high_precision(_unwrap_hybrid_B(b, layout)) for b in B] + A = [_unwrap_identity_tensor(_unwrap_hybrid_A(a, layout)) for a in A] + B = [_unwrap_identity_tensor(_unwrap_hybrid_B(b, layout)) for b in B] transa = layout[0] == "T" transb = layout[1] == "T" diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 20514b68e76..c2eb11aafe6 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1059,7 +1059,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # Return early if recipe state matches recipe if self.fp8_meta_tensors_initialized: recipe_state = self.fp8_meta[fp8_meta_tensor_key] - # Follow-up: Match built-in recipes by full config, not just RecipeState type, so + # Follow-up (#3157): Match built-in recipes by full config, not just RecipeState type, so # same-class mid-training changes rebuild quantizers/workspaces correctly. if recipe.delayed() and isinstance(recipe_state, DelayedScalingRecipeState): self.adjust_amax_history_length(recipe.amax_history_len, fwd=fwd) @@ -1077,7 +1077,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return if recipe.custom() and isinstance(recipe_state, CustomRecipeState): - # Follow-up: Compare CustomRecipe/qfactory config here. qfactory changes made + # Follow-up (#3157): Compare CustomRecipe/qfactory config here. qfactory changes made # mid-training on the same recipe object currently do not take effect because # stale quantizers are reused. if recipe_state.recipe is recipe: diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a43c60269b7..9c457053543 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1583,7 +1583,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP # activations are gathered in high precision and re-quantized whole, so # every rank already sees the same global amax. - # TODO(negvet): once native quantized all-gather lands (see + # TODO(#3158): once native quantized all-gather lands (see # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path # quantizes per-shard, needing a hybrid branch here that mirrors the # current-scaling / NVFP4 SP reduction above: diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 96cec2c0215..6c5ae2e776a 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2188,7 +2188,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP # activations are gathered in high precision and re-quantized whole, so # every rank already sees the same global amax. - # TODO(negvet): once native quantized all-gather lands (see + # TODO(#3158): once native quantized all-gather lands (see # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path # quantizes per-shard, needing a hybrid branch here that mirrors the # current-scaling / NVFP4 SP reduction above: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 10878bac32a..722cf324859 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1786,7 +1786,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP # activations are gathered in high precision and re-quantized whole, so # every rank already sees the same global amax. - # TODO(negvet): once native quantized all-gather lands (see + # TODO(#3158): once native quantized all-gather lands (see # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path # quantizes per-shard, needing a hybrid branch here that mirrors the # current-scaling / NVFP4 SP reduction above: diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 56d0581cd07..5db53a9734f 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -58,6 +58,13 @@ class HybridQuantizer(Quantizer): columnwise_source="rowwise_dequantized", ) + In a ``CustomRecipe`` factory this can be combined with role-based + specialization. For example, a factory can return a ``HybridQuantizer`` only + for ``role.tensor_type == "weight"`` and use regular quantizers for inputs + and gradients. See + ``custom_recipes.quantization_factory_zoo.nvfp4_1d_double_quantized_weight_quantizer_factory`` + for a weight-only double-quantization example. + """ _COLUMNWISE_SOURCES = ("original", "rowwise_dequantized") @@ -81,6 +88,7 @@ def __init__( ("columnwise", columnwise_quantizer), ): if isinstance(quantizer, QuantizerRequest): + # TODO(#3158): Support delayed-scaling requests inside hybrid sub-quantizers. raise TypeError( "HybridQuantizer does not support nested QuantizerRequest " f"objects yet; got {type(quantizer).__name__} for the {role} " @@ -158,11 +166,19 @@ def with_amax_reduction(self, value: bool) -> None: @property def amax_reduction_group(self): """Amax-reduction group of the sub-quantizers, or ``None`` if unset.""" + result = None for sub in (self.rowwise_quantizer, self.columnwise_quantizer): group = getattr(sub, "amax_reduction_group", None) - if group is not None: - return group - return None + if group is None: + continue + if result is None: + result = group + elif group is not result: + raise RuntimeError( + "HybridQuantizer sub-quantizers have inconsistent " + "amax_reduction_group values." + ) + return result @amax_reduction_group.setter def amax_reduction_group(self, value) -> None: @@ -311,7 +327,7 @@ def supports_only_rowwise_all_gather(self) -> bool: (``False``) keeps the smaller, wgrad-ready columnwise shard saved — which is the more efficient memory choice. - TODO(negvet): Add native hybrid dispatch to + TODO(#3158): Add native hybrid dispatch to ``gather_along_first_dim`` to remove the BF16 detour. * **Scope.** Branch at the top of ``gather_along_first_dim`` that @@ -351,7 +367,7 @@ def supports_only_rowwise_all_gather(self) -> bool: return False def allows_save_original_input_for_backward(self) -> bool: - # TODO: Add an explicit recompute-from-original policy for deterministic + # TODO(#3158): Add an explicit recompute-from-original policy for deterministic # quantizers that can trade backward compute for saved activation memory. return self.columnwise_source == "original" @@ -363,7 +379,7 @@ def _get_compatible_recipe(self): # We trust that users who write a CustomRecipe know what they're doing # with regard to per-operand scaling mode compatibility. # - # TODO(negvet): validate per-operand scaling-mode compatibility at + # TODO(#3158): validate per-operand scaling-mode compatibility at # recipe-build time instead of at cuBLAS-dispatch time. Concretely: # 1. Walk the qfactory outputs for a given module_type (``linear``, # ``grouped_linear``, ``dpa``) — call the factory for each @@ -626,7 +642,7 @@ def fsdp_pre_all_gather( # pylint: disable=unused-argument format-specific padding (e.g. MXFP8 block-scale alignment) before the gather so concatenation along dim-0 is well-defined. - TODO(negvet): bandwidth optimization — pack both directions into a + TODO(#3158): bandwidth optimization — pack both directions into a single flat buffer sized ``max(row_bytes, col_bytes)`` (not ``row_bytes + col_bytes``) to halve comm volume for asymmetric format pairs. Planned implementation: a new per-sub-storage diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 1d9adbe3403..8bc4e6a15d2 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1162,7 +1162,7 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # Not supported (raise NotImplementedError per-direction + TODO): # # - MXFP8Quantizer as a hybrid sub-quantizer (any direction) -# TODO(hybrid-mxfp8-distopt): the distopt cast kernels +# TODO(#3158, hybrid-mxfp8-distopt): the distopt cast kernels # (`tex.mxfp8_scaling_compute_partial_amax`, `tex.mxfp8_scaling_partial_cast`) # are bidirectional — both rowwise and colwise outputs required — so they # cannot ingest a single-direction hybrid sub-storage. (Unrelated to the @@ -1173,7 +1173,7 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # Also unlocks cross-format MXFP8 row + col. # # - NVFP4Quantizer as a hybrid sub-quantizer (any direction) -# TODO(hybrid-nvfp4-distopt): load-bearing blocker is the kernel assertion +# TODO(#3158, hybrid-nvfp4-distopt): load-bearing blocker is the kernel assertion # `return_identity || !use_2d_quantization` in # `quantize_transpose_vector_blockwise_fp4.cu`, which rejects exactly the # columnwise-only 2D configuration that `HybridQuantizer.__init__` produces @@ -1190,7 +1190,7 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # `row_sub`'s gathered rowwise). # # - Float8BlockQuantizer as a hybrid sub-quantizer -# TODO(hybrid-fp8-blockwise): same shape as the NVFP4 secondary blocker — +# TODO(#3158, hybrid-fp8-blockwise): same shape as the NVFP4 secondary blocker — # `_cast_master_weights_to_fp8_blockwise_scaling` writes only `_rowwise_data` # with per-tensor post-AG `_create_columnwise()` that doesn't reach hybrid's # separate col sub-storage. Unlike NVFP4, there is no kernel-level @@ -1261,7 +1261,7 @@ def _route_hybrid_to_buckets( elif isinstance(sub_q, IdentityQuantizer): identity_params.append(entry) elif isinstance(sub_q, MXFP8Quantizer): - # TODO(hybrid-mxfp8-distopt): the distopt cast kernels are + # TODO(#3158, hybrid-mxfp8-distopt): the distopt cast kernels are # bidirectional, so a single-direction hybrid sub-storage cannot be # fed in. See top-of-file TODO block for the unblocker (single- # direction variants of the two distopt kernels). @@ -1271,7 +1271,7 @@ def _route_hybrid_to_buckets( "block above _route_hybrid_to_buckets for the unblocker shape." ) elif isinstance(sub_q, NVFP4Quantizer): - # TODO(hybrid-nvfp4-distopt): load-bearing blocker is the kernel + # TODO(#3158, hybrid-nvfp4-distopt): load-bearing blocker is the kernel # assertion that rejects columnwise-only 2D NVFP4 — which is # exactly what hybrid's col sub-quantizer pin produces. Secondary # blocker (gated on the kernel fix) is the per-tensor post-AG @@ -1283,7 +1283,7 @@ def _route_hybrid_to_buckets( "block above _route_hybrid_to_buckets for details." ) elif isinstance(sub_q, Float8BlockQuantizer): - # Pending hybrid-fp8-blockwise work: same shape as the NVFP4 + # Pending hybrid-fp8-blockwise work (#3158): same shape as the NVFP4 # secondary blocker (and only that one — no kernel-level construction # blocker for Block FP8). Python-side post-AG fix. See the # _route_hybrid_to_buckets design note above for details. From 3a8ab2272f3983a49aa6e1bd3b9ea21377b1eaea Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 30 Jun 2026 13:49:25 +0000 Subject: [PATCH 40/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 11 ++++------- .../dot_product_attention/dot_product_attention.py | 1 + .../custom_recipes/quantization_factory_zoo.py | 14 +++++++++----- transformer_engine/pytorch/tensor/hybrid_tensor.py | 3 +-- 4 files changed, 15 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 083c445952a..5d7f953853f 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -1617,10 +1617,9 @@ def _clear_native_dpa_recipe(monkeypatch): @staticmethod def _assert_equal(actual, expected, label): - assert torch.equal(actual, expected), ( - f"{label} mismatch: max diff = " - f"{(actual.float() - expected.float()).abs().max().item()}" - ) + assert torch.equal( + actual, expected + ), f"{label} mismatch: max diff = {(actual.float() - expected.float()).abs().max().item()}" def _run_model(self, model, inp, grad, fp8_recipe, seed): run_inp = inp.clone().detach().requires_grad_(True) @@ -1629,9 +1628,7 @@ def _run_model(self, model, inp, grad, fp8_recipe, seed): with autocast(enabled=True, recipe=fp8_recipe): out = model(run_inp) out.backward(grad) - local_recipes = [ - type(r).__name__ for r in model.dpa.fp8_meta.get("local_recipes", []) - ] + local_recipes = [type(r).__name__ for r in model.dpa.fp8_meta.get("local_recipes", [])] return ( out.detach().clone(), run_inp.grad.detach().clone(), diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index cbe2a88281b..156b021a214 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -154,6 +154,7 @@ def _delayed_scaling_recipe() -> Optional[DelayedScaling]: return None + """ This feature is **experimental** and subject to change. diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index 3c25f3cd9aa..3eb02ddcbf1 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -15,7 +15,7 @@ Factories are ordered from conservative to more aggressive quantization. Organization: - * Linear / grouped-linear recipes (pre-training). Favor more precision + * Linear / grouped-linear recipes (pre-training). Favor more precision on the forward pass. * RL-oriented recipes: Favor more precision in backward GEMMs. * Linear + attention recipes: factories that also cover ``DotProductAttention`` @@ -378,8 +378,10 @@ def nvfp4_linear_fp8_dpa_factory( from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer is_dpa = role is not None and role.module_type == "dpa" - is_dpa_boundary = role is not None and not role.module_type and ( - "dpa_output" in role.name or "dpa_grad_input" in role.name + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) ) # Native NVFP4 + FP8 attention uses delayed scaling for S/dP. @@ -456,8 +458,10 @@ def nvfp4_linear_mxfp8_dpa_factory( from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer is_dpa = role is not None and role.module_type == "dpa" - is_dpa_boundary = role is not None and not role.module_type and ( - "dpa_output" in role.name or "dpa_grad_input" in role.name + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) ) if is_dpa or is_dpa_boundary: diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 5db53a9734f..58cc878bed3 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -175,8 +175,7 @@ def amax_reduction_group(self): result = group elif group is not result: raise RuntimeError( - "HybridQuantizer sub-quantizers have inconsistent " - "amax_reduction_group values." + "HybridQuantizer sub-quantizers have inconsistent amax_reduction_group values." ) return result From 81ffa7ea5b4e68815b0ef6e3bd419926a833ee03 Mon Sep 17 00:00:00 2001 From: root Date: Tue, 7 Jul 2026 06:20:35 -0700 Subject: [PATCH 41/97] Add MHA boudary quantization tests Signed-off-by: root --- tests/pytorch/test_hybrid_quantization.py | 159 ++++++++++++++++++ .../pytorch/attention/multi_head_attention.py | 8 +- 2 files changed, 163 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 5d7f953853f..90024c5ef20 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -1874,6 +1874,165 @@ def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) f"{case_name} MHA param grad {name}", ) + def test_update_output_quantizer_roles_wires_independent_boundaries(self): + from transformer_engine.pytorch.quantization import QuantizerRole + + model = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + qkv = model.layernorm_qkv if model.input_layernorm else model.qkv + + expected_qkv = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=model.core_attention.name or "", + ) + expected_do = QuantizerRole( + module_type="dpa", + tensor_type="do", + name=model.core_attention.name or "", + ) + expected_o = QuantizerRole( + module_type="linear", + tensor_type="input", + name=model.proj.name or "", + ) + expected_dqkv = QuantizerRole( + module_type="linear", + tensor_type="grad_output", + name=qkv.name or "", + ) + + def boundary_roles(): + return ( + qkv.output_quantizer_role, + model.proj.grad_input_quantizer_role, + model.core_attention.output_quantizer_role, + model.core_attention.grad_input_quantizer_role, + ) + + model._update_output_quantizer_roles(True, False, False) + assert boundary_roles() == (expected_qkv, None, None, None) + + model._update_output_quantizer_roles(False, True, False) + assert boundary_roles() == (None, expected_do, None, None) + + model._update_output_quantizer_roles(False, False, True) + assert boundary_roles() == (None, None, expected_o, expected_dqkv) + + model._update_output_quantizer_roles(False, False, False) + assert boundary_roles() == (None, None, None, None) + + def test_mxfp8_qfactory_uses_plain_bf16_mha_boundaries(self, monkeypatch): + """MXFP8 DPA stays internal; MHA boundary tensors remain plain BF16.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.attention.dot_product_attention import ( + backends as dpa_backends, + ) + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_linear_mxfp8_dpa_factory, + ) + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + self._clear_native_dpa_recipe(monkeypatch) + torch.manual_seed(2401) + model = te.MultiheadAttention( + self.hidden_size, + self.num_heads, + kv_channels=self.kv_channels, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="mha", + ).cuda() + + boundary_tensors = {} + saved_dpa_tensors = {} + orig_prepare_for_saving = dpa_backends.prepare_for_saving + + def _record_prepare_for_saving(*tensors, **kwargs): + saved_dpa_tensors["fp8_o"] = tensors[3] + saved_dpa_tensors["f16_o"] = tensors[7] + return orig_prepare_for_saving(*tensors, **kwargs) + + monkeypatch.setattr(dpa_backends, "prepare_for_saving", _record_prepare_for_saving) + + def _record_grad(name): + def _hook(grad): + boundary_tensors[name] = grad + return grad + + return _hook + + def _record_dpa_inputs(_module, inputs): + for name, tensor in zip(("q", "k", "v"), inputs[:3]): + boundary_tensors[name] = tensor + tensor.register_hook(_record_grad(f"d{name}")) + + def _record_dpa_output(_module, _inputs, output): + boundary_tensors["o"] = output + output.register_hook(_record_grad("do")) + + def _record_projection_input(_module, inputs): + boundary_tensors["proj_input"] = inputs[0] + boundary_tensors["o_is_proj_input"] = boundary_tensors["o"] is inputs[0] + + handles = ( + model.core_attention.register_forward_pre_hook(_record_dpa_inputs), + model.core_attention.register_forward_hook(_record_dpa_output), + model.proj.register_forward_pre_hook(_record_projection_input), + ) + + torch.manual_seed(2402) + inp = torch.randn( + self.seq_len, + self.batch, + self.hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + grad = torch.randn_like(inp) + qfactory_recipe = recipe.CustomRecipe( + qfactory=nvfp4_linear_mxfp8_dpa_factory, + fp8_dpa=True, + ) + try: + self._run_mha_model(model, inp, grad, qfactory_recipe, seed=2403) + finally: + for handle in handles: + handle.remove() + + assert boundary_tensors["o_is_proj_input"] + assert saved_dpa_tensors["fp8_o"] is None + saved_o = saved_dpa_tensors["f16_o"] + assert type(saved_o) is torch.Tensor + assert saved_o.dtype is torch.bfloat16 + assert ( + saved_o.untyped_storage().data_ptr() + == boundary_tensors["o"].untyped_storage().data_ptr() + ) + + for name in ("q", "k", "v", "o", "proj_input", "dq", "dk", "dv", "do"): + tensor = boundary_tensors[name] + assert type(tensor) is torch.Tensor, f"{name} is {type(tensor)}" + assert tensor.dtype is torch.bfloat16, f"{name} has dtype {tensor.dtype}" + @pytest.mark.skipif(not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling) class TestHybridGemmBitwiseIdenticalBlockFP8: diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index e495c3831ab..2343033fee1 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -526,11 +526,11 @@ def _update_output_quantizer_roles( 1. ``qkv_fp8_output`` — **QKV linear → DPA (fwd)**: the QKV linear's ``output_quantizer_role`` is told its consumer is DPA. - 2. ``proj_fp8_grad`` — **Proj linear ← DPA (bwd)**: proj's - ``grad_input_quantizer_role`` is told its producer is DPA. + 2. ``proj_fp8_grad`` — **Proj linear → DPA (bwd)**: proj's + ``grad_input_quantizer_role`` is told its consumer is DPA. 3. ``dpa_fp8_output`` — **DPA → Proj linear (fwd)**: DPA's ``output_quantizer_role`` is told its consumer is the proj linear. - 4. ``dpa_fp8_output`` — **DPA ← QKV linear (bwd)**: DPA's + 4. ``dpa_fp8_output`` — **DPA → QKV linear (bwd)**: DPA's ``grad_input_quantizer_role`` is told its consumer is QKV linear. When a flag is ``False`` the corresponding role is reset to ``None`` @@ -556,7 +556,7 @@ def _update_output_quantizer_roles( self.query_layer.output_quantizer_role = qkv_output_role self.key_value.output_quantizer_role = qkv_output_role - # ── Boundary 2 (bwd): Proj grad-input ← produced by DPA ────────── + # ── Boundary 2 (bwd): Proj grad-input (dO) → consumed by DPA ───── proj_grad_input_role = ( QuantizerRole(module_type="dpa", tensor_type="do", name=dpa_name) if proj_fp8_grad From 4e54f0203ba2bbf47abad647dc28022011f7e007 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 13 Jul 2026 15:15:34 +0000 Subject: [PATCH 42/97] Improve numerical correctness of the tests Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/conftest.py | 14 - .../fsdp2_tests/run_fsdp2_fused_adam.py | 785 ++++++++-- tests/pytorch/distributed/run_hybrid_tp_sp.py | 190 ++- .../pytorch/distributed/test_hybrid_tp_sp.py | 6 +- tests/pytorch/test_cpu_offloading.py | 59 +- tests/pytorch/test_cpu_offloading_v1.py | 94 ++ tests/pytorch/test_hybrid_quantization.py | 1307 +++++++++++------ tests/pytorch/test_identity_quantizer.py | 81 +- 8 files changed, 1906 insertions(+), 630 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index bbc5bd246de..91627ab7397 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -52,13 +52,6 @@ def _check_nvfp4_support(): ("HybridMixed_MXFP8_FP8", fp8.check_mxfp8_support), ] -_HYBRID_FLOAT8_BLOCK_FSDP2_XFAIL_REASON = ( - "Tracked by #3158: HybridFloat8BlockScaling + FSDP2 is not supported when dim-0 shards split " - "128-row Float8Block scale tiles. Each shard stores independently rounded " - "scale rows, but the gathered tensor expects scale rows rounded from the " - "global row count, so naively all-gathered scale buffers have the wrong shape." -) - def _parametrize_recipes(): params = [] @@ -75,13 +68,6 @@ def _parametrize_hybrid_recipes(): for name, check_fn in _HYBRID_RECIPE_CONFIGS: supported, reason = check_fn() marks = [pytest.mark.skipif(not supported, reason=reason)] - if name == "HybridFloat8BlockScaling": - marks.append( - pytest.mark.xfail( - raises=RuntimeError, - reason=_HYBRID_FLOAT8_BLOCK_FSDP2_XFAIL_REASON, - ) - ) params.append(pytest.param(name, id=name, marks=marks)) return params diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 4fa4d232029..1169215e2f0 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -44,12 +44,16 @@ from torch.distributed.tensor import DTensor import transformer_engine.pytorch as te -from transformer_engine.pytorch import QuantizedTensor +from transformer_engine.pytorch import HybridQuantizedTensor, QuantizedTensor +from transformer_engine.pytorch.tensor import ( + Float8BlockwiseQTensorStorage, + Float8TensorStorage, + MXFP8TensorStorage, +) import transformer_engine.common.recipe from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs - HIDDEN_SIZE = 256 FFN_HIDDEN_SIZE = 1024 NUM_ATTENTION_HEADS = 8 @@ -160,6 +164,185 @@ def _get_dist_info(): return world_size, device +def _collective_assert(condition, message): + """Raise on every rank only after all ranks report structural status.""" + failed = torch.tensor([not condition], dtype=torch.uint8, device="cuda") + if not condition: + print(f"[rank {dist.get_rank()}] {message}", flush=True) + dist.all_reduce(failed, dist.ReduceOp.MAX) + assert not bool(failed.item()), f"{message}: failed on at least one rank" + + +def _collective_assert_same_shape(tensor, message): + """Check tensor ranks/shapes before entering a shape-sensitive collective.""" + # Quantized parameter buffers in these tests are at most 2-D, but leave + # room for future formats without introducing object collectives. + layout = torch.full((9,), -1, dtype=torch.int64, device=tensor.device) + _collective_assert(tensor.ndim < layout.numel(), f"{message}: ndim={tensor.ndim}") + layout[0] = tensor.ndim + layout[1 : tensor.ndim + 1] = torch.tensor( + tensor.shape, dtype=torch.int64, device=tensor.device + ) + min_layout = layout.clone() + max_layout = layout.clone() + dist.all_reduce(min_layout, dist.ReduceOp.MIN) + dist.all_reduce(max_layout, dist.ReduceOp.MAX) + _collective_assert( + torch.equal(min_layout, max_layout), + f"{message}: tensor shape differs across ranks", + ) + + +def _record_exact(errors, actual, expected, label): + """Record an exact mismatch without interrupting later distributed work.""" + if actual is None or expected is None: + if actual is not expected: + errors.append(f"{label}: one tensor is None") + return + try: + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + except (AssertionError, TypeError) as exc: + errors.append(f"{label}: {exc}") + + +def _raise_collective_errors(errors, context): + """Aggregate delayed exact failures after every rank finishes collectives.""" + if errors: + print(f"[rank {dist.get_rank()}] {context}:\n" + "\n".join(errors), flush=True) + failed = torch.tensor([bool(errors)], dtype=torch.uint8, device="cuda") + dist.all_reduce(failed, dist.ReduceOp.MAX) + assert not bool(failed.item()), f"{context}: exact check failed on at least one rank" + + +def _check_hybrid_direction_buffers( + local_sub, + full_sub, + expected_type, + *, + direction, + param_name, + world_size, + errors, +): + """Compare one direction's gathered raw data, scale buffers, and metadata.""" + _collective_assert( + isinstance(local_sub, expected_type), + f"{param_name}: {direction} local storage is {type(local_sub).__name__}, " + f"expected {expected_type.__name__}", + ) + _collective_assert( + isinstance(full_sub, expected_type), + f"{param_name}: {direction} full storage is {type(full_sub).__name__}, " + f"expected {expected_type.__name__}", + ) + + local_buffers, local_meta = local_sub.fsdp_extract_buffers() + full_buffers, full_meta = full_sub.fsdp_extract_buffers() + local_count = torch.tensor([len(local_buffers)], dtype=torch.int64, device="cuda") + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() == len(full_buffers), + f"{param_name}: {direction} buffer count mismatch: local={len(local_buffers)}, " + f"full={len(full_buffers)}, rank range=({min_count.item()}, {max_count.item()})", + ) + + expected_buffers = [] + for buffer in local_buffers: + _collective_assert( + buffer is not None, + f"{param_name}: {direction} unexpectedly exposed a None FSDP buffer", + ) + _collective_assert_same_shape(buffer, f"{param_name}: {direction} FSDP buffer") + gathered = [torch.zeros_like(buffer) for _ in range(world_size)] + dist.all_gather(gathered, buffer) + expected_buffers.append(torch.cat(gathered, dim=0)) + + for key in ("field_names", "direction"): + if key in local_meta or key in full_meta: + if local_meta.get(key) != full_meta.get(key): + errors.append( + f"{param_name}: {direction} metadata {key!r} differs: " + f"{full_meta.get(key)!r} != {local_meta.get(key)!r}" + ) + for index, (actual, expected) in enumerate(zip(full_buffers, expected_buffers)): + field_names = local_meta.get("field_names", ()) + field = field_names[index] if index < len(field_names) else f"buffer[{index}]" + _record_exact(errors, actual, expected, f"{param_name}: {direction} {field}") + + # Per-tensor FP8 scale is metadata rather than an FSDP buffer. It must be + # identical on every shard and preserved on the reconstructed full tensor. + if isinstance(local_sub, Float8TensorStorage): + local_scale_inv = getattr(local_sub, "_scale_inv", None) + _collective_assert( + isinstance(local_scale_inv, torch.Tensor), + f"{param_name}: {direction} Float8 storage has no tensor _scale_inv", + ) + local_scale = local_scale_inv.detach().clone() + _collective_assert_same_shape(local_scale, f"{param_name}: {direction} Float8 scale") + gathered_scales = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered_scales, local_scale) + for rank, scale in enumerate(gathered_scales[1:], start=1): + _record_exact( + errors, + scale, + gathered_scales[0], + f"{param_name}: {direction} scale rank {rank}", + ) + _record_exact( + errors, + full_sub._scale_inv, + gathered_scales[0], + f"{param_name}: {direction} reconstructed scale", + ) + + +def _manual_reconstruct_hybrid(local, *, param_name, world_size): + """Run the Hybrid FSDP pre/post protocol on manually gathered raw buffers.""" + _collective_assert( + isinstance(local, HybridQuantizedTensor), + f"{param_name}: local shard is {type(local).__name__}, expected HybridQuantizedTensor", + ) + sharded_tensors, metadata = local.fsdp_pre_all_gather( + mesh=None, + orig_size=local.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + local_count = torch.tensor([len(sharded_tensors)], dtype=torch.int64, device=local.device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{param_name}: Hybrid FSDP buffer count range is ({min_count.item()}, {max_count.item()})", + ) + + gathered_outputs = [] + for index, shard in enumerate(sharded_tensors): + _collective_assert( + shard is not None, + f"{param_name}: Hybrid FSDP buffer {index} is None", + ) + _collective_assert_same_shape(shard, f"{param_name}: Hybrid FSDP buffer {index}") + gathered = [torch.zeros_like(shard) for _ in range(world_size)] + dist.all_gather(gathered, shard) + gathered_outputs.append(torch.cat(gathered, dim=0)) + + reconstructed, _ = local.fsdp_post_all_gather( + tuple(gathered_outputs), metadata, local.dtype, out=None + ) + _collective_assert( + isinstance(reconstructed, HybridQuantizedTensor), + f"{param_name}: Hybrid FSDP reconstruction returned {type(reconstructed).__name__}", + ) + return reconstructed + + def test_fused_adam_fp8_master_weights(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init (FP8 params). @@ -1150,7 +1333,8 @@ def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name): invoking ``fsdp_post_all_gather(out=...)`` twice per step -- while ``False`` keeps the gathered copy alive through backward (one gather per step). The gathered quantized bytes are identical either way, so both schedules must - produce **bitwise-identical** loss trajectories. + produce bitwise-identical outputs, losses, input/parameter gradients, + master parameters, moments, and step counters. Strictly stronger than "loss decreased": it locks in that the hybrid all-gather hooks are schedule-invariant across both FSDP2 passes, and @@ -1165,6 +1349,26 @@ def test_fused_adam_hybrid_reshard_variants(hybrid_recipe_name): x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) target = torch.randn_like(x) + def assert_exact(actual, expected, path): + if torch.is_tensor(expected): + torch.testing.assert_close( + actual, + expected, + rtol=0.0, + atol=0.0, + msg=lambda m: f"reshard schedule changed {path}: {m}", + ) + elif isinstance(expected, dict): + assert actual.keys() == expected.keys(), f"{path}: keys differ" + for key in expected: + assert_exact(actual[key], expected[key], f"{path}.{key}") + elif isinstance(expected, (list, tuple)): + assert len(actual) == len(expected), f"{path}: lengths differ" + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + assert_exact(actual_item, expected_item, f"{path}[{index}]") + else: + assert actual == expected, f"{path}: {actual!r} != {expected!r}" + def run(reshard_after_forward): # Re-seed so both schedules get identical weight init from reset_parameters(). torch.manual_seed(42) @@ -1180,29 +1384,53 @@ def run(reshard_after_forward): master_weights=True, master_weight_dtype=torch.float32, ) - losses = [] + run_x = x.detach().clone().requires_grad_() + artifacts = { + "outputs": [], + "losses": [], + "input_grads": [], + "param_grads": [], + "optimizer": [], + } for _ in range(NUM_STEPS): optimizer.zero_grad(set_to_none=True) + run_x.grad = None with te.autocast(enabled=True, recipe=hybrid_recipe): - output = model(x) + output = model(run_x) + artifacts["outputs"].append(output.detach().clone()) loss = F.mse_loss(output, target) - losses.append(loss.item()) + artifacts["losses"].append(loss.detach().clone()) loss.backward() + artifacts["input_grads"].append(run_x.grad.detach().clone()) + step_grads = [] + for param in model.parameters(): + grad = param.grad + if grad is not None: + grad = grad.to_local() if isinstance(grad, DTensor) else grad + grad = grad.detach().clone() + step_grads.append(grad) + artifacts["param_grads"].append(step_grads) optimizer.step() - return losses - - losses_resharded = run(reshard_after_forward=True) # re-gather in backward - losses_kept = run(reshard_after_forward=False) # keep gathered weight through backward - - # Both schedules must train (strictly monotonic decrease) ... + step_state = [] + for param in model.parameters(): + param_state = {} + for key, value in optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + param_state[key] = value + step_state.append(param_state) + artifacts["optimizer"].append(step_state) + return artifacts + + artifacts_resharded = run(reshard_after_forward=True) # re-gather in backward + artifacts_kept = run(reshard_after_forward=False) # keep gathered weight through backward + + losses_resharded = [loss.item() for loss in artifacts_resharded["losses"]] assert all( losses_resharded[i + 1] < losses_resharded[i] for i in range(NUM_STEPS - 1) ), f"reshard_after_forward=True loss not strictly decreasing: {losses_resharded}" - # ... and be numerically identical, since reshard is a schedule, not a math change. - assert losses_resharded == losses_kept, ( - "reshard_after_forward changed numerics (must be schedule-invariant): " - f"True={losses_resharded} vs False={losses_kept}" - ) + assert_exact(artifacts_resharded, artifacts_kept, "training") def test_fused_adam_hybrid_bf16_vs_hybrid_parity(hybrid_recipe_name): @@ -1288,10 +1516,11 @@ def test_fused_adam_hybrid_vs_base_recipe_parity(hybrid_recipe_name): """Same-format hybrid must match its vanilla recipe bitwise through the full FSDP2 + FusedAdam loop. - Forward output, weight gradients, and per-step loss are all asserted - bitwise-identical every iteration -- regression guard for the amax-reduction - fix. Uses a bare ``te.Linear`` stack (see ``_build_linear_parity_stack``) to - isolate GEMM-operand quantization. + Every output, loss, input/parameter gradient, master parameter, optimizer + moment, and step counter is asserted bitwise-identical -- a regression guard + for both amax reduction and master-weight requantization. Uses a bare + ``te.Linear`` stack (see ``_build_linear_parity_stack``) to isolate + GEMM-operand quantization. """ if hybrid_recipe_name not in _HYBRID_TO_BASE_RECIPE: pytest.skip( @@ -1310,6 +1539,18 @@ def test_fused_adam_hybrid_vs_base_recipe_parity(hybrid_recipe_name): target = torch.randn_like(x) def run_training(build_fn, recipe_for_autocast): + def snapshot_optimizer(model, optimizer): + snapshots = [] + for param in model.parameters(): + state = {} + for key, value in optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + state[key] = value + snapshots.append(state) + return snapshots + # Re-seed so both models get identical init from reset_parameters() (run # after sharding); with same-format quantization and a dropout-free loop # the full trajectory is then deterministic. @@ -1322,18 +1563,22 @@ def run_training(build_fn, recipe_for_autocast): master_weights=True, master_weight_dtype=torch.float32, ) - first_output = None + run_x = x.detach().clone().requires_grad_() + outputs = [] losses = [] + input_grads = [] grads_per_step = [] + optimizer_states = [] for step in range(NUM_STEPS): optimizer.zero_grad(set_to_none=True) + run_x.grad = None with te.autocast(enabled=True, recipe=recipe_for_autocast): - output = model(x) - if step == 0: - first_output = output.detach().clone() + output = model(run_x) + outputs.append(output.detach().clone()) loss = F.mse_loss(output, target) losses.append(loss.detach().clone()) loss.backward() + input_grads.append(run_x.grad.detach().clone()) # Snapshot grad local shards before the optimizer consumes them # (p.grad is a DTensor under FSDP2) to assert backward parity directly. step_grads = [] @@ -1346,26 +1591,32 @@ def run_training(build_fn, recipe_for_autocast): step_grads.append(g.detach().clone()) grads_per_step.append(step_grads) optimizer.step() - return first_output, losses, grads_per_step + optimizer_states.append(snapshot_optimizer(model, optimizer)) + return outputs, losses, input_grads, grads_per_step, optimizer_states base_recipe = get_recipe_from_string(base_recipe_name) hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) - base_first, base_losses, base_grads = run_training( + base_outputs, base_losses, base_input_grads, base_grads, base_opt_states = run_training( lambda: _build_linear_parity_stack(base_recipe), base_recipe ) - hybrid_first, hybrid_losses, hybrid_grads = run_training( - lambda: _build_linear_parity_stack(hybrid_recipe), hybrid_recipe - ) - - # (1) First forward: bitwise-identical (the core operand-equivalence check). - torch.testing.assert_close( - hybrid_first, - base_first, - rtol=0.0, - atol=0.0, - msg=lambda m: f"[{hybrid_recipe_name} vs {base_recipe_name}] first forward output not bitwise-identical (a same-format hybrid must match its vanilla recipe before any optimizer step): {m}", - ) + ( + hybrid_outputs, + hybrid_losses, + hybrid_input_grads, + hybrid_grads, + hybrid_opt_states, + ) = run_training(lambda: _build_linear_parity_stack(hybrid_recipe), hybrid_recipe) + + # (1) Every forward: bitwise-identical before and after optimizer updates. + for step, (base_output, hybrid_output) in enumerate(zip(base_outputs, hybrid_outputs)): + torch.testing.assert_close( + hybrid_output, + base_output, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} forward output not bitwise-identical: {m}", + ) # (2) Every per-step loss: bitwise-identical across the whole optimizer loop. for step, (b_loss, h_loss) in enumerate(zip(base_losses, hybrid_losses)): @@ -1378,6 +1629,15 @@ def run_training(build_fn, recipe_for_autocast): ) # (3) Backward: every weight-gradient shard at every step bitwise-identical + for step, (base_grad, hybrid_grad) in enumerate(zip(base_input_grads, hybrid_input_grads)): + torch.testing.assert_close( + hybrid_grad, + base_grad, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} input gradient not bitwise-identical: {m}", + ) + # (implied by the loss trajectory, but asserted directly to be explicit). for step, (b_step, h_step) in enumerate(zip(base_grads, hybrid_grads)): for i, (b_grad, h_grad) in enumerate(zip(b_step, h_step)): @@ -1395,6 +1655,33 @@ def run_training(build_fn, recipe_for_autocast): msg=lambda m, s=step, i=i: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i} gradient not bitwise-identical to the vanilla recipe: {m}", ) + # Optimizer continuation state: FP32 master params, moments, and counters. + assert len(base_opt_states) == len(hybrid_opt_states) + for step, (base_step, hybrid_step) in enumerate(zip(base_opt_states, hybrid_opt_states)): + assert len(base_step) == len(hybrid_step) + for param_idx, (base_state, hybrid_state) in enumerate(zip(base_step, hybrid_step)): + assert base_state.keys() == hybrid_state.keys(), ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} param {param_idx} " + f"optimizer state keys differ: {base_state.keys()} != {hybrid_state.keys()}" + ) + for key in base_state: + base_value = base_state[key] + hybrid_value = hybrid_state[key] + if torch.is_tensor(base_value): + torch.testing.assert_close( + hybrid_value, + base_value, + rtol=0.0, + atol=0.0, + msg=lambda m, s=step, i=param_idx, k=key: f"[{hybrid_recipe_name} vs {base_recipe_name}] step {s} param {i} optimizer state {k!r} not bitwise-identical: {m}", + ) + else: + assert hybrid_value == base_value, ( + f"[{hybrid_recipe_name} vs {base_recipe_name}] step {step} " + f"param {param_idx} optimizer state {key!r} differs: " + f"{hybrid_value!r} != {base_value!r}" + ) + def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): """Per-tensor hybrid weights must share ONE amax-reduced scale across FSDP2 @@ -1418,29 +1705,32 @@ def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name): model = _build_hybrid_model(hybrid_recipe) model = _shard_model(model, world_size) - checked = 0 + checked = {"rowwise": 0, "columnwise": 0} for name, param in model.named_parameters(): if not ( isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) ): continue - row_sub = param._local_tensor._rowwise_storage - scale_inv = getattr(row_sub, "_scale_inv", None) if row_sub is not None else None - if scale_inv is None: - continue - local_scale = scale_inv.detach().reshape(-1).clone() - gathered = [torch.zeros_like(local_scale) for _ in range(world_size)] - dist.all_gather(gathered, local_scale) - for r in range(1, world_size): - torch.testing.assert_close( - gathered[r], - gathered[0], - rtol=0.0, - atol=0.0, - msg=lambda m, n=name, r=r: f"{n}: rank {r} rowwise _scale_inv differs from rank 0 -- cross-shard amax reduction was not applied to the hybrid current-scaling weight: {m}", - ) - checked += 1 - assert checked > 0, "no hybrid current-scaling weights found to check" + for direction in checked: + sub_storage = getattr(param._local_tensor, f"_{direction}_storage") + scale_inv = getattr(sub_storage, "_scale_inv", None) + if scale_inv is None: + continue + local_scale = scale_inv.detach().reshape(-1).clone() + gathered = [torch.zeros_like(local_scale) for _ in range(world_size)] + dist.all_gather(gathered, local_scale) + for r in range(1, world_size): + torch.testing.assert_close( + gathered[r], + gathered[0], + rtol=0.0, + atol=0.0, + msg=lambda m, n=name, r=r, d=direction: f"{n}: rank {r} {d} _scale_inv differs from rank 0 -- cross-shard amax reduction was not applied to the hybrid current-scaling weight: {m}", + ) + checked[direction] += 1 + assert all( + count > 0 for count in checked.values() + ), f"missing hybrid current-scaling directions: {checked}" def test_fused_adam_hybrid_identity_fp8_master_weights(): @@ -1551,19 +1841,24 @@ def test_fused_adam_hybrid_identity_fp8_master_weights(): def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): - """Validate that FSDP2 all-gather + post-gather reconstruction produces - correct results by comparing ``unshard(param).dequantize()`` with a manual - all-gather of dequantized local shards. - - For the stateless formats supported here (per-tensor FP8, MXFP8), FSDP2's - all-gather concatenates rowwise bytes along dim-0 and the per-block / - per-tensor scales follow the same concatenation. Dequantizing the gathered - bytes is therefore bitwise-identical to concatenating the dequantized - shards — the tolerance is effectively ``assert_equal``. + """Validate both Hybrid directions through FSDP2 at raw-buffer precision. + + The rowwise and columnwise sub-storages are checked independently. Every + gathered data buffer, scale buffer, field layout, and per-tensor scale + metadata value must exactly match a manual dim-0 all-gather. """ - from transformer_engine.pytorch import HybridQuantizedTensor from fsdp2_utils import get_hybrid_recipe_from_string + expected_types = { + "HybridFP8CurrentScaling": (Float8TensorStorage, Float8TensorStorage), + "HybridMXFP8": (MXFP8TensorStorage, MXFP8TensorStorage), + "HybridFloat8BlockScaling": ( + Float8BlockwiseQTensorStorage, + Float8BlockwiseQTensorStorage, + ), + "HybridMixed_MXFP8_FP8": (MXFP8TensorStorage, Float8TensorStorage), + } + row_type, col_type = expected_types[hybrid_recipe_name] hybrid_recipe = get_hybrid_recipe_from_string(hybrid_recipe_name) world_size, device = _get_dist_info() @@ -1574,34 +1869,53 @@ def test_fused_adam_hybrid_allgather_correctness(hybrid_recipe_name): with te.autocast(enabled=True, recipe=hybrid_recipe): _ = model(x) - checked = 0 - for name, param in model.named_parameters(): - if not (isinstance(param, DTensor) and isinstance(param._local_tensor, QuantizedTensor)): - continue - local_shard = param._local_tensor - - local_deq = local_shard.dequantize().contiguous() - gathered_list = [torch.zeros_like(local_deq) for _ in range(world_size)] - dist.all_gather(gathered_list, local_deq) - manual_full = torch.cat(gathered_list, dim=0) - - full_param = param.full_tensor() - if isinstance(full_param, QuantizedTensor): - fsdp_full_deq = full_param.dequantize() - else: - fsdp_full_deq = full_param.float() + params = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ] + local_count = torch.tensor([len(params)], dtype=torch.int64, device=device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{hybrid_recipe_name}: Hybrid parameter count range is " + f"({min_count.item()}, {max_count.item()})", + ) - # Stateless formats gather identical bytes -> exact match. - torch.testing.assert_close( - manual_full.float(), - fsdp_full_deq[: manual_full.shape[0]].float(), - msg=lambda m, n=name: f"Allgather mismatch for {n}: {m}", - rtol=0.0, - atol=0.0, + errors = [] + for name, param in params: + local = param._local_tensor + reconstructed = _manual_reconstruct_hybrid(local, param_name=name, world_size=world_size) + public_full = param.full_tensor() + _record_exact( + errors, + public_full, + reconstructed.dequantize(), + f"{name}: public full_tensor vs Hybrid reconstruction", + ) + _check_hybrid_direction_buffers( + local._rowwise_storage, + reconstructed._rowwise_storage, + row_type, + direction="rowwise", + param_name=name, + world_size=world_size, + errors=errors, + ) + _check_hybrid_direction_buffers( + local._columnwise_storage, + reconstructed._columnwise_storage, + col_type, + direction="columnwise", + param_name=name, + world_size=world_size, + errors=errors, ) - checked += 1 - assert checked > 0, "No quantized DTensor params found to validate" + _raise_collective_errors(errors, f"{hybrid_recipe_name} raw Hybrid all-gather") def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): @@ -1665,39 +1979,149 @@ def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): out = model(x) out.sum().backward() - # Compare dim-0 all-gather (bytes) with FSDP2's reconstruction. - for name, param in model.named_parameters(): - if not ( - isinstance(param, DTensor) and isinstance(param._local_tensor, QuantizedTensor) - ): - continue - local_shard = param._local_tensor - local_deq = local_shard.dequantize().contiguous() - gathered_list = [torch.zeros_like(local_deq) for _ in range(world_size)] - dist.all_gather(gathered_list, local_deq) - manual_full = torch.cat(gathered_list, dim=0) - - full_param = param.full_tensor() - fsdp_full_deq = ( - full_param.dequantize() - if isinstance(full_param, QuantizedTensor) - else full_param.float() + col_type = MXFP8TensorStorage if recipe_name == "HybridMXFP8" else Float8TensorStorage + params = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ] + local_count = torch.tensor([len(params)], dtype=torch.int64, device=device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + f"{recipe_name}: awkward-shape Hybrid parameter count range is " + f"({min_count.item()}, {max_count.item()})", + ) + + errors = [] + for name, param in params: + local = param._local_tensor + label = f"{recipe_name}:{name}" + reconstructed = _manual_reconstruct_hybrid( + local, param_name=label, world_size=world_size + ) + public_full = param.full_tensor() + _record_exact( + errors, + public_full, + reconstructed.dequantize(), + f"{label}: public full_tensor vs Hybrid reconstruction", ) + _check_hybrid_direction_buffers( + local._rowwise_storage, + reconstructed._rowwise_storage, + MXFP8TensorStorage, + direction="rowwise", + param_name=label, + world_size=world_size, + errors=errors, + ) + _check_hybrid_direction_buffers( + local._columnwise_storage, + reconstructed._columnwise_storage, + col_type, + direction="columnwise", + param_name=label, + world_size=world_size, + errors=errors, + ) + _raise_collective_errors(errors, f"{recipe_name} awkward raw Hybrid all-gather") - torch.testing.assert_close( - manual_full.float(), - fsdp_full_deq[: manual_full.shape[0]].float(), - rtol=0.0, - atol=0.0, - msg=lambda m, n=name, r=recipe_name: f"[{r}] Allgather mismatch for {n} at awkward shard shape: {m}", + +@pytest.mark.xfail( + strict=True, + raises=RuntimeError, + reason=( + "Tracked by #3158: gathering dim-0 shards that split a 128-row " + "Float8Block scale tile produces too many scale rows." + ), +) +def test_fused_adam_hybrid_float8_block_unaligned_shard_shape(): + """Positive FSDP2 roundtrip for a shard that splits a 128-row scale tile.""" + from transformer_engine.pytorch import fp8 + from fsdp2_utils import get_hybrid_recipe_from_string + + supported, reason = fp8.check_fp8_block_scaling_support() + if not supported: + pytest.skip(reason) + + world_size, device = _get_dist_info() + if world_size != 2: + pytest.skip("test shape is defined for exactly two FSDP ranks") + + in_features = 256 + per_rank_out = 192 + out_features = per_rank_out * world_size + assert out_features % 128 == 0 + assert per_rank_out % 128 != 0 + + hybrid_recipe = get_hybrid_recipe_from_string("HybridFloat8BlockScaling") + with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = torch.nn.Sequential( + te.Linear( + in_features, + out_features, + params_dtype=torch.bfloat16, + device="meta", + ) + ) + model = _shard_model(model, world_size) + + x = torch.randn(128, in_features, dtype=torch.bfloat16, device=device) + with te.autocast(enabled=True, recipe=hybrid_recipe): + output = model(x) + output.sum().backward() + + params = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) + ] + local_count = torch.tensor([len(params)], dtype=torch.int64, device=device) + min_count = local_count.clone() + max_count = local_count.clone() + dist.all_reduce(min_count, dist.ReduceOp.MIN) + dist.all_reduce(max_count, dist.ReduceOp.MAX) + _collective_assert( + min_count.item() == max_count.item() and min_count.item() > 0, + "Float8Block unaligned Hybrid parameter count range is " + f"({min_count.item()}, {max_count.item()})", + ) + + errors = [] + for name, param in params: + local = param._local_tensor + reconstructed = _manual_reconstruct_hybrid(local, param_name=name, world_size=world_size) + public_full = param.full_tensor() + _record_exact( + errors, + public_full, + reconstructed.dequantize(), + f"{name}: public full_tensor vs Hybrid reconstruction", + ) + for direction in ("rowwise", "columnwise"): + _check_hybrid_direction_buffers( + getattr(local, f"_{direction}_storage"), + getattr(reconstructed, f"_{direction}_storage"), + Float8BlockwiseQTensorStorage, + direction=direction, + param_name=name, + world_size=world_size, + errors=errors, ) + _raise_collective_errors(errors, "Float8Block unaligned raw Hybrid all-gather") def test_hybrid_dcp_output_parity(hybrid_recipe_name): - """DCP save+load roundtrip: output after load matches output before save. + """DCP roundtrip and exact forked optimizer continuation. - Trains a hybrid model, saves with DCP, loads into fresh model, - and asserts forward output parity. + Trains and checkpoints a hybrid model plus FusedAdam, loads both into fresh + objects, and compares the identical next step against uninterrupted training: + output, loss, input/parameter gradients, master params, moments, counters, + and post-step output must all match bitwise. """ import torch.distributed.checkpoint as dcp @@ -1726,6 +2150,72 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) target = torch.randn_like(x) + failures = [] + + def snapshot_optimizer(current_model, current_optimizer): + snapshots = [] + for param in current_model.parameters(): + state = {} + for key, value in current_optimizer.state[param].items(): + if torch.is_tensor(value): + value = value.to_local() if isinstance(value, DTensor) else value + value = value.detach().clone() + state[key] = value + snapshots.append(state) + return snapshots + + def check_optimizer_state(actual, expected, label): + if len(actual) != len(expected): + failures.append( + f"{label}: parameter count differs: {len(actual)} != {len(expected)}" + ) + for param_idx, (actual_state, expected_state) in enumerate(zip(actual, expected)): + if actual_state.keys() != expected_state.keys(): + failures.append( + f"{label}: param {param_idx} state keys differ: " + f"{actual_state.keys()} != {expected_state.keys()}" + ) + for key in actual_state.keys() & expected_state.keys(): + actual_value = actual_state[key] + expected_value = expected_state[key] + if torch.is_tensor(expected_value): + _record_exact( + failures, + actual_value, + expected_value, + f"{label}: param {param_idx} optimizer state {key!r}", + ) + elif actual_value != expected_value: + failures.append( + f"{label}: param {param_idx} optimizer state {key!r} differs: " + f"{actual_value!r} != {expected_value!r}" + ) + + def run_continuation_step(current_model, current_optimizer): + step_x = x.detach().clone().requires_grad_() + current_optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=hybrid_recipe): + step_output = current_model(step_x) + step_loss = F.mse_loss(step_output, target) + step_loss.backward() + step_grads = [] + for param in current_model.parameters(): + grad = param.grad + if grad is not None: + grad = grad.to_local() if isinstance(grad, DTensor) else grad + grad = grad.detach().clone() + step_grads.append(grad) + current_optimizer.step() + with torch.no_grad(), te.autocast(enabled=True, recipe=hybrid_recipe): + post_step_output = current_model(x).detach().clone() + return { + "output": step_output.detach().clone(), + "loss": step_loss.detach().clone(), + "input_grad": None if step_x.grad is None else step_x.grad.detach().clone(), + "param_grads": step_grads, + "optimizer": snapshot_optimizer(current_model, current_optimizer), + "post_step_output": post_step_output, + } for _ in range(NUM_STEPS): optimizer.zero_grad(set_to_none=True) @@ -1743,6 +2233,7 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): "optimizer": optimizer.state_dict(), } dcp.save(save_state, checkpoint_id=checkpoint_dir) + saved_optimizer_state = snapshot_optimizer(model, optimizer) model2 = _build_hybrid_model(hybrid_recipe) model2 = _shard_model(model2, world_size) @@ -1769,14 +2260,56 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): with torch.no_grad(): with te.autocast(enabled=True, recipe=hybrid_recipe): loaded_output = model2(x) + check_optimizer_state( + snapshot_optimizer(model2, optimizer2), + saved_optimizer_state, + "DCP optimizer state immediately after load", + ) + _record_exact(failures, loaded_output, ref_output, "DCP roundtrip output") + _raise_collective_errors(failures, "DCP state immediately after load") + failures.clear() + + # Fork from the checkpoint and execute the identical next step. This + # catches an ignored optimizer checkpoint even when model-only output + # parity succeeds immediately after load. + torch.manual_seed(9876) + torch.cuda.manual_seed(9876) + reference_step = run_continuation_step(model, optimizer) + torch.manual_seed(9876) + torch.cuda.manual_seed(9876) + resumed_step = run_continuation_step(model2, optimizer2) + + for key in ("output", "loss", "input_grad", "post_step_output"): + _record_exact( + failures, + resumed_step[key], + reference_step[key], + f"DCP continuation {key}", + ) - torch.testing.assert_close( - loaded_output, - ref_output, - rtol=0, - atol=0, - msg=lambda m: f"DCP roundtrip output mismatch: {m}", + resumed_grads = resumed_step["param_grads"] + reference_grads = reference_step["param_grads"] + if len(resumed_grads) != len(reference_grads): + failures.append( + "DCP continuation parameter-gradient count differs: " + f"{len(resumed_grads)} != {len(reference_grads)}" + ) + for param_idx, (resumed_grad, reference_grad) in enumerate( + zip(resumed_grads, reference_grads) + ): + _record_exact( + failures, + resumed_grad, + reference_grad, + f"DCP continuation param {param_idx} gradient", + ) + + check_optimizer_state( + resumed_step["optimizer"], + reference_step["optimizer"], + "DCP optimizer state after continuation", ) + _raise_collective_errors(failures, "DCP forked continuation") finally: dist.barrier() if rank == 0: @@ -1806,9 +2339,7 @@ def test_hybrid_dcp_output_parity(hybrid_recipe_name): test_fused_adam_hybrid_identity_fp8_master_weights ), "fused_adam_hybrid_allgather_correctness": test_fused_adam_hybrid_allgather_correctness, - "fused_adam_hybrid_mxfp8_awkward_shard_shape": ( - test_fused_adam_hybrid_mxfp8_awkward_shard_shape - ), + "fused_adam_hybrid_mxfp8_awkward_shard_shape": test_fused_adam_hybrid_mxfp8_awkward_shard_shape, "hybrid_dcp_output_parity": test_hybrid_dcp_output_parity, } diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index dad5e619278..ec94b6fc462 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -34,7 +34,6 @@ sys.path.insert(0, str(TEST_ROOT)) from run_layer_with_overlap import _compare_tensors # noqa: E402 - # ── Global state ───────────────────────────────────────────────────── SEQ_LEN = 32 @@ -236,9 +235,11 @@ def _match_param_sizes(dist_param, single_param): return single_param[tuple(indices)] -def _check_outputs(output_single, output_dist, label="outputs"): +def _check_outputs(output_single, output_dist, label="outputs", *, tolerances=None): failed = torch.tensor([0], dtype=torch.uint8, device="cuda") - f, info = _compare_tensors(label, output_dist, output_single, **_get_tolerances()) + f, info = _compare_tensors( + label, output_dist, output_single, **(tolerances or _get_tolerances()) + ) if f: dist_print(info, src=WORLD_RANK, error=True) failed[0] = int(f) @@ -246,20 +247,112 @@ def _check_outputs(output_single, output_dist, label="outputs"): assert not bool(failed.item()), f"{label}: numerical check failed on at least one rank" -def _check_gradients(model_dist, model_single): - for i, ((name, pd), ps) in enumerate( - zip(model_dist.named_parameters(), model_single.parameters()) - ): - if pd.grad is None or ps.grad is None: +def _collective_assert(condition, label): + """Raise on every rank only after all ranks report structural status.""" + failed = torch.tensor([not condition], dtype=torch.uint8, device="cuda") + if not condition: + dist_print(label, src=WORLD_RANK, error=True) + dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) + assert not bool(failed.item()), f"{label}: failed on at least one rank" + + +def _check_gradients( + model_dist, model_single, *, bitwise=False, reduce_replicated=False, tolerances=None +): + """Compare every parameter gradient, including presence and parameter count. + + Sequence-parallel replicated parameters accumulate a partial gradient on + each rank. ``reduce_replicated`` reconstructs their full reference gradient + without mutating the gradient held by the model. + """ + dist_params = list(model_dist.named_parameters()) + single_params = list(model_single.named_parameters()) + local_counts = torch.tensor( + [len(dist_params), len(single_params)], dtype=torch.int64, device="cuda" + ) + min_counts = local_counts.clone() + max_counts = local_counts.clone() + dist.all_reduce(min_counts, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_counts, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_counts, max_counts) and len(dist_params) == len(single_params), + "parameter count mismatch: " + f"local distributed={len(dist_params)}, local single={len(single_params)}, " + f"global min={min_counts.tolist()}, global max={max_counts.tolist()}", + ) + for i, ((name, pd), (single_name, ps)) in enumerate(zip(dist_params, single_params)): + _collective_assert( + name == single_name, + f"parameter {i} name mismatch: {name!r} != {single_name!r}", + ) + local_presence = torch.tensor( + [pd.grad is not None, ps.grad is not None], dtype=torch.uint8, device="cuda" + ) + min_presence = local_presence.clone() + max_presence = local_presence.clone() + dist.all_reduce(min_presence, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_presence, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_presence, max_presence) + and local_presence[0].item() == local_presence[1].item(), + f"grad[{i}].{name}: gradient presence differs locally or across ranks: " + f"local={local_presence.tolist()}, global min={min_presence.tolist()}, " + f"global max={max_presence.tolist()}", + ) + if pd.grad is None: continue - failed = torch.tensor([0], dtype=torch.uint8, device="cuda") - ps_grad = _match_param_sizes(pd.grad, ps.grad) - f, info = _compare_tensors(f"grad[{i}].{name}", pd.grad, ps_grad, **_get_tolerances()) - if f: - dist_print(info, src=WORLD_RANK, error=True) - failed[0] = int(f) - dist.all_reduce(failed, dist.ReduceOp.MAX, NCCL_WORLD) - assert not bool(failed.item()), f"grad[{i}].{name}: failed on at least one rank" + pd_grad = pd.grad + local_reduce = torch.tensor( + [reduce_replicated and pd_grad.shape == ps.grad.shape], + dtype=torch.uint8, + device="cuda", + ) + min_reduce = local_reduce.clone() + max_reduce = local_reduce.clone() + dist.all_reduce(min_reduce, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_reduce, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_reduce, max_reduce), + f"grad[{i}].{name}: replicated-gradient reduction branch differs across ranks", + ) + if local_reduce.item(): + pd_grad = pd_grad.detach().clone() + dist.all_reduce(pd_grad, dist.ReduceOp.SUM, NCCL_WORLD) + ps_grad = _match_param_sizes(pd_grad, ps.grad) + label = f"grad[{i}].{name}" + if bitwise: + _check_bitwise(pd_grad, ps_grad, label) + else: + _check_outputs(ps_grad, pd_grad, label=label, tolerances=tolerances) + + +def _check_input_gradient(inp_dist, inp_single, label, *, bitwise=False): + """Compare the local TP/SP input-gradient shard with its full reference.""" + local_presence = torch.tensor( + [inp_dist.grad is not None, inp_single.grad is not None], + dtype=torch.uint8, + device="cuda", + ) + min_presence = local_presence.clone() + max_presence = local_presence.clone() + dist.all_reduce(min_presence, dist.ReduceOp.MIN, NCCL_WORLD) + dist.all_reduce(max_presence, dist.ReduceOp.MAX, NCCL_WORLD) + _collective_assert( + torch.equal(min_presence, max_presence) + and local_presence[0].item() == local_presence[1].item(), + f"{label}: input-gradient presence differs locally or across ranks: " + f"local={local_presence.tolist()}, global min={min_presence.tolist()}, " + f"global max={max_presence.tolist()}", + ) + _collective_assert( + inp_dist.grad is not None, + f"{label}: both input gradients are unexpectedly absent", + ) + expected = _match_param_sizes(inp_dist.grad, inp_single.grad) + if bitwise: + _check_bitwise(inp_dist.grad, expected, label) + else: + _check_outputs(expected, inp_dist.grad, label=label) def _apply_models(model_single, model_dist, inp_single, inp_dist, **kwargs): @@ -338,9 +431,12 @@ def _test_linear(parallel_mode, sequence_parallel, params_dtype=torch.bfloat16, label=f"linear[{parallel_mode},sp={sequence_parallel},amax_stress={amax_stress}]", ) - # Other configs need extra grad sync, matching run_numerics.py. - if parallel_mode == "column" or not sequence_parallel: - _check_gradients(model_dist, model_single) + _check_input_gradient( + inp_dist, + inp_single, + label=f"linear[{parallel_mode},sp={sequence_parallel},amax_stress={amax_stress}] dgrad", + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) def test_linear(): @@ -363,6 +459,8 @@ def vanilla_recipe(): return te_recipe.MXFP8BlockScaling() if QUANTIZATION == "hybrid_nvfp4": return te_recipe.NVFP4BlockScaling() + if QUANTIZATION == "identity": + return None raise ValueError(f"No vanilla recipe for QUANTIZATION={QUANTIZATION!r}") @@ -413,7 +511,7 @@ def run(recipe): inp = inp[:, WORLD_RANK * split : (WORLD_RANK + 1) * split].clone() inp.requires_grad_() - with te.autocast(enabled=True, recipe=recipe): + with te.autocast(enabled=recipe is not None, recipe=recipe): out = model(inp) # Fixed, recipe-independent target so both backward graphs match. torch.manual_seed(54321) @@ -431,9 +529,8 @@ def run(recipe): # Same-topology fprop operands should match bitwise. _check_bitwise(out_h, out_v, f"{tag} forward") - # Skip backward bitwise when SP changes the gather path or NVFP4 consumes - # columnwise SR RNG differently. Loose TP/SP checks cover those cases. - if not sequence_parallel and not _backward_not_bitwise_comparable(): + # NVFP4 backward consumes columnwise stochastic-rounding RNG differently. + if not _backward_not_bitwise_comparable(): _check_bitwise(dinp_h, dinp_v, f"{tag} dgrad") assert len(wgrads_h) == len(wgrads_v), f"{tag}: weight-grad count mismatch" for i, (gh, gv) in enumerate(zip(wgrads_h, wgrads_v)): @@ -443,7 +540,6 @@ def run(recipe): def test_linear_vs_vanilla(): # These recipes have no same-format vanilla bitwise target. if QUANTIZATION in ( - "identity", "hybrid_mxfp8_nvfp4", "hybrid_fp8_identity", "hybrid_mxfp8_identity", @@ -456,17 +552,20 @@ def test_linear_vs_vanilla(): def _same_format_parity_supported(): - return QUANTIZATION in ("hybrid_fp8", "hybrid_mxfp8") + return QUANTIZATION in ("hybrid_fp8", "hybrid_mxfp8", "identity") def _check_same_topology_parity( - out_h, dinp_h, model_h, out_v, dinp_v, model_v, tag, *, check_grads + out_h, dinp_h, model_h, out_v, dinp_v, model_v, tag, *, tolerances=None ): - # Larger modules may use different fused/unfused norm paths. - _check_outputs(out_v, out_h, label=f"{tag} forward") - if check_grads: - _check_outputs(dinp_v, dinp_h, label=f"{tag} dgrad") - _check_gradients(model_h, model_v) + check = ( + _check_bitwise + if tolerances is None + else lambda a, e, label: _check_outputs(e, a, label=label, tolerances=tolerances) + ) + check(out_h, out_v, f"{tag} forward") + check(dinp_h, dinp_v, f"{tag} dgrad") + _check_gradients(model_h, model_v, bitwise=tolerances is None, tolerances=tolerances) def _test_layernorm_linear_vs_vanilla(sequence_parallel, params_dtype=torch.bfloat16): @@ -491,7 +590,7 @@ def run(recipe_obj): torch.cuda.manual_seed(45670) inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) inp.requires_grad_() - with te.autocast(enabled=True, recipe=recipe_obj): + with te.autocast(enabled=recipe_obj is not None, recipe=recipe_obj): out = model(inp) torch.manual_seed(45671) torch.cuda.manual_seed(45671) @@ -508,7 +607,6 @@ def run(recipe_obj): dinp_v, model_v, f"layernorm_linear_vs_vanilla[sp={sequence_parallel}]", - check_grads=not sequence_parallel, ) @@ -539,7 +637,7 @@ def run(recipe_obj): torch.cuda.manual_seed(56780) inp = torch.randn((BATCH_SIZE, HIDDEN_SIZE)).cuda().to(params_dtype) inp.requires_grad_() - with te.autocast(enabled=True, recipe=recipe_obj): + with te.autocast(enabled=recipe_obj is not None, recipe=recipe_obj): out = model(inp) torch.manual_seed(56781) torch.cuda.manual_seed(56781) @@ -556,7 +654,10 @@ def run(recipe_obj): dinp_v, model_v, f"layernorm_mlp_vs_vanilla[sp={sequence_parallel}]", - check_grads=not sequence_parallel, + # Hybrid CustomRecipe disables quantized-norm fusion while the native + # recipe uses it. Both paths consume identical quantized GEMM operands, + # but the BF16 norm boundary can differ by one rounding step. + tolerances=(None if QUANTIZATION == "identity" else {"rtol": 2**-7, "atol": 2**-10}), ) @@ -603,6 +704,11 @@ def _test_layernorm_linear(sequence_parallel, params_dtype=torch.bfloat16): _loss_backward(out_single, out_dist) _check_outputs(out_single, out_dist, label=f"layernorm_linear[sp={sequence_parallel}]") + _check_input_gradient( + inp_dist, inp_single, label=f"layernorm_linear[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) + def test_layernorm_linear(): for sequence_parallel in [False, True]: @@ -649,9 +755,10 @@ def _test_layernorm_mlp(sequence_parallel, params_dtype=torch.bfloat16): _loss_backward(out_single, out_dist) _check_outputs(out_single, out_dist, label=f"layernorm_mlp[sp={sequence_parallel}]") - # SP grads need extra sync, matching run_numerics.py. - if not sequence_parallel: - _check_gradients(model_dist, model_single) + _check_input_gradient( + inp_dist, inp_single, label=f"layernorm_mlp[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) def test_layernorm_mlp(): @@ -711,9 +818,10 @@ def _test_transformer_layer(sequence_parallel, params_dtype=torch.bfloat16): _loss_backward(out_single, out_dist) _check_outputs(out_single, out_dist, label=f"transformer_layer[sp={sequence_parallel}]") - # SP grads need extra sync, matching run_numerics.py. - if not sequence_parallel: - _check_gradients(model_dist, model_single) + _check_input_gradient( + inp_dist, inp_single, label=f"transformer_layer[sp={sequence_parallel}] dgrad" + ) + _check_gradients(model_dist, model_single, reduce_replicated=sequence_parallel) def test_transformer_layer(): diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py index 8cc78d793a1..27305aee7aa 100644 --- a/tests/pytorch/distributed/test_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -12,7 +12,6 @@ import torch import transformer_engine.pytorch as te - if torch.cuda.device_count() < 2: pytest.skip( "Distributed TP/SP tests need at least 2 GPUs.", @@ -56,11 +55,13 @@ def test_hybrid_fp8_linear_vs_vanilla(): _run_test("hybrid_fp8", "linear_vs_vanilla") +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_layernorm_linear_vs_vanilla(): """Same-topology LayerNormLinear parity against vanilla FP8.""" _run_test("hybrid_fp8", "layernorm_linear_vs_vanilla") +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_layernorm_mlp_vs_vanilla(): """Same-topology LayerNormMLP parity against vanilla FP8.""" _run_test("hybrid_fp8", "layernorm_mlp_vs_vanilla") @@ -90,6 +91,7 @@ def test_hybrid_fp8_identity_linear(): _run_test("hybrid_fp8_identity", "linear") +@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_identity_all_modules(): """All-Identity TP/SP end-to-end coverage for every supported TE module.""" _run_test("identity", "all") @@ -111,10 +113,12 @@ def test_hybrid_mxfp8_linear_vs_vanilla(): _run_test("hybrid_mxfp8", "linear_vs_vanilla") +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") def test_hybrid_mxfp8_layernorm_linear_vs_vanilla(): _run_test("hybrid_mxfp8", "layernorm_linear_vs_vanilla") +@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") def test_hybrid_mxfp8_layernorm_mlp_vs_vanilla(): _run_test("hybrid_mxfp8", "layernorm_mlp_vs_vanilla") diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 697ee06b81b..773ced21e7a 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -866,44 +866,79 @@ def forward(self, x): ): param_offload.data.copy_(param_no_offload.data) - x = Utils.create_tensor(None) + x = Utils.create_tensor(None, requires_grad=True) if use_cuda_graphs: callable_offload = te.make_graphed_callables( callable_offload, (x,), enabled=recipe is not None, - recipe=(Utils.create_recipe_ctx(recipe) if recipe is not None else None), + recipe=recipe, ) # warm up (for example to compute sf for delayed scaling) for _ in range(4): + callable_offload.zero_grad(set_to_none=True) + callable_no_offload.zero_grad(set_to_none=True) + x.grad = None out = callable_offload(x) out.sum().backward() + x.grad = None out = callable_no_offload(x) out.sum().backward() callable_offload.zero_grad(set_to_none=True) + callable_no_offload.zero_grad(set_to_none=True) + x.grad = None + rng_state = torch.cuda.get_rng_state() out_offload = callable_offload(x) out_offload.sum().backward() - # save out and gradients - offload_outs = [out_offload] + # Clone before the no-offload pass. CUDA graphs may reuse static output + # buffers, and the comparison must cover computed values rather than + # aliases or unchanged parameters. + offload_out = out_offload.detach().clone() + assert x.grad is not None + offload_input_grad = x.grad.detach().clone() + offload_param_grads = [] for param in callable_offload.parameters(): - offload_outs.append(param.detach().clone()) + assert param.grad is not None + offload_param_grads.append(param.grad.detach().clone()) torch.cuda.reset_peak_memory_stats() + torch.cuda.set_rng_state(rng_state) + x.grad = None out_no_offload = callable_no_offload(x) out_no_offload.sum().backward() - # collect gradients - no_offload_outs = [out_no_offload] + no_offload_out = out_no_offload.detach().clone() + assert x.grad is not None + no_offload_input_grad = x.grad.detach().clone() + no_offload_param_grads = [] for param in callable_no_offload.parameters(): - no_offload_outs.append(param.detach().clone()) - - # check if tensors are the same - for i in range(len(offload_outs)): - assert torch.allclose(offload_outs[i], no_offload_outs[i]), f"Error in tensor {i}." + assert param.grad is not None + no_offload_param_grads.append(param.grad.detach().clone()) + + # CPU offload is byte-preserving transport. It must not alter forward + # results, input gradients, or any parameter gradient. + torch.testing.assert_close(offload_out, no_offload_out, rtol=0.0, atol=0.0) + torch.testing.assert_close( + offload_input_grad, + no_offload_input_grad, + rtol=0.0, + atol=0.0, + ) + assert len(offload_param_grads) == len(no_offload_param_grads) + for index, (offload_grad, no_offload_grad) in enumerate( + zip(offload_param_grads, no_offload_param_grads) + ): + torch.testing.assert_close( + offload_grad, + no_offload_grad, + rtol=0.0, + atol=0.0, + msg=f"Parameter gradient {index} differs with CPU offload", + ) torch.cuda.synchronize() diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index bf39cfa3e77..b942fc8ab09 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -73,6 +73,10 @@ def _hybrid_mxfp8_nvfp4_qfactory(role): if mxfp8_available and nvfp4_available: quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_mxfp8_nvfp4_qfactory)) +hybrid_quantization_recipes = [ + item for item in quantization_recipes if item is not None and item.custom() +] + model_config = { "small": ModelConfig(8, 512, 8, 64, num_layers=5, eps=0.1), } @@ -288,3 +292,93 @@ def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: s # in that case the memory-savings assertion above is the only check. if memory_from_cached_weights is not None: assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON + + +@pytest.mark.parametrize( + "quantization_recipe", + hybrid_quantization_recipes, + ids=[item.qfactory.__name__ for item in hybrid_quantization_recipes], +) +def test_hybrid_cpu_offload_numerics_exact(quantization_recipe: recipe.Recipe) -> None: + """V1 offload must preserve hybrid forward and backward values bitwise.""" + + def run(modules, inp, grad_output, *, cpu_offload): + if cpu_offload: + offload_context, sync_function = te.get_cpu_offload_context( + enabled=True, + num_layers=len(modules), + model_layers=len(modules) + 1, + offload_activations=True, + offload_weights=False, + ) + else: + offload_context = contextlib.nullcontext() + sync_function = lambda tensor: tensor + + out = inp + for module in modules: + with te.autocast(enabled=True, recipe=quantization_recipe), offload_context: + out = module(out) + out = sync_function(out) + # Commit the final offload group, mirroring the memory test above. + with offload_context: + out = out.clone() + out = sync_function(out) + out.backward(grad_output) + + grads = {} + for module_index, module in enumerate(modules): + for name, param in module.named_parameters(): + assert param.grad is not None + grads[f"{module_index}.{name}"] = param.grad.detach().clone() + assert inp.grad is not None + return out.detach().clone(), inp.grad.detach().clone(), grads + + torch.manual_seed(9000) + reference_modules = [model_types["linear"]() for _ in range(2)] + torch.manual_seed(9001) + offload_modules = [model_types["linear"]() for _ in range(2)] + with torch.no_grad(): + for offload_module, reference_module in zip(offload_modules, reference_modules): + for offload_param, reference_param in zip( + offload_module.parameters(), reference_module.parameters() + ): + offload_param.copy_(reference_param) + + # V1 requires one pass to initialize cached quantized weight workspaces. + warmup_rng_state = torch.cuda.get_rng_state() + _warmup_model(offload_modules, quantization_recipe) + for module in offload_modules: + module.zero_grad(set_to_none=True) + torch.cuda.set_rng_state(warmup_rng_state) + _warmup_model(reference_modules, quantization_recipe) + for module in reference_modules: + module.zero_grad(set_to_none=True) + + torch.manual_seed(9002) + x = torch.randn((8, SIZE, SIZE), device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn_like(x) + x_offload = x.detach().clone().requires_grad_(True) + x_reference = x.detach().clone().requires_grad_(True) + + # Replay identical stochastic-rounding/RHT randomness in both paths. + rng_state = torch.cuda.get_rng_state() + offload_out, offload_input_grad, offload_param_grads = run( + offload_modules, x_offload, grad_output, cpu_offload=True + ) + torch.cuda.set_rng_state(rng_state) + reference_out, reference_input_grad, reference_param_grads = run( + reference_modules, x_reference, grad_output, cpu_offload=False + ) + + torch.testing.assert_close(offload_out, reference_out, rtol=0.0, atol=0.0) + torch.testing.assert_close(offload_input_grad, reference_input_grad, rtol=0.0, atol=0.0) + assert offload_param_grads.keys() == reference_param_grads.keys() + for name, reference_grad in reference_param_grads.items(): + torch.testing.assert_close( + offload_param_grads[name], + reference_grad, + rtol=0.0, + atol=0.0, + msg=f"V1 CPU-offload parameter gradient mismatch for {name}", + ) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 90024c5ef20..e949e6ec392 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -106,6 +106,144 @@ def _make_hybrid_quantizer_fp4_row_fp8_col(): ) +def _as_data_tensor_tuple(storage): + """Return a storage's raw buffers as a tuple without copying them.""" + if storage is None: + return () + tensors = storage.get_data_tensors() + return tensors if isinstance(tensors, tuple) else (tensors,) + + +def _snapshot_storage_tensor_metadata(storage, *, clone=False): + """Capture every tensor-valued concrete-storage metadata field without mutation.""" + if storage is None: + return None + metadata = storage.get_metadata() + tensor_metadata = {} + for name, value in metadata.items(): + # Optional tensor fields are represented by ``None``. Including them + # makes missing rowwise/columnwise data, scales, and amax explicit. + if isinstance(value, torch.Tensor) or value is None: + snapshot_value = value.detach().clone() if clone and value is not None else value + if ( + value is not None + and getattr(storage, "_is_2D_scaled", False) + and name in ("rowwise_scale_inv", "columnwise_scale_inv") + ): + # Float8Block 2D scales pad one tile dimension to a multiple + # of four. Kernels never initialize or consume that padding, + # and FSDP deliberately strips then zero-repads it. Canonicalize + # only those non-logical elements while still comparing the + # complete tensor field, shape, presence, and every valid scale. + snapshot_value = value.detach().clone() + m, n = storage._fsdp_logical_mn() + block_len = storage._FSDP_BLOCK_LEN + m_tiles = (m + block_len - 1) // block_len + n_tiles = (n + block_len - 1) // block_len + if name == "rowwise_scale_inv": + snapshot_value[:, n_tiles:] = 0 + else: + snapshot_value[:, m_tiles:] = 0 + tensor_metadata[name] = snapshot_value + return { + "storage_type": type(storage).__name__, + "tensor_metadata": tensor_metadata, + } + + +def _assert_storage_data_exact(actual, expected, *, context): + """Assert every tensor-valued data/scale/amax metadata field exactly.""" + _assert_nested_state_exact( + _snapshot_storage_tensor_metadata(actual), + _snapshot_storage_tensor_metadata(expected), + path=context, + ) + + +def _clone_nested_state(value): + """Clone tensors in nested checkpoint state so later steps cannot mutate snapshots.""" + if isinstance(value, torch.Tensor): + return value.detach().clone() + if isinstance(value, dict): + return {key: _clone_nested_state(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_nested_state(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_nested_state(item) for item in value) + return value + + +def _assert_nested_state_exact(actual, expected, *, path="state"): + """Recursively compare checkpoint state, using zero tolerance for tensors.""" + if isinstance(expected, torch.Tensor): + assert isinstance(actual, torch.Tensor), f"{path}: expected Tensor, got {type(actual)}" + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0, msg=path) + return + if isinstance(expected, dict): + assert isinstance(actual, dict), f"{path}: expected dict, got {type(actual)}" + assert actual.keys() == expected.keys(), f"{path}: dictionary keys differ" + for key in expected: + _assert_nested_state_exact(actual[key], expected[key], path=f"{path}.{key}") + return + if isinstance(expected, (list, tuple)): + assert isinstance( + actual, type(expected) + ), f"{path}: expected {type(expected).__name__}, got {type(actual).__name__}" + assert len(actual) == len(expected), f"{path}: sequence lengths differ" + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + _assert_nested_state_exact( + actual_item, + expected_item, + path=f"{path}[{index}]", + ) + return + assert actual == expected, f"{path}: {actual!r} != {expected!r}" + + +def _assert_hybrid_tensor_exact(actual, expected, *, context): + """Compare both hybrid directions, including all metadata and native dequantization.""" + for direction in ("rowwise", "columnwise"): + actual_storage = getattr(actual, f"{direction}_sub_storage") + expected_storage = getattr(expected, f"{direction}_sub_storage") + _assert_storage_data_exact( + actual_storage, + expected_storage, + context=f"{context} {direction}", + ) + if expected_storage is None: + continue + try: + expected_dequantized = expected_storage.dequantize() + except NotImplementedError: + # Some direction-only formats (notably NVFP4 columnwise) expose + # bytes/scales for GEMM but intentionally do not dequantize. + continue + actual_dequantized = actual_storage.dequantize() + torch.testing.assert_close( + actual_dequantized, + expected_dequantized, + rtol=0.0, + atol=0.0, + msg=f"{context} {direction} dequantized value differs", + ) + + +def _snapshot_model_parameters(model): + """Capture normal values and all tensor metadata in both hybrid directions.""" + snapshot = {} + for name, param in model.named_parameters(): + if isinstance(param, HybridQuantizedTensor): + snapshot[name] = { + "rowwise": _snapshot_storage_tensor_metadata(param.rowwise_sub_storage, clone=True), + "columnwise": _snapshot_storage_tensor_metadata( + param.columnwise_sub_storage, clone=True + ), + } + else: + snapshot[name] = param.detach().clone() + return snapshot + + class _CountingIdentityQuantizer(IdentityQuantizer): """Identity quantizer that counts quantize calls for regression tests.""" @@ -954,6 +1092,15 @@ def hybrid_tensor(self): def test_save_restore_roundtrip(self, hybrid_tensor): dq_before = hybrid_tensor.dequantize() + expected_metadata = { + "rowwise": _snapshot_storage_tensor_metadata( + hybrid_tensor.rowwise_sub_storage, clone=True + ), + "columnwise": _snapshot_storage_tensor_metadata( + hybrid_tensor.columnwise_sub_storage, clone=True + ), + } + buffers_before = _as_data_tensor_tuple(hybrid_tensor) tensors, obj = hybrid_tensor.prepare_for_saving() assert isinstance(tensors, list) assert all(t is None or isinstance(t, torch.Tensor) for t in tensors) @@ -963,7 +1110,20 @@ def test_save_restore_roundtrip(self, hybrid_tensor): assert len(remainder) == 0 dq_after = hybrid_tensor.dequantize() - torch.testing.assert_close(dq_before, dq_after) + buffers_after = _as_data_tensor_tuple(hybrid_tensor) + assert len(buffers_after) == len(buffers_before) + for before, after in zip(buffers_before, buffers_after): + assert after is before, "Direct save/restore must reattach the exact saved buffer" + torch.testing.assert_close(dq_before, dq_after, rtol=0.0, atol=0.0) + for direction in ("rowwise", "columnwise"): + actual_metadata = _snapshot_storage_tensor_metadata( + getattr(hybrid_tensor, f"{direction}_sub_storage"), clone=False + ) + _assert_nested_state_exact( + actual_metadata, + expected_metadata[direction], + path=f"save/restore {direction}", + ) def test_save_clears_data(self, hybrid_tensor): tensors, obj = hybrid_tensor.prepare_for_saving() @@ -1289,10 +1449,15 @@ def test_get_data_tensors(self): hq = _make_hybrid_quantizer_fp8_row_fp4_col() result = hq.quantize(inp) data_tensors = result.get_data_tensors() + row_tensors = _as_data_tensor_tuple(result.rowwise_sub_storage) + col_tensors = _as_data_tensor_tuple(result.columnwise_sub_storage) + assert isinstance(data_tensors, tuple) - assert len(data_tensors) > 0 - has_non_none = any(t is not None for t in data_tensors) - assert has_non_none + assert row_tensors and col_tensors + assert len(data_tensors) == len(row_tensors) + len(col_tensors) + assert all( + actual is expected for actual, expected in zip(data_tensors, row_tensors + col_tensors) + ), "Hybrid get_data_tensors must concatenate both sub-storages in direction order" @requires_fp8_and_nvfp4 @@ -2634,94 +2799,62 @@ class TestHybridMixedWithNonHybrid: wgrad (NT): input columnwise must match grad_output columnwise """ - def test_hybrid_input_plain_weight_fwd_bwd(self): - """Input is hybrid (FP8 row / FP8 col), weight + grad_output plain FP8. - - Wgrad columnwise: FP8 (input.col) × FP8 (grad_output.col) → compatible. - """ - torch.manual_seed(77) + @staticmethod + def _assert_native_fp8_parity(hybrid_role, *, seed): + torch.manual_seed(seed) in_features, out_features, batch = 128, 128, 32 + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_native = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_native.load_state_dict(model_hybrid.state_dict()) + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) - model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn( - batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True - ) - - def factory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type == "input": + def mixed_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) + if is_linear and role.tensor_type == hybrid_role: return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_fp8_quantizer(), ) - if is_linear and role.tensor_type == "weight": - return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, - device="cuda", - ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, - device="cuda", - ) + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - mixed_recipe = recipe.CustomRecipe(qfactory=factory) - with autocast(enabled=True, recipe=mixed_recipe): - out = model(inp) - - assert not torch.isnan(out).any() + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=mixed_factory), + seed=seed + 100, + ) + native_result = _run_linear_forward_backward( + model_native, + base_input, + grad_output, + recipe.Float8CurrentScaling(), + seed=seed + 100, + ) + _assert_linear_results_exact( + hybrid_result, + native_result, + output=True, + input_grad=True, + param_grads=True, + ) - loss = out.float().sum() - loss.backward() + def test_hybrid_input_plain_weight_fwd_bwd(self): + """Input is hybrid (FP8 row / FP8 col), weight + grad_output plain FP8. - assert inp.grad is not None - assert not torch.isnan(inp.grad).any() - for name, p in model.named_parameters(): - assert p.grad is not None, f"Gradient for '{name}' is None" + Wgrad columnwise: FP8 (input.col) × FP8 (grad_output.col) → compatible. + """ + self._assert_native_fp8_parity("input", seed=77) def test_plain_input_hybrid_weight_fwd_bwd(self): """Input is plain FP8, weight is hybrid (FP8 row / FP8 col).""" - torch.manual_seed(88) - in_features, out_features, batch = 128, 128, 32 - - model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn( - batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True - ) - - def factory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type == "input": - return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, - device="cuda", - ) - if is_linear and role.tensor_type == "weight": - return HybridQuantizer( - rowwise_quantizer=_make_fp8_quantizer(), - columnwise_quantizer=_make_fp8_quantizer(), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return Float8CurrentScalingQuantizer( - tex.DType.kFloat8E5M2, - device="cuda", - ) - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - - mixed_recipe = recipe.CustomRecipe(qfactory=factory) - with autocast(enabled=True, recipe=mixed_recipe): - out = model(inp) - - assert not torch.isnan(out).any() - - loss = out.float().sum() - loss.backward() - - assert inp.grad is not None - assert not torch.isnan(inp.grad).any() - for name, p in model.named_parameters(): - assert p.grad is not None, f"Gradient for '{name}' is None" + self._assert_native_fp8_parity("weight", seed=88) # --------------------------------------------------------------------------- @@ -2823,6 +2956,70 @@ def _build_cross_format_params(): return params +def _set_quantization_test_seed(seed): + """Reset every RNG that a quantizer or GEMM helper may consume.""" + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _run_linear_forward(model, base_input, fp8_recipe, *, seed): + """Run a deterministic forward pass and return a detached result.""" + _set_quantization_test_seed(seed) + with torch.no_grad(): + with autocast(enabled=True, recipe=fp8_recipe): + output = model(base_input) + return output.detach().clone() + + +def _run_linear_forward_backward(model, base_input, grad_output, fp8_recipe, *, seed): + """Run Linear with an external gradient and capture every numerical result.""" + model.zero_grad(set_to_none=True) + run_input = base_input.clone().detach().requires_grad_(True) + _set_quantization_test_seed(seed) + with autocast(enabled=True, recipe=fp8_recipe): + output = model(run_input) + output.backward(grad_output) + return ( + output.detach().clone(), + run_input.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + }, + ) + + +def _plain_linear_qfactory(operand_factory, grad_factory): + """Build a non-hybrid factory with role-correct operand and grad dtypes.""" + + def factory(role): + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type in ("grad_output", "grad_input"): + return grad_factory() + return operand_factory() + + return factory + + +def _assert_linear_results_exact(actual, expected, *, output, input_grad, param_grads): + """Compare selected Linear results with zero tolerance.""" + if output: + torch.testing.assert_close(actual[0], expected[0], rtol=0.0, atol=0.0) + if input_grad: + torch.testing.assert_close(actual[1], expected[1], rtol=0.0, atol=0.0) + if param_grads: + assert actual[2].keys() == expected[2].keys() + for name in actual[2]: + torch.testing.assert_close( + actual[2][name], + expected[2][name], + rtol=0.0, + atol=0.0, + msg=f"Parameter gradient {name!r} differs", + ) + + class TestHybridCrossFormatParametrized: """Parametrized fwd+bwd over all stateless quantizer cross-format pairs.""" @@ -2831,53 +3028,73 @@ def test_fwd_bwd(self, row_name, col_name): torch.manual_seed(42) in_features, out_features, batch = 128, 128, 32 - model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn( - batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True - ) + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_bwd_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref.load_state_dict(model_hybrid.state_dict()) + model_bwd_ref.load_state_dict(model_hybrid.state_dict()) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) row_cfg = _QUANTIZER_CONFIGS[row_name] col_cfg = _QUANTIZER_CONFIGS[col_name] - make_row_e4m3 = row_cfg[0] - make_col_e4m3 = col_cfg[0] + make_row_operand = row_cfg[0] + make_row_grad = row_cfg[1] if row_cfg[1] is not None else row_cfg[0] + make_col_operand = col_cfg[0] make_col_grad = col_cfg[1] if col_cfg[1] is not None else col_cfg[0] - def factory(role): - if ( - role is not None - and role.module_type in ("linear", "grouped_linear") - and role.tensor_type in ("input", "weight") - ): + def hybrid_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) + if is_linear and role.tensor_type in ("input", "weight"): return HybridQuantizer( - rowwise_quantizer=make_row_e4m3(), - columnwise_quantizer=make_col_e4m3(), + rowwise_quantizer=make_row_operand(), + columnwise_quantizer=make_col_operand(), ) - if ( - role is not None - and role.module_type in ("linear", "grouped_linear") - and role.tensor_type in ("grad_output", "grad_input") - ): + if is_linear and role.tensor_type in ("grad_output", "grad_input"): return make_col_grad() - return make_row_e4m3() - - mixed_recipe = recipe.CustomRecipe(qfactory=factory) - with autocast(enabled=True, recipe=mixed_recipe): - out = model(inp) + return make_row_operand() - assert out.shape == (batch, out_features) - assert not torch.isnan(out).any(), f"Output NaN ({row_name} row × {col_name} col)" - assert not torch.isinf(out).any(), f"Output Inf ({row_name} row × {col_name} col)" + hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_factory) + fprop_ref_recipe = recipe.CustomRecipe( + qfactory=_plain_linear_qfactory(make_row_operand, make_row_grad) + ) + bwd_ref_recipe = recipe.CustomRecipe( + qfactory=_plain_linear_qfactory(make_col_operand, make_col_grad) + ) - loss = out.float().sum() - loss.backward() + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + hybrid_recipe, + seed=1234, + ) + fprop_ref = _run_linear_forward( + model_fprop_ref, + base_input, + fprop_ref_recipe, + seed=1234, + ) + bwd_ref = _run_linear_forward_backward( + model_bwd_ref, + base_input, + grad_output, + bwd_ref_recipe, + seed=1234, + ) - assert inp.grad is not None, "Input gradient is None" - assert not torch.isnan(inp.grad).any(), f"Input grad NaN ({row_name} row × {col_name} col)" - for name, p in model.named_parameters(): - assert p.grad is not None, f"Gradient for '{name}' is None" - assert not torch.isnan( - p.grad - ).any(), f"Gradient for '{name}' NaN ({row_name} row × {col_name} col)" + torch.testing.assert_close(hybrid_result[0], fprop_ref, rtol=0.0, atol=0.0) + _assert_linear_results_exact( + hybrid_result, + bwd_ref, + output=False, + input_grad=True, + param_grads=True, + ) # --------------------------------------------------------------------------- @@ -2947,8 +3164,9 @@ def test_push_pop_roundtrip(self, row_name, col_name): reloaded = self._run_roundtrip(hybrid) + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload roundtrip") assert isinstance(reloaded, HybridQuantizedTensor) - torch.testing.assert_close(reloaded.dequantize(), expected) + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) @requires_fp8_and_nvfp4 def test_push_pop_preserves_sub_storage_types(self): @@ -2964,6 +3182,7 @@ def test_push_pop_preserves_sub_storage_types(self): reloaded = self._run_roundtrip(hybrid) assert isinstance(reloaded.rowwise_sub_storage, row_type) + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload storage types") assert isinstance(reloaded.columnwise_sub_storage, col_type) @requires_fp8_and_nvfp4 @@ -2982,8 +3201,9 @@ def test_push_pop_with_rowwise_only(self): assert isinstance(reloaded, HybridQuantizedTensor) assert reloaded.columnwise_sub_storage is None + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload rowwise-only") assert reloaded.rowwise_sub_storage is not None - torch.testing.assert_close(reloaded.dequantize(), expected) + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) @requires_fp8_and_nvfp4 def test_push_pop_with_columnwise_only(self): @@ -3008,8 +3228,9 @@ def test_push_pop_with_columnwise_only(self): assert isinstance(reloaded, HybridQuantizedTensor) assert reloaded.rowwise_sub_storage is None + _assert_hybrid_tensor_exact(reloaded, hybrid, context="CPU offload columnwise-only") assert reloaded.columnwise_sub_storage is not None - torch.testing.assert_close(reloaded.dequantize(), expected) + torch.testing.assert_close(reloaded.dequantize(), expected, rtol=0.0, atol=0.0) @requires_fp8_and_nvfp4 def test_push_pop_roundtrip_does_not_leak_intermediate_buffers(self): @@ -3063,131 +3284,119 @@ class TestHybridThreeFormats: grad_output is itself hybrid (FormatB row + FormatC col) when B ≠ C. """ - def test_fp8_fprop_mxfp8_dgrad_nvfp4_wgrad(self): - """FP8 current (fprop) + MXFP8 (dgrad) + NVFP4 (wgrad).""" - torch.manual_seed(300) + @staticmethod + def _assert_three_format_routing( + make_fprop, make_dgrad, make_wgrad, *, seed, plain_grad_output=False + ): in_features, out_features, batch = 128, 128, 32 - - model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn( - batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True - ) - - def factory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + torch.manual_seed(seed) + model_hybrid = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_fprop_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_dgrad_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + model_wgrad_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + state_dict = model_hybrid.state_dict() + model_fprop_ref.load_state_dict(state_dict) + model_dgrad_ref.load_state_dict(state_dict) + model_wgrad_ref.load_state_dict(state_dict) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + grad_output = torch.randn(batch, out_features, device="cuda", dtype=torch.bfloat16) + + def hybrid_factory(role): + is_linear = role is not None and role.module_type in ( + "linear", + "grouped_linear", + ) if is_linear and role.tensor_type == "weight": return HybridQuantizer( - rowwise_quantizer=_make_fp8_quantizer(), - columnwise_quantizer=_make_mxfp8_quantizer(), + rowwise_quantizer=make_fprop(), + columnwise_quantizer=make_dgrad(), ) if is_linear and role.tensor_type == "input": return HybridQuantizer( - rowwise_quantizer=_make_fp8_quantizer(), - columnwise_quantizer=_make_nvfp4_quantizer(), + rowwise_quantizer=make_fprop(), + columnwise_quantizer=make_wgrad(), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): + if plain_grad_output: + return make_dgrad() return HybridQuantizer( - rowwise_quantizer=_make_mxfp8_quantizer(), - columnwise_quantizer=_make_nvfp4_quantizer(), + rowwise_quantizer=make_dgrad(), + columnwise_quantizer=make_wgrad(), ) - return _make_fp8_quantizer() - - with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): - out = model(inp) - - assert out.shape == (batch, out_features) - assert not torch.isnan(out).any(), "Output contains NaN" - - loss = out.float().sum() - loss.backward() + return make_fprop() + + hybrid_result = _run_linear_forward_backward( + model_hybrid, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=hybrid_factory), + seed=seed + 100, + ) + fprop_ref = _run_linear_forward( + model_fprop_ref, + base_input, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_fprop, make_fprop)), + seed=seed + 100, + ) + dgrad_ref = _run_linear_forward_backward( + model_dgrad_ref, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_dgrad, make_dgrad)), + seed=seed + 100, + ) + wgrad_ref = _run_linear_forward_backward( + model_wgrad_ref, + base_input, + grad_output, + recipe.CustomRecipe(qfactory=_plain_linear_qfactory(make_wgrad, make_wgrad)), + seed=seed + 100, + ) + + torch.testing.assert_close(hybrid_result[0], fprop_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(hybrid_result[1], dgrad_ref[1], rtol=0.0, atol=0.0) + torch.testing.assert_close( + hybrid_result[2]["weight"], + wgrad_ref[2]["weight"], + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + hybrid_result[2]["bias"], + dgrad_ref[2]["bias"], + rtol=0.0, + atol=0.0, + ) - assert inp.grad is not None, "Input gradient is None" - assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" - for name, p in model.named_parameters(): - assert p.grad is not None, f"Gradient for '{name}' is None" - assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + def test_fp8_fprop_mxfp8_dgrad_nvfp4_wgrad(self): + """FP8 current (fprop) + MXFP8 (dgrad) + NVFP4 (wgrad).""" + self._assert_three_format_routing( + _make_fp8_quantizer, + _make_mxfp8_quantizer, + _make_nvfp4_quantizer, + seed=300, + ) def test_nvfp4_fprop_fp8_dgrad_mxfp8_wgrad(self): """NVFP4 (fprop) + FP8 current (dgrad) + MXFP8 (wgrad).""" - torch.manual_seed(301) - in_features, out_features, batch = 128, 128, 32 - - model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn( - batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + self._assert_three_format_routing( + _make_nvfp4_quantizer, + _make_fp8_quantizer, + _make_mxfp8_quantizer, + seed=301, ) - def factory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type == "weight": - return HybridQuantizer( - rowwise_quantizer=_make_nvfp4_quantizer(), - columnwise_quantizer=_make_fp8_quantizer(), - ) - if is_linear and role.tensor_type == "input": - return HybridQuantizer( - rowwise_quantizer=_make_nvfp4_quantizer(), - columnwise_quantizer=_make_mxfp8_quantizer(), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return HybridQuantizer( - rowwise_quantizer=_make_fp8_quantizer(), - columnwise_quantizer=_make_mxfp8_quantizer(), - ) - return _make_nvfp4_quantizer() - - with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): - out = model(inp) - - assert out.shape == (batch, out_features) - assert not torch.isnan(out).any(), "Output contains NaN" - - loss = out.float().sum() - loss.backward() - - assert inp.grad is not None, "Input gradient is None" - assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" - for name, p in model.named_parameters(): - assert p.grad is not None, f"Gradient for '{name}' is None" - assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" - def test_same_dgrad_wgrad_reduces_to_plain_grad(self): """When dgrad format == wgrad format, grad_output can be a plain quantizer.""" - torch.manual_seed(302) - in_features, out_features, batch = 128, 128, 32 - - model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - inp = torch.randn( - batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + self._assert_three_format_routing( + _make_nvfp4_quantizer, + _make_mxfp8_quantizer, + _make_mxfp8_quantizer, + plain_grad_output=True, + seed=302, ) - def factory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type == "weight": - return HybridQuantizer( - rowwise_quantizer=_make_nvfp4_quantizer(), - columnwise_quantizer=_make_mxfp8_quantizer(), - ) - if is_linear and role.tensor_type == "input": - return HybridQuantizer( - rowwise_quantizer=_make_nvfp4_quantizer(), - columnwise_quantizer=_make_mxfp8_quantizer(), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_mxfp8_quantizer() - return _make_nvfp4_quantizer() - - with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): - out = model(inp) - - loss = out.float().sum() - loss.backward() - - assert inp.grad is not None - assert not torch.isnan(inp.grad).any() - for name, p in model.named_parameters(): - assert p.grad is not None, f"Gradient for '{name}' is None" - # --------------------------------------------------------------------------- # All-modules test: hybrid quantization through every TE module type @@ -3238,21 +3447,48 @@ class TestHybridAllModules: batch = 16 seq_len = 8 - def _run_fwd_bwd(self, model, inp): - hybrid_recipe = recipe.CustomRecipe(qfactory=_make_hybrid_fp8_factory()) - with autocast(enabled=True, recipe=hybrid_recipe): - out = model(inp) - loss = out.float().sum() - loss.backward() + def _run_fwd_bwd(self, model, inp, *model_args, output_atol=0.0, param_atols=None): + """Compare same-format hybrid numerics against the native FP8 recipe.""" + import copy + + model_native = copy.deepcopy(model) + param_atols = {} if param_atols is None else param_atols + grad_output = torch.randn_like(inp) + + def run(run_model, run_recipe): + run_model.zero_grad(set_to_none=True) + run_input = inp.detach().clone().requires_grad_(True) + _set_quantization_test_seed(3370) + with autocast(enabled=True, recipe=run_recipe): + output = run_model(run_input, *model_args) + output.backward(grad_output) + return ( + output.detach().clone(), + run_input.grad.detach().clone(), + { + name: param.grad.detach().clone() + for name, param in run_model.named_parameters() + if param.grad is not None + }, + ) - assert not torch.isnan(out).any(), "Output contains NaN" - assert not torch.isinf(out).any(), "Output contains Inf" - assert inp.grad is not None, "Input gradient is None" - assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" - for name, p in model.named_parameters(): - if p.requires_grad: - assert p.grad is not None, f"Gradient for '{name}' is None" - assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + native_result = run(model_native, recipe.Float8CurrentScaling()) + hybrid_result = run( + model, + recipe.CustomRecipe(qfactory=_make_hybrid_fp8_factory()), + ) + + torch.testing.assert_close(hybrid_result[0], native_result[0], rtol=0.0, atol=output_atol) + torch.testing.assert_close(hybrid_result[1], native_result[1], rtol=0.0, atol=0.0) + assert hybrid_result[2].keys() == native_result[2].keys() + for name in hybrid_result[2]: + torch.testing.assert_close( + hybrid_result[2][name], + native_result[2][name], + rtol=0.0, + atol=param_atols.get(name, 0.0), + msg=f"Parameter gradient {name!r} differs", + ) def test_linear(self): torch.manual_seed(500) @@ -3300,7 +3536,15 @@ def test_layernorm_mlp(self): dtype=torch.bfloat16, requires_grad=True, ) - self._run_fwd_bwd(model, inp) + # Native fuses LN+quantize while Hybrid takes the explicit quantize path. + # The only non-bitwise results are output (1 BF16 quantum here) and + # fc2_weight grad; all other gradients remain zero-tolerance checks. + self._run_fwd_bwd( + model, + inp, + output_atol=0.0009765625, + param_atols={"fc2_weight": 0.0234375}, + ) def test_grouped_linear(self): torch.manual_seed(504) @@ -3322,20 +3566,7 @@ def test_grouped_linear(self): rem = self.batch % num_gemms m_splits = [base + (1 if i < rem else 0) for i in range(num_gemms)] - hybrid_recipe = recipe.CustomRecipe(qfactory=_make_hybrid_fp8_factory()) - with autocast(enabled=True, recipe=hybrid_recipe): - out = model(inp, m_splits) - loss = out.float().sum() - loss.backward() - - assert not torch.isnan(out).any(), "Output contains NaN" - assert not torch.isinf(out).any(), "Output contains Inf" - assert inp.grad is not None, "Input gradient is None" - assert not torch.isnan(inp.grad).any(), "Input gradient contains NaN" - for name, p in model.named_parameters(): - if p.requires_grad: - assert p.grad is not None, f"Gradient for '{name}' is None" - assert not torch.isnan(p.grad).any(), f"Gradient for '{name}' contains NaN" + self._run_fwd_bwd(model, inp, m_splits) def test_transformer_layer(self): torch.manual_seed(503) @@ -3356,7 +3587,15 @@ def test_transformer_layer(self): dtype=torch.bfloat16, requires_grad=True, ) - self._run_fwd_bwd(model, inp) + # The LayerNormMLP submodule has the same fused-vs-explicit ordering. + # Keep dgrad and every other parameter exact; bound only the observed + # BF16 output and final projection-weight accumulation roundoff. + self._run_fwd_bwd( + model, + inp, + output_atol=0.015625, + param_atols={"layernorm_mlp.fc2_weight": 0.140625}, + ) @requires_fp8 @@ -3523,6 +3762,49 @@ def test_hybrid_split_quantize_rejects_mixed_parent_usage_flags(self): with pytest.raises(ValueError, match="mixed parent usage flags"): _hybrid_split_quantize(tensor, [16, 16], quantizers) + @requires_fp8_and_nvfp4 + @pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason=( + "Production follow-up: GroupedLinear bulk split-quantize ignores " + "HybridQuantizer.columnwise_source" + ), + ) + def test_hybrid_split_quantize_honors_rowwise_dequantized_source(self): + """Bulk split must match standalone HybridQuantizer provenance semantics.""" + from transformer_engine.pytorch.module.grouped_linear import ( + _hybrid_split_quantize, + ) + + torch.manual_seed(3598) + # NVFP4 grouped split-quantize requires each M split to be a multiple + # of 64; valid geometry ensures the xfail reaches provenance checking. + tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + + def make_quantizer(): + return HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + + actual = _hybrid_split_quantize( + tensor, + [64, 64], + [make_quantizer(), make_quantizer()], + ) + expected = [ + make_quantizer().quantize(part) for part in torch.split(tensor, [64, 64], dim=0) + ] + + for index, (actual_part, expected_part) in enumerate(zip(actual, expected)): + _assert_storage_data_exact( + actual_part.columnwise_sub_storage, + expected_part.columnwise_sub_storage, + context=f"GroupedLinear split {index} columnwise provenance", + ) + # =========================================================================== # Quantized Parameters (quantized_model_init) tests for hybrid quantization @@ -3691,10 +3973,10 @@ def test_quantized_param_skips_workspace(self): inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) with autocast(enabled=True, recipe=hybrid_recipe): - out = model(inp) + out = model(inp, is_first_microbatch=True) assert out.shape == (32, 128) - assert not torch.isnan(out).any() + assert "weight" not in model._fp8_workspaces def test_bf16_weight_creates_hybrid_workspace(self): """When weight is BF16 and recipe produces HybridQuantizer, the workspace @@ -3704,10 +3986,13 @@ def test_bf16_weight_creates_hybrid_workspace(self): inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) with autocast(enabled=True, recipe=hybrid_recipe): - out = model(inp) + out = model(inp, is_first_microbatch=True) assert out.shape == (32, 128) - assert not torch.isnan(out).any() + workspace = model._fp8_workspaces.get("weight") + assert isinstance(workspace, HybridQuantizedTensorStorage) + assert workspace.rowwise_sub_storage is not None + assert workspace.columnwise_sub_storage is not None def test_workspace_cache_reuse_across_microbatches(self): """Cached hybrid workspace should be reused on 2nd+ microbatches.""" @@ -3718,12 +4003,16 @@ def test_workspace_cache_reuse_across_microbatches(self): with autocast(enabled=True, recipe=hybrid_recipe): with torch.no_grad(): out1 = model(inp, is_first_microbatch=True) + workspace = model._fp8_workspaces["weight"] + buffers = _as_data_tensor_tuple(workspace) out2 = model(inp, is_first_microbatch=False) - # Both should produce valid, identical results (same weight, same input) - assert not torch.isnan(out1).any() - assert not torch.isnan(out2).any() - torch.testing.assert_close(out1, out2) + assert isinstance(workspace, HybridQuantizedTensorStorage) + assert model._fp8_workspaces["weight"] is workspace + current_buffers = _as_data_tensor_tuple(workspace) + assert len(current_buffers) == len(buffers) + assert all(current is cached for current, cached in zip(current_buffers, buffers)) + torch.testing.assert_close(out1, out2, rtol=0.0, atol=0.0) def test_workspace_cache_invalidation_on_usage_change(self): """If usage requirements change (e.g. inference→training), the cache @@ -3736,17 +4025,26 @@ def test_workspace_cache_invalidation_on_usage_change(self): # First pass: inference (no columnwise needed) with torch.no_grad(): with autocast(enabled=True, recipe=hybrid_recipe): - out_infer = model(inp, is_first_microbatch=True) - assert not torch.isnan(out_infer).any() + model(inp, is_first_microbatch=True) + inference_workspace = model._fp8_workspaces["weight"] + assert isinstance(inference_workspace, HybridQuantizedTensorStorage) + assert inference_workspace.rowwise_sub_storage is not None + assert inference_workspace.columnwise_sub_storage is None # Second pass: training (columnwise now needed for backward) with autocast(enabled=True, recipe=hybrid_recipe): out_train = model(inp, is_first_microbatch=True) + training_workspace = model._fp8_workspaces["weight"] + + assert isinstance(training_workspace, HybridQuantizedTensorStorage) + assert training_workspace is not inference_workspace + assert training_workspace.rowwise_sub_storage is not None + assert training_workspace.columnwise_sub_storage is not None + loss = out_train.float().sum() loss.backward() assert inp.grad is not None - assert not torch.isnan(inp.grad).any() # --------------------------------------------------------------------------- @@ -4034,6 +4332,14 @@ def _reference_fp8_bytes(master, scale, dtype=torch.bfloat16): quantizer.update_quantized(master.reshape(1, -1), temp) return temp._data.reshape(-1) + @staticmethod + def _logical_float8_bytes(storage): + """Return FP8 payload in the tensor's logical row-major order.""" + if storage._data is not None: + return storage._data.reshape(-1) + assert storage._transpose is not None + return storage._transpose.transpose(-2, -1).contiguous().reshape(-1) + def test_fp8_current_transpose_only_nonzero_offset(self): """Current-scaling distopt update handles Hopper-style columnwise storage. @@ -4114,6 +4420,19 @@ def test_fp8_current_same_format_full_master(self): assert weight._rowwise_storage is not None assert weight._columnwise_storage is not None + master_bf16 = master_flat.to(weight.dtype).reshape(weight.shape) + expected_row = weight._quantizer.rowwise_quantizer.copy().quantize(master_bf16) + expected_column = weight._quantizer.columnwise_quantizer.copy().quantize(master_bf16) + _assert_storage_data_exact( + weight._rowwise_storage, + expected_row, + context="full-master rowwise independent quantization", + ) + _assert_storage_data_exact( + weight._columnwise_storage, + expected_column, + context="full-master columnwise independent quantization", + ) dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) # FP8 E4M3 round-trip; matches the loose tolerance the equivalent @@ -4137,20 +4456,58 @@ def test_fp8_current_nonzero_start_offset(self): weight, hp_master_full = _build_hybrid_linear_weight(64, 64, hybrid_recipe) half = hp_master_full.numel() // 2 - hp_master_shard = hp_master_full.view(-1)[half:].contiguous() + # Negation preserves the shard amax while guaranteeing that this update + # is observable. The fixed seed puts the full-tensor amax in this shard, + # so the current-scaling factor must remain bitwise stable. + hp_master_shard = -hp_master_full.view(-1)[half:].contiguous() start_offset = half + before = {} + for direction, storage in ( + ("rowwise", weight._rowwise_storage), + ("columnwise", weight._columnwise_storage), + ): + assert storage is not None + before[direction] = { + "bytes": self._logical_float8_bytes(storage).clone(), + "scale_inv": storage._scale_inv.clone(), + "dequantized": storage.dequantize(dtype=torch.float32).reshape(-1).clone(), + } + quantize_master_weights([weight], [hp_master_shard], [start_offset], group=group) post_all_gather_processing([weight]) - # The second-half slice (which the master shard covered) should match. - # The first-half slice was already written at quantized_model_init time - # from the same high-precision init val, so it should also match (modulo - # the second amax all-reduce shifting the per-tensor scale). - dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32).view(-1) - dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32).view(-1) - torch.testing.assert_close(dq_row[start_offset:], hp_master_shard, rtol=0.125, atol=0.1) - torch.testing.assert_close(dq_col[start_offset:], hp_master_shard, rtol=0.125, atol=0.1) + for direction, storage in ( + ("rowwise", weight._rowwise_storage), + ("columnwise", weight._columnwise_storage), + ): + previous = before[direction] + actual_bytes = self._logical_float8_bytes(storage) + torch.testing.assert_close(storage._scale_inv, previous["scale_inv"], rtol=0, atol=0) + torch.testing.assert_close( + actual_bytes[:start_offset], previous["bytes"][:start_offset], rtol=0, atol=0 + ) + expected_shard_bytes = self._reference_fp8_bytes( + hp_master_shard.to(weight.dtype), + torch.reciprocal(storage._scale_inv), + weight.dtype, + ) + torch.testing.assert_close( + actual_bytes[start_offset:], expected_shard_bytes, rtol=0, atol=0 + ) + assert not torch.equal( + actual_bytes[start_offset:], previous["bytes"][start_offset:] + ), f"{direction} updated shard unexpectedly retained every FP8 byte" + dequantized = storage.dequantize(dtype=torch.float32).reshape(-1) + torch.testing.assert_close( + dequantized[:start_offset], + previous["dequantized"][:start_offset], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + dequantized[start_offset:], hp_master_shard, rtol=0.125, atol=0.1 + ) def test_fp8_delayed_same_format_full_master(self): """Same-format delayed scaling on both directions. Both sub-storages @@ -4329,6 +4686,27 @@ def test_nvfp4_rowwise_raises(self): with pytest.raises(NotImplementedError, match="NVFP4Quantizer rowwise"): quantize_master_weights([weight], [master_flat], [0], group=group) + @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") + def test_nvfp4_columnwise_raises(self): + """NVFP4 in only the columnwise slot is rejected per-direction.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=True + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="NVFP4Quantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + @pytest.mark.skipif( not fp8_block_scaling_available, reason=f"Float8 block scaling: {reason_for_no_fp8_block_scaling}", @@ -4347,6 +4725,33 @@ def test_blockwise_rowwise_raises(self): with pytest.raises(NotImplementedError, match="Float8BlockQuantizer rowwise"): quantize_master_weights([weight], [master_flat], [0], group=group) + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8 block scaling: {reason_for_no_fp8_block_scaling}", + ) + def test_blockwise_columnwise_raises(self): + """Float8BlockQuantizer in only the columnwise slot is rejected.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + col_factory=lambda: Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + grad_factory=lambda: Float8CurrentScalingQuantizer( + tex.DType.kFloat8E5M2, device="cuda" + ), + ) + weight, hp_master = _build_hybrid_linear_weight(128, 128, hybrid_recipe) + master_flat = hp_master.view(-1).contiguous() + + with pytest.raises(NotImplementedError, match="Float8BlockQuantizer columnwise"): + quantize_master_weights([weight], [master_flat], [0], group=group) + def test_rowwise_only_fp8_current_full_master(self): """Single-direction hybrid: columnwise dropped via update_usage. @@ -4373,6 +4778,14 @@ def test_rowwise_only_fp8_current_full_master(self): # Columnwise stays dropped (the cast must not silently revive it). assert weight._columnwise_storage is None + expected_row = weight._quantizer.rowwise_quantizer.copy().quantize( + master_flat.to(weight.dtype).reshape(weight.shape) + ) + _assert_storage_data_exact( + weight._rowwise_storage, + expected_row, + context="rowwise-only full-master independent quantization", + ) # Rowwise is populated and dequantizes close to the master. dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) @@ -4401,6 +4814,14 @@ def test_columnwise_only_fp8_current_full_master(self): # Rowwise stays dropped (the cast must not silently revive it). assert weight._rowwise_storage is None + expected_column = weight._quantizer.columnwise_quantizer.copy().quantize( + master_flat.to(weight.dtype).reshape(weight.shape) + ) + _assert_storage_data_exact( + weight._columnwise_storage, + expected_column, + context="columnwise-only full-master independent quantization", + ) # Columnwise is populated and dequantizes close to the master. dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) @@ -4530,20 +4951,13 @@ def test_quantize_inplace_updates_data(self): original = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") tensor = hq.quantize(original) - dq_before = tensor.dequantize().clone() - # Update with different data new_data = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") - tensor.quantize_(new_data) - - dq_after = tensor.dequantize() + expected = hq.quantize(new_data) + result = tensor.quantize_(new_data) - # Should be close to new data, not old data - diff_new = (dq_after.float() - new_data.float()).abs().mean() - diff_old = (dq_after.float() - original.float()).abs().mean() - assert ( - diff_new < diff_old - ), f"After quantize_(), data is closer to old ({diff_old:.4f}) than new ({diff_new:.4f})" + assert result is tensor + _assert_hybrid_tensor_exact(tensor, expected, context="quantize_") def test_quantize_inplace_preserves_tensor_identity(self): """quantize_() should update in-place, not create a new tensor.""" @@ -4556,11 +4970,10 @@ def test_quantize_inplace_preserves_tensor_identity(self): original = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") tensor = hq.quantize(original) - tensor_id = id(tensor) new_data = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") result = tensor.quantize_(new_data) - assert id(tensor) == tensor_id, "quantize_() should return same object" + assert result is tensor, "quantize_() must return the object it updated" # noop_flag is a delayed-scaling feature; not tested here since # delayed scaling is out of scope for hybrid quantization. @@ -5014,21 +5427,86 @@ def _run_training_loop(self, model, train_recipe, x, target, num_steps): master_weights=True, master_weight_dtype=torch.float32, ) - losses = [] - for _ in range(num_steps): + trajectory = [] + hybrid_metadata_trajectory = [] + for step in range(num_steps): optimizer.zero_grad(set_to_none=True) + step_input = x.detach().clone().requires_grad_(True) + _set_quantization_test_seed(199 + step) with autocast(enabled=True, recipe=train_recipe): - output = model(x) + output = model(step_input) loss = torch.nn.functional.mse_loss(output, target) - losses.append(loss.item()) loss.backward() + gradients = { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + } + pre_step_cuda_rng_state = torch.cuda.get_rng_state() optimizer.step() - master_params = [ - optimizer.get_unscaled_state(p, "master_param") - for p in model.parameters() - if p.requires_grad - ] - return losses, master_params + post_step_cuda_rng_state = torch.cuda.get_rng_state() + + hybrid_parameters = [ + (name, param) + for name, param in model.named_parameters() + if isinstance(param, HybridQuantizedTensor) + ] + expected_storages = {} + torch.cuda.set_rng_state(pre_step_cuda_rng_state) + try: + for name, param in hybrid_parameters: + master = optimizer.get_unscaled_state(param, "master_param") + expected_storages[name] = ( + param._quantizer.rowwise_quantizer.copy().quantize(master), + param._quantizer.columnwise_quantizer.copy().quantize(master), + ) + finally: + # The oracle must be observational only. Restore the state left + # by the real optimizer step after replaying NVFP4 stochastic + # rounding from its pre-step RNG state. + torch.cuda.set_rng_state(post_step_cuda_rng_state) + + hybrid_metadata = {} + for name, param in hybrid_parameters: + expected_row, expected_column = expected_storages[name] + _assert_storage_data_exact( + param.rowwise_sub_storage, + expected_row, + context=f"step {step} {name} rowwise writeback", + ) + _assert_storage_data_exact( + param.columnwise_sub_storage, + expected_column, + context=f"step {step} {name} columnwise writeback", + ) + hybrid_metadata[name] = { + "rowwise": _snapshot_storage_tensor_metadata( + param.rowwise_sub_storage, clone=True + ), + "columnwise": _snapshot_storage_tensor_metadata( + param.columnwise_sub_storage, clone=True + ), + } + hybrid_metadata_trajectory.append(hybrid_metadata) + logical_parameters = { + name: ( + param.dequantize(dtype=torch.float32).detach().clone() + if isinstance(param, QuantizedTensor) + else param.detach().float().clone() + ) + for name, param in model.named_parameters() + } + trajectory.append( + { + "output": output.detach().clone(), + "loss": loss.detach().clone(), + "input_gradient": step_input.grad.detach().clone(), + "parameter_gradients": gradients, + "logical_parameters": logical_parameters, + "optimizer": _clone_nested_state(optimizer.state_dict()), + } + ) + return trajectory, hybrid_metadata_trajectory def _test_equivalence(self): model_ref, model_hyb = self._build_models() @@ -5037,51 +5515,36 @@ def _test_equivalence(self): x = torch.randn(4, 32, self.hidden_size, dtype=torch.bfloat16, device="cuda") target = torch.randn_like(x) - torch.manual_seed(199) - torch.cuda.manual_seed_all(199) - losses_ref, masters_ref = self._run_training_loop( + ref_trajectory, ref_hybrid_metadata = self._run_training_loop( model_ref, self._vanilla_recipe(), x, target, self.num_steps, ) - torch.manual_seed(199) - torch.cuda.manual_seed_all(199) - losses_hyb, masters_hyb = self._run_training_loop( + hybrid_trajectory, hybrid_metadata_trajectory = self._run_training_loop( model_hyb, self._hybrid_recipe(), x, target, self.num_steps, ) - - # Losses should be very close (same quantization, same training dynamics) - for i, (lr, lh) in enumerate(zip(losses_ref, losses_hyb)): - assert abs(lr - lh) < 0.1 * max( - abs(lr), 1e-6 - ), f"Step {i}: loss diverged — vanilla={lr:.6f}, hybrid={lh:.6f}" - - # Master weights should be close after training - for i, (mr, mh) in enumerate(zip(masters_ref, masters_hyb)): - torch.testing.assert_close( - mr, - mh, - rtol=1e-3, - atol=1e-3, - msg=f"Master weight {i} diverged after {self.num_steps} steps", - ) + _assert_nested_state_exact( + hybrid_trajectory, + ref_trajectory, + path="same-format quantized-parameter trajectory", + ) + assert all(not step_metadata for step_metadata in ref_hybrid_metadata) + assert len(hybrid_metadata_trajectory) == self.num_steps + for step_metadata in hybrid_metadata_trajectory: + assert step_metadata.keys() == {"weight"} + assert step_metadata["weight"]["rowwise"] is not None + assert step_metadata["weight"]["columnwise"] is not None @requires_fp8 class TestQuantizedParamsEquivalenceFP8CurrentScaling(_QuantizedParamsEquivalenceBase): - """Vanilla Float8CurrentScaling vs hybrid FP8 current (same format both dirs). - - Note: vanilla Float8Tensor params use the fused multi_tensor_adam_fp8 - kernel in FusedAdam, while HybridQuantizedTensor falls to the FP32 - master + quantize_() writeback path. These are numerically different - codepaths, so we use relaxed tolerances. - """ + """Vanilla versus same-format hybrid FP8-current training parity.""" def _vanilla_recipe(self): return recipe.Float8CurrentScaling() @@ -5090,37 +5553,7 @@ def _hybrid_recipe(self): return recipe.CustomRecipe(qfactory=_hybrid_fp8_current_qfactory) def test_equivalence(self): - model_ref, model_hyb = self._build_models() - - torch.manual_seed(99) - x = torch.randn(4, 32, self.hidden_size, dtype=torch.bfloat16, device="cuda") - target = torch.randn_like(x) - - losses_ref, _ = self._run_training_loop( - model_ref, - self._vanilla_recipe(), - x, - target, - self.num_steps, - ) - losses_hyb, _ = self._run_training_loop( - model_hyb, - self._hybrid_recipe(), - x, - target, - self.num_steps, - ) - - # Both should decrease (training works in both paths) - assert losses_ref[-1] < losses_ref[0], f"Vanilla loss didn't decrease: {losses_ref}" - assert losses_hyb[-1] < losses_hyb[0], f"Hybrid loss didn't decrease: {losses_hyb}" - - # Losses should be in the same ballpark (different optimizer kernels - # cause small divergence that compounds over steps) - for i, (lr, lh) in enumerate(zip(losses_ref, losses_hyb)): - assert ( - abs(lr - lh) / max(abs(lr), 1e-6) < 0.5 - ), f"Step {i}: losses diverged too much — vanilla={lr:.6f}, hybrid={lh:.6f}" + self._test_equivalence() @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") @@ -5208,7 +5641,7 @@ def test_state_dict_save_load_roundtrip(self): with autocast(enabled=True, recipe=hybrid_recipe): out_after = model2(inp) - torch.testing.assert_close(out_before, out_after) + torch.testing.assert_close(out_before, out_after, rtol=0.0, atol=0.0) def test_state_dict_contains_weight(self): """state_dict should contain the weight key.""" @@ -5276,79 +5709,104 @@ def test_state_dict_torch_save_load(self): with autocast(enabled=True, recipe=hybrid_recipe): out_after = model2(inp) - torch.testing.assert_close(out_before, out_after) + torch.testing.assert_close(out_before, out_after, rtol=0.0, atol=0.0) finally: os.unlink(tmp_path) + @staticmethod + def _checkpoint_training_step(model, optimizer, x, target, hybrid_recipe): + optimizer.zero_grad(set_to_none=True) + step_input = x.detach().clone().requires_grad_(True) + with autocast(enabled=True, recipe=hybrid_recipe): + output = model(step_input) + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + gradients = { + name: param.grad.detach().clone() + for name, param in model.named_parameters() + if param.grad is not None + } + optimizer.step() + return { + "output": output.detach().clone(), + "loss": loss.detach().clone(), + "input_gradient": step_input.grad.detach().clone(), + "gradients": gradients, + "parameters": _snapshot_model_parameters(model), + "optimizer": _clone_nested_state(optimizer.state_dict()), + } + def test_checkpoint_resume_training(self): """Save mid-training, load into new model+optimizer, verify training continues.""" - import tempfile import os + import tempfile - torch.manual_seed(42) + _set_quantization_test_seed(42) hybrid_recipe = self._hybrid_checkpoint_recipe() with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() - optimizer = te.optimizers.FusedAdam( model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, ) - x = torch.randn(4, 32, 256, dtype=torch.bfloat16, device="cuda") target = torch.randn_like(x) - # Train for a few steps for _ in range(3): - optimizer.zero_grad(set_to_none=True) - with autocast(enabled=True, recipe=hybrid_recipe): - output = model(x) - loss = torch.nn.functional.mse_loss(output, target) - loss.backward() - optimizer.step() + self._checkpoint_training_step(model, optimizer, x, target, hybrid_recipe) - loss_before_save = loss.item() - - # Save checkpoint - with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as f: + with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as checkpoint_file: torch.save( { "model": model.state_dict(), "optimizer": optimizer.state_dict(), + "cpu_rng_state": torch.get_rng_state(), + "cuda_rng_state": torch.cuda.get_rng_state(), }, - f.name, + checkpoint_file.name, ) - tmp_path = f.name + checkpoint_path = checkpoint_file.name try: - # Load into fresh model + uninterrupted = self._checkpoint_training_step( + model, + optimizer, + x, + target, + hybrid_recipe, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): - model2 = Linear(256, 256, params_dtype=torch.bfloat16).cuda() - optimizer2 = te.optimizers.FusedAdam( - model2.parameters(), + resumed_model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + resumed_optimizer = te.optimizers.FusedAdam( + resumed_model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, ) + checkpoint = torch.load(checkpoint_path, weights_only=False) + resumed_model.load_state_dict(checkpoint["model"]) + resumed_optimizer.load_state_dict(checkpoint["optimizer"]) + torch.set_rng_state(checkpoint["cpu_rng_state"]) + torch.cuda.set_rng_state(checkpoint["cuda_rng_state"]) + + resumed = self._checkpoint_training_step( + resumed_model, + resumed_optimizer, + x, + target, + hybrid_recipe, + ) - checkpoint = torch.load(tmp_path, weights_only=False) - model2.load_state_dict(checkpoint["model"]) - optimizer2.load_state_dict(checkpoint["optimizer"]) - - # Continue training -- loss should not spike - optimizer2.zero_grad(set_to_none=True) - with autocast(enabled=True, recipe=hybrid_recipe): - output2 = model2(x) - loss_after_load = torch.nn.functional.mse_loss(output2, target).item() - - assert loss_after_load <= loss_before_save * 1.5, ( - f"Loss spiked after checkpoint resume: {loss_before_save:.4f} →" - f" {loss_after_load:.4f}" + _assert_nested_state_exact( + resumed, + uninterrupted, + path="checkpoint continuation", ) finally: - os.unlink(tmp_path) + os.unlink(checkpoint_path) # --------------------------------------------------------------------------- @@ -5553,6 +6011,10 @@ def test_split_preserves_hybrid_type(self, hybrid_param): dim0 = hybrid_param.shape[0] chunk_size = dim0 // 2 pieces = torch.split(hybrid_param, chunk_size, dim=0) + expected_rowwise = torch.split(hybrid_param.rowwise_sub_storage, chunk_size, dim=0) + expected_columnwise = torch.split(hybrid_param.columnwise_sub_storage, chunk_size, dim=0) + assert len(pieces) == len(expected_rowwise) == len(expected_columnwise) + assert len(pieces) >= 2 for piece in pieces: assert isinstance( @@ -5560,13 +6022,26 @@ def test_split_preserves_hybrid_type(self, hybrid_param): ), f"Expected HybridQuantizedTensor, got {type(piece).__name__}" assert piece.rowwise_sub_storage is not None assert piece.columnwise_sub_storage is not None + for index, (piece, expected_row, expected_column) in enumerate( + zip(pieces, expected_rowwise, expected_columnwise) + ): + _assert_storage_data_exact( + piece.rowwise_sub_storage, + expected_row, + context=f"split {index} rowwise", + ) + _assert_storage_data_exact( + piece.columnwise_sub_storage, + expected_column, + context=f"split {index} columnwise", + ) total_rows = sum(p.shape[0] for p in pieces) assert total_rows == dim0 orig_deq = hybrid_param.dequantize() reassembled = torch.cat([p.dequantize() for p in pieces], dim=0) - torch.testing.assert_close(orig_deq, reassembled) + torch.testing.assert_close(orig_deq, reassembled, rtol=0.0, atol=0.0) def test_split_sub_storage_types_preserved(self, hybrid_param): """After split, sub-storage types must match the original.""" @@ -5628,16 +6103,18 @@ def test_copy_between_hybrid_tensors(self, hybrid_param): aten.copy_.default(dst, hybrid_param) dst_deq = dst.dequantize() - torch.testing.assert_close(src_deq, dst_deq) + torch.testing.assert_close(src_deq, dst_deq, rtol=0.0, atol=0.0) + _assert_hybrid_tensor_exact(dst, hybrid_param, context="hybrid copy_") def test_copy_from_bf16_to_hybrid(self, hybrid_param): """copy_ from BF16 into HybridQuantizedTensor triggers quantize_.""" param = hybrid_param.detach() bf16_data = torch.randn_like(param.dequantize()) + expected = param._quantizer.quantize(bf16_data) aten.copy_.default(param, bf16_data) - result_deq = param.dequantize() assert isinstance(param, HybridQuantizedTensor) - assert result_deq.shape == bf16_data.shape + _assert_hybrid_tensor_exact(param, expected, context="BF16 copy_") + torch.testing.assert_close(param.dequantize(), expected.dequantize(), rtol=0.0, atol=0.0) def test_new_zeros_returns_hybrid(self, hybrid_param): """new_zeros must return a usable HybridQuantizedTensor container. @@ -5668,7 +6145,10 @@ def test_new_zeros_returns_hybrid(self, hybrid_param): # Functional contract: the container must be writable via copy_ from # another hybrid (how FSDP2 populates the buffer post-gather). aten.copy_.default(result, hybrid_param) - torch.testing.assert_close(result.dequantize(), hybrid_param.dequantize()) + torch.testing.assert_close( + result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="new_zeros then copy_") def test_empty_like_returns_hybrid(self, hybrid_param): """empty_like must return a HybridQuantizedTensor.""" @@ -5686,7 +6166,10 @@ def test_clone_returns_hybrid(self, hybrid_param): result, HybridQuantizedTensor ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" assert result is not hybrid_param - torch.testing.assert_close(result.dequantize(), hybrid_param.dequantize()) + torch.testing.assert_close( + result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="clone") # --------------------------------------------------------------------------- @@ -5862,6 +6345,7 @@ def test_post_all_gather_buffer_reuse(self, hybrid_param): assert ( second_result is first_result ), "Buffer reuse: post_all_gather(out=prev) should return the same object" + _assert_hybrid_tensor_exact(second_result, hybrid_param, context="post-all-gather reuse") def test_post_all_gather_dequantize_matches_original(self, hybrid_param): """Reconstructed tensor should dequantize close to the original.""" @@ -5881,7 +6365,8 @@ def test_post_all_gather_dequantize_matches_original(self, hybrid_param): out=None, ) result_deq = result.dequantize() - torch.testing.assert_close(orig_deq, result_deq) + _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather") + torch.testing.assert_close(orig_deq, result_deq, rtol=0.0, atol=0.0) def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): """Reconstructed tensor's sub-storages match the original types.""" @@ -5902,6 +6387,7 @@ def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): out=None, ) assert type(result.rowwise_sub_storage) is orig_row_type + _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather storage types") assert type(result.columnwise_sub_storage) is orig_col_type @@ -5936,7 +6422,8 @@ def test_pre_post_roundtrip_preserves_data(self, hybrid_param): hybrid_param.dtype, out=None, ) - torch.testing.assert_close(orig_deq, result.dequantize()) + _assert_hybrid_tensor_exact(result, hybrid_param, context="FSDP roundtrip") + torch.testing.assert_close(orig_deq, result.dequantize(), rtol=0.0, atol=0.0) def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): """Second roundtrip with out=previous preserves data (iteration 2+ simulation).""" @@ -5968,7 +6455,10 @@ def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): out=first_result, ) assert second_result is first_result - torch.testing.assert_close(hybrid_param.dequantize(), second_result.dequantize()) + torch.testing.assert_close( + hybrid_param.dequantize(), second_result.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(second_result, hybrid_param, context="FSDP roundtrip reuse") def test_scale_refresh_across_iterations(self): """After a sharded optimizer-style requantize, iter-2+ gathers see the new scale. @@ -6039,7 +6529,10 @@ def test_scale_refresh_across_iterations(self): torch.testing.assert_close( hybrid_param.dequantize(), gathered_refreshed.dequantize(), + rtol=0.0, + atol=0.0, ) + _assert_hybrid_tensor_exact(gathered_refreshed, hybrid_param, context="FSDP scale refresh") # And the magnitude really did change (sanity: this test would pass # vacuously if update_quantized didn't actually change anything). assert gathered_refreshed.dequantize().abs().max() > 10.0, ( @@ -6047,12 +6540,6 @@ def test_scale_refresh_across_iterations(self): "weight; the scale-refresh invariant is not being exercised" ) - @pytest.mark.xfail( - reason=( - "Tracked by #3158: Hybrid FSDP2 does not support NVFP4 sub-storages yet; NVFP4 uses " - "dedicated tensor hooks and does not implement the hybrid fsdp_buffer_fields protocol." - ) - ) def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): """Hybrid FSDP2 with an NVFP4 sub-storage must raise a clear error. @@ -6094,7 +6581,8 @@ def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): ) msg = str(exc_info.value) assert "NVFP4Tensor" in msg - assert "hybrid_quantization_fsdp.md" in msg + assert "rowwise sub-storage" in msg + assert "Use a supported sub-quantizer" in msg assert "fsdp_buffer_fields" in msg @@ -6127,7 +6615,8 @@ def test_make_like_preserves_sub_storages(self): assert copy.shape == param.shape assert type(copy.rowwise_sub_storage) is type(param.rowwise_sub_storage) assert type(copy.columnwise_sub_storage) is type(param.columnwise_sub_storage) - torch.testing.assert_close(copy.dequantize(), param.dequantize()) + _assert_hybrid_tensor_exact(copy, param, context="make_like") + torch.testing.assert_close(copy.dequantize(), param.dequantize(), rtol=0.0, atol=0.0) def test_make_like_is_independent(self): """make_like result should not share the same tensor identity.""" diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 7d96d0720a7..2d1220e3370 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -256,7 +256,15 @@ def test_grouped_split_all_identity_uses_plain_tensor_views(self): x, m_splits, cast_quantizers, activation_dtype=torch.bfloat16 ) assert all(isinstance(t, IdentityTensorStorage) for t in cast_out) - assert all(t.dequantize().dtype == torch.float32 for t in cast_out) + for actual, expected in zip(cast_out, torch.split(x, m_splits)): + dequantized = actual.dequantize() + assert dequantized.dtype == torch.float32 + torch.testing.assert_close( + dequantized, + expected.to(torch.float32), + rtol=0.0, + atol=0.0, + ) def test_grouped_split_rejects_mixed_identity_and_quantized_operands(self): from transformer_engine.pytorch.module.grouped_linear import ( @@ -465,17 +473,34 @@ def qfactory(role): # pylint: disable=unused-argument columnwise_quantizer=IdentityQuantizer(), ) + torch.manual_seed(1701) + ref = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) custom_recipe = CustomRecipe(qfactory=qfactory) + torch.manual_seed(1702) with te.quantized_model_init(enabled=True, recipe=custom_recipe): - op = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) + test = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) - x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with torch.no_grad(): + test.weight.copy_(ref.weight) + + torch.manual_seed(1703) + x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + x_ref = x.detach().clone().requires_grad_(True) + x_test = x.detach().clone().requires_grad_(True) + + y_ref = ref(x_ref) with te.autocast(enabled=True, recipe=custom_recipe): - y = op(x) - y.sum().backward() + y_test = test(x_test) + + torch.manual_seed(1704) + grad_output = torch.randn_like(y_ref) + y_ref.backward(grad_output) + y_test.backward(grad_output) - assert isinstance(op.weight, HybridQuantizedTensor) - assert x.grad is not None + assert isinstance(test.weight, HybridQuantizedTensor) + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0.0, atol=0.0) + torch.testing.assert_close(test.weight.grad, ref.weight.grad, rtol=0.0, atol=0.0) @pytest.mark.parametrize( "qfactory", @@ -493,14 +518,23 @@ def qfactory(role): # pylint: disable=unused-argument def test_te_ops_quantize_then_gelu_accepts_identity_backed_tensors(self, qfactory): import transformer_engine.pytorch.ops as te_ops - model = te_ops.Sequential(te_ops.Quantize(forward=True), te_ops.GELU()) - x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + ref = te_ops.GELU() + test = te_ops.Sequential(te_ops.Quantize(forward=True), te_ops.GELU()) + x = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + x_ref = x.detach().clone().requires_grad_(True) + x_test = x.detach().clone().requires_grad_(True) + y_ref = ref(x_ref) with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): - y = model(x) + y_test = test(x_test) - assert isinstance(y, torch.Tensor) - assert y.shape == x.shape + torch.manual_seed(1801) + grad_output = torch.randn_like(y_ref) + y_ref.backward(grad_output) + y_test.backward(grad_output) + + torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0.0, atol=0.0) def test_hybrid_fsdp_rejects_storage_only_sub_storages(self): row_quantizer = IdentityQuantizer() @@ -560,7 +594,9 @@ def test_dequantize_bitwise_identical(self): def test_dtype_cast(self): x = torch.randn(4, 8, device="cuda", dtype=torch.float32) out = IdentityQuantizer(dtype=torch.bfloat16)(x) - assert out.dequantize().dtype == torch.bfloat16 + dequantized = out.dequantize() + assert dequantized.dtype == torch.bfloat16 + torch.testing.assert_close(dequantized, x.to(torch.bfloat16), rtol=0.0, atol=0.0) def test_update_usage_is_noop(self): x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) @@ -931,21 +967,13 @@ def test_identity_recipe_matches_bf16_bitwise(self, module_name, qfactory): y_ref, dx_ref, wg_ref = _fwd_bwd_module(module_name, ref, x, recipe=None) y_id, dx_id, wg_id = _fwd_bwd_module(module_name, test, x, recipe=recipe) - # Linear / LayerNormLinear / GroupedLinear route through the same HP - # math with Identity and should stay bitwise exact. Composite modules - # can select different fused/unfused BF16 kernel paths after prior FP8 - # tests have warmed TE/CUDA state, so require tight BF16 numerical - # parity instead of order-dependent bitwise identity. - kwargs = ( - {"rtol": 0.0, "atol": 0.0} - if module_name in ("Linear", "LayerNormLinear", "GroupedLinear") - else {"rtol": 2.0e-2, "atol": 8.0e-3} - ) - torch.testing.assert_close(y_id, y_ref, **kwargs) - torch.testing.assert_close(dx_id, dx_ref, **kwargs) + # Identity is a high-precision passthrough. Entering the recipe context + # must not change any module's BF16 kernel selection or arithmetic. + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) assert len(wg_id) == len(wg_ref) for g_id, g_ref in zip(wg_id, wg_ref): - torch.testing.assert_close(g_id, g_ref, **kwargs) + torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") def test_grouped_linear_mxfp8_forward_identity_backward_matches_override(self): @@ -1089,6 +1117,7 @@ def test_activation_recompute_matches_no_checkpoint(self, format_name, use_reent "LayerNormMLP", marks=pytest.mark.xfail( reason="LayerNormMLP does not support built-in backward_override=dequantized", + raises=AssertionError, strict=True, ), id="layernorm_mlp", From c7e11d2c4c37392a7a2fa217698bc7d475ec79b0 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 15 Jul 2026 10:36:52 +0000 Subject: [PATCH 43/97] Resolve comments: fix rowwise_dequantized and improve quantizer validation in GroupedLinear, etc Signed-off-by: Evgeny --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/distributed/run_hybrid_tp_sp.py | 7 +- tests/pytorch/test_hybrid_quantization.py | 703 ++++++++++++++---- tests/pytorch/test_identity_quantizer.py | 119 ++- .../dot_product_attention.py | 29 +- transformer_engine/pytorch/module/base.py | 4 +- .../pytorch/module/grouped_linear.py | 354 +++++---- .../pytorch/quantized_tensor.py | 3 +- .../pytorch/tensor/_quantization_helpers.py | 57 +- .../pytorch/tensor/float8_tensor.py | 15 +- .../pytorch/tensor/hybrid_tensor.py | 131 +++- .../pytorch/tensor/identity_tensor.py | 84 +-- .../tensor/storage/float8_tensor_storage.py | 15 +- .../tensor/storage/hybrid_tensor_storage.py | 13 + .../tensor/storage/identity_tensor_storage.py | 5 +- 15 files changed, 1184 insertions(+), 356 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index d31a0ce494b..cc4b0b0ec18 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -40,6 +40,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_torch_compile.xm python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwisetensor.xml $TE_PATH/tests/pytorch/test_float8blockwisetensor.py || test_fail "test_float8blockwisetensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_view_shape.xml $TE_PATH/tests/pytorch/test_float8_view_shape.py || test_fail "test_float8_view_shape.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_PATH/tests/pytorch/test_grouped_tensor.py || test_fail "test_grouped_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH/tests/pytorch/test_gqa.py || test_fail "test_gqa.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py" diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index ec94b6fc462..a8c20ed0276 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -17,7 +17,6 @@ from torch import nn import transformer_engine.pytorch as te -import transformer_engine_torch as tex from transformer_engine.common import recipe as te_recipe from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, @@ -55,11 +54,11 @@ # Factories stay small; these tests target TP/SP plumbing. -def _make_fp8_current_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): +def _make_fp8_current_quantizer(*, fp8_dtype=te.DType.kFloat8E4M3): return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") -def _make_mxfp8_quantizer(*, fp8_dtype=tex.DType.kFloat8E4M3): +def _make_mxfp8_quantizer(*, fp8_dtype=te.DType.kFloat8E4M3): return MXFP8Quantizer(fp8_dtype=fp8_dtype) @@ -72,7 +71,7 @@ def _hybrid_fp8_qfactory(role): columnwise_quantizer=_make_fp8_current_quantizer(), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_fp8_current_quantizer(fp8_dtype=tex.DType.kFloat8E5M2) + return _make_fp8_current_quantizer(fp8_dtype=te.DType.kFloat8E5M2) return _make_fp8_current_quantizer() diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index e949e6ec392..aa46a56c51b 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -796,15 +796,16 @@ def test_storage_dequantize(self, input_tensor): hq.internal = False -@requires_fp8_and_nvfp4 class TestHybridUpdateUsage: """Test update_usage semantics and sub-storage cleanup.""" @pytest.fixture def hybrid_tensor(self): - torch.manual_seed(42) - inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") - hq = _make_hybrid_quantizer_fp8_row_fp4_col() + inp = torch.randn(4, 8) + hq = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) return hq.quantize(inp) def test_initial_usages(self, hybrid_tensor): @@ -841,6 +842,43 @@ def test_request_true_is_noop(self, hybrid_tensor): assert hybrid_tensor.rowwise_sub_storage is row_before assert hybrid_tensor.columnwise_sub_storage is col_before + def test_request_missing_columnwise_raises(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + + with pytest.raises(RuntimeError, match="no columnwise sub-storage"): + hybrid_tensor.update_usage(columnwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is not None + assert hybrid_tensor.columnwise_sub_storage is None + + def test_request_missing_rowwise_raises(self, hybrid_tensor): + hybrid_tensor.update_usage(rowwise_usage=False) + + with pytest.raises(RuntimeError, match="no rowwise sub-storage"): + hybrid_tensor.update_usage(rowwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is None + assert hybrid_tensor.columnwise_sub_storage is not None + + def test_missing_direction_request_is_atomic(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + row_before = hybrid_tensor.rowwise_sub_storage + + with pytest.raises(RuntimeError, match="no columnwise sub-storage"): + hybrid_tensor.update_usage(rowwise_usage=False, columnwise_usage=True) + + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is None + + def test_none_preserves_missing_direction(self, hybrid_tensor): + hybrid_tensor.update_usage(columnwise_usage=False) + row_before = hybrid_tensor.rowwise_sub_storage + + hybrid_tensor.update_usage(rowwise_usage=None, columnwise_usage=None) + + assert hybrid_tensor.rowwise_sub_storage is row_before + assert hybrid_tensor.columnwise_sub_storage is None + def test_repr_after_drop(self, hybrid_tensor): hybrid_tensor.update_usage(rowwise_usage=False) r = repr(hybrid_tensor) @@ -1686,6 +1724,82 @@ def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwis ) +@requires_fp8 +class TestCustomDPALocalRecipeCache: + """Custom-DPA native recipe labels track the quantizer rebuild.""" + + def test_inference_runs_once_per_quantizer_state_and_clears_stale_labels(self, monkeypatch): + from transformer_engine.pytorch.attention.dot_product_attention import ( + dot_product_attention as dpa_module, + ) + from transformer_engine.pytorch.module.base import TransformerEngineBaseModule + from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + custom_recipe = recipe.CustomRecipe(qfactory=lambda _role: IdentityQuantizer()) + monkeypatch.setattr( + FP8GlobalStateManager, + "get_fp8_recipe", + classmethod(lambda _cls: custom_recipe), + ) + + state = [object()] + quantizer = [object()] + + def fake_base_init(module, num_gemms=1): # pylint: disable=unused-argument + module.fp8_meta["scaling_fwd"] = state[0] + module.quantizers["scaling_fwd"] = [quantizer[0]] + + monkeypatch.setattr( + TransformerEngineBaseModule, + "init_fp8_metadata", + fake_base_init, + ) + + inferred_labels = [recipe.MXFP8BlockScaling()] + mutated_recipe_labels = [recipe.MXFP8BlockScaling(fp8_mha=True)] + inference_results = iter((inferred_labels, mutated_recipe_labels, None)) + inference_calls = 0 + + def fake_infer(*_args, **_kwargs): + nonlocal inference_calls + inference_calls += 1 + return next(inference_results) + + monkeypatch.setattr(dpa_module, "_infer_custom_dpa_local_recipes", fake_infer) + + dpa = te.DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + ) + dpa.init_fp8_metadata() + assert dpa.fp8_meta["local_recipes"] is inferred_labels + dpa.init_fp8_metadata() + assert inference_calls == 1 + assert dpa.fp8_meta["local_recipes"] is inferred_labels + + # Native labels also copy these mutable fields from CustomRecipe. They + # must refresh even when the quantizer generation itself is unchanged. + custom_recipe.fp8_mha = True + dpa.init_fp8_metadata() + assert inference_calls == 2 + assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels + dpa.init_fp8_metadata() + assert inference_calls == 2 + assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels + + # A rebuilt recipe state/quantizer list invalidates the cache. If the + # new family has no native label, the old label must not survive. + state[0] = object() + quantizer[0] = object() + dpa.init_fp8_metadata() + assert inference_calls == 3 + assert "local_recipes" not in dpa.fp8_meta + dpa.init_fp8_metadata() + assert inference_calls == 3 + assert "local_recipes" not in dpa.fp8_meta + + @requires_fp8_and_nvfp4 class TestAttentionFactoryNativeRecipeParity: """Linear + DPA qfactories should match native DPA recipe switches bitwise.""" @@ -3599,52 +3713,33 @@ def test_transformer_layer(self): @requires_fp8 -class TestHybridGroupedLinearClassifier: - """Unit tests for ``grouped_linear._is_hybrid_quantizer_list``. - - ``GroupedLinear`` dispatches its split-quantize between two mutually- - exclusive backends: ``tex.split_quantize`` (plain) and - ``_hybrid_split_quantize`` (all-hybrid). Neither can consume a mixed - list — ``tex.split_quantize`` doesn't recognise ``HybridQuantizer``, - and ``_hybrid_split_quantize`` calls ``q.rowwise_quantizer`` on every - element. Before the classifier was tightened, ``_has_hybrid_quantizer`` - used ``any(...)`` semantics: a single hybrid entry in a mixed list - would route to ``_hybrid_split_quantize`` and raise ``AttributeError`` - deep inside a grouped C++ call. These tests pin the new strict - classifier contract.""" - - def test_all_hybrid_returns_true(self): - from transformer_engine.pytorch.module.grouped_linear import ( - _is_hybrid_quantizer_list, - ) - - quantizers = [_make_hybrid_quantizer_fp8_row_fp4_col() for _ in range(3)] - assert _is_hybrid_quantizer_list(quantizers) is True +class TestHybridGroupedLinearValidation: + """GroupedLinear generation-validation and split-dispatch coverage. - def test_all_plain_returns_false(self): - from transformer_engine.pytorch.module.grouped_linear import ( - _is_hybrid_quantizer_list, - ) + Structural compatibility is validated once per real quantizer generation. + Steady-state dispatch reads the first expert after that uniformity check.""" - quantizers = [_make_fp8_quantizer() for _ in range(3)] - assert _is_hybrid_quantizer_list(quantizers) is False - - def test_all_none_returns_false(self): - """No quantizers at all (BF16 path) — classifier returns False so - the caller takes the non-fp8 branch.""" + @pytest.mark.parametrize( + "quantizers", + [ + pytest.param( + [_make_hybrid_quantizer_fp8_row_fp4_col() for _ in range(3)], + id="hybrid", + ), + pytest.param([_make_fp8_quantizer() for _ in range(3)], id="plain"), + pytest.param([None, None, None], id="none"), + ], + ) + def test_uniform_lists_validate(self, quantizers): from transformer_engine.pytorch.module.grouped_linear import ( - _is_hybrid_quantizer_list, + _validate_grouped_quantizer_list, ) - assert _is_hybrid_quantizer_list([None, None, None]) is False + _validate_grouped_quantizer_list(quantizers, operand_name="input") def test_mixed_hybrid_and_plain_raises(self): - """The actual bug: a mixed list used to silently route to - ``_hybrid_split_quantize`` and crash with ``AttributeError`` on - ``plain_q.rowwise_quantizer``. Now it fails fast at the - classifier with a user-actionable error.""" from transformer_engine.pytorch.module.grouped_linear import ( - _is_hybrid_quantizer_list, + _validate_grouped_quantizer_list, ) quantizers = [ @@ -3652,32 +3747,12 @@ def test_mixed_hybrid_and_plain_raises(self): _make_fp8_quantizer(), _make_hybrid_quantizer_fp8_row_fp4_col(), ] - with pytest.raises(ValueError) as exc_info: - _is_hybrid_quantizer_list(quantizers) - msg = str(exc_info.value) - # Error names both counts and points at the root cause so users - # can fix their ``qfactory`` without digging. - assert "mixes HybridQuantizer" in msg - assert "2 hybrid" in msg - assert "1 non-hybrid" in msg - assert "CustomRecipe" in msg and "qfactory" in msg + with pytest.raises(ValueError, match="mix HybridQuantizer and non-hybrid"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") def test_none_plus_hybrid_raises(self): - """None entries mixed with HybridQuantizer must NOT classify as - all-hybrid: ``_hybrid_split_quantize`` would later iterate the - full list with ``isinstance(q, HybridQuantizer)`` and raise - ``TypeError`` on the ``None`` entry. The classifier rejects - upfront with a clear ValueError so users see a single, - actionable error. - - In current TE flows ``CustomRecipeState.make_quantizers`` - rejects ``None`` returns from ``qfactory``, so this combination - shouldn't actually arise — but if a future "intentional no-op" - ``IdentityQuantizer`` ever loosens that contract, this guard - prevents the silent crash. - """ from transformer_engine.pytorch.module.grouped_linear import ( - _is_hybrid_quantizer_list, + _validate_grouped_quantizer_list, ) quantizers = [ @@ -3685,33 +3760,30 @@ def test_none_plus_hybrid_raises(self): None, _make_hybrid_quantizer_fp8_row_fp4_col(), ] - with pytest.raises(ValueError) as exc_info: - _is_hybrid_quantizer_list(quantizers) - msg = str(exc_info.value) - assert "mixes HybridQuantizer" in msg - assert "2 hybrid" in msg - assert "1 None" in msg - - def test_hybrid_split_quantize_rejects_plain_element(self): - """Defense-in-depth: even if a caller bypasses the classifier, - ``_hybrid_split_quantize`` itself asserts uniformity and raises - ``TypeError`` with a list of received types, rather than the - opaque ``AttributeError: 'Float8CurrentScalingQuantizer' object - has no attribute 'rowwise_quantizer'`` from the old code.""" + with pytest.raises(ValueError, match="mix None and concrete quantizers"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_mixed_identity_dtype_raises(self): from transformer_engine.pytorch.module.grouped_linear import ( - _hybrid_split_quantize, + _validate_grouped_quantizer_list, ) - tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") quantizers = [ - _make_hybrid_quantizer_fp8_row_fp4_col(), - _make_fp8_quantizer(), # Not hybrid — should trigger TypeError + IdentityQuantizer(dtype=torch.bfloat16), + IdentityQuantizer(dtype=torch.float16), ] - with pytest.raises(TypeError) as exc_info: - _hybrid_split_quantize(tensor, [16, 16], quantizers) - msg = str(exc_info.value) - assert "HybridQuantizer" in msg - assert "Float8CurrentScalingQuantizer" in msg + with pytest.raises(ValueError, match="incompatible plain backend configurations"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_distinct_delayed_scaling_state_is_allowed(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [_make_delayed_quantizer(), _make_delayed_quantizer()] + quantizers[1].scale.fill_(2.0) + quantizers[1].amax.fill_(3.0) + _validate_grouped_quantizer_list(quantizers, operand_name="input") @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize( @@ -3742,44 +3814,215 @@ def test_hybrid_split_quantize_respects_parent_usage_flags(self, usage, expected assert [storage.get_usages() for storage in out] == [expected, expected] - @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) - def test_hybrid_split_quantize_rejects_mixed_parent_usage_flags(self): - from transformer_engine.pytorch.module.grouped_linear import ( - _hybrid_split_quantize, + def test_columnwise_only_rowwise_dequantized_uses_transient_grouped_row(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + result = real_split_quantize(tensor, m_splits, quantizers, **kwargs) + calls.append((tensor, result)) + return result + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=False, columnwise=True) + + out = grouped_linear._hybrid_split_quantize(tensor, [16, 16], quantizers) + + assert len(calls) == 2 + expected_columnwise_source = torch.cat( + [result.dequantize() for result in calls[0][1]], dim=0 ) + torch.testing.assert_close(calls[1][0], expected_columnwise_source, rtol=0.0, atol=0.0) + assert calls[1][0] is not tensor + assert all(storage.rowwise_sub_storage is None for storage in out) + assert all(storage.columnwise_sub_storage is not None for storage in out) + + def test_rowwise_only_skips_columnwise_quantization(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + calls.append(tensor) + return real_split_quantize(tensor, m_splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in range(2) + ] + for quantizer in quantizers: + quantizer.set_usage(rowwise=True, columnwise=False) + + out = grouped_linear._hybrid_split_quantize(tensor, [16, 16], quantizers) + + assert calls == [tensor] + assert all(storage.rowwise_sub_storage is not None for storage in out) + assert all(storage.columnwise_sub_storage is None for storage in out) + def test_original_source_preserves_two_bulk_call_fast_path(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + calls = [] + + def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): + calls.append(tensor) + return real_split_quantize(tensor, m_splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") quantizers = [ HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source="original", ) for _ in range(2) ] - quantizers[0].set_usage(rowwise=True, columnwise=False) - quantizers[1].set_usage(rowwise=True, columnwise=True) - with pytest.raises(ValueError, match="mixed parent usage flags"): - _hybrid_split_quantize(tensor, [16, 16], quantizers) + out = grouped_linear._hybrid_split_quantize(tensor, [16, 16], quantizers) + + assert calls == [tensor, tensor] + assert all(storage.rowwise_sub_storage is not None for storage in out) + assert all(storage.columnwise_sub_storage is not None for storage in out) + + def test_validation_rejects_mixed_columnwise_source_policies(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=source, + ) + for source in ("original", "rowwise_dequantized") + ] + + with pytest.raises(ValueError, match="mixed columnwise source policies"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_validation_rejects_same_family_config_mismatch(self): + from transformer_engine.pytorch.module.grouped_linear import ( + _validate_grouped_quantizer_list, + ) + + quantizers = [_make_fp8_quantizer(), _make_fp8_quantizer()] + quantizers[1].force_pow_2_scales = True + + with pytest.raises( + ValueError, + match="incompatible plain backend configurations", + ): + _validate_grouped_quantizer_list(quantizers, operand_name="input") + + def test_validation_runs_only_with_quantizer_generation(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + def make_qfactory(columnwise_source): + def qfactory(_role): + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=columnwise_source, + ) + + return qfactory + + model = GroupedLinear(2, 128, 128, bias=False, params_dtype=torch.bfloat16).cuda() + tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + m_splits = torch.tensor([64, 64], dtype=torch.int64) + original_recipe = recipe.CustomRecipe(qfactory=make_qfactory("original")) + + real_validate = grouped_linear._validate_grouped_quantizer_list + validation_calls = [] + + def tracked_validate(quantizers, *, operand_name="operand"): + validation_calls.append((operand_name, id(quantizers[0]))) + return real_validate(quantizers, operand_name=operand_name) + + monkeypatch.setattr( + grouped_linear, + "_validate_grouped_quantizer_list", + tracked_validate, + ) + + with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): + model(tensor, m_splits) + first_call_count = len(validation_calls) + first_generation = model._validated_quantizer_generations["scaling_fwd"] + assert first_call_count > 0 + + with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): + model(tensor, m_splits) + assert len(validation_calls) == first_call_count + assert model._validated_quantizer_generations["scaling_fwd"] is first_generation + + rebuilt_recipe = recipe.CustomRecipe(qfactory=make_qfactory("rowwise_dequantized")) + with torch.no_grad(), autocast(enabled=True, recipe=rebuilt_recipe): + model(tensor, m_splits) + rebuilt_generation = model._validated_quantizer_generations["scaling_fwd"] + assert len(validation_calls) > first_call_count + assert rebuilt_generation is not first_generation + assert rebuilt_generation[0].columnwise_source == "rowwise_dequantized" + + input_count = 0 + + def mixed_source_qfactory(role): + nonlocal input_count + source = "original" + if role is not None and role.tensor_type == "input": + source = "original" if input_count == 0 else "rowwise_dequantized" + input_count += 1 + return HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + columnwise_source=source, + ) + + mixed_recipe = recipe.CustomRecipe(qfactory=mixed_source_qfactory) + # A failed generation is never marked validated. Base metadata can then + # early-return on retry, so the O(1) guard must validate it again. + for _ in range(2): + with pytest.raises(ValueError, match="mixed columnwise source policies"): + with torch.no_grad(), autocast(enabled=True, recipe=mixed_recipe): + model(tensor, m_splits) + assert model._validated_quantizer_generations["scaling_fwd"] is rebuilt_generation + + # Stale invalid recipe metadata must not affect the non-quantized path. + with torch.no_grad(): + model(tensor, m_splits) @requires_fp8_and_nvfp4 - @pytest.mark.xfail( - strict=True, - raises=AssertionError, - reason=( - "Production follow-up: GroupedLinear bulk split-quantize ignores " - "HybridQuantizer.columnwise_source" - ), - ) def test_hybrid_split_quantize_honors_rowwise_dequantized_source(self): - """Bulk split must match standalone HybridQuantizer provenance semantics.""" + """NVFP4 column data must derive from the actual grouped row result.""" from transformer_engine.pytorch.module.grouped_linear import ( _hybrid_split_quantize, ) torch.manual_seed(3598) # NVFP4 grouped split-quantize requires each M split to be a multiple - # of 64; valid geometry ensures the xfail reaches provenance checking. + # of 64. tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") def make_quantizer(): @@ -3789,19 +4032,20 @@ def make_quantizer(): columnwise_source="rowwise_dequantized", ) + quantizers = [make_quantizer(), make_quantizer()] actual = _hybrid_split_quantize( tensor, [64, 64], - [make_quantizer(), make_quantizer()], + quantizers, ) - expected = [ - make_quantizer().quantize(part) for part in torch.split(tensor, [64, 64], dim=0) - ] - for index, (actual_part, expected_part) in enumerate(zip(actual, expected)): + for index, (actual_part, quantizer) in enumerate(zip(actual, quantizers)): + expected_columnwise = quantizer.columnwise_quantizer.quantize( + actual_part.rowwise_sub_storage.dequantize() + ) _assert_storage_data_exact( actual_part.columnwise_sub_storage, - expected_part.columnwise_sub_storage, + expected_columnwise, context=f"GroupedLinear split {index} columnwise provenance", ) @@ -5993,6 +6237,211 @@ def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self) ] +class TestHybridNewZeros: + """Public new_zeros semantics independent of FSDP-specific allocation.""" + + @staticmethod + def _make_identity_hybrid(): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) + return quantizer, source + + @pytest.mark.parametrize( + "keep_rowwise,keep_columnwise", + [(True, True), (True, False), (False, True)], + ) + def test_initializes_every_present_direction_and_preserves_kwargs( + self, + keep_rowwise, + keep_columnwise, + ): + quantizer, source = self._make_identity_hybrid() + if not keep_rowwise: + source.update_usage(rowwise_usage=False) + if not keep_columnwise: + source.update_usage(columnwise_usage=False) + + # Direction selection must come from source storage, not mutable parent + # usage flags, and new_zeros must not change those flags. + quantizer.set_usage(rowwise=False, columnwise=False) + result = torch.ops.aten.new_zeros.default( + source, + [3, 5], + dtype=torch.float32, + device=source.device, + ) + + assert isinstance(result, HybridQuantizedTensor) + assert result.shape == torch.Size((3, 5)) + assert result.dtype == torch.float32 + assert result.device == source.device + assert (result.rowwise_sub_storage is not None) is keep_rowwise + assert (result.columnwise_sub_storage is not None) is keep_columnwise + assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} + assert result._quantizer is not quantizer + assert result._quantizer.get_usages() == { + "rowwise": keep_rowwise, + "columnwise": keep_columnwise, + } + + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + if sub_storage is not None: + torch.testing.assert_close( + sub_storage.dequantize(), + torch.zeros((3, 5), dtype=torch.float32, device=source.device), + rtol=0.0, + atol=0.0, + ) + + # A plain-source copy routes through the result's parent quantizer. It + # must update every allocated direction even though the source parent + # had both of its mutable usage flags disabled before new_zeros. + plain_source = torch.full((3, 5), 7.0, dtype=torch.float32, device=source.device) + result.copy_(plain_source) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + if sub_storage is not None: + torch.testing.assert_close( + sub_storage.dequantize(), plain_source, rtol=0.0, atol=0.0 + ) + assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} + + def test_identity_substorages_allow_integer_dtype(self): + _, source = self._make_identity_hybrid() + result = source.new_zeros((2, 3), dtype=torch.int32) + + assert result.dtype == torch.int32 + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + assert sub_storage.dequantize().dtype == torch.int32 + assert torch.count_nonzero(sub_storage.dequantize()).item() == 0 + + @requires_fp8 + @pytest.mark.parametrize("unsupported_dtype", (torch.int32, torch.float64, torch.bool)) + def test_non_identity_substorage_rejects_unsupported_dtype(self, unsupported_dtype): + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) + + with pytest.raises(TypeError, match="new_zeros only supports"): + source.new_zeros((2, 3), dtype=unsupported_dtype) + + def test_does_not_invoke_live_quantizers_or_consume_rng(self, monkeypatch): + quantizer, source = self._make_identity_hybrid() + cpu_rng_before = torch.get_rng_state().clone() + cuda_rng_before = torch.cuda.get_rng_state(source.device).clone() + + def fail_live_make_empty(*args, **kwargs): + raise AssertionError("new_zeros invoked a live quantizer") + + monkeypatch.setattr(quantizer, "make_empty", fail_live_make_empty) + monkeypatch.setattr( + source.rowwise_sub_storage._quantizer, + "make_empty", + fail_live_make_empty, + ) + monkeypatch.setattr( + source.columnwise_sub_storage._quantizer, + "make_empty", + fail_live_make_empty, + ) + + result = source.new_zeros((2, 6)) + + assert isinstance(result, HybridQuantizedTensor) + assert torch.equal(torch.get_rng_state(), cpu_rng_before) + assert torch.equal(torch.cuda.get_rng_state(source.device), cuda_rng_before) + + def test_rejects_empty_hybrid(self): + _, source = self._make_identity_hybrid() + source.update_usage(rowwise_usage=False, columnwise_usage=False) + + with pytest.raises(RuntimeError, match="at least one present sub-storage"): + source.new_zeros((2, 6)) + + @requires_fp8_and_nvfp4 + def test_initializes_nvfp4_data_scale_and_amax_buffers(self): + quantizer = _make_hybrid_quantizer_fp8_row_fp4_col() + source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) + + result = source.new_zeros(source.shape) + + assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + buffers, storage = sub_storage.prepare_for_saving() + try: + assert all( + buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers + ) + finally: + assert storage.restore_from_saved(buffers) == [] + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8Blockwise: {reason_for_no_fp8_block_scaling}", + ) + def test_initializes_float8_block_data_and_scale_buffers(self): + quantizer = HybridQuantizer( + rowwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + columnwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + ) + source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) + + result = source.new_zeros(source.shape) + + assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + buffers, storage = sub_storage.prepare_for_saving() + try: + assert all( + buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers + ) + finally: + assert storage.restore_from_saved(buffers) == [] + + @requires_fp8 class TestHybridTorchDispatchFSDP2Ops: """Test aten ops that FSDP2 relies on to preserve the HybridQuantizedTensor type. @@ -6117,22 +6566,10 @@ def test_copy_from_bf16_to_hybrid(self, hybrid_param): torch.testing.assert_close(param.dequantize(), expected.dequantize(), rtol=0.0, atol=0.0) def test_new_zeros_returns_hybrid(self, hybrid_param): - """new_zeros must return a usable HybridQuantizedTensor container. - - FSDP2 calls ``new_zeros`` only to allocate an all-gather destination - buffer that is immediately overwritten by ``copy_``; the initial - contents are never observed. The hybrid dispatch therefore delegates - to ``HybridQuantizer.make_empty`` (uninitialized bytes) rather than - quantizing a BF16 zeros temporary. This test asserts the contract - we actually depend on — correct container type / shape / sub-storage - presence, and the ability to copy into it and read back — NOT that - the raw dequantize value happens to be zero. - """ + """new_zeros returns initialized storage that remains FSDP-copyable.""" new_shape = list(hybrid_param.shape) result = aten.new_zeros.default(hybrid_param, new_shape) - # Structural contract: FSDP2 needs a HybridQuantizedTensor with both - # sub-storages populated so the gathered buffers have a destination. assert isinstance( result, HybridQuantizedTensor ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" @@ -6141,9 +6578,14 @@ def test_new_zeros_returns_hybrid(self, hybrid_param): assert result.columnwise_sub_storage is not None assert type(result.rowwise_sub_storage) is type(hybrid_param.rowwise_sub_storage) assert type(result.columnwise_sub_storage) is type(hybrid_param.columnwise_sub_storage) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) - # Functional contract: the container must be writable via copy_ from - # another hybrid (how FSDP2 populates the buffer post-gather). + # FSDP2 overwrites the initialized destination via copy_ after gather. aten.copy_.default(result, hybrid_param) torch.testing.assert_close( result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 @@ -6543,10 +6985,9 @@ def test_scale_refresh_across_iterations(self): def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): """Hybrid FSDP2 with an NVFP4 sub-storage must raise a clear error. - Per the hybrid FSDP2 design (see ``hybrid_quantization_fsdp.md`` §9 - Gap 5), NVFP4 FSDP2 support is not implemented yet — packed FP4 data - alignment for dim-0 splitting, columnwise dequant, and RHT cache - handling all need work. Until that lands, hybrid pre-all-gather must + NVIDIA/TransformerEngine#3158 tracks the missing NVFP4 FSDP2 support: + packed FP4 dim-0 alignment, columnwise dequantization, and RHT-cache + handling are not implemented. Until that lands, hybrid pre-all-gather must refuse an NVFP4 sub-storage cleanly via the ``fsdp_buffer_fields`` protocol rather than silently producing wrong data. @@ -6568,7 +7009,7 @@ def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): param = model.weight # Clean refusal: hybrid's pre_all_gather raises an NVFP4-specific - # message pointing to the design doc, not a generic + # message identifying the unsupported sub-storage protocol, not a generic # "NVFP4Tensor does not implement fsdp_buffer_fields" from deep inside # the base class. with pytest.raises(NotImplementedError) as exc_info: diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 2d1220e3370..758cbcfc6b0 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -268,11 +268,9 @@ def test_grouped_split_all_identity_uses_plain_tensor_views(self): def test_grouped_split_rejects_mixed_identity_and_quantized_operands(self): from transformer_engine.pytorch.module.grouped_linear import ( - _split_quantize_with_identity_fallback, + _validate_grouped_quantizer_list, ) - x = torch.empty(64, 64, dtype=torch.bfloat16) - m_splits = [32, 32] cases = [ [IdentityQuantizer(), _mxfp8(tex.DType.kFloat8E4M3)], [ @@ -288,13 +286,8 @@ def test_grouped_split_rejects_mixed_identity_and_quantized_operands(self): ] for quantizers in cases: - with pytest.raises(ValueError, match="mixes Identity-backed and non-Identity-backed"): - _split_quantize_with_identity_fallback( - x, - m_splits, - quantizers, - activation_dtype=torch.bfloat16, - ) + with pytest.raises(ValueError, match="mix Identity-backed and quantized"): + _validate_grouped_quantizer_list(quantizers, operand_name="input") def test_hybrid_split_forwards_disable_bulk_allocation_to_both_directions(self, monkeypatch): import transformer_engine.pytorch.module.grouped_linear as grouped_linear @@ -350,9 +343,9 @@ def qfactory(role): calls = [] def fake_hybrid_split_quantize( - tensor, m_splits, quantizers, *, disable_bulk_allocation=False + tensor, m_splits, quantizers, *, disable_bulk_allocation=False, **kwargs ): - del tensor, m_splits, quantizers + del tensor, m_splits, quantizers, kwargs calls.append(disable_bulk_allocation) raise StopAfterFlagCapture("captured hybrid split kwargs") @@ -387,7 +380,7 @@ def qfactory(role): x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) - with pytest.raises(ValueError, match="mixes Identity-backed and non-Identity-backed"): + with pytest.raises(ValueError, match="mix Identity-backed and quantized"): with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): model(x, m_splits=m_splits) @@ -598,6 +591,27 @@ def test_dtype_cast(self): assert dequantized.dtype == torch.bfloat16 torch.testing.assert_close(dequantized, x.to(torch.bfloat16), rtol=0.0, atol=0.0) + def test_storage_dequantize_defaults_to_nominal_dtype(self): + payload = torch.randn(4, 8, dtype=torch.float32) + storage = IdentityTensorStorage( + hp_data=payload, + fake_dtype=torch.bfloat16, + ) + + default = storage.dequantize() + assert default.dtype == torch.bfloat16 + torch.testing.assert_close(default, payload.to(torch.bfloat16), rtol=0.0, atol=0.0) + + explicit = storage.dequantize(dtype=torch.float16) + assert explicit.dtype == torch.float16 + torch.testing.assert_close(explicit, payload.to(torch.float16), rtol=0.0, atol=0.0) + + same_dtype_storage = IdentityTensorStorage( + hp_data=payload, + fake_dtype=torch.float32, + ) + assert same_dtype_storage.dequantize() is payload + def test_update_usage_is_noop(self): x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) q = IdentityQuantizer() @@ -663,6 +677,85 @@ def test_tensor_ops_preserve_identity_and_values(self): dst.copy_(t) torch.testing.assert_close(dst.dequantize(), x, rtol=0.0, atol=0.0) + def test_view_ops_preserve_offsets_strides_and_aliasing(self): + base = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor = IdentityQuantizer()(base) + + sliced = torch.ops.aten.slice.Tensor(tensor, 0, 2, 6, 1) + omitted_offset = torch.ops.aten.as_strided.default( + sliced, + [2, 4], + [4, 1], + ) + explicit_offset = torch.ops.aten.as_strided.default( + sliced, + [2, 4], + [4, 1], + 0, + ) + torch.testing.assert_close(omitted_offset.dequantize(), base[2:4], rtol=0.0, atol=0.0) + torch.testing.assert_close(explicit_offset.dequantize(), base[:2], rtol=0.0, atol=0.0) + assert sliced.storage_offset() == base[2:].storage_offset() + assert omitted_offset.storage_offset() == base[2:].storage_offset() + assert explicit_offset.storage_offset() == 0 + + noncontiguous = torch.ops.aten.slice.Tensor(tensor, 1, 0, 4, 2) + assert noncontiguous.stride() == base[:, ::2].stride() + assert noncontiguous.storage_offset() == base[:, ::2].storage_offset() + replacement = torch.full_like(base[:, ::2], -3) + noncontiguous.copy_(replacement) + torch.testing.assert_close(base[:, ::2], replacement, rtol=0.0, atol=0.0) + + pieces = torch.split(tensor, 2, dim=0) + piece_value = torch.full_like(base[:2], 7) + pieces[0].copy_(piece_value) + torch.testing.assert_close(base[:2], piece_value, rtol=0.0, atol=0.0) + + def test_detach_and_clone_preserve_view_metadata(self): + base = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor = IdentityQuantizer()(base) + + offset_view = torch.ops.aten.slice.Tensor(tensor, 0, 2, 6, 1) + detached = offset_view.detach() + assert detached.stride() == detached._hp_data.stride() == offset_view.stride() + assert detached.storage_offset() == detached._hp_data.storage_offset() == 8 + replacement = torch.full_like(base[2:6], -5) + detached.copy_(replacement) + torch.testing.assert_close(base[2:6], replacement, rtol=0.0, atol=0.0) + + base2 = torch.arange(24, device="cuda", dtype=torch.float32).reshape(6, 4) + tensor2 = IdentityQuantizer()(base2) + noncontiguous = torch.ops.aten.slice.Tensor(tensor2, 1, 0, 4, 2) + detached_noncontiguous = noncontiguous.detach() + assert detached_noncontiguous.stride() == detached_noncontiguous._hp_data.stride() + assert detached_noncontiguous.storage_offset() == 0 + + transposed = torch.ops.aten.as_strided.default( + tensor2, + [4, 6], + [1, 4], + ) + cloned = transposed.clone() + assert cloned.stride() == cloned._hp_data.stride() + assert cloned.storage_offset() == cloned._hp_data.storage_offset() + torch.testing.assert_close(cloned.dequantize(), base2.as_strided((4, 6), (1, 4))) + original_snapshot = base2.clone() + cloned.copy_(torch.zeros_like(cloned.dequantize())) + torch.testing.assert_close(base2, original_snapshot, rtol=0.0, atol=0.0) + + def test_copy_from_plain_and_identity_sources(self): + dst_data = torch.zeros(3, 4, device="cuda", dtype=torch.float32) + dst = IdentityQuantizer()(dst_data) + + plain_src = torch.arange(12, device="cuda", dtype=torch.float32).reshape(3, 4) + assert dst.copy_(plain_src) is dst + torch.testing.assert_close(dst.dequantize(), plain_src, rtol=0.0, atol=0.0) + + identity_src_data = plain_src.neg() + identity_src = IdentityQuantizer()(identity_src_data) + assert dst.copy_(identity_src) is dst + torch.testing.assert_close(dst.dequantize(), identity_src_data, rtol=0.0, atol=0.0) + def test_fsdp_pre_post_all_gather_roundtrip(self): x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) t = IdentityQuantizer()(x) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index f3ed4c2e220..dad3bc5b942 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -464,6 +464,12 @@ def __init__( ) -> None: super().__init__(name=name) + # Cache the native recipe labels inferred from custom DPA quantizers. + # ``init_fp8_metadata`` runs on every forward, while the quantizers only + # change when their recipe state is rebuilt. + self._custom_dpa_local_recipes_cache_key: Optional[Tuple[Any, ...]] = None + self._custom_dpa_local_recipes_cache: Optional[List[Recipe]] = None + self.logger = logging.getLogger("DotProductAttention") self.logger.setLevel(attn_log._log_level) if not self.logger.hasHandlers(): @@ -703,11 +709,26 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): super().init_fp8_metadata(num_gemms=num_gemms) - local_recipes = _infer_custom_dpa_local_recipes( - fp8_recipe, self.fp8_meta, self.quantizers + fwd_quantizers = self.quantizers.get("scaling_fwd", ()) + cache_key = ( + id(self.fp8_meta.get("scaling_fwd")), + tuple(id(quantizer) for quantizer in fwd_quantizers), + fp8_recipe.fp8_format, + fp8_recipe.fp8_dpa, + fp8_recipe.fp8_mha, ) - if local_recipes is not None: - self.fp8_meta["local_recipes"] = local_recipes + if cache_key != self._custom_dpa_local_recipes_cache_key: + self._custom_dpa_local_recipes_cache = _infer_custom_dpa_local_recipes( + fp8_recipe, self.fp8_meta, self.quantizers + ) + self._custom_dpa_local_recipes_cache_key = cache_key + + if self._custom_dpa_local_recipes_cache is None: + # Do not leave labels from an earlier supported quantizer + # family attached after a rebuild to an unsupported family. + self.fp8_meta.pop("local_recipes", None) + else: + self.fp8_meta["local_recipes"] = self._custom_dpa_local_recipes_cache return # switch/append recipe: fp8_recipe stays unchanged, but DPA.fp8_meta["recipe"] may be set to diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index c2eb11aafe6..ebb802b2ae7 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1059,7 +1059,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # Return early if recipe state matches recipe if self.fp8_meta_tensors_initialized: recipe_state = self.fp8_meta[fp8_meta_tensor_key] - # Follow-up (#3157): Match built-in recipes by full config, not just RecipeState type, so + # TODO(#3157): Match built-in recipes by full config, not just RecipeState type, so # same-class mid-training changes rebuild quantizers/workspaces correctly. if recipe.delayed() and isinstance(recipe_state, DelayedScalingRecipeState): self.adjust_amax_history_length(recipe.amax_history_len, fwd=fwd) @@ -1077,7 +1077,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return if recipe.custom() and isinstance(recipe_state, CustomRecipeState): - # Follow-up (#3157): Compare CustomRecipe/qfactory config here. qfactory changes made + # TODO(#3157): Compare CustomRecipe/qfactory config here. qfactory changes made # mid-training on the same recipe object currently do not take effect because # stale quantizers are reused. if recipe_state.recipe is recipe: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 88d4837dec0..a883b40a517 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -86,11 +86,6 @@ def _uses_identity_quantizer(quantizer): return False -def _has_identity_quantizer_list(quantizers): - """Whether any quantizer in a grouped list uses Identity.""" - return any(_uses_identity_quantizer(q) for q in quantizers) - - def _identity_quantizer_signature(quantizer): """Identity usage per GEMM direction: (rowwise, columnwise).""" if isinstance(quantizer, HybridQuantizer): @@ -102,29 +97,130 @@ def _identity_quantizer_signature(quantizer): return (identity, identity) -def _check_uniform_identity_quantizer_list(quantizers): - """Reject grouped lists that mix Identity-backed and quantized directions.""" - signatures = [_identity_quantizer_signature(q) for q in quantizers] - if not any(rowwise or columnwise for rowwise, columnwise in signatures): - return - if all(signature == signatures[0] for signature in signatures): +_DYNAMIC_QUANTIZER_SIGNATURE_FIELDS = frozenset( + { + "rowwise_usage", + "columnwise_usage", + "internal", + "optimize_for_gemm", + } +) + + +def _backend_quantizer_signature(quantizer): + """Return backend configuration that grouped kernels require to be uniform.""" + if quantizer is None: + return None + + # Identity is not registered as a torch.compile value quantizer, but its + # dtype changes the grouped GEMM input type and therefore must be uniform. + if isinstance(quantizer, IdentityQuantizer): + return (type(quantizer), (("dtype", quantizer.dtype),)) + + fields = quantizer._value_fields() + if fields is None: + # Delayed-scaling Float8Quantizer carries per-expert scale/amax tensors, + # which are intentionally different, but its emitted FP8 dtype is a + # group-wide backend choice. Other unregistered/custom quantizers retain + # the conservative exact-family behavior until they expose value fields. + fields = ("dtype",) if isinstance(quantizer, Float8Quantizer) else () + + config = [] + for name in fields: + if name in _DYNAMIC_QUANTIZER_SIGNATURE_FIELDS: + continue + value = getattr(quantizer, name) + if name == "dtype": + value = int(value) + config.append((name, value)) + return (type(quantizer), tuple(config)) + + +def _validate_backend_match(reference, quantizer, operand_name, direction, expert_index): + """Validate one expert against the group's reference backend.""" + if type(quantizer) is not type(reference): + raise ValueError( + f"GroupedLinear {operand_name} quantizers use incompatible {direction} backend" + f" families across experts: expert 0 uses {type(reference).__name__}, but expert" + f" {expert_index} uses {type(quantizer).__name__}. Grouped operands require one" + " quantizer family per direction." + ) + reference_signature = _backend_quantizer_signature(reference) + quantizer_signature = _backend_quantizer_signature(quantizer) + if quantizer_signature != reference_signature: + raise ValueError( + f"GroupedLinear {operand_name} quantizers use incompatible {direction} backend" + f" configurations across experts: expert 0 uses {reference_signature}, but expert" + f" {expert_index} uses {quantizer_signature}. Grouped operands require the same" + " backend-relevant configuration per direction." + ) + + +def _validate_grouped_quantizer_list(quantizers, *, operand_name="operand") -> None: + """Validate one grouped operand once when its quantizer generation changes.""" + if not quantizers: return - raise ValueError( - "GroupedLinear quantizer list mixes Identity-backed and non-Identity-backed" - f" directions: {signatures}. This combination is not supported because" - " grouped GEMM requires a uniform scaling mode for every tensor in each" - " operand list. Make the CustomRecipe `qfactory` return Identity" - " consistently for every expert in the grouped operand, or return no" - " Identity quantizers for that operand." - ) + reference = quantizers[0] + reference_is_hybrid = isinstance(reference, HybridQuantizer) + reference_identity = _identity_quantizer_signature(reference) -def _is_plain_identity_passthrough_list(quantizers, activation_dtype): - """Whether every quantizer is a plain Identity passthrough.""" - return bool(quantizers) and all( - isinstance(q, IdentityQuantizer) and (q.dtype is None or q.dtype == activation_dtype) - for q in quantizers - ) + for expert_index, quantizer in enumerate(quantizers[1:], start=1): + if (quantizer is None) != (reference is None): + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix None and concrete quantizers" + f" across experts: expert 0 is {type(reference).__name__}, but expert" + f" {expert_index} is {type(quantizer).__name__}." + ) + if reference is None: + continue + + quantizer_is_hybrid = isinstance(quantizer, HybridQuantizer) + if quantizer_is_hybrid != reference_is_hybrid: + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix HybridQuantizer and non-hybrid" + f" quantizers across experts: expert 0 is {type(reference).__name__}, but expert" + f" {expert_index} is {type(quantizer).__name__}." + ) + + identity = _identity_quantizer_signature(quantizer) + if identity != reference_identity: + raise ValueError( + f"GroupedLinear {operand_name} quantizers mix Identity-backed and quantized" + f" directions across experts: expert 0 uses {reference_identity}, but expert" + f" {expert_index} uses {identity}." + ) + + if reference_is_hybrid: + _validate_backend_match( + reference.rowwise_quantizer, + quantizer.rowwise_quantizer, + operand_name, + "rowwise", + expert_index, + ) + _validate_backend_match( + reference.columnwise_quantizer, + quantizer.columnwise_quantizer, + operand_name, + "columnwise", + expert_index, + ) + if quantizer.columnwise_source != reference.columnwise_source: + raise ValueError( + f"GroupedLinear {operand_name} HybridQuantizer list has mixed columnwise" + " source policies across experts: expert 0 uses" + f" {reference.columnwise_source!r}, but expert {expert_index} uses" + f" {quantizer.columnwise_source!r}." + ) + else: + _validate_backend_match( + reference, + quantizer, + operand_name, + "plain", + expert_index, + ) def _split_quantize_with_identity_fallback( @@ -135,9 +231,9 @@ def _split_quantize_with_identity_fallback( *, disable_bulk_allocation=False, ): - """Split+quantize, avoiding native C++ split kernels for Identity quantizers.""" - # No Identity anywhere: use the native grouped split+quantize kernel. - if not _has_identity_quantizer_list(quantizers): + """Split+quantize a list validated as uniform for this recipe generation.""" + reference = quantizers[0] + if not _uses_identity_quantizer(reference): return tex.split_quantize( tensor, m_splits, @@ -145,114 +241,59 @@ def _split_quantize_with_identity_fallback( disable_bulk_allocation=disable_bulk_allocation, ) - _check_uniform_identity_quantizer_list(quantizers) tensor = cast_if_needed(tensor, activation_dtype) - # Plain all-Identity passthrough: match the native high-precision BF16 split path. - if _is_plain_identity_passthrough_list(quantizers, activation_dtype): + if isinstance(reference, IdentityQuantizer) and ( + reference.dtype is None or reference.dtype == activation_dtype + ): return torch.split(tensor, m_splits) - # Uniform Identity-backed wrappers still need per-split Python quantizer calls. return [ quantizer(tensor_part) if quantizer is not None else tensor_part for tensor_part, quantizer in zip(torch.split(tensor, m_splits), quantizers) ] -def _is_hybrid_quantizer_list(quantizers): - """Classify a GroupedLinear quantizer list as hybrid-uniform or plain-uniform. - - Returns ``True`` when every non-``None`` entry is a ``HybridQuantizer``, - ``False`` when none are. Raises :class:`ValueError` on a mixed list. - - Why it's a hard "or": neither dispatch branch at the call sites - supports a mixed list. ``tex.split_quantize`` (the plain path) does - not recognize ``HybridQuantizer``. :func:`_hybrid_split_quantize` - (the hybrid path) treats every entry as hybrid — it calls - ``q.rowwise_quantizer`` / ``q.columnwise_quantizer`` unconditionally, - which would ``AttributeError`` on a plain quantizer. Rejecting at - the classifier gives a clear actionable error instead of failing - deep inside a grouped C++ call. - - Supporting mixed would require a per-element grouped kernel that - accepts a heterogeneous quantizer vector; no such kernel exists - today. If that becomes a real requirement (e.g. per-expert hybrid - recipes in MoE), implement element-wise fallback here rather than - silently masking the mismatch. - """ - non_none = [q for q in quantizers if q is not None] - if not non_none: - return False - hybrid_count = sum(1 for q in non_none if isinstance(q, HybridQuantizer)) - if hybrid_count == 0: - return False - # Reject a list mixing HybridQuantizer with None entries: ``_hybrid_split_quantize`` - # subsequently iterates the *full* list with ``isinstance(q, HybridQuantizer)`` which - # would raise ``TypeError`` on the ``None`` entries. Forcing the list to be entirely - # non-``None`` before claiming "all hybrid" matches the dispatch's actual capability. - if hybrid_count == len(quantizers): - return True - raise ValueError( - "GroupedLinear quantizer list mixes HybridQuantizer and non-hybrid" - f" quantizers ({hybrid_count} hybrid," - f" {len(non_none) - hybrid_count} non-hybrid," - f" {len(quantizers) - len(non_none)} None). This combination is not" - " supported: neither `tex.split_quantize` nor `_hybrid_split_quantize`" - " can consume a heterogeneous list. Make the CustomRecipe `qfactory`" - " return a consistent type (all HybridQuantizer or all non-hybrid)" - " across every GEMM for the same role." - ) - - -def _hybrid_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocation=False): - """Grouped split+quantize for an **all-hybrid** quantizer list. - - Precondition: every ``q`` in ``quantizers`` is a ``HybridQuantizer``. - Enforce via :func:`_is_hybrid_quantizer_list` at the call site (or - the explicit assert below as a defense-in-depth). - - Runs ``tex.split_quantize`` twice — once over the rowwise sub- - quantizers, once over the columnwise sub-quantizers — then zips - the two per-split results back into ``HybridQuantizedTensorStorage`` - per GEMM. Two grouped C++ calls instead of ``2 * num_gemms`` - ungrouped Python calls. - """ +def _hybrid_split_quantize( + tensor, + m_splits, + quantizers, + *, + disable_bulk_allocation=False, +): + """Grouped split+quantize for an all-hybrid, generation-validated operand.""" from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage as HybridStorage - if not all(isinstance(q, HybridQuantizer) for q in quantizers): - raise TypeError( - "_hybrid_split_quantize requires every quantizer to be a" - " HybridQuantizer; callers must gate on _is_hybrid_quantizer_list." - f" Got types: {[type(q).__name__ for q in quantizers]}" - ) - - usage_signatures = [(q.rowwise_usage, q.columnwise_usage) for q in quantizers] - if not all(signature == usage_signatures[0] for signature in usage_signatures): - raise ValueError( - "GroupedLinear HybridQuantizer list has mixed parent usage flags " - f"{usage_signatures}. This is not supported by the grouped " - "split-quantize path; all experts for a grouped operand must " - "request the same rowwise/columnwise directions." - ) - rowwise_enabled, columnwise_enabled = usage_signatures[0] - - row_quantizers = [q.rowwise_quantizer for q in quantizers] - col_quantizers = [q.columnwise_quantizer for q in quantizers] + reference = quantizers[0] + rowwise_enabled = reference.rowwise_usage + columnwise_enabled = reference.columnwise_usage + columnwise_source = reference.columnwise_source + rowwise_quantizers = [quantizer.rowwise_quantizer for quantizer in quantizers] + columnwise_quantizers = [quantizer.columnwise_quantizer for quantizer in quantizers] + needs_rowwise_result = rowwise_enabled or ( + columnwise_enabled and columnwise_source == "rowwise_dequantized" + ) row_results = ( tex.split_quantize( tensor, m_splits, - row_quantizers, + rowwise_quantizers, disable_bulk_allocation=disable_bulk_allocation, ) - if rowwise_enabled + if needs_rowwise_result else [None] * len(quantizers) ) + + columnwise_src = tensor + if columnwise_enabled and columnwise_source == "rowwise_dequantized": + # Assemble the exact grouped row results in split order. NVFP4 padding + # and scale layout can differ from independently quantizing each split. + columnwise_src = torch.cat([result.dequantize() for result in row_results], dim=0) col_results = ( tex.split_quantize( - tensor, + columnwise_src, m_splits, - col_quantizers, + columnwise_quantizers, disable_bulk_allocation=disable_bulk_allocation, ) if columnwise_enabled @@ -261,7 +302,7 @@ def _hybrid_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocat return [ HybridStorage( - rowwise_storage=row, + rowwise_storage=row if rowwise_enabled else None, columnwise_storage=col, quantizer=q, fake_dtype=tensor.dtype, @@ -731,11 +772,6 @@ def forward( for output_quantizer in output_quantizers: output_quantizer.set_usage(rowwise=True, columnwise=False) - if fp8 and not debug: - _check_uniform_identity_quantizer_list(input_quantizers) - _check_uniform_identity_quantizer_list(weight_quantizers) - _check_uniform_identity_quantizer_list(grad_output_quantizers) - # Initialize input tensors in_features = weights[0].size(-1) if inp.size(-1) != in_features: @@ -783,9 +819,11 @@ def forward( inp_view = inp.reshape(-1, in_features) inputmats: list - identity = _has_identity_quantizer_list(input_quantizers) - hybrid = False if identity else _is_hybrid_quantizer_list(input_quantizers) - if fp8 and not debug and not hybrid: + input_reference = input_quantizers[0] + input_hybrid = isinstance( + input_reference, HybridQuantizer + ) and not _uses_identity_quantizer(input_reference) + if fp8 and not debug and not input_hybrid: # Disable bulk allocation when CPU offloading is active: offloading skips small # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. @@ -796,7 +834,7 @@ def forward( activation_dtype, disable_bulk_allocation=cpu_offloading, ) - elif fp8 and hybrid: + elif fp8 and input_hybrid: inputmats = _hybrid_split_quantize( inp_view, m_splits, @@ -1201,11 +1239,19 @@ def backward( grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) grad_output = [None] * ctx.num_gemms grad_biases = [None] * ctx.num_gemms - grad_output_identity = _has_identity_quantizer_list(ctx.grad_output_quantizers) + grad_output_reference = ctx.grad_output_quantizers[0] + if ctx.fp8 and isinstance(grad_output_reference, HybridQuantizer): + # Usage is a runtime decision, not part of generation validation. + # Apply it uniformly so dispatch can read the first parent without + # rescanning every expert. + for grad_output_quantizer in ctx.grad_output_quantizers: + grad_output_quantizer.set_usage( + rowwise=ctx.requires_dgrad, + columnwise=ctx.weights_requires_grad, + ) + grad_output_identity = _uses_identity_quantizer(grad_output_reference) grad_output_hybrid = ( - False - if grad_output_identity - else _is_hybrid_quantizer_list(ctx.grad_output_quantizers) + isinstance(grad_output_reference, HybridQuantizer) and not grad_output_identity ) if ctx.fp8 and not ctx.debug and not grad_output_hybrid: if ctx.use_bias: @@ -1357,10 +1403,10 @@ def backward( else: input_quantizer.set_usage(rowwise=False, columnwise=True) inputmats: list - input_identity = _has_identity_quantizer_list(ctx.input_quantizers) - input_hybrid = ( - False if input_identity else _is_hybrid_quantizer_list(ctx.input_quantizers) - ) + input_reference = ctx.input_quantizers[0] + input_hybrid = isinstance( + input_reference, HybridQuantizer + ) and not _uses_identity_quantizer(input_reference) if ctx.fp8 and not ctx.debug and not input_hybrid: inputmats = _split_quantize_with_identity_fallback( inp_view, @@ -1611,6 +1657,7 @@ def __init__( "fwd": 3, "bwd": 2, } + self._validated_quantizer_generations = {} if tp_group is None: self.tp_size = tp_size @@ -1699,6 +1746,43 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) + self._validate_quantizer_generation(fwd) + + def _validate_quantizer_generation(self, fwd: bool) -> None: + """Validate grouped-kernel invariants once per quantizer generation.""" + # Recipe state replaces this list object only when it constructs a new + # quantizer generation. The O(1) identity guard keeps validation off the + # steady-state forward path. Record a generation only after all of its + # operand roles pass, so a failed recipe transition is retried. + meta_key = "scaling_fwd" if fwd else "scaling_bwd" + generation = self.quantizers.get(meta_key) + if generation is None: + return + if self._validated_quantizer_generations.get(meta_key) is generation: + return + + if fwd: + stride = self._num_fp8_tensors_per_gemm["fwd"] + input_quantizers = tuple( + generation[self._offsets["input"] + i * stride] for i in range(self.num_gemms) + ) + weight_quantizers = tuple( + generation[self._offsets["weight"] + i * stride] for i in range(self.num_gemms) + ) + _validate_grouped_quantizer_list(input_quantizers, operand_name="input") + _validate_grouped_quantizer_list(weight_quantizers, operand_name="weight") + else: + stride = self._num_fp8_tensors_per_gemm["bwd"] + grad_output_quantizers = tuple( + generation[self._offsets["grad_output"] + i * stride] for i in range(self.num_gemms) + ) + _validate_grouped_quantizer_list( + grad_output_quantizers, + operand_name="grad_output", + ) + + self._validated_quantizer_generations[meta_key] = generation + def get_quantizer_roles( self, *, @@ -2229,6 +2313,14 @@ def _get_weight_quantizers(self) -> List[Quantizer]: return weight_quantizers def _get_quantizers(self): + if self.fp8: + # Normally validated while installing recipe metadata. Keep this + # O(1) generation guard so failed transitions cannot reuse stale + # validation state if base metadata takes an early return on retry. + self._validate_quantizer_generation(True) + if torch.is_grad_enabled(): + self._validate_quantizer_generation(False) + weight_quantizers = self._get_weight_quantizers() input_quantizers, output_quantizers = ( [None] * self.num_gemms, diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index a37d44e7c75..f8a447ae4b9 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -570,6 +570,7 @@ def __new__( requires_grad: bool = False, device: Optional[torch.device] = None, stride: Optional[Iterable[int]] = None, + storage_offset: int = 0, ): if fake_dtype is not None and fake_dtype != dtype: raise ValueError(f"fake_dtype ({fake_dtype}) does not match dtype ({dtype})") @@ -582,7 +583,7 @@ def __new__( cls, shape, strides=stride, - storage_offset=0, + storage_offset=storage_offset, dtype=dtype, layout=torch.strided, requires_grad=requires_grad, diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 1b08039dda2..afb25ac279c 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -9,13 +9,68 @@ """ from __future__ import annotations -from typing import Callable, Optional, Tuple, Any, Dict, TYPE_CHECKING +from typing import Callable, Optional, Tuple, Any, Dict, Iterable, TYPE_CHECKING import torch if TYPE_CHECKING: from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +def _resolve_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: + """Resolve a requested view shape with PyTorch-compatible semantics. + + The concrete-integer path avoids constructing a temporary meta tensor. If + either shape contains symbolic dimensions, retain the previous meta-tensor + path so that PyTorch remains responsible for symbolic shape handling. + """ + input_shape = tuple(input_shape) + shape = tuple(shape) + if len(shape) == 1 and isinstance(shape[0], (list, tuple, torch.Size)): + shape = tuple(shape[0]) + + # Avoid comparisons that specialize or guard SymInts. The meta fallback is + # also useful for preserving PyTorch's type checking of non-integer dims. + if any( + not isinstance(dim, int) or isinstance(dim, bool) for dim in (*input_shape, *shape) + ): + return torch.empty(input_shape, device="meta").view(shape).shape + + input_numel = 1 + for dim in input_shape: + input_numel *= dim + + inferred_dim = None + known_numel = 1 + for index, dim in enumerate(shape): + if dim == -1: + if inferred_dim is not None: + raise RuntimeError("only one dimension can be inferred") + inferred_dim = index + elif dim < 0: + raise RuntimeError( + f"invalid shape dimension {dim} at index {index} of shape {list(shape)}" + ) + else: + known_numel *= dim + + resolved_shape = list(shape) + if inferred_dim is not None: + if known_numel == 0: + if input_numel == 0: + raise RuntimeError( + f"cannot reshape tensor of 0 elements into shape {list(shape)} because " + "the unspecified dimension size -1 can be any value and is ambiguous" + ) + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + if input_numel % known_numel != 0: + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + resolved_shape[inferred_dim] = input_numel // known_numel + elif known_numel != input_numel: + raise RuntimeError(f"shape '{list(shape)}' is invalid for input of size {input_numel}") + + return torch.Size(resolved_shape) + + class _QuantizeFunc(torch.autograd.Function): """Quantize tensor""" diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index d8362670712..2c962fd5986 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -19,7 +19,11 @@ from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer from ..dynamo import register_value_opaque_quantizer -from ._quantization_helpers import _IdentityFunc, safe_quantized_repr +from ._quantization_helpers import ( + _IdentityFunc, + _resolve_view_shape, + safe_quantized_repr, +) from ..constants import dist_group_type, DType aten = torch.ops.aten @@ -38,11 +42,6 @@ } -def _canonical_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> torch.Size: - """Resolve PyTorch view shape syntax, including ``-1`` inference.""" - return torch.empty(tuple(input_shape), device="meta").view(*shape).shape - - def _columnwise_shape_for(rowwise_shape: Iterable[int]) -> torch.Size: """Physical columnwise FP8 shape for a logical rowwise shape.""" shape = torch.Size(rowwise_shape) @@ -614,7 +613,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ) out_shape = out_data.size() else: - out_shape = _canonical_view_shape(tensor.shape, args[1:]) + out_shape = _resolve_view_shape(tensor.shape, args[1:]) out_transpose = None if tensor._transpose_invalid else tensor._transpose if out_transpose is not None: @@ -1117,7 +1116,7 @@ def forward( out_data = tensor._data.view(*shape) out_shape = out_data.size() else: - out_shape = _canonical_view_shape(tensor.shape, shape) + out_shape = _resolve_view_shape(tensor.shape, shape) out_transpose = None if tensor._transpose_invalid else tensor._transpose if out_transpose is not None: view_shape_for_transpose = _columnwise_shape_for(out_shape) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 58cc878bed3..ce3892ef10f 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -10,6 +10,7 @@ import torch from .storage.hybrid_tensor_storage import HybridQuantizedTensorStorage +from .storage.identity_tensor_storage import IdentityTensorStorage from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer aten = torch.ops.aten @@ -846,6 +847,52 @@ def _delegate(sub): quantizer=tensor._quantizer, ) + @staticmethod + def _new_zeroed_sub_storage( + sub_storage: QuantizedTensorStorage, + shape: torch.Size, + *, + dtype: torch.dtype, + device: torch.device, + pin_memory: bool, + ) -> QuantizedTensorStorage: + """Allocate and initialize a zero sub-storage without touching live state.""" + + try: + quantizer = sub_storage._get_quantizer().copy() + out = quantizer.make_empty( + shape, + dtype=dtype, + device=device, + requires_grad=False, + pin_memory=pin_memory, + ) + except NotImplementedError as exc: + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros cannot construct zero-initialized " + f"{type(sub_storage).__name__} storage" + ) from exc + + if not isinstance(out, type(sub_storage)): + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros did not preserve sub-storage format: " + f"expected {type(sub_storage).__name__}, got {type(out).__name__}" + ) + + buffers, storage = out.prepare_for_saving() + try: + with torch.no_grad(): + for buffer in buffers: + if buffer is not None: + buffer.zero_() + finally: + leftover = storage.restore_from_saved(buffers) + if leftover: + raise RuntimeError( + f"{type(out).__name__}.restore_from_saved did not consume all buffers" + ) + return out + @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): if kwargs is None: @@ -1006,14 +1053,82 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): # ── FSDP2: new_zeros ───────────────────────────────────────── if func == aten.new_zeros.default: tensor = args[0] - new_shape = args[1] - # FSDP2 allocates new_zeros buffers as all-gather destinations - # that are immediately overwritten by copy_. Use make_empty - # (uninitialized storage with the right container shape/fields). - return tensor._quantizer.make_empty( - new_shape, - dtype=tensor.dtype, - device=tensor.device, + new_shape = torch.Size(args[1]) + if tensor._rowwise_storage is None and tensor._columnwise_storage is None: + raise RuntimeError( + "HybridQuantizedTensor.new_zeros requires at least one present sub-storage" + ) + dtype = kwargs.get("dtype") or tensor.dtype + device = torch.device(kwargs.get("device") or tensor.device) + layout = kwargs.get("layout") + pin_memory = bool(kwargs.get("pin_memory")) + if layout is not None and layout != torch.strided: + raise NotImplementedError( + "HybridQuantizedTensor.new_zeros only supports torch.strided layout" + ) + + supported_quantized_dtypes = (torch.float32, torch.float16, torch.bfloat16) + for direction, sub_storage in ( + ("rowwise", tensor._rowwise_storage), + ("columnwise", tensor._columnwise_storage), + ): + if sub_storage is None or isinstance(sub_storage, IdentityTensorStorage): + continue + if dtype not in supported_quantized_dtypes: + raise TypeError( + "HybridQuantizedTensor.new_zeros only supports float32, float16, or " + f"bfloat16 with a {direction} {type(sub_storage).__name__}; got {dtype}." + ) + + row = ( + cls._new_zeroed_sub_storage( + tensor._rowwise_storage, + new_shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + if tensor._rowwise_storage is not None + else None + ) + col = ( + cls._new_zeroed_sub_storage( + tensor._columnwise_storage, + new_shape, + dtype=dtype, + device=device, + pin_memory=pin_memory, + ) + if tensor._columnwise_storage is not None + else None + ) + + # The source's parent usage flags are mutable and may no longer + # describe its surviving sub-storages. Give the result an isolated + # parent whose dynamic usage matches what was actually allocated; + # otherwise a later copy_ from a plain tensor can silently skip a + # present direction. Bind its children to copies of the newly + # allocated storage quantizers so their format/state stays coherent. + quantizer = tensor._quantizer.copy() + if row is not None: + quantizer.rowwise_quantizer = row._get_quantizer().copy() + quantizer.rowwise_quantizer.set_usage(rowwise=True, columnwise=False) + if col is not None: + quantizer.columnwise_quantizer = col._get_quantizer().copy() + quantizer.columnwise_quantizer.set_usage(rowwise=False, columnwise=True) + quantizer.set_usage( + rowwise=row is not None, + columnwise=col is not None, + ) + + return HybridQuantizedTensor( + shape=new_shape, + dtype=dtype, + rowwise_storage=row, + columnwise_storage=col, + quantizer=quantizer, + requires_grad=False, + device=device, ) # ── FSDP2: clone ───────────────────────────────────────────── diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index cb9ea2e0b0c..5d80fcbd0cd 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -170,7 +170,7 @@ def view(self, *shape) -> "IdentityTensor": def detach(self) -> "IdentityTensor": # pylint: disable=missing-function-docstring - return IdentityTensor.make_like(self) + return self._wrap_data_view(self._hp_data.detach(), requires_grad=False) def clone(self) -> "IdentityTensor": # pylint: disable=missing-function-docstring @@ -182,6 +182,8 @@ def clone(self) -> "IdentityTensor": quantizer=self._quantizer, requires_grad=self.requires_grad, device=self.device, + stride=data.stride(), + storage_offset=data.storage_offset(), ) def contiguous( @@ -230,16 +232,38 @@ def fsdp_post_all_gather( ) return out, all_gather_outputs - def _wrap_data_view(self, data: torch.Tensor) -> "IdentityTensor": + def _wrap_data_view( + self, data: torch.Tensor, *, requires_grad: Optional[bool] = None + ) -> "IdentityTensor": + requires_grad = self.requires_grad if requires_grad is None else requires_grad return IdentityTensor( shape=data.shape, dtype=self.dtype, hp_data=data, quantizer=self._quantizer, - requires_grad=self.requires_grad, + requires_grad=requires_grad, device=data.device, + stride=data.stride(), + storage_offset=data.storage_offset(), ) + @classmethod + def _delegate_view_op(cls, func, tensor, args, kwargs): + """Apply an alias-preserving view op to the held tensor and rewrap it.""" + + result = func(tensor._hp_data, *args[1:], **kwargs) + + def _wrap(value): + if isinstance(value, torch.Tensor): + return tensor._wrap_data_view(value) + return value + + if isinstance(result, tuple): + return tuple(_wrap(value) for value in result) + if isinstance(result, list): + return [_wrap(value) for value in result] + return _wrap(result) + @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): if kwargs is None: @@ -251,51 +275,21 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.clone.default: return args[0].clone() - if func == aten.view.default: - tensor = args[0] - shape = args[1] - if list(shape) == list(tensor.shape): - return IdentityTensor.make_like(tensor) - return tensor._wrap_data_view(tensor._hp_data.view(*shape)) - - if func == aten.split.Tensor: - tensor = args[0] - split_size = args[1] - dim = kwargs.get("dim", args[2] if len(args) > 2 else 0) - return [ - tensor._wrap_data_view(piece) - for piece in torch.split(tensor._hp_data, split_size, dim=dim) - ] - - if func == aten.as_strided.default: - tensor = args[0] - shape = args[1] - strides = args[2] - storage_offset = kwargs.get("storage_offset", args[3] if len(args) > 3 else 0) - if ( - tuple(shape) == tuple(tensor.shape) - and tuple(strides) == tuple(tensor.stride()) - and storage_offset == tensor.storage_offset() - ): - return IdentityTensor.make_like(tensor) - return tensor._wrap_data_view( - torch.as_strided(tensor._hp_data, shape, strides, storage_offset) - ) - - if func == aten.slice.Tensor: - tensor = args[0] - dim = args[1] - start = args[2] - end = args[3] - step = args[4] if len(args) > 4 else 1 - if start == 0 and end == tensor.size(dim) and step == 1: - return IdentityTensor.make_like(tensor) - return tensor._wrap_data_view(aten.slice.Tensor(tensor._hp_data, dim, start, end, step)) + if func in ( + aten.view.default, + aten.split.Tensor, + aten.as_strided.default, + aten.slice.Tensor, + ): + # Preserve optional arguments exactly; omitted as_strided offset + # means reuse the input view's current storage offset, not zero. + return cls._delegate_view_op(func, args[0], args, kwargs) if func == aten.copy_.default: dst, src = args[0], args[1] - if isinstance(dst, IdentityTensor) and isinstance(src, IdentityTensor): - dst._hp_data.copy_(src._hp_data, *args[2:], **kwargs) + if isinstance(dst, IdentityTensor): + src_data = src._hp_data if isinstance(src, IdentityTensor) else src + dst._hp_data.copy_(src_data, *args[2:], **kwargs) return dst if func == aten.new_zeros.default: diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index eaa7e48b017..e0bd5613dcd 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -12,7 +12,7 @@ import transformer_engine_torch as tex from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from .._quantization_helpers import safe_quantized_repr +from .._quantization_helpers import _resolve_view_shape, safe_quantized_repr from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType @@ -175,12 +175,15 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: dtype = self._dtype return _FromFloat8Func.forward(None, self, dtype) - def size(self, *args, **kwargs): + def size(self, dim: Optional[int] = None) -> Union[torch.Size, int]: # pylint: disable=missing-function-docstring if self._data is not None: - return self._data.size(*args, **kwargs) - size = self._transpose.size(*args, **kwargs) - return torch.Size([size[-1], math.prod(size[:-1])]) + return self._data.size(dim) + transpose_size = self._transpose.size() + size = torch.Size([transpose_size[-1], math.prod(transpose_size[:-1])]) + if dim is None: + return size + return size[dim] @property def device(self): @@ -197,7 +200,7 @@ def view(self, shape: torch.Size): if out_data is not None: out_shape = out_data.size() else: - out_shape = torch.empty(tuple(self.size()), device="meta").view(shape).shape + out_shape = _resolve_view_shape(self.size(), shape) out_transpose = None if self._transpose_invalid else self._transpose if out_transpose is not None: if len(out_shape) == 0: diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index e376591cdbb..86a9a7534c9 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -75,6 +75,19 @@ def update_usage( rowwise_usage: Optional[bool] = None, columnwise_usage: Optional[bool] = None, ): + # A storage object cannot reconstruct a representation that has already + # been dropped. Validate every requested direction before mutating + # either one so mixed drop/enable requests are atomic. + if rowwise_usage and self._rowwise_storage is None: + raise RuntimeError( + "Requested rowwise usage, but HybridQuantizedTensorStorage " + "has no rowwise sub-storage" + ) + if columnwise_usage and self._columnwise_storage is None: + raise RuntimeError( + "Requested columnwise usage, but HybridQuantizedTensorStorage " + "has no columnwise sub-storage" + ) if rowwise_usage is not None and not rowwise_usage: self._rowwise_storage = None if columnwise_usage is not None and not columnwise_usage: diff --git a/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py index 4cb7fc092a1..bcf606b7e66 100644 --- a/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py @@ -30,7 +30,6 @@ class IdentityTensorStorage(QuantizedTensorStorage): """ _hp_data: Optional[torch.Tensor] - _quantizer: Optional[Quantizer] def __new__( cls, @@ -103,7 +102,9 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """Return the held high-precision tensor (no-op dequantization).""" if self._hp_data is None: raise RuntimeError("IdentityTensorStorage has no data to dequantize") - if dtype is not None and self._hp_data.dtype != dtype: + if dtype is None: + dtype = self._dtype + if self._hp_data.dtype != dtype: return self._hp_data.to(dtype) return self._hp_data From 05f4d1e6d365917805962c738119337424610aa5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:40:04 +0000 Subject: [PATCH 44/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/tensor/_quantization_helpers.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index afb25ac279c..10672bbcfbd 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -30,9 +30,7 @@ def _resolve_view_shape(input_shape: Iterable[int], shape: Iterable[int]) -> tor # Avoid comparisons that specialize or guard SymInts. The meta fallback is # also useful for preserving PyTorch's type checking of non-integer dims. - if any( - not isinstance(dim, int) or isinstance(dim, bool) for dim in (*input_shape, *shape) - ): + if any(not isinstance(dim, int) or isinstance(dim, bool) for dim in (*input_shape, *shape)): return torch.empty(input_shape, device="meta").view(shape).shape input_numel = 1 From 0dd5a5f59c305ace4aa869841a57ef526a3e13a5 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 17 Jul 2026 10:44:50 +0000 Subject: [PATCH 45/97] Finish comment resolving: materialize quantizer once, + minor fixes Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/fsdp2_utils.py | 10 +- tests/pytorch/test_hybrid_quantization.py | 122 ++++++++++++++++++ .../dot_product_attention.py | 55 ++++++++ .../pytorch/attention/multi_head_attention.py | 51 +++----- .../quantization_factory_base.py | 6 - .../quantization_factory_zoo.py | 2 - 6 files changed, 200 insertions(+), 46 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 39b7b0ad088..fff55471025 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -79,11 +79,11 @@ def _identity_qfactory(role): # pylint: disable=unused-argument return IdentityQuantizer() -# The qfactories above are registered here as module-level functions (not -# lambdas or closures) on purpose: DCP serializes ``CustomRecipe`` via -# ``pickle``, and closure-based qfactories (or inner functions capturing state) -# are not picklable. Keeping them at module scope lets them pickle by reference. -# See ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. +# CustomRecipe has dynamic TE extra-state handling. Once FP8 state is +# initialized, TE's get_extra_state() pickles the recipe on save, so +# checkpoint-test qfactories must be module-level and picklable. On load, +# payloads without delayed-scaling state are identified and ignored without +# unpickling. See ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. _HYBRID_QFACTORIES = { "HybridFP8CurrentScaling": _hybrid_fp8_current_qfactory, "HybridMXFP8": _hybrid_mxfp8_qfactory, diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index aa46a56c51b..4949708a1b2 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -1799,6 +1799,128 @@ def fake_infer(*_args, **_kwargs): assert inference_calls == 3 assert "local_recipes" not in dpa.fp8_meta + @pytest.mark.parametrize( + "factory_name,expected", + [ + ("current_scaling_quantizer_factory", (True, False)), + ("delayed_scaling_quantizer_factory", (False, False)), + ], + ) + def test_qkv_capabilities_reuse_canonical_quantizer(self, factory_name, expected): + """Capability queries must not call qfactory outside recipe-state setup.""" + from transformer_engine.pytorch.custom_recipes import quantization_factory_base + + base_qfactory = getattr(quantization_factory_base, factory_name) + calls = [] + + def counting_qfactory(role): + calls.append(role) + # Model a factory with observable RNG state. An extra classification + # probe would consume another value and change subsequent results. + _ = torch.rand((), device="cuda") + return base_qfactory(role) + + custom_recipe = recipe.CustomRecipe( + qfactory=counting_qfactory, + fp8_dpa=True, + fp8_mha=True, + ) + dpa = te.DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + name="counted_dpa", + ).cuda() + + with autocast(enabled=True, recipe=custom_recipe): + first = dpa.get_qkv_quantization_capabilities() + canonical_qkv = dpa._qkv_capabilities_quantizer + calls_after_first = len(calls) + second = dpa.get_qkv_quantization_capabilities() + + assert first == expected + assert second == first + assert dpa._qkv_capabilities_quantizer is canonical_qkv + assert len(calls) == calls_after_first + + # A recipe-state rebuild creates a new canonical slot and must + # invalidate the capability cache automatically. + def rebuilt_qfactory(role): + return counting_qfactory(role) + + rebuilt_recipe = recipe.CustomRecipe( + qfactory=rebuilt_qfactory, + fp8_dpa=True, + fp8_mha=True, + ) + with autocast(enabled=True, recipe=rebuilt_recipe): + rebuilt = dpa.get_qkv_quantization_capabilities() + rebuilt_qkv = dpa._qkv_capabilities_quantizer + + assert rebuilt == first + assert rebuilt_qkv is not canonical_qkv + assert len(calls) > calls_after_first + + @pytest.mark.parametrize("fp8_mha", [False, True]) + def test_mha_forward_does_not_probe_qfactory(self, monkeypatch, fp8_mha): + """Repeated MHA forwards must not create discarded DPA quantizers.""" + from transformer_engine.pytorch.attention import multi_head_attention as mha_module + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + current_scaling_quantizer_factory, + ) + from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") + + monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", "") + calls = [] + + def valid_fp8_dpa_qfactory(role): + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + if is_dpa or is_dpa_boundary: + return nvfp4_linear_fp8_dpa_factory(role) + return current_scaling_quantizer_factory(role) + + def counting_qfactory(role): + calls.append(role) + _ = torch.rand((), device="cuda") + return valid_fp8_dpa_qfactory(role) + + custom_recipe = recipe.CustomRecipe( + qfactory=counting_qfactory, + fp8_dpa=True, + fp8_mha=fp8_mha, + ) + model = te.MultiheadAttention( + hidden_size=32, + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + attn_mask_type="no_mask", + params_dtype=torch.bfloat16, + bias=False, + qkv_format="sbhd", + name="counted_mha", + ).cuda() + inp = torch.randn(128, 2, 32, device="cuda", dtype=torch.bfloat16) + + with torch.no_grad(), autocast(enabled=True, recipe=custom_recipe): + model(inp) + calls_after_first = len(calls) + model(inp) + + assert len(calls) == calls_after_first + @requires_fp8_and_nvfp4 class TestAttentionFactoryNativeRecipeParity: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index f03248eff67..6c43e2725f9 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -468,6 +468,8 @@ def __init__( # change when their recipe state is rebuilt. self._custom_dpa_local_recipes_cache_key: Optional[Tuple[Any, ...]] = None self._custom_dpa_local_recipes_cache: Optional[List[Recipe]] = None + self._qkv_capabilities_quantizer: Optional[Any] = None + self._qkv_capabilities_cache: Optional[Tuple[bool, bool]] = None self.logger = logging.getLogger("DotProductAttention") self.logger.setLevel(attn_log._log_level) @@ -915,6 +917,59 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # Clear cached workspaces as they were created with the old recipe/quantizer type self._fp8_workspaces.clear() + def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: + """Return MHA boundary capabilities from the canonical QKV quantizer. + + The returned flags are ``(float8_current_scaling, mxfp8_scaling)``. + """ + self.init_fp8_metadata(num_gemms=3) + try: + qkv_quantizer = self.quantizers["scaling_fwd"][dpa_utils.META_QKV] + except (KeyError, IndexError, TypeError) as exc: + role = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=self.name or "", + ) + raise RuntimeError( + "DotProductAttention did not materialize the canonical QKV " + f"quantizer for {role}." + ) from exc + + if qkv_quantizer is self._qkv_capabilities_quantizer: + assert self._qkv_capabilities_cache is not None + return self._qkv_capabilities_cache + + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, + ) + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): + capabilities = (True, False) + elif isinstance(qkv_quantizer, MXFP8Quantizer): + capabilities = (False, True) + elif isinstance(qkv_quantizer, Float8Quantizer): + capabilities = (False, False) + else: + capabilities = None + + if capabilities is not None: + self._qkv_capabilities_quantizer = qkv_quantizer + self._qkv_capabilities_cache = capabilities + return capabilities + + role = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=self.name or "", + ) + raise TypeError( + f"Unsupported CustomRecipe quantizer for {role}: " + f"{type(qkv_quantizer).__name__}." + ) + def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> None: """Override to allow multiple recipes. Init scales and amaxes for fwd | bwd.""" if isinstance(recipe, Recipe) and recipe.custom(): diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 2343033fee1..22afe996b27 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -43,33 +43,6 @@ _dpa_fp8_recipe_mha = os.getenv("NVTE_DPA_FP8_RECIPE_MHA", "0") == "1" -def _infer_custom_dpa_scaling(recipe, dpa_name: str) -> Tuple[bool, bool]: - """Infer DPA quantizer family for CustomRecipe boundary decisions. - - MultiheadAttention must decide boundary behavior before - DotProductAttention builds its per-slot quantizers. Built-in recipes expose - the DPA quantizer family through recipe predicates; CustomRecipe exposes it - only through the qfactory role contract, so probe the DPA QKV role. - """ - if not recipe.custom() or not getattr(recipe, "fp8_dpa", False): - return False, False - - try: - qkv_quantizer = recipe.qfactory( - QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) - ) - except Exception: # pragma: no cover, pylint: disable=broad-exception-caught - return False, False - - from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - return ( - isinstance(qkv_quantizer, Float8CurrentScalingQuantizer), - isinstance(qkv_quantizer, MXFP8Quantizer), - ) - - class MultiheadAttention(torch.nn.Module): r""" Multi-head Attention (MHA), including Query, @@ -904,12 +877,19 @@ def forward( fp8_mha = fp8_recipe.fp8_mha float8_current_scaling = fp8_recipe.float8_current_scaling() mxfp8_scaling = fp8_recipe.mxfp8() - if custom_recipe and fp8_dpa: - custom_float8_cs, custom_mxfp8 = _infer_custom_dpa_scaling( - fp8_recipe, self.core_attention.name or "" + if fp8 and custom_recipe and fp8_mha: + # Wire every boundary this CustomRecipe may quantize before + # DPA materializes its recipe state. Some quantizer families + # disable the corresponding output below, but pre-wiring avoids + # rebuilding that state after inspecting its canonical QKV slot. + self._update_output_quantizer_roles( + rotary_pos_emb is None, + True, + True, + ) + float8_current_scaling, mxfp8_scaling = ( + self.core_attention.get_qkv_quantization_capabilities() ) - float8_current_scaling = custom_float8_cs - mxfp8_scaling = custom_mxfp8 else: fp8_dpa = _dpa_fp8_recipe_dpa fp8_mha = _dpa_fp8_recipe_mha @@ -939,7 +919,12 @@ def forward( # 1. FP8CS recipe: produce F16 grads; again, due to cuBLAS limitation proj_fp8_grad = dpa_fp8_output and not float8_current_scaling - self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + # Custom fp8_mha boundaries were wired before DPA recipe-state setup so + # querying its canonical QKV quantizer cannot trigger a second build. + if not (fp8 and custom_recipe and fp8_mha and _dpa_fp8_recipe == ""): + self._update_output_quantizer_roles( + qkv_fp8_output, proj_fp8_grad, dpa_fp8_output + ) layernorm_output = None if self.attention_type == "self": diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py index 0a2b8f99334..7dae0891777 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py @@ -164,8 +164,6 @@ def nvfp4_quantizer_factory( if is_weight: return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, with_rht=False, with_post_rht_amax=False, with_2d_quantization=True, @@ -176,8 +174,6 @@ def nvfp4_quantizer_factory( if is_grad: return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, with_rht=True, with_post_rht_amax=True, with_2d_quantization=False, @@ -188,8 +184,6 @@ def nvfp4_quantizer_factory( # For input and unknown roles return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, with_rht=True, with_post_rht_amax=True, with_2d_quantization=False, diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index 3eb02ddcbf1..aca031d39c0 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -127,8 +127,6 @@ def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): return NVFP4Quantizer( fp4_dtype=DType.kFloat4E2M1, - rowwise=True, - columnwise=True, with_rht=False, with_post_rht_amax=False, with_2d_quantization=False, From 77ce9ffea1eebb94c634eccf570a51777a3d30bf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 10:46:05 +0000 Subject: [PATCH 46/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../dot_product_attention/dot_product_attention.py | 6 ++---- .../pytorch/attention/multi_head_attention.py | 4 +--- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 6c43e2725f9..787efa9fa8a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -932,8 +932,7 @@ def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: name=self.name or "", ) raise RuntimeError( - "DotProductAttention did not materialize the canonical QKV " - f"quantizer for {role}." + f"DotProductAttention did not materialize the canonical QKV quantizer for {role}." ) from exc if qkv_quantizer is self._qkv_capabilities_quantizer: @@ -966,8 +965,7 @@ def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: name=self.name or "", ) raise TypeError( - f"Unsupported CustomRecipe quantizer for {role}: " - f"{type(qkv_quantizer).__name__}." + f"Unsupported CustomRecipe quantizer for {role}: {type(qkv_quantizer).__name__}." ) def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> None: diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 22afe996b27..09df941b4a4 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -922,9 +922,7 @@ def forward( # Custom fp8_mha boundaries were wired before DPA recipe-state setup so # querying its canonical QKV quantizer cannot trigger a second build. if not (fp8 and custom_recipe and fp8_mha and _dpa_fp8_recipe == ""): - self._update_output_quantizer_roles( - qkv_fp8_output, proj_fp8_grad, dpa_fp8_output - ) + self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) layernorm_output = None if self.attention_type == "self": From d0f604e51fa5ca79be5625f0a3cbdcc3000c3fb3 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 17 Jul 2026 12:59:43 +0000 Subject: [PATCH 47/97] Add automatic alignment contract for CustomRecipe Signed-off-by: Evgeny --- tests/pytorch/test_custom_recipe.py | 101 +++++++++++++++++- transformer_engine/common/recipe/__init__.py | 14 ++- .../pytorch/module/fp8_padding.py | 7 +- .../pytorch/module/fp8_unpadding.py | 7 +- transformer_engine/pytorch/quantization.py | 12 ++- 5 files changed, 132 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 79ca5fe6126..e88ef93c958 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -11,13 +11,18 @@ from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx from transformer_engine.pytorch import ( autocast, + Fp8Padding, + Fp8Unpadding, Linear, LayerNormLinear, LayerNormMLP, GroupedLinear, Float8CurrentScalingQuantizer, ) -from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.quantization import ( + QuantizerRole, + get_align_size_for_quantization, +) import transformer_engine.pytorch.ops as te_ops from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( current_scaling_quantizer_factory, @@ -27,6 +32,9 @@ delayed_scaling_quantizer_factory, high_precision_factory, ) +from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_quantizer_factory, +) from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import ( nvfp4_ref_rht_2d_quantizer_factory, ) @@ -1884,3 +1892,94 @@ def test_slot_role_supports_module_type_only_role(): assert resolved.name == "" # Tensor-type-only recipes fall back to positional for this slot. assert state._slot_tensor_type(0) == "input" + + +_alignment_fp8_available, _alignment_fp8_reason = te.is_fp8_available(return_reason=True) +_alignment_mxfp8_available, _alignment_mxfp8_reason = te.is_mxfp8_available(return_reason=True) +_alignment_nvfp4_available, _alignment_nvfp4_reason = te.is_nvfp4_available(return_reason=True) + + +def test_custom_recipe_quantization_alignment_contract(): + """The alignment contract is safe, configurable, and does not invoke qfactory.""" + calls = [] + + def qfactory(role): + calls.append(role) + return current_scaling_quantizer_factory(role) + + custom_recipe = recipe.CustomRecipe(qfactory=qfactory) + assert get_align_size_for_quantization(custom_recipe) == 128 + assert calls == [] + + custom_recipe = recipe.CustomRecipe(qfactory=qfactory, quantization_alignment=32) + assert get_align_size_for_quantization(custom_recipe) == 32 + assert calls == [] + + with pytest.raises(ValueError, match="quantization_alignment must be positive"): + recipe.CustomRecipe(qfactory=qfactory, quantization_alignment=0) + + +@pytest.mark.parametrize( + "qfactory", + [ + pytest.param( + current_scaling_quantizer_factory, + marks=pytest.mark.skipif( + not _alignment_fp8_available, + reason=_alignment_fp8_reason, + ), + id="fp8_current_scaling", + ), + pytest.param( + mxfp8_quantizer_factory, + marks=pytest.mark.skipif( + not _alignment_mxfp8_available, + reason=_alignment_mxfp8_reason, + ), + id="mxfp8", + ), + pytest.param( + nvfp4_quantizer_factory, + marks=pytest.mark.skipif( + not _alignment_nvfp4_available, + reason=_alignment_nvfp4_reason, + ), + id="nvfp4", + ), + pytest.param( + mxfp8_fwd_nvfp4_bwd_quantizer_factory, + marks=pytest.mark.skipif( + not (_alignment_mxfp8_available and _alignment_nvfp4_available), + reason=f"MXFP8: {_alignment_mxfp8_reason}; NVFP4: {_alignment_nvfp4_reason}", + ), + id="hybrid_mxfp8_nvfp4", + ), + ], +) +def test_custom_recipe_automatic_padding_bad_splits(qfactory): + """Automatic alignment makes misaligned CustomRecipe grouped GEMMs valid.""" + custom_recipe = recipe.CustomRecipe(qfactory=qfactory) + padding = Fp8Padding(2) + unpadding = Fp8Unpadding(2) + grouped_linear = GroupedLinear( + 2, + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + ) + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + original_splits = [17, 15] + + with autocast(enabled=True, recipe=custom_recipe): + padded_inp, padded_splits = padding(inp, original_splits) + assert padding.align_size == get_align_size_for_quantization(custom_recipe) == 128 + assert padded_splits == [128, 128] + padded_out = grouped_linear(padded_inp, padded_splits) + out = unpadding(padded_out, original_splits) + + out.sum().backward() + assert out.shape == (sum(original_splits), 128) + assert torch.isfinite(out).all() + assert torch.isfinite(inp.grad).all() diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 8a03f2f51a2..9f492572957 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -660,6 +660,14 @@ class CustomRecipe(Recipe): `high_precision` keeps original high-precision operands for backward, and `dequantized` dequantizes saved operands to the active high-precision compute dtype (e.g. BF16/FP16/FP32) for backward. + quantization_alignment : int, default = 128 + Conservative recipe-wide fallback used by automatic padding for grouped operations. + This must be at least the largest alignment required by any quantizer + that ``qfactory`` may return for a grouped operation, across all roles + and module names. The default of 128 safely supports all current TE + formats. It can be lowered when the factory's full output space is known; + for example, a factory restricted to MXFP8 may use 32. + Automatic padding reads this value without invoking ``qfactory``. """ qfactory: Callable[..., Any] @@ -672,15 +680,19 @@ class CustomRecipe(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + quantization_alignment: int = 128 def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + if self.quantization_alignment <= 0: + raise ValueError("CustomRecipe quantization_alignment must be positive.") def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"qfactory={self.qfactory}, " - f"backward_override={self.backward_override}" + f"backward_override={self.backward_override}, " + f"quantization_alignment={self.quantization_alignment}" ) diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index 8ac49c9bae7..5c46230b47d 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -80,9 +80,10 @@ class Fp8Padding(torch.nn.Module): num_gemms : int number of GEMMs to be performed simultaneously. align_size : int, optional - the alignment size for the input tensor. If not provided, the alignment size will - be determined by the FP8/FP4 recipe (32 for MXFP8/NVFP4 and 16 for others) in the first - forward pass. + Alignment size for each grouped input. If not provided, it is + determined from the active recipe on the first forward pass: + 32 for MXFP8, 128 for NVFP4, 16 for other built-in recipes, + and ``CustomRecipe.quantization_alignment`` for custom recipes. """ def __init__( diff --git a/transformer_engine/pytorch/module/fp8_unpadding.py b/transformer_engine/pytorch/module/fp8_unpadding.py index c5d396837f2..d34605e1951 100644 --- a/transformer_engine/pytorch/module/fp8_unpadding.py +++ b/transformer_engine/pytorch/module/fp8_unpadding.py @@ -78,9 +78,10 @@ class Fp8Unpadding(torch.nn.Module): num_gemms : int number of GEMMs to be performed simultaneously. align_size : int, optional - The alignment size for the input tensor. If not provided, the alignment size will - be automatically determined based on the FP8/FP4 recipe in the first forward pass: - 32 for MXFP8 or NVFP4, otherwise 16. + Alignment size used to pad each grouped input. If not provided, + it is determined from the active recipe on the first forward + pass: 32 for MXFP8, 128 for NVFP4, 16 for other built-in + recipes, and ``CustomRecipe.quantization_alignment`` for custom recipes. """ def __init__( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 3c3158da96f..07a3b804830 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -260,7 +260,17 @@ def get_default_recipe() -> Recipe: def get_align_size_for_quantization(recipe: Recipe) -> int: - """Get the alignment size for quantization.""" + """Get the alignment used to pad grouped quantized operations. + + Built-in recipes use their format requirement. Custom recipes use their + declarative ``quantization_alignment`` contract rather than invoking + ``qfactory``, since a factory may be stateful or role-dependent. + """ + # TODO(#3158): Prefer module/role-specific alignment derived from canonical + # cached quantizers when that context is available. Keep the recipe-wide + # alignment as the conservative fallback for context-free callers. + if recipe.custom(): + return recipe.quantization_alignment if recipe.mxfp8(): return 32 if recipe.nvfp4(): From cf92c3d5f74826b6630b6d5243f0520e9c4ebe14 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 17 Jul 2026 17:57:02 +0000 Subject: [PATCH 48/97] Misc fixes: float8 fsdp, column-only float8, checks, etc. Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 394 +++++++++++++++--- tests/pytorch/test_identity_quantizer.py | 70 ++++ transformer_engine/pytorch/module/base.py | 8 + .../pytorch/module/grouped_linear.py | 20 +- .../pytorch/ops/basic/layer_norm.py | 11 +- .../pytorch/tensor/float8_tensor.py | 142 ++++++- .../pytorch/tensor/hybrid_tensor.py | 31 +- .../pytorch/tensor/identity_tensor.py | 28 +- .../tensor/storage/float8_tensor_storage.py | 79 +++- .../tensor/storage/grouped_tensor_storage.py | 12 + 10 files changed, 693 insertions(+), 102 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 4949708a1b2..c193012f2de 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -4,6 +4,7 @@ """Tests for hybrid quantization (mixed rowwise/columnwise formats).""" +import io import pytest import torch @@ -606,6 +607,15 @@ def test_rowwise_dequantized_identity_columnwise_matches_rowwise(self, input_ten columnwise_dq = out.columnwise_sub_storage.dequantize() torch.testing.assert_close(columnwise_dq, rowwise_dq, rtol=0.0, atol=0.0) + def test_internal_rowwise_storage_preserves_input_dtype(self, input_tensor): + hq = self._make_quantizer(columnwise_source="rowwise_dequantized") + hq.rowwise_quantizer.internal = True + + columnwise_src = hq._columnwise_src_from_rowwise(input_tensor, None) + + assert columnwise_src.dtype == input_tensor.dtype + + def test_columnwise_only_uses_transient_rowwise_source(self, input_tensor): hq = self._make_quantizer(columnwise_source="rowwise_dequantized") hq.set_usage(rowwise=False, columnwise=True) @@ -3964,13 +3974,46 @@ def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): assert len(calls) == 2 expected_columnwise_source = torch.cat( - [result.dequantize() for result in calls[0][1]], dim=0 + [result.dequantize(dtype=tensor.dtype) for result in calls[0][1]], dim=0 ) torch.testing.assert_close(calls[1][0], expected_columnwise_source, rtol=0.0, atol=0.0) assert calls[1][0] is not tensor assert all(storage.rowwise_sub_storage is None for storage in out) assert all(storage.columnwise_sub_storage is not None for storage in out) + @requires_nvfp4 + @pytest.mark.parametrize("m_splits", ([128, 128], [0, 128])) + def test_nvfp4_rowwise_dequantized_preserves_source_dtype(self, monkeypatch, m_splits): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + real_split_quantize = grouped_linear.tex.split_quantize + input_dtypes = [] + + def tracked_split_quantize(tensor, splits, quantizers, **kwargs): + input_dtypes.append(tensor.dtype) + return real_split_quantize(tensor, splits, quantizers, **kwargs) + + monkeypatch.setattr(grouped_linear.tex, "split_quantize", tracked_split_quantize) + tensor = torch.randn( + sum(m_splits), + 128, + dtype=torch.bfloat16, + device="cuda", + ) + quantizers = [ + HybridQuantizer( + rowwise_quantizer=_make_nvfp4_quantizer(), + columnwise_quantizer=_make_nvfp4_quantizer(), + columnwise_source="rowwise_dequantized", + ) + for _ in m_splits + ] + + out = grouped_linear._hybrid_split_quantize(tensor, m_splits, quantizers) + + assert input_dtypes == [tensor.dtype, tensor.dtype] + assert len(out) == len(m_splits) + def test_rowwise_only_skips_columnwise_quantization(self, monkeypatch): import transformer_engine.pytorch.module.grouped_linear as grouped_linear @@ -6358,6 +6401,113 @@ def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self) (16, 2), ] + @staticmethod + def _make_valid_transpose_only_float8_tensor(shape=(12, 16)): + """Build transpose-only storage with known, numerically valid FP8 bytes.""" + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + rowwise_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + rowwise = rowwise_quantizer(source) + columnwise_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + tensor = Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=rowwise._data.movedim(-1, 0).contiguous(), + fp8_scale_inv=rowwise._scale_inv.detach().clone(), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=columnwise_quantizer, + requires_grad=False, + device="cuda", + ) + return tensor, rowwise.dequantize() + + @pytest.mark.parametrize("weights_only", (False, True)) + def test_serialization_preserves_transpose_only_payload(self, weights_only): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + buffer = io.BytesIO() + + torch.save(tensor, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=weights_only) + + assert loaded._data is None + assert loaded._transpose_invalid is False + assert torch.equal(loaded._transpose, tensor._transpose) + torch.testing.assert_close(loaded.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_dequantize_from_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + torch.testing.assert_close(tensor.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_slice_and_select_preserve_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + sliced = tensor[2:8] + selected = torch.select(tensor, 1, 3) + + assert isinstance(sliced, Float8Tensor) + assert sliced._data is None + assert sliced._transpose.shape == torch.Size((16, 6)) + torch.testing.assert_close( + sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0 + ) + assert isinstance(selected, Float8Tensor) + assert selected._data is None + assert selected._transpose.shape == torch.Size((12,)) + torch.testing.assert_close( + selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0 + ) + + def test_as_strided_row_shard_preserves_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + shard = torch.as_strided(tensor, (5, 16), (16, 1), 16) + + assert isinstance(shard, Float8Tensor) + assert shard._data is None + assert shard._transpose.shape == torch.Size((16, 5)) + torch.testing.assert_close( + shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0 + ) + + def test_as_strided_falls_back_for_nonrepresentable_layout(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + output = torch.as_strided(tensor, (16, 12), (1, 16), 0) + + assert type(output) is torch.Tensor + reference = torch.as_strided(expected, (16, 12), (1, 16), 0) + torch.testing.assert_close(output, reference, rtol=0.0, atol=0.0) + + def test_new_zeros_preserves_transpose_only_layout(self): + tensor, _ = self._make_valid_transpose_only_float8_tensor() + + output = tensor.new_zeros((5, 16)) + + assert isinstance(output, Float8Tensor) + assert output.shape == torch.Size((5, 16)) + assert output._data is None + assert output._transpose.shape == torch.Size((16, 5)) + assert output._transpose.dtype == torch.uint8 + assert torch.count_nonzero(output._transpose).item() == 0 + torch.testing.assert_close( + output.dequantize(), + torch.zeros((5, 16), dtype=output.dtype, device=output.device), + rtol=0.0, + atol=0.0, + ) + class TestHybridNewZeros: """Public new_zeros semantics independent of FSDP-specific allocation.""" @@ -6625,6 +6775,18 @@ def test_split_sub_storage_types_preserved(self, hybrid_param): assert type(piece.rowwise_sub_storage) is orig_row_type assert type(piece.columnwise_sub_storage) is orig_col_type + @requires_mxfp8 + def test_split_rejects_mxfp8_high_precision_fallback(self): + """Fail before wrapping unquantizable MXFP8 shards for Hybrid FSDP2.""" + quantizer = HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + tensor = quantizer(torch.randn(64, 64, dtype=torch.bfloat16, device="cuda")) + + with pytest.raises(NotImplementedError, match="local shape.*divisible by 32"): + torch.split(tensor, 16, dim=0) + def test_view_preserves_hybrid_type(self, hybrid_param): """view must return a HybridQuantizedTensor (used by FSDP2 reset_sharded_param).""" shape_2d = hybrid_param.shape @@ -6677,6 +6839,30 @@ def test_copy_between_hybrid_tensors(self, hybrid_param): torch.testing.assert_close(src_deq, dst_deq, rtol=0.0, atol=0.0) _assert_hybrid_tensor_exact(dst, hybrid_param, context="hybrid copy_") + def test_copy_between_mismatched_usages_raises_atomically(self, hybrid_param): + quantizer = hybrid_param._quantizer.copy() + src = quantizer.quantize(torch.ones_like(hybrid_param.dequantize())) + dst = quantizer.quantize(torch.zeros_like(hybrid_param.dequantize())) + src.update_usage(columnwise_usage=False) + dst_before = tuple( + None if tensor is None else tensor.clone() + for tensor in _as_data_tensor_tuple(dst) + ) + + with pytest.raises( + NotImplementedError, + match="requires matching rowwise/columnwise usages", + ): + aten.copy_.default(dst, src) + + dst_after = _as_data_tensor_tuple(dst) + assert len(dst_after) == len(dst_before) + for before, after in zip(dst_before, dst_after): + if before is None: + assert after is None + else: + torch.testing.assert_close(after, before, rtol=0.0, atol=0.0) + def test_copy_from_bf16_to_hybrid(self, hybrid_param): """copy_ from BF16 into HybridQuantizedTensor triggers quantize_.""" param = hybrid_param.detach() @@ -7223,77 +7409,159 @@ def test_make_like_is_independent(self): @requires_hopper_fp8 class TestHybridFloat8ColumnwiseOnlyHopperPath: - """Float8TensorStorage columnwise-only sub-storage exercises the - ``_transpose`` field on Hopper. The FSDP2 buffer protocol must - recognize this layout. - """ + """Hopper transpose-only Float8 uses an M-major FSDP transport layout.""" - def _make_columnwise_only_float8_storage(self): - """Build a Float8TensorStorage in the layout a columnwise-only - hybrid sub-storage would have on Hopper: ``_data=None`` and - the actual quantized bytes in ``_transpose``. - """ - q = Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, + @staticmethod + def _make_columnwise_only_float8_storage(shape=(32, 64), value=1.0, quantizer=None): + source = torch.full(shape, value, device="cuda", dtype=torch.bfloat16) + byte_quantizer = Float8Quantizer( + scale=torch.ones(1, device="cuda", dtype=torch.float32), + amax=torch.zeros(1, device="cuda", dtype=torch.float32), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + rowwise = byte_quantizer(source) + if quantizer is None: + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + storage = Float8Tensor( + shape=shape, + dtype=source.dtype, + data=None, + data_transpose=rowwise._data.movedim(-1, 0).contiguous(), + fp8_scale_inv=rowwise._scale_inv.clone(), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + requires_grad=False, device="cuda", - rowwise=False, - columnwise=True, ) - src = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda") - out = q(src) - # Columnwise-only Float8 on Hopper: _data is None, _transpose holds data - assert ( - out._data is None - ), f"Test precondition failed: expected _data is None on Hopper, got {out._data}" - assert out._transpose is not None, "Test precondition failed: _transpose is None" - return out - - def test_fsdp_buffer_fields_returns_transpose(self): - """``fsdp_buffer_fields`` must return ``("_transpose",)`` when - ``_data`` is ``None`` and ``_transpose`` is populated. The - unconditional ``("_data",)`` would have FSDP2 all-gather a - ``None`` tensor on Hopper hybrid + FSDP2. - """ - storage = self._make_columnwise_only_float8_storage() - assert storage.fsdp_buffer_fields() == ("_transpose",) + return storage, rowwise.dequantize() + + @classmethod + def _make_hybrid_shard(cls, value): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + source = torch.full((32, 64), value, device="cuda", dtype=torch.bfloat16) + rowwise = quantizer.rowwise_quantizer(source) + columnwise, expected = cls._make_columnwise_only_float8_storage( + value=value, quantizer=quantizer.columnwise_quantizer + ) + hybrid = HybridQuantizedTensor( + shape=source.shape, + dtype=source.dtype, + rowwise_storage=rowwise, + columnwise_storage=columnwise, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + return hybrid, expected + + def test_columnwise_only_float8_fsdp_buffer_fields_returns_transpose(self): + out, _ = self._make_columnwise_only_float8_storage() + + assert out._data is None + assert out._transpose is not None + assert not out._transpose_invalid + assert out.fsdp_buffer_fields() == ("_transpose",) + + def test_columnwise_only_float8_fsdp_extract_uses_m_major_transport(self): + out, _ = self._make_columnwise_only_float8_storage() + + buffers, metadata = out.fsdp_extract_buffers() - def test_fsdp_extract_buffers_returns_transpose_data(self): - """``fsdp_extract_buffers`` (default impl, reads named fields) - must return the actual ``_transpose`` tensor, not ``None``. - """ - storage = self._make_columnwise_only_float8_storage() - buffers, meta = storage.fsdp_extract_buffers() assert len(buffers) == 1 - assert buffers[0] is not None - assert buffers[0] is storage._transpose - assert meta["field_names"] == ("_transpose",) - - def test_fsdp_assign_gathered_resets_transpose_invalid(self): - """After the gathered transpose buffer is written back via - ``fsdp_assign_gathered``, ``_transpose_invalid`` must be False - — otherwise ``update_usage`` / ``get_usages`` would treat the - freshly gathered transpose as stale on first use. - """ - storage = self._make_columnwise_only_float8_storage() - # Simulate stale state pre-gather - storage._transpose_invalid = True - new_transpose = torch.zeros_like(storage._transpose) - storage.fsdp_assign_gathered((new_transpose,), {"field_names": ("_transpose",)}) - assert storage._transpose is new_transpose - assert storage._transpose_invalid is False - # And ``get_usages`` correctly reports columnwise-available - assert storage.get_usages()["columnwise"] is True + assert buffers[0].shape == out.shape + assert buffers[0].is_contiguous() + torch.testing.assert_close( + buffers[0], out._transpose.movedim(0, -1).contiguous() + ) + assert metadata == { + "field_names": ("_transpose",), + "transport_layout": "columnwise_m_major", + } + + def test_columnwise_only_float8_fsdp_assign_restores_columnwise_layout(self): + out, expected = self._make_columnwise_only_float8_storage() + buffers, metadata = out.fsdp_extract_buffers() + gathered = torch.cat((buffers[0], buffers[0]), dim=0) + rebuilt = Float8Tensor.make_like(out, shape=gathered.shape) + + rebuilt.fsdp_assign_gathered((gathered,), metadata) + + assert rebuilt.shape == torch.Size((64, 64)) + assert rebuilt._data is None + assert rebuilt._transpose.shape == torch.Size((64, 64)) + assert not rebuilt._transpose_invalid + torch.testing.assert_close( + rebuilt._transpose, gathered.movedim(-1, 0).contiguous() + ) + torch.testing.assert_close( + rebuilt.dequantize(), torch.cat((expected, expected), dim=0) + ) + + def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): + shards = [self._make_hybrid_shard(value) for value in (1.0, 2.0)] + extracted = [ + shard.fsdp_pre_all_gather( + mesh=None, + orig_size=shard.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + for shard, _ in shards + ] + gathered = tuple( + torch.cat((extracted[0][0][i], extracted[1][0][i]), dim=0) + for i in range(len(extracted[0][0])) + ) + expected = torch.cat((shards[0][1], shards[1][1]), dim=0) + + rebuilt, _ = shards[0][0].fsdp_post_all_gather( + gathered, extracted[0][1], torch.bfloat16 + ) + + assert rebuilt.shape == torch.Size((64, 64)) + assert rebuilt._columnwise_storage.shape == torch.Size((64, 64)) + assert rebuilt._columnwise_storage._data is None + assert rebuilt._columnwise_storage._transpose.shape == torch.Size((64, 64)) + torch.testing.assert_close(rebuilt._columnwise_storage.dequantize(), expected) + + reused, _ = shards[0][0].fsdp_post_all_gather( + gathered, extracted[0][1], torch.bfloat16, out=rebuilt + ) + assert reused is rebuilt + assert reused._columnwise_storage._transpose.shape == torch.Size((64, 64)) + torch.testing.assert_close(reused._columnwise_storage.dequantize(), expected) def test_fsdp_buffer_fields_falls_back_to_data_when_both_present(self): - """A normally-constructed Float8TensorStorage has ``_data`` - populated; ``fsdp_buffer_fields`` should still prefer ``_data`` - — direction-aware logic must not regress the vanilla path. - """ - q = Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - src = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda") - out = q(src) + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ) + out = quantizer(torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)) + assert out._data is not None + assert out._transpose is not None assert out.fsdp_buffer_fields() == ("_data",) + buffers, metadata = out.fsdp_extract_buffers() + assert len(buffers) == 1 + assert buffers[0] is out._data + assert metadata == { + "field_names": ("_data",), + "transport_layout": "native", + } @requires_hopper_fp8 diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 758cbcfc6b0..ad4bd66b54b 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -591,6 +591,76 @@ def test_dtype_cast(self): assert dequantized.dtype == torch.bfloat16 torch.testing.assert_close(dequantized, x.to(torch.bfloat16), rtol=0.0, atol=0.0) + def test_make_empty_honors_configured_dtype(self): + tensor = IdentityQuantizer(dtype=torch.float32).make_empty( + (4, 8), + dtype=torch.bfloat16, + device="cuda", + ) + + assert tensor.dtype == torch.float32 + assert tensor._hp_data.dtype == torch.float32 + assert tensor.dequantize().dtype == torch.float32 + + def test_update_quantized_synchronizes_dtype(self): + dst = IdentityQuantizer().make_empty( + (4, 8), + dtype=torch.bfloat16, + device="cuda", + ) + quantizer = IdentityQuantizer(dtype=torch.float32) + + quantizer.update_quantized(torch.ones_like(dst._hp_data), dst) + + assert dst.dtype == torch.float32 + assert dst._hp_data.dtype == torch.float32 + assert dst.dequantize().dtype == torch.float32 + + @pytest.mark.parametrize("noop", [0, 1]) + def test_update_quantized_honors_noop(self, noop): + quantizer = IdentityQuantizer() + dst = quantizer(torch.full((4, 8), 3.0, device="cuda")) + src = torch.full((4, 8), 7.0, device="cuda") + + quantizer.update_quantized( + src, dst, noop_flag=torch.tensor(noop, dtype=torch.float32, device="cuda") + ) + + expected = 3.0 if noop else 7.0 + torch.testing.assert_close(dst.dequantize(), torch.full_like(src, expected)) + + @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") + def test_hybrid_update_honors_shared_noop(self): + quantizer = HybridQuantizer( + rowwise_quantizer=_fp8_cs(), + columnwise_quantizer=IdentityQuantizer(), + ) + old = torch.full((32, 32), 3.0, dtype=torch.bfloat16, device="cuda") + dst = quantizer(old) + expected_row = dst._rowwise_storage.dequantize().clone() + + quantizer.update_quantized( + torch.full_like(old, 7.0), + dst, + noop_flag=torch.ones(1, dtype=torch.float32, device="cuda"), + ) + + torch.testing.assert_close(dst._rowwise_storage.dequantize(), expected_row) + torch.testing.assert_close(dst._columnwise_storage.dequantize(), old) + + def test_hybrid_new_zeros_preserves_identity_dtypes(self): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(dtype=torch.bfloat16), + columnwise_quantizer=IdentityQuantizer(dtype=torch.float32), + ) + tensor = quantizer(torch.ones(4, 8, dtype=torch.bfloat16, device="cuda")) + + zeros = tensor.new_zeros(tensor.shape) + + assert zeros.rowwise_sub_storage._hp_data.dtype == torch.bfloat16 + assert zeros.columnwise_sub_storage._hp_data.dtype == torch.float32 + + def test_storage_dequantize_defaults_to_nominal_dtype(self): payload = torch.randn(4, 8, dtype=torch.float32) storage = IdentityTensorStorage( diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index ebb802b2ae7..f2c2d77e68b 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -610,6 +610,14 @@ def fill_userbuffers_buffer_for_all_gather( local_tensor = quantizer(local_tensor) return local_tensor, local_tensor + # TODO(#3158): Support HybridQuantizer with standard Userbuffers. + if isinstance(quantizer, HybridQuantizer): + raise NotImplementedError( + "Standard Userbuffers does not support HybridQuantizer yet. " + "Disable Userbuffers all-gather overlap or use a supported " + "single-format quantizer. See #3158." + ) + # Tensor dimensions local_shape = local_tensor.size() if not local_shape: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 890606e64f0..e6bb5852d28 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -288,7 +288,10 @@ def _hybrid_split_quantize( if columnwise_enabled and columnwise_source == "rowwise_dequantized": # Assemble the exact grouped row results in split order. NVFP4 padding # and scale layout can differ from independently quantizing each split. - columnwise_src = torch.cat([result.dequantize() for result in row_results], dim=0) + columnwise_src = torch.cat( + [result.dequantize(dtype=tensor.dtype) for result in row_results], + dim=0, + ) col_results = ( tex.split_quantize( columnwise_src, @@ -1874,6 +1877,21 @@ def make_grouped_weights(self, defer_init=False) -> None: return weight_quantizers = self._get_weight_quantizers() + # TODO(#3158): Support Identity/Hybrid single grouped weights. + unsupported_quantizers = tuple( + type(quantizer).__name__ + for quantizer in weight_quantizers + if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer)) + ) + if unsupported_quantizers: + quantizer_names = ", ".join(dict.fromkeys(unsupported_quantizers)) + raise NotImplementedError( + "GroupedLinear(single_grouped_weight=True) does not support " + f"{quantizer_names} weight quantizers yet. Set " + "single_grouped_weight=False or unset " + "NVTE_GROUPED_LINEAR_SINGLE_PARAM. See #3158." + ) + recipe = ( weight_quantizers[0]._get_compatible_recipe() if weight_quantizers and weight_quantizers[0] is not None diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index 3fda5145c62..a85ad483ad6 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -16,7 +16,7 @@ from ...constants import TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...export import is_in_onnx_export_mode -from ...tensor import Quantizer +from ...tensor import HybridQuantizer, Quantizer from ...utils import ( canonicalize_device, canonicalize_dtype, @@ -183,6 +183,15 @@ def op_forward( if is_in_onnx_export_mode(): return self.op_onnx_forward(input_) + # TODO(#3158): Support producing both directional representations at + # the fused normalization boundary. + if isinstance(next_op_input_quantizer, HybridQuantizer): + raise NotImplementedError( + "te.ops.LayerNorm does not support HybridQuantizer output yet. " + "Use transformer_engine.pytorch.LayerNormMLP or a single-format " + "quantizer instead. See #3158." + ) + # Check tensor dims weight = self.weight weight_dims = tuple(weight.size()) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 2c962fd5986..f12de50aa2d 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -646,17 +646,22 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func in (aten.slice.Tensor, aten.select.int): tensor = args[0] data = tensor._data - data_slice = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) + data_slice = None + if data is not None: + data_slice = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) transpose_slice = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose - ndim = data.dim() + ndim = tensor.dim() + if ndim == 0: + return super().__torch_dispatch__(func, types, args, kwargs) dim = args[1] if len(args) > 1 else 0 + dim %= ndim t_dim = 0 if dim == ndim - 1 else dim + 1 transpose_slice = transpose.__torch_dispatch__( func, @@ -664,11 +669,34 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [transpose, t_dim] + list(args[2:]), kwargs, ) + if ( + func == aten.select.int + and dim == ndim - 1 + and transpose_slice.dim() > 1 + ): + transpose_slice = transpose_slice.movedim(-1, 0).contiguous() + + if data_slice is not None: + out_shape = data_slice.shape + else: + logical = torch.empty(tensor.shape, device="meta") + out_shape = func(logical, *args[1:], **(kwargs or {})).shape + if transpose_slice is None: + raise RuntimeError( + "Float8Tensor slice/select requires rowwise or columnwise data" + ) + expected_transpose_shape = _columnwise_shape_for(out_shape) + if transpose_slice.shape != expected_transpose_shape: + raise RuntimeError( + "Float8Tensor slice/select produced incompatible columnwise storage: " + f"expected {tuple(expected_transpose_shape)}, got " + f"{tuple(transpose_slice.shape)}" + ) return Float8Tensor.make_like( tensor, data=data_slice, data_transpose=transpose_slice, - shape=data_slice.shape, + shape=out_shape, ) # Related to FSDP2 @@ -734,12 +762,18 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): # create fresh new tensor with zeros. tensor = args[0] data = tensor._data - func_out = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) + storage_kwargs = dict(kwargs or {}) + output_dtype = storage_kwargs.pop("dtype", None) or tensor.dtype + storage_kwargs.pop("layout", None) + storage_kwargs.pop("requires_grad", None) + func_out = None + if data is not None: + func_out = data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + storage_kwargs, + ) func_transposed_out = None if tensor._transpose is not None and not tensor._transpose_invalid: transpose = tensor._transpose @@ -749,25 +783,71 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): func, types, [transpose, t_shape] + list(args[2:]), - kwargs, + storage_kwargs, + ) + if func_out is None and func_transposed_out is None: + raise RuntimeError( + "Float8Tensor.new_zeros requires rowwise or columnwise data" ) scale_inv = tensor._scale_inv.detach().clone() + reference = func_out if func_out is not None else func_transposed_out + if scale_inv.device != reference.device: + scale_inv = scale_inv.to(reference.device) quantizer = tensor._quantizer # Deep-copied in constructor out_tensor = Float8Tensor( data=func_out, - shape=func_out.shape, - dtype=tensor.dtype, + shape=torch.Size(args[1]), + dtype=output_dtype, fp8_dtype=tensor._fp8_dtype, fp8_scale_inv=scale_inv, data_transpose=func_transposed_out, quantizer=quantizer, - device=tensor.device, + device=reference.device, ) return out_tensor if func == torch.ops.aten.as_strided.default: tensor = args[0] data = tensor._data + if data is None: + size = torch.Size(args[1]) + stride = tuple(args[2]) + storage_offset = (kwargs or {}).get( + "storage_offset", + args[3] if len(args) > 3 else tensor.storage_offset(), + ) + # A contiguous 2D row shard maps to a column slice in the + # persistent transpose and can remain a Float8Tensor. + has_valid_transpose = ( + tensor._transpose is not None and not tensor._transpose_invalid + ) + is_contiguous_row_shard = ( + tensor.dim() == 2 + and tensor.shape[1] > 0 + and len(size) == 2 + and size[1] == tensor.shape[1] + and stride == tuple(tensor.stride()) + ) + if ( + has_valid_transpose + and is_contiguous_row_shard + and storage_offset % tensor.shape[1] == 0 + ): + row_start = storage_offset // tensor.shape[1] + row_end = row_start + size[0] + if 0 <= row_start and row_end <= tensor.shape[0]: + transpose_out = tensor._transpose[:, row_start:row_end] + return Float8Tensor.make_like( + tensor, + data=None, + data_transpose=transpose_out, + shape=size, + ) + + # Arbitrary logical strides are not generally affine views of + # transposed storage. Fall back explicitly to high precision. + return func(tensor.dequantize(), *args[1:], **(kwargs or {})) + # Apply as_strided to the primary uint8 data func_out = data.__torch_dispatch__( func, @@ -979,9 +1059,23 @@ def __reduce_ex__(self, protocol: int) -> tuple: with ``torch.serialization.add_safe_globals`` to keep ``torch.load(weights_only=True)`` compatibility. """ + data_transpose = None + if ( + self._data is None + and self._transpose is not None + and not self._transpose_invalid + ): + data_transpose = self._transpose return ( _make_float8_tensor_in_reduce_ex, - (self._data, self._fp8_dtype, self._scale_inv, self.dtype, self.shape), + ( + self._data, + self._fp8_dtype, + self._scale_inv, + self.dtype, + self.shape, + data_transpose, + ), ) @classmethod @@ -1077,11 +1171,12 @@ def _set_data(self, tensor: torch.Tensor) -> None: def _make_float8_tensor_in_reduce_ex( - data: torch.Tensor, + data: Optional[torch.Tensor], fp8_dtype: DType, fp8_scale_inv: torch.Tensor, dtype: torch.dtype, shape: torch.Size, + data_transpose: Optional[torch.Tensor] = None, ) -> Float8Tensor: """Reconstruct a ``Float8Tensor`` from its ``__reduce_ex__`` payload.""" return Float8Tensor( @@ -1090,7 +1185,12 @@ def _make_float8_tensor_in_reduce_ex( fp8_scale_inv=fp8_scale_inv, dtype=dtype, shape=shape, - device=data.device if data is not None else None, + data_transpose=data_transpose, + device=( + data.device + if data is not None + else data_transpose.device if data_transpose is not None else None + ), ) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index ce3892ef10f..7ab84f32d44 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -193,7 +193,7 @@ def _columnwise_src_from_rowwise( ) -> torch.Tensor: if rowwise_result is None: rowwise_result = self.rowwise_quantizer.quantize(tensor) - return rowwise_result.dequantize() + return rowwise_result.dequantize(dtype=tensor.dtype) def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: # Gate each sub-quantizer call on the parent usage flag. Sub-quantizers @@ -984,6 +984,26 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if row_pieces is None and col_pieces is None: return super().__torch_dispatch__(func, types, args, kwargs) + # TODO(#3158): Support Hybrid sub-storages that fall back to a + # high-precision tensor for an unquantizable local shard. + for direction, pieces in ( + ("rowwise", row_pieces), + ("columnwise", col_pieces), + ): + if pieces is None: + continue + for piece in pieces: + if not isinstance(piece, QuantizedTensor): + raise NotImplementedError( + "HybridQuantizedTensor split produced a high-precision " + f"{type(piece).__name__} for the {direction} sub-storage " + f"with local shape {tuple(piece.shape)}. Hybrid FSDP2 does " + "not support high-precision fallback children. MXFP8 FSDP " + "shards must have a first dimension divisible by 32; adjust " + "the sharding topology or disable Hybrid MXFP8 for this " + "parameter. See #3158." + ) + num_pieces = len(row_pieces) if row_pieces is not None else len(col_pieces) return [ HybridQuantizedTensor( @@ -1044,6 +1064,15 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.copy_.default: dst, src = args[0], args[1] if isinstance(dst, HybridQuantizedTensor) and isinstance(src, HybridQuantizedTensor): + dst_usages = dst.get_usages() + src_usages = src.get_usages() + if dst_usages != src_usages: + raise NotImplementedError( + "HybridQuantizedTensor.copy_ requires matching rowwise/columnwise " + f"usages, but source has {src_usages} and destination has " + f"{dst_usages}. Copy from a high-precision tensor instead." + ) + if dst._rowwise_storage is not None and src._rowwise_storage is not None: aten.copy_.default(dst._rowwise_storage, src._rowwise_storage) if dst._columnwise_storage is not None and src._columnwise_storage is not None: diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index 5d80fcbd0cd..cae581fe98e 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -106,10 +106,11 @@ def make_empty( if device is None: device = torch.device("cuda") device = torch.device(device) - data = torch.empty(tuple(shape), dtype=dtype, device=device, pin_memory=pin_memory) + data_dtype = self.dtype if self.dtype is not None else dtype + data = torch.empty(tuple(shape), dtype=data_dtype, device=device, pin_memory=pin_memory) return IdentityTensor( data.shape, - dtype, + data_dtype, hp_data=data, quantizer=self, requires_grad=requires_grad, @@ -121,7 +122,7 @@ def update_quantized( src: torch.Tensor, dst: QuantizedTensorStorage, *, - noop_flag: Optional[torch.Tensor] = None, # pylint: disable=unused-argument + noop_flag: Optional[torch.Tensor] = None, ) -> QuantizedTensorStorage: if not isinstance(dst, IdentityTensorStorage): raise ValueError( @@ -134,9 +135,15 @@ def update_quantized( and dst._hp_data.dtype == data.dtype and dst._hp_data.device == data.device ): - dst._hp_data.copy_(data) + if noop_flag is None: + dst._hp_data.copy_(data) + else: + torch.where(noop_flag == 0, data, dst._hp_data, out=dst._hp_data) else: + if noop_flag is not None and noop_flag.item() != 0: + return dst dst._hp_data = data.detach() + dst._dtype = data.dtype return dst def calibrate(self, tensor: torch.Tensor) -> None: @@ -219,12 +226,23 @@ def fsdp_post_all_gather( """Rebuild IdentityTensor from the gathered high-precision buffer.""" (data,) = all_gather_outputs (quantizer,) = metadata + logical_dtype = ( + quantizer.dtype + if quantizer is not None and quantizer.dtype is not None + else param_dtype + ) + if data.dtype != logical_dtype: + raise RuntimeError( + "IdentityTensor FSDP payload dtype does not match its logical dtype: " + f"payload={data.dtype}, logical={logical_dtype}." + ) if out is not None: out._hp_data = data + out._dtype = logical_dtype else: out = IdentityTensor( shape=data.shape, - dtype=param_dtype, + dtype=logical_dtype, hp_data=data, quantizer=quantizer, requires_grad=False, diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index e0bd5613dcd..874b60ca96c 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -45,7 +45,25 @@ def forward( # Cast from FP8 return tex.dequantize(tensor, te_dtype) - raise NotImplementedError("Casting back from the transpose not implemented yet!") + if tensor._transpose is not None and not tensor._transpose_invalid: + # A columnwise-only tensor stores the logical last dimension first + # in its physical FP8 buffer. Dequantize that buffer as ordinary + # FP8 data, then restore logical row-major dimension order. + transpose_tensor = Float8TensorStorage( + data=tensor._transpose, + data_transpose=None, + fp8_scale_inv=tensor._scale_inv, + fp8_dtype=tensor._fp8_dtype, + fake_dtype=tensor._dtype, + ) + output = _FromFloat8Func.forward(None, transpose_tensor, dtype) + if output.dim() > 0: + output = output.movedim(0, -1).contiguous() + return output + + raise RuntimeError( + "Float8TensorStorage has neither rowwise nor columnwise data" + ) @staticmethod def backward( @@ -332,20 +350,61 @@ def fsdp_buffer_fields(self) -> Tuple[str, ...]: # the caller would have seen pre-direction-aware logic. return ("_data",) + def fsdp_extract_buffers( + self, + ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]: + """Extract dim-0 gather-safe buffers for per-tensor Float8. + + Rowwise ``_data`` is already M-major. A transpose-only Hopper/L40 + payload is physically ``[N, *M]`` and cannot be gathered directly + because FSDP2 grows physical dimension 0. Move N back to the logical + last dimension and make the transport buffer contiguous so FSDP2 sees + ``[*M, N]``. :meth:`fsdp_assign_gathered` restores the persistent + columnwise layout after the collective. + """ + names = self.fsdp_buffer_fields() + if names == ("_transpose",): + if self._transpose is None or self._transpose_invalid: + raise RuntimeError( + "Float8TensorStorage cannot extract an invalid columnwise transpose" + ) + transport = self._transpose + if transport.dim() > 0: + transport = transport.movedim(0, -1).contiguous() + return (transport,), { + "field_names": names, + "transport_layout": "columnwise_m_major", + } + buffers = tuple(getattr(self, name) for name in names) + return buffers, {"field_names": names, "transport_layout": "native"} + def fsdp_assign_gathered( self, gathered: Tuple[Optional[torch.Tensor], ...], meta: Dict[str, Any], ) -> None: - """Write gathered Float8 buffers back, refreshing ``_transpose_invalid``. - - The base implementation just ``setattr``s the gathered tensors into - the named fields. For Float8 we additionally need to clear - ``_transpose_invalid`` when the gathered field is ``_transpose`` — - otherwise a freshly-gathered transpose buffer is treated as stale - on first use (see :attr:`_transpose_invalid` semantics in - ``update_usage`` / ``get_usages``). - """ + """Restore gathered Float8 buffers to their persistent layout.""" + transport_layout = meta.get("transport_layout", "native") + if transport_layout == "columnwise_m_major": + if meta.get("field_names") != ("_transpose",) or len(gathered) != 1: + raise RuntimeError( + "Float8TensorStorage got invalid metadata for columnwise FSDP transport" + ) + (transport,) = gathered + if transport is None: + raise RuntimeError( + "Float8TensorStorage got a None columnwise FSDP transport buffer" + ) + self._data = None + if transport.dim() > 0: + transport = transport.movedim(-1, 0).contiguous() + self._transpose = transport + self._transpose_invalid = False + return + if transport_layout != "native": + raise RuntimeError( + f"Float8TensorStorage got unknown FSDP transport layout {transport_layout!r}" + ) super().fsdp_assign_gathered(gathered, meta) if "_transpose" in meta["field_names"]: self._transpose_invalid = False diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 7ede75943ae..3473024c03f 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -654,6 +654,18 @@ def make_grouped_tensor( A GroupedTensor. """ + if quantizer is not None: + # TODO(#3158): Support Identity/Hybrid packed grouped storage. + from ..hybrid_tensor import HybridQuantizer + from ..identity_tensor import IdentityQuantizer + + if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer)): + raise NotImplementedError( + "GroupedTensorStorage does not support IdentityQuantizer or " + "HybridQuantizer yet. Use separate tensors, or set " + "GroupedLinear(single_grouped_weight=False). See #3158." + ) + # Set device if device is None: device = torch.cuda.current_device() From 6fa849a4de065cfcfd65cb8fa89cca9a484f8a80 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:02:41 +0000 Subject: [PATCH 49/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 32 +++++-------------- tests/pytorch/test_identity_quantizer.py | 1 - .../pytorch/tensor/float8_tensor.py | 16 ++-------- .../tensor/storage/float8_tensor_storage.py | 4 +-- 4 files changed, 12 insertions(+), 41 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index c193012f2de..bc956c7c910 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -615,7 +615,6 @@ def test_internal_rowwise_storage_preserves_input_dtype(self, input_tensor): assert columnwise_src.dtype == input_tensor.dtype - def test_columnwise_only_uses_transient_rowwise_source(self, input_tensor): hq = self._make_quantizer(columnwise_source="rowwise_dequantized") hq.set_usage(rowwise=False, columnwise=True) @@ -6459,15 +6458,11 @@ def test_slice_and_select_preserve_transpose_only_payload(self): assert isinstance(sliced, Float8Tensor) assert sliced._data is None assert sliced._transpose.shape == torch.Size((16, 6)) - torch.testing.assert_close( - sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0) assert isinstance(selected, Float8Tensor) assert selected._data is None assert selected._transpose.shape == torch.Size((12,)) - torch.testing.assert_close( - selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0) def test_as_strided_row_shard_preserves_transpose_only_payload(self): tensor, expected = self._make_valid_transpose_only_float8_tensor() @@ -6477,9 +6472,7 @@ def test_as_strided_row_shard_preserves_transpose_only_payload(self): assert isinstance(shard, Float8Tensor) assert shard._data is None assert shard._transpose.shape == torch.Size((16, 5)) - torch.testing.assert_close( - shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0) def test_as_strided_falls_back_for_nonrepresentable_layout(self): tensor, expected = self._make_valid_transpose_only_float8_tensor() @@ -6845,8 +6838,7 @@ def test_copy_between_mismatched_usages_raises_atomically(self, hybrid_param): dst = quantizer.quantize(torch.zeros_like(hybrid_param.dequantize())) src.update_usage(columnwise_usage=False) dst_before = tuple( - None if tensor is None else tensor.clone() - for tensor in _as_data_tensor_tuple(dst) + None if tensor is None else tensor.clone() for tensor in _as_data_tensor_tuple(dst) ) with pytest.raises( @@ -7483,9 +7475,7 @@ def test_columnwise_only_float8_fsdp_extract_uses_m_major_transport(self): assert len(buffers) == 1 assert buffers[0].shape == out.shape assert buffers[0].is_contiguous() - torch.testing.assert_close( - buffers[0], out._transpose.movedim(0, -1).contiguous() - ) + torch.testing.assert_close(buffers[0], out._transpose.movedim(0, -1).contiguous()) assert metadata == { "field_names": ("_transpose",), "transport_layout": "columnwise_m_major", @@ -7503,12 +7493,8 @@ def test_columnwise_only_float8_fsdp_assign_restores_columnwise_layout(self): assert rebuilt._data is None assert rebuilt._transpose.shape == torch.Size((64, 64)) assert not rebuilt._transpose_invalid - torch.testing.assert_close( - rebuilt._transpose, gathered.movedim(-1, 0).contiguous() - ) - torch.testing.assert_close( - rebuilt.dequantize(), torch.cat((expected, expected), dim=0) - ) + torch.testing.assert_close(rebuilt._transpose, gathered.movedim(-1, 0).contiguous()) + torch.testing.assert_close(rebuilt.dequantize(), torch.cat((expected, expected), dim=0)) def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): shards = [self._make_hybrid_shard(value) for value in (1.0, 2.0)] @@ -7528,9 +7514,7 @@ def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): ) expected = torch.cat((shards[0][1], shards[1][1]), dim=0) - rebuilt, _ = shards[0][0].fsdp_post_all_gather( - gathered, extracted[0][1], torch.bfloat16 - ) + rebuilt, _ = shards[0][0].fsdp_post_all_gather(gathered, extracted[0][1], torch.bfloat16) assert rebuilt.shape == torch.Size((64, 64)) assert rebuilt._columnwise_storage.shape == torch.Size((64, 64)) diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index ad4bd66b54b..92484b172ce 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -660,7 +660,6 @@ def test_hybrid_new_zeros_preserves_identity_dtypes(self): assert zeros.rowwise_sub_storage._hp_data.dtype == torch.bfloat16 assert zeros.columnwise_sub_storage._hp_data.dtype == torch.float32 - def test_storage_dequantize_defaults_to_nominal_dtype(self): payload = torch.randn(4, 8, dtype=torch.float32) storage = IdentityTensorStorage( diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index f12de50aa2d..e21d283eb48 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -669,11 +669,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [transpose, t_dim] + list(args[2:]), kwargs, ) - if ( - func == aten.select.int - and dim == ndim - 1 - and transpose_slice.dim() > 1 - ): + if func == aten.select.int and dim == ndim - 1 and transpose_slice.dim() > 1: transpose_slice = transpose_slice.movedim(-1, 0).contiguous() if data_slice is not None: @@ -786,9 +782,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): storage_kwargs, ) if func_out is None and func_transposed_out is None: - raise RuntimeError( - "Float8Tensor.new_zeros requires rowwise or columnwise data" - ) + raise RuntimeError("Float8Tensor.new_zeros requires rowwise or columnwise data") scale_inv = tensor._scale_inv.detach().clone() reference = func_out if func_out is not None else func_transposed_out if scale_inv.device != reference.device: @@ -1060,11 +1054,7 @@ def __reduce_ex__(self, protocol: int) -> tuple: ``torch.load(weights_only=True)`` compatibility. """ data_transpose = None - if ( - self._data is None - and self._transpose is not None - and not self._transpose_invalid - ): + if self._data is None and self._transpose is not None and not self._transpose_invalid: data_transpose = self._transpose return ( _make_float8_tensor_in_reduce_ex, diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 874b60ca96c..37b669f707c 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -61,9 +61,7 @@ def forward( output = output.movedim(0, -1).contiguous() return output - raise RuntimeError( - "Float8TensorStorage has neither rowwise nor columnwise data" - ) + raise RuntimeError("Float8TensorStorage has neither rowwise nor columnwise data") @staticmethod def backward( From 495abd5bbbda853afbeb4fa7747d93a0bc95b6dd Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 20 Jul 2026 18:30:02 +0000 Subject: [PATCH 50/97] Hopper fixes Signed-off-by: Evgeny --- qa/L0_pytorch_unittest/test.sh | 1 - .../distributed/fsdp2_tests/conftest.py | 15 ++ .../fsdp2_tests/run_fsdp2_fused_adam.py | 62 +---- .../pytorch/distributed/test_hybrid_tp_sp.py | 17 ++ tests/pytorch/test_float8blockwisetensor.py | 44 ++++ tests/pytorch/test_hybrid_quantization.py | 238 +++++++++++++++--- tests/pytorch/test_identity_quantizer.py | 20 ++ transformer_engine/pytorch/csrc/quantizer.cpp | 22 ++ .../float8_blockwise_tensor_storage.py | 28 +++ 9 files changed, 364 insertions(+), 83 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index cc4b0b0ec18..d31a0ce494b 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -40,7 +40,6 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_torch_compile.xm python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwisetensor.xml $TE_PATH/tests/pytorch/test_float8blockwisetensor.py || test_fail "test_float8blockwisetensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_view_shape.xml $TE_PATH/tests/pytorch/test_float8_view_shape.py || test_fail "test_float8_view_shape.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_PATH/tests/pytorch/test_grouped_tensor.py || test_fail "test_grouped_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH/tests/pytorch/test_gqa.py || test_fail "test_gqa.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py" diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index 91627ab7397..32ea3e1d65d 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -14,6 +14,7 @@ import torch import torch.distributed as dist from transformer_engine.pytorch import fp8 +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported # Ensure the correct CUDA device is active before _parametrize_recipes() # runs at collection time, since the session-scoped dist_init fixture @@ -68,6 +69,20 @@ def _parametrize_hybrid_recipes(): for name, check_fn in _HYBRID_RECIPE_CONFIGS: supported, reason = check_fn() marks = [pytest.mark.skipif(not supported, reason=reason)] + if name == "HybridFP8CurrentScaling" and not is_non_tn_fp8_gemm_supported(): + marks.append( + pytest.mark.xfail( + raises=pytest.RaisesExc( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ), + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 " + "quantization; tracked by NVIDIA/TransformerEngine#3158" + ), + ) + ) params.append(pytest.param(name, id=name, marks=marks)) return params diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 1169215e2f0..08e762045ae 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -1935,6 +1935,10 @@ def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): """ from fsdp2_utils import get_hybrid_recipe_from_string + supported, reason = te.is_mxfp8_available(return_reason=True) + if not supported: + pytest.skip(f"MXFP8: {reason}") + world_size, device = _get_dist_info() # FSDP2 shards a Linear weight of shape (out_features, in_features) along @@ -2031,16 +2035,8 @@ def test_fused_adam_hybrid_mxfp8_awkward_shard_shape(): _raise_collective_errors(errors, f"{recipe_name} awkward raw Hybrid all-gather") -@pytest.mark.xfail( - strict=True, - raises=RuntimeError, - reason=( - "Tracked by #3158: gathering dim-0 shards that split a 128-row " - "Float8Block scale tile produces too many scale rows." - ), -) def test_fused_adam_hybrid_float8_block_unaligned_shard_shape(): - """Positive FSDP2 roundtrip for a shard that splits a 128-row scale tile.""" + """Unaligned local shards are rejected before gathering incompatible scale tiles.""" from transformer_engine.pytorch import fp8 from fsdp2_utils import get_hybrid_recipe_from_string @@ -2071,48 +2067,12 @@ def test_fused_adam_hybrid_float8_block_unaligned_shard_shape(): model = _shard_model(model, world_size) x = torch.randn(128, in_features, dtype=torch.bfloat16, device=device) - with te.autocast(enabled=True, recipe=hybrid_recipe): - output = model(x) - output.sum().backward() - - params = [ - (name, param) - for name, param in model.named_parameters() - if isinstance(param, DTensor) and isinstance(param._local_tensor, HybridQuantizedTensor) - ] - local_count = torch.tensor([len(params)], dtype=torch.int64, device=device) - min_count = local_count.clone() - max_count = local_count.clone() - dist.all_reduce(min_count, dist.ReduceOp.MIN) - dist.all_reduce(max_count, dist.ReduceOp.MAX) - _collective_assert( - min_count.item() == max_count.item() and min_count.item() > 0, - "Float8Block unaligned Hybrid parameter count range is " - f"({min_count.item()}, {max_count.item()})", - ) - - errors = [] - for name, param in params: - local = param._local_tensor - reconstructed = _manual_reconstruct_hybrid(local, param_name=name, world_size=world_size) - public_full = param.full_tensor() - _record_exact( - errors, - public_full, - reconstructed.dequantize(), - f"{name}: public full_tensor vs Hybrid reconstruction", - ) - for direction in ("rowwise", "columnwise"): - _check_hybrid_direction_buffers( - getattr(local, f"_{direction}_storage"), - getattr(reconstructed, f"_{direction}_storage"), - Float8BlockwiseQTensorStorage, - direction=direction, - param_name=name, - world_size=world_size, - errors=errors, - ) - _raise_collective_errors(errors, "Float8Block unaligned raw Hybrid all-gather") + with pytest.raises( + RuntimeError, + match="local flattened M dimension.*not a multiple of 128", + ): + with te.autocast(enabled=True, recipe=hybrid_recipe): + model(x) def test_hybrid_dcp_output_parity(hybrid_recipe_name): diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py index 27305aee7aa..320acd0b7a3 100644 --- a/tests/pytorch/distributed/test_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -11,6 +11,7 @@ import pytest import torch import transformer_engine.pytorch as te +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported if torch.cuda.device_count() < 2: pytest.skip( @@ -26,6 +27,15 @@ NUM_PROCS = min(2, torch.cuda.device_count()) LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] +xfail_hopper_columnwise_per_tensor_fp8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + def _run_test(quantization: str, test: str = "all"): script = TEST_ROOT / "run_hybrid_tp_sp.py" @@ -44,42 +54,49 @@ def _run_test(quantization: str, test: str = "all"): @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 def test_hybrid_fp8_linear(): """Linear TP/SP coverage for hybrid FP8 current scaling.""" _run_test("hybrid_fp8", "linear") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 def test_hybrid_fp8_linear_vs_vanilla(): """Same-topology Linear parity against Float8CurrentScaling.""" _run_test("hybrid_fp8", "linear_vs_vanilla") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 def test_hybrid_fp8_layernorm_linear_vs_vanilla(): """Same-topology LayerNormLinear parity against vanilla FP8.""" _run_test("hybrid_fp8", "layernorm_linear_vs_vanilla") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 def test_hybrid_fp8_layernorm_mlp_vs_vanilla(): """Same-topology LayerNormMLP parity against vanilla FP8.""" _run_test("hybrid_fp8", "layernorm_mlp_vs_vanilla") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 def test_hybrid_fp8_layernorm_linear(): """LayerNormLinear TP/SP coverage for hybrid FP8.""" _run_test("hybrid_fp8", "layernorm_linear") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 def test_hybrid_fp8_layernorm_mlp(): """LayerNormMLP TP/SP coverage for hybrid FP8.""" _run_test("hybrid_fp8", "layernorm_mlp") @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") +@xfail_hopper_columnwise_per_tensor_fp8 def test_hybrid_fp8_transformer_layer(): """TransformerLayer TP/SP coverage for hybrid FP8.""" _run_test("hybrid_fp8", "transformer_layer") diff --git a/tests/pytorch/test_float8blockwisetensor.py b/tests/pytorch/test_float8blockwisetensor.py index 1b5e17f4635..ab00786e4fe 100644 --- a/tests/pytorch/test_float8blockwisetensor.py +++ b/tests/pytorch/test_float8blockwisetensor.py @@ -96,6 +96,50 @@ def test_constructor( assert tensor.dtype == dtype, "Incorrect nominal dtype" assert tensor.is_cuda, "Incorrect device" + def test_fsdp_assign_gathered_rejects_invalid_scale_geometry(self): + """Reject gathered scales that cannot describe the gathered data tiles.""" + shape = (128, 256) + quantizer = Float8BlockQuantizer( + fp8_dtype=DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ) + tensor = Float8BlockwiseQTensor( + shape=shape, + dtype=torch.bfloat16, + rowwise_data=torch.zeros(shape, device="cuda", dtype=torch.uint8), + rowwise_scale_inv=torch.ones( + quantizer.get_scale_shape(shape, columnwise=False), + device="cuda", + dtype=torch.float32, + ), + columnwise_data=torch.zeros( + quantizer.get_columnwise_shape(shape), + device="cuda", + dtype=torch.uint8, + ), + columnwise_scale_inv=torch.ones( + quantizer.get_scale_shape(shape, columnwise=True), + device="cuda", + dtype=torch.float32, + ), + fp8_dtype=DType.kFloat8E4M3, + is_2D_scaled=True, + quantizer=quantizer, + ) + original_data = tensor._rowwise_data + gathered_data = torch.zeros((256, 256), device="cuda", dtype=torch.uint8) + invalid_scale = torch.ones((3, 2), device="cuda", dtype=torch.float32) + + with pytest.raises(RuntimeError, match="gathered scale geometry"): + tensor.fsdp_assign_gathered( + (gathered_data, invalid_scale), + {"direction": "rowwise"}, + ) + + assert tensor._rowwise_data is original_data + def _test_quantize_dequantize( self, quantizer: Float8BlockQuantizer, diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index bc956c7c910..972f0e4534b 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -45,6 +45,7 @@ _unwrap_hybrid_A, _unwrap_hybrid_B, ) +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) @@ -53,6 +54,23 @@ return_reason=True ) +_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR = ( + "Columnwise-only per-tensor FP8 quantization is not implemented" +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=pytest.RaisesExc( + NotImplementedError, + match=_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR, + ), + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + requires_fp8 = pytest.mark.skipif( not fp8_available, reason=f"FP8: {reason_for_no_fp8}", @@ -1008,6 +1026,7 @@ def test_clear_is_idempotent(self, input_tensor): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridTensorShapeOps: """Shape ops that preserve supported Hybrid sub-storages.""" @@ -1530,6 +1549,7 @@ def test_size_from_storage(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridGemmBitwiseIdentical: """Hybrid quantizer with same FP8 format in both directions must produce bitwise-identical results to the vanilla Float8CurrentScaling recipe.""" @@ -1876,6 +1896,7 @@ def test_mha_forward_does_not_probe_qfactory(self, monkeypatch, fp8_mha): from transformer_engine.pytorch.attention import multi_head_attention as mha_module from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( current_scaling_quantizer_factory, + delayed_scaling_quantizer_factory, ) from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, @@ -1887,9 +1908,15 @@ def test_mha_forward_does_not_probe_qfactory(self, monkeypatch, fp8_mha): pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") monkeypatch.setattr(mha_module, "_dpa_fp8_recipe", "") + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") calls = [] def valid_fp8_dpa_qfactory(role): + # Hopper FP8 attention supports DelayedScaling. Use it for the + # complete MHA recipe so that all requests share one coherent + # delayed-scaling state, including the DPA boundary tensors. + if cc < (10, 0): + return delayed_scaling_quantizer_factory(role) is_dpa = role is not None and role.module_type == "dpa" is_dpa_boundary = ( role is not None @@ -1911,9 +1938,9 @@ def counting_qfactory(role): fp8_mha=fp8_mha, ) model = te.MultiheadAttention( - hidden_size=32, + hidden_size=128, num_attention_heads=2, - kv_channels=16, + kv_channels=64, attention_dropout=0.0, attn_mask_type="no_mask", params_dtype=torch.bfloat16, @@ -1921,7 +1948,7 @@ def counting_qfactory(role): qkv_format="sbhd", name="counted_mha", ).cuda() - inp = torch.randn(128, 2, 32, device="cuda", dtype=torch.bfloat16) + inp = torch.randn(128, 2, 128, device="cuda", dtype=torch.bfloat16) with torch.no_grad(), autocast(enabled=True, recipe=custom_recipe): model(inp) @@ -2792,6 +2819,7 @@ def test_fp8_tensor_passthrough(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridBiasGradient: """Verify bias gradients are computed correctly with HybridQuantizer. @@ -3034,6 +3062,7 @@ def factory(role): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridMixedWithNonHybrid: """Only one operand is hybrid; the other uses a plain TE quantizer. @@ -3168,6 +3197,22 @@ def _make_block_quantizer_e5m2(*, rowwise=True, columnwise=True): } +def _make_role_aware_quantizer(factory, role): + """Construct a quantizer with the standard Block-FP8 GEMM geometry.""" + quantizer = factory() + if isinstance(quantizer, Float8BlockQuantizer): + is_weight = ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type == "weight" + ) + # Match Float8BlockScaling: 2D blocks for weights, 1D blocks for + # activations and gradients. A GEMM cannot consume 2D-scaled inputs + # on both sides. + quantizer.block_scaling_dim = 2 if is_weight else 1 + return quantizer + + def _build_cross_format_params(): """Build parametrize list for all stateless cross-format hybrid combos.""" combos = [ @@ -3197,6 +3242,8 @@ def _build_cross_format_params(): marks = [] if hw_skip: marks.append(pytest.mark.skipif(True, reason=hw_reason or "N/A")) + if col == "fp8_current": + marks.append(_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8) params.append(pytest.param(row, col, id=f"{row}_row_x_{col}_col", marks=marks)) return params @@ -3241,8 +3288,8 @@ def _plain_linear_qfactory(operand_factory, grad_factory): def factory(role): is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return grad_factory() - return operand_factory() + return _make_role_aware_quantizer(grad_factory, role) + return _make_role_aware_quantizer(operand_factory, role) return factory @@ -3296,12 +3343,12 @@ def hybrid_factory(role): ) if is_linear and role.tensor_type in ("input", "weight"): return HybridQuantizer( - rowwise_quantizer=make_row_operand(), - columnwise_quantizer=make_col_operand(), + rowwise_quantizer=_make_role_aware_quantizer(make_row_operand, role), + columnwise_quantizer=_make_role_aware_quantizer(make_col_operand, role), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return make_col_grad() - return make_row_operand() + return _make_role_aware_quantizer(make_col_grad, role) + return _make_role_aware_quantizer(make_row_operand, role) hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_factory) fprop_ref_recipe = recipe.CustomRecipe( @@ -3679,6 +3726,7 @@ def factory(role): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridAllModules: """Hybrid quantization through all TE module types (not just Linear). @@ -3800,16 +3848,16 @@ def test_grouped_linear(self): self.ffn_hidden_size, params_dtype=torch.bfloat16, ).cuda() + # Hopper FP8 grouped wgrad uses each expert's token count as a + # leading dimension, which cuBLAS requires to be divisible by 16. + m_splits = [16] * num_gemms inp = torch.randn( - self.batch, + sum(m_splits), self.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True, ) - base = self.batch // num_gemms - rem = self.batch % num_gemms - m_splits = [base + (1 if i < rem else 0) for i in range(num_gemms)] self._run_fwd_bwd(model, inp, m_splits) @@ -3921,8 +3969,18 @@ def test_distinct_delayed_scaling_state_is_allowed(self): ("usage", "expected"), [ pytest.param((True, False), {"rowwise": True, "columnwise": False}, id="rowwise"), - pytest.param((False, True), {"rowwise": False, "columnwise": True}, id="columnwise"), - pytest.param((True, True), {"rowwise": True, "columnwise": True}, id="both"), + pytest.param( + (False, True), + {"rowwise": False, "columnwise": True}, + id="columnwise", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + pytest.param( + (True, True), + {"rowwise": True, "columnwise": True}, + id="both", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), ], ) def test_hybrid_split_quantize_respects_parent_usage_flags(self, usage, expected): @@ -3945,6 +4003,7 @@ def test_hybrid_split_quantize_respects_parent_usage_flags(self, usage, expected assert [storage.get_usages() for storage in out] == [expected, expected] + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_columnwise_only_rowwise_dequantized_uses_transient_grouped_row(self, monkeypatch): import transformer_engine.pytorch.module.grouped_linear as grouped_linear @@ -4042,6 +4101,7 @@ def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): assert all(storage.rowwise_sub_storage is not None for storage in out) assert all(storage.columnwise_sub_storage is None for storage in out) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_original_source_preserves_two_bulk_call_fast_path(self, monkeypatch): import transformer_engine.pytorch.module.grouped_linear as grouped_linear @@ -4240,12 +4300,12 @@ def qfactory(role): is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( - rowwise_quantizer=row_factory(), - columnwise_quantizer=col_factory(), + rowwise_quantizer=_make_role_aware_quantizer(row_factory, role), + columnwise_quantizer=_make_role_aware_quantizer(col_factory, role), ) if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return grad_factory() - return row_factory() + return _make_role_aware_quantizer(grad_factory, role) + return _make_role_aware_quantizer(row_factory, role) return recipe.CustomRecipe(qfactory=qfactory) @@ -4256,6 +4316,7 @@ def qfactory(role): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridQuantizedModelInit: """Verify that quantized_model_init with a hybrid CustomRecipe produces HybridQuantizedTensor parameters.""" @@ -4372,6 +4433,7 @@ def _hybrid_fp8_recipe(self): ), ) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_quantized_param_skips_workspace(self): """When weight is already a HybridQuantizedTensor (quantized params), get_weight_workspace should return it directly without creating a workspace.""" @@ -4386,6 +4448,7 @@ def test_quantized_param_skips_workspace(self): assert out.shape == (32, 128) assert "weight" not in model._fp8_workspaces + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_bf16_weight_creates_hybrid_workspace(self): """When weight is BF16 and recipe produces HybridQuantizer, the workspace should be a HybridQuantizedTensor.""" @@ -4422,6 +4485,7 @@ def test_workspace_cache_reuse_across_microbatches(self): assert all(current is cached for current, cached in zip(current_buffers, buffers)) torch.testing.assert_close(out1, out2, rtol=0.0, atol=0.0) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_workspace_cache_invalidation_on_usage_change(self): """If usage requirements change (e.g. inference→training), the cache should be invalidated and a fresh workspace created.""" @@ -4461,6 +4525,7 @@ def test_workspace_cache_invalidation_on_usage_change(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridUpdateWeightQuantizers: """Test that quantizer refresh propagates correctly to hybrid sub-quantizers.""" @@ -4806,6 +4871,7 @@ def test_fp8_delayed_transpose_only_nonzero_offset(self): # ---------- Positive tests (same-format) ---------- + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_fp8_current_same_format_full_master(self): """Full master (start_offset=0) routes both sub-storages through the existing per-format current-scaling helper. Verifies both directions @@ -4848,6 +4914,7 @@ def test_fp8_current_same_format_full_master(self): torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_fp8_current_nonzero_start_offset(self): """Mimic DP-sharded master: master covers logical elements [start_offset, start_offset + master.numel()) of the full model weight. @@ -4917,6 +4984,7 @@ def test_fp8_current_nonzero_start_offset(self): dequantized[start_offset:], hp_master_shard, rtol=0.125, atol=0.1 ) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_fp8_delayed_same_format_full_master(self): """Same-format delayed scaling on both directions. Both sub-storages route into the delayed-scaling bucket as independent entries; the @@ -4943,6 +5011,7 @@ def test_fp8_delayed_same_format_full_master(self): torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_fp8_delayed_row_current_col_full_master(self): """Cross-format per-tensor Float8: delayed row + current col. @@ -4973,6 +5042,7 @@ def test_fp8_delayed_row_current_col_full_master(self): torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_fp8_current_row_delayed_col_full_master(self): """Cross-format per-tensor Float8: current row + delayed col. @@ -5160,6 +5230,7 @@ def test_blockwise_columnwise_raises(self): with pytest.raises(NotImplementedError, match="Float8BlockQuantizer columnwise"): quantize_master_weights([weight], [master_flat], [0], group=group) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_rowwise_only_fp8_current_full_master(self): """Single-direction hybrid: columnwise dropped via update_usage. @@ -5198,6 +5269,7 @@ def test_rowwise_only_fp8_current_full_master(self): dq_row = weight._rowwise_storage.dequantize(dtype=torch.float32) torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_columnwise_only_fp8_current_full_master(self): """Single-direction hybrid: rowwise dropped via update_usage. @@ -5234,6 +5306,7 @@ def test_columnwise_only_fp8_current_full_master(self): dq_col = weight._columnwise_storage.dequantize(dtype=torch.float32) torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_both_sub_storages_none_raises(self): """Both sub-storages dropped via update_usage — nothing left to cast. @@ -5257,6 +5330,7 @@ def test_both_sub_storages_none_raises(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridPostAllGatherProcessing: """Hybrid branch of `post_all_gather_processing` is exercised indirectly by the positive `TestHybridQuantizeMasterWeights` tests; the case below pins @@ -5296,6 +5370,7 @@ def test_post_ag_idempotent_for_fp8_current_hybrid(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridRecipeCorrespondence: """Test _check_weight_tensor_recipe_correspondence with hybrid params.""" @@ -5340,6 +5415,7 @@ def test_hybrid_param_with_mismatched_recipe_raises(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridQuantizeInPlace: """Test in-place re-quantization (quantize_) for HybridQuantizedTensor. @@ -5393,6 +5469,7 @@ def test_quantize_inplace_preserves_tensor_identity(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridFusedAdam: """Test FusedAdam optimizer with HybridQuantizedTensor parameters.""" @@ -5495,6 +5572,7 @@ def test_fused_adam_requires_master_weights(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridQuantizedParamsEndToEnd: """Full training loop: quantized_model_init + autocast fwd + bwd + FusedAdam.step().""" @@ -5951,6 +6029,7 @@ def _test_equivalence(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestQuantizedParamsEquivalenceFP8CurrentScaling(_QuantizedParamsEquivalenceBase): """Vanilla versus same-format hybrid FP8-current training parity.""" @@ -6251,7 +6330,11 @@ def _mxfp8_factory(): _dispatch_configs = [ - pytest.param("fp8_fp8", id="same-format-fp8"), + pytest.param( + "fp8_fp8", + id="same-format-fp8", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), ] if mxfp8_available: _dispatch_configs.append(pytest.param("mxfp8_mxfp8", id="same-format-mxfp8")) @@ -6592,6 +6675,7 @@ def test_identity_substorages_allow_integer_dtype(self): assert torch.count_nonzero(sub_storage.dequantize()).item() == 0 @requires_fp8 + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 @pytest.mark.parametrize("unsupported_dtype", (torch.int32, torch.float64, torch.bool)) def test_non_identity_substorage_rejects_unsupported_dtype(self, unsupported_dtype): quantizer = HybridQuantizer( @@ -6934,7 +7018,13 @@ def _make_fsdp_protocol_param(config_name): return model.weight -_fsdp_protocol_configs = [pytest.param("fp8_fp8", id="same-format")] +_fsdp_protocol_configs = [ + pytest.param( + "fp8_fp8", + id="same-format", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ) +] if mxfp8_available: _fsdp_protocol_configs.append(pytest.param("mxfp8_fp8", id="mixed-mxfp8-fp8")) if fp8_block_scaling_available: @@ -7202,6 +7292,7 @@ def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): ) _assert_hybrid_tensor_exact(second_result, hybrid_param, context="FSDP roundtrip reuse") + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_scale_refresh_across_iterations(self): """After a sharded optimizer-style requantize, iter-2+ gathers see the new scale. @@ -7333,6 +7424,7 @@ def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): @requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridMakeLike: """Test that make_like produces correct copies for __torch_dispatch__ usage.""" @@ -7386,8 +7478,6 @@ def test_make_like_is_independent(self): # Skip on Blackwell where the C++ kernel always populates ``_data`` and # the columnwise-only Float8 path doesn't exercise ``_transpose``. -from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported # noqa: E402 - requires_hopper_fp8 = pytest.mark.skipif( is_non_tn_fp8_gemm_supported() or not fp8_available, reason=( @@ -7549,6 +7639,58 @@ def test_fsdp_buffer_fields_falls_back_to_data_when_both_present(self): @requires_hopper_fp8 +class TestPerTensorFloat8ColumnwiseOnlyGuard: + """Per-tensor FP8 cannot yet emit only the physical transpose.""" + + @staticmethod + def _make_quantizer(scaling): + if scaling == "current": + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + return Float8Quantizer( + scale=torch.ones(1, device="cuda", dtype=torch.float32), + amax=torch.zeros(1, device="cuda", dtype=torch.float32), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=False, + columnwise=True, + ) + + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_initial_quantize_raises_not_implemented(self, scaling): + quantizer = self._make_quantizer(scaling) + source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantizer(source) + + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_update_quantized_raises_not_implemented(self, scaling): + quantizer = self._make_quantizer(scaling) + source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + output = quantizer.make_empty( + source.shape, + dtype=source.dtype, + device=source.device, + ) + assert output._data is None + assert output._transpose is not None + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantizer.update_quantized(source, output) + + +@requires_hopper_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridFsdpPostAllGatherUpdateUsage: """``HybridQuantizedTensor.fsdp_post_all_gather`` must call ``update_usage`` on each sub-storage after writing gathered data @@ -7774,6 +7916,7 @@ def _run_transformer_layer(self, recipe_obj, *, checkpoint_fn=None): # ----- te.checkpoint, reentrant --------------------------------- + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_te_checkpoint_reentrant_linear_fp8_bitwise(self): """te.checkpoint(use_reentrant=True) around te.Linear with same-format FP8 hybrid → bitwise parity with non-recompute. @@ -7806,6 +7949,7 @@ def fn(model, inp): # ----- te.checkpoint, non-reentrant ----------------------------- + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_te_checkpoint_non_reentrant_linear_fp8_bitwise(self): """te.checkpoint(use_reentrant=False) — the saved-tensors-hooks path. Different recompute infra (``_checkpoint_hook``) than the @@ -7868,19 +8012,34 @@ def fn(model, inp): # weight-workspace caching hits it under vanilla torch checkpoint), # but hybrid amplifies and surfaces it via the 2x sub-storage count. + _TORCH_CHECKPOINT_FP8_XFAIL = pytest.mark.xfail( + raises=( + pytest.RaisesExc( + NotImplementedError, + match=_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR, + ) + if not is_non_tn_fp8_gemm_supported() + else torch.utils.checkpoint.CheckpointError + ), + strict=True, + reason=( + "On Hopper, same-format FP8 hybrid reaches the unsupported columnwise-only " + "per-tensor quantization guard before checkpointing. On architectures that " + "support non-TN FP8 GEMMs, vanilla torch.utils.checkpoint(use_reentrant=False) " + "is incompatible with TE's weight-workspace cache. Tracked by #3158." + ), + ) + _TORCH_CHECKPOINT_CACHE_XFAIL = pytest.mark.xfail( raises=torch.utils.checkpoint.CheckpointError, strict=True, reason=( - "Vanilla torch.utils.checkpoint(use_reentrant=False) is" - " incompatible with TE's weight-workspace cache: cache-miss" - " on the first forward saves a different tensor count than" - " cache-hit on recompute. Use te.checkpoint instead (tested" - " above). Tracked by #3158." + "Vanilla torch.utils.checkpoint(use_reentrant=False) is incompatible " + "with TE's weight-workspace cache. Tracked by #3158." ), ) - @_TORCH_CHECKPOINT_CACHE_XFAIL + @_TORCH_CHECKPOINT_FP8_XFAIL def test_torch_checkpoint_non_reentrant_linear_fp8_bitwise(self): """Vanilla ``torch.utils.checkpoint.checkpoint`` without TE wrapper around a hybrid-quantized te.Linear. @@ -7931,6 +8090,7 @@ def fn(model, inp): # ----- TransformerLayer ----------------------------------------- + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_te_checkpoint_reentrant_transformer_layer_fp8(self): """te.checkpoint(reentrant) around a full TransformerLayer under hybrid FP8. Exercises LayerNormLinear + DPA + LayerNormMLP in one @@ -7953,6 +8113,7 @@ def fn(model, inp): test = self._run_transformer_layer(self._same_format_fp8_recipe(), checkpoint_fn=fn) _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) TransformerLayer FP8") + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_te_checkpoint_non_reentrant_transformer_layer_fp8(self): """Same TransformerLayer setup but through the non-reentrant saved-tensors-hooks recompute path. Same bitwise-equality @@ -7970,6 +8131,7 @@ def fn(model, inp): # ----- quantized_model_init + recompute ------------------------- + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_te_checkpoint_reentrant_quantized_model_init_fp8_bitwise(self): """Combine ``quantized_model_init`` (persistent HybridQuantizedTensor weights) with activation recompute — @@ -8046,6 +8208,7 @@ def _run_grouped_linear(self, recipe_obj, *, checkpoint_fn=None): out.float().sum().backward() return _collect_outputs(out, inp, model) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_te_checkpoint_reentrant_grouped_linear_fp8_bitwise(self): """GroupedLinear + te.checkpoint(reentrant) under same-format FP8 hybrid. Exercises the MoE ``_hybrid_split_quantize`` + list-of- @@ -8059,6 +8222,7 @@ def fn(model, inp, m_splits): test = self._run_grouped_linear(self._same_format_fp8_recipe(), checkpoint_fn=fn) _assert_outputs_bitwise_equal(ref, test, "te.checkpoint(reentrant) GroupedLinear FP8") + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_te_checkpoint_non_reentrant_grouped_linear_fp8_bitwise(self): """Same GroupedLinear recompute setup but through the non- reentrant saved-tensors-hooks path — verifies that the list of @@ -8075,6 +8239,7 @@ def fn(model, inp, m_splits): # ----- Selective attention recompute ---------------------------- + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_selective_attention_recompute_transformer_layer_fp8_bitwise(self): """``TransformerLayer(..., checkpoint_core_attention=True)`` — the Megatron default memory-savings pattern. @@ -8132,8 +8297,18 @@ def _run(checkpoint_core_attention): @pytest.mark.parametrize( "format_name,reentrant", [ - pytest.param("fp8_current", True, id="fp8_current-reentrant"), - pytest.param("fp8_current", False, id="fp8_current-nonreentrant"), + pytest.param( + "fp8_current", + True, + id="fp8_current-reentrant", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), + pytest.param( + "fp8_current", + False, + id="fp8_current-nonreentrant", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), pytest.param( "mxfp8", True, @@ -8228,6 +8403,7 @@ def fn(model, inp): # ----- save_for_backward round-trip (unit-level) ---------------- + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_prepare_restore_roundtrip_is_identity(self): """Unit-level guarantee: the ``prepare_for_saving`` / ``restore_from_saved`` chain used by diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 92484b172ce..4900521efd6 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -29,6 +29,7 @@ from transformer_engine.pytorch.tensor.storage.identity_tensor_storage import ( IdentityTensorStorage, ) +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) @@ -37,6 +38,23 @@ return_reason=True ) +_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR = ( + "Columnwise-only per-tensor FP8 quantization is not implemented" +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=pytest.RaisesExc( + NotImplementedError, + match=_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR, + ), + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + # ── Module-level qfactories (picklable / autocast-friendly) ────────── @@ -935,6 +953,7 @@ def test_quantize_master_weights_identity_exact_nonzero_offset(self): torch.testing.assert_close(weight.dequantize(), expected, rtol=0.0, atol=0.0) @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_quantize_master_weights_hybrid_identity_fp8_current(self): from transformer_engine.pytorch.tensor.utils import ( post_all_gather_processing, @@ -1529,6 +1548,7 @@ def test_whole_layer_hp_matches_bf16_bitwise(self): for g_id, g_ref in zip(wg_id, wg_ref): torch.testing.assert_close(g_id, g_ref, rtol=0.0, atol=0.0) + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_fwd_hp_bwd_fp8_forward_bitwise(self): """High-precision forward must be bitwise-equal to BF16 forward; the backward runs in FP8 (finite, close to BF16 within a loose tolerance).""" diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4308464c527..1c6f85638eb 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -17,6 +17,26 @@ namespace transformer_engine::pytorch { namespace { +/*! @brief Reject unsupported columnwise-only per-tensor FP8 quantization + * + * Per-tensor FP8 uses the legacy cast-transpose kernel whenever columnwise + * output is requested. That kernel requires both rowwise and columnwise + * output buffers, so a transpose-only output cannot be produced without an + * otherwise-unused rowwise allocation and write. Keep this limitation + * explicit until the kernel supports independently selecting its outputs. + */ +void check_per_tensor_fp8_quantize_output(const TensorWrapper& output) { + const auto rowwise_data = output.get_rowwise_data(); + const auto columnwise_data = output.get_columnwise_data(); + if (rowwise_data.data_ptr == nullptr && columnwise_data.data_ptr != nullptr) { + PyErr_SetString( + PyExc_NotImplementedError, + "Columnwise-only per-tensor FP8 quantization is not implemented; " + "the cast-transpose kernel requires rowwise output storage."); + throw py::error_already_set(); + } +} + /*! @brief Resolve an optional device to a concrete CUDA device * * If no device is provided, uses the current CUDA device. @@ -615,6 +635,7 @@ void Float8Quantizer::quantize(const TensorWrapper& input, TensorWrapper& out, if (input.numel() == 0) { return; } + check_per_tensor_fp8_quantize_output(out); QuantizationConfigWrapper quant_config; if (noop_flag) { quant_config.set_noop_tensor(noop_flag->data()); @@ -975,6 +996,7 @@ void Float8CurrentScalingQuantizer::quantize_impl(const TensorWrapper& input, Te out.set_scale(nullptr, DType::kFloat32, out.defaultShape); return; } + check_per_tensor_fp8_quantize_output(out); // Quantization configs QuantizationConfigWrapper quant_config; diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 37245b36d16..368b83031de 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -508,6 +508,13 @@ def fsdp_extract_buffers( names = self.fsdp_buffer_fields() block_len = self._FSDP_BLOCK_LEN m, n = self._fsdp_logical_mn() + if m % block_len != 0: + raise RuntimeError( + "FSDP2 cannot all-gather a 2-D Float8BlockwiseQTensor whose " + f"local flattened M dimension ({m}) is not a multiple of {block_len}; " + "the shard boundary splits a block-scale tile. Choose aligned shard " + "boundaries or use a supported non-blockwise recipe." + ) m_tiles = (m + block_len - 1) // block_len last_tiles = (n + block_len - 1) // block_len @@ -552,6 +559,27 @@ def fsdp_assign_gathered( block_len = self._FSDP_BLOCK_LEN direction = meta["direction"] data, scale = gathered + if direction not in ("rowwise", "columnwise"): + raise RuntimeError(f"Invalid Float8Block FSDP gather direction: {direction!r}") + if data is None or scale is None: + raise RuntimeError( + "Float8Block FSDP gathered data and scale buffers must both be present." + ) + + m_full = 1 + for dim in data.shape[:-1]: + m_full *= dim + last_dim = data.size(-1) if data.dim() > 0 else 1 + expected_scale_shape = ( + (m_full + block_len - 1) // block_len, + (last_dim + block_len - 1) // block_len, + ) + if tuple(scale.shape) != expected_scale_shape: + raise RuntimeError( + "Float8Block FSDP gathered scale geometry does not match gathered data: " + f"got scale shape {tuple(scale.shape)}, expected {expected_scale_shape} " + f"for gathered data shape {tuple(data.shape)}." + ) if direction == "rowwise": last_dim = data.size(-1) From 416288bd4b04642915ab6c667688e5295b954192 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:31:43 +0000 Subject: [PATCH 51/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/quantizer.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 1c6f85638eb..ddea0f429dc 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -29,10 +29,9 @@ void check_per_tensor_fp8_quantize_output(const TensorWrapper& output) { const auto rowwise_data = output.get_rowwise_data(); const auto columnwise_data = output.get_columnwise_data(); if (rowwise_data.data_ptr == nullptr && columnwise_data.data_ptr != nullptr) { - PyErr_SetString( - PyExc_NotImplementedError, - "Columnwise-only per-tensor FP8 quantization is not implemented; " - "the cast-transpose kernel requires rowwise output storage."); + PyErr_SetString(PyExc_NotImplementedError, + "Columnwise-only per-tensor FP8 quantization is not implemented; " + "the cast-transpose kernel requires rowwise output storage."); throw py::error_already_set(); } } From a3f52586fe2c72a899a001f9ac214c4384c0032b Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 21 Jul 2026 13:08:02 +0000 Subject: [PATCH 52/97] More Hopper fixes Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/conftest.py | 5 +- tests/pytorch/test_hybrid_quantization.py | 79 +++++++++++++++++-- tests/pytorch/test_identity_quantizer.py | 5 +- tests/pytorch/test_quantized_tensor.py | 26 ++++++ .../tensor/storage/mxfp8_tensor_storage.py | 2 +- transformer_engine/pytorch/tensor/utils.py | 43 ++++++++++ 6 files changed, 143 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index 32ea3e1d65d..0d1953e4c92 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -72,10 +72,7 @@ def _parametrize_hybrid_recipes(): if name == "HybridFP8CurrentScaling" and not is_non_tn_fp8_gemm_supported(): marks.append( pytest.mark.xfail( - raises=pytest.RaisesExc( - NotImplementedError, - match="Columnwise-only per-tensor FP8 quantization is not implemented", - ), + raises=NotImplementedError, strict=True, reason=( "Hopper does not yet support columnwise-only per-tensor FP8 " diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 972f0e4534b..c67359a6baf 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -60,10 +60,7 @@ _XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( condition=not is_non_tn_fp8_gemm_supported(), - raises=pytest.RaisesExc( - NotImplementedError, - match=_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR, - ), + raises=NotImplementedError, strict=True, reason=( "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " @@ -4805,6 +4802,75 @@ def _reference_fp8_bytes(master, scale, dtype=torch.bfloat16): quantizer.update_quantized(master.reshape(1, -1), temp) return temp._data.reshape(-1) + def test_fp8_current_fsdp_transpose_only_raises_before_mutation(self): + """Reject Hopper FSDP updates that would materialize rowwise column storage.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + shape = (4, 8) + shard_shape = (2, 8) + row_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + col_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + quantizer = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + shard_source = source[: shard_shape[0]].contiguous() + weight = HybridQuantizedTensor( + shape=shape, + dtype=source.dtype, + rowwise_storage=row_quantizer(source), + columnwise_storage=self._make_transpose_only_float8_weight(shape, col_quantizer), + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + shard = HybridQuantizedTensor( + shape=shard_shape, + dtype=source.dtype, + rowwise_storage=row_quantizer(shard_source), + columnwise_storage=self._make_transpose_only_float8_weight(shard_shape, col_quantizer), + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + weight_col = weight._columnwise_storage + shard_col = shard._columnwise_storage + weight_scale_before = weight_col._scale_inv.clone() + weight_transpose_before = weight_col._transpose.clone() + shard_scale_before = shard_col._scale_inv.clone() + shard_transpose_before = shard_col._transpose.clone() + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantize_master_weights( + [weight], + [shard_source.float()], + [0], + group=None, + fsdp_shard_model_weights=[shard], + ) + + assert weight_col._data is None + assert shard_col._data is None + torch.testing.assert_close(weight_col._scale_inv, weight_scale_before, rtol=0, atol=0) + torch.testing.assert_close(shard_col._scale_inv, shard_scale_before, rtol=0, atol=0) + assert torch.equal(weight_col._transpose, weight_transpose_before) + assert torch.equal(shard_col._transpose, shard_transpose_before) + @staticmethod def _logical_float8_bytes(storage): """Return FP8 payload in the tensor's logical row-major order.""" @@ -8014,10 +8080,7 @@ def fn(model, inp): _TORCH_CHECKPOINT_FP8_XFAIL = pytest.mark.xfail( raises=( - pytest.RaisesExc( - NotImplementedError, - match=_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR, - ) + NotImplementedError if not is_non_tn_fp8_gemm_supported() else torch.utils.checkpoint.CheckpointError ), diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 4900521efd6..48f5f04c71a 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -44,10 +44,7 @@ _XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( condition=not is_non_tn_fp8_gemm_supported(), - raises=pytest.RaisesExc( - NotImplementedError, - match=_COLUMNWISE_ONLY_PER_TENSOR_FP8_ERROR, - ), + raises=NotImplementedError, strict=True, reason=( "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index b2f77ecd665..a07d54d22b5 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -20,6 +20,7 @@ Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, + MXFP8TensorStorage, NVFP4Tensor, QuantizedTensor, ) @@ -670,6 +671,31 @@ def test_cpu_dequantize( assert y_cpu.shape == ref_cpu.shape torch.testing.assert_close(y_cpu, ref_cpu, rtol=0, atol=0) + def test_mxfp8_fsdp_extract_keeps_partial_columnwise_scale_block(self) -> None: + """FSDP extraction must retain the scale for a partial final 32-row block.""" + rows, cols = 48, 64 + columnwise_scale_inv = torch.arange(4 * 128, dtype=torch.float32).reshape(4, 128) + storage = MXFP8TensorStorage( + rowwise_data=None, + rowwise_scale_inv=None, + columnwise_data=torch.zeros((rows, cols), dtype=torch.uint8), + columnwise_scale_inv=columnwise_scale_inv, + fp8_dtype=te.DType.kFloat8E4M3, + quantizer=None, + with_gemm_swizzled_scales=False, + fake_dtype=torch.bfloat16, + ) + + buffers, metadata = storage.fsdp_extract_buffers() + + assert metadata == {"field_names": ("_columnwise_data", "_columnwise_scale_inv")} + assert len(buffers) == 2 + assert buffers[0] is storage._columnwise_data + extracted_scale_inv = buffers[1] + assert extracted_scale_inv is not None + assert extracted_scale_inv.shape == (2, 128) + torch.testing.assert_close(extracted_scale_inv, columnwise_scale_inv[:2], rtol=0, atol=0) + @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("dim", [0, 1]) def test_chunk( diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 9c8327626f5..da5d6c4e6ba 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -367,7 +367,7 @@ def fsdp_extract_buffers( if t.size(0) != flattened_in_shape0: t = t[:flattened_in_shape0] elif name == "_columnwise_scale_inv" and t is not None: - expected = flattened_in_shape0 // MXFP8_BLOCK_SCALING_SIZE + expected = math.ceil(flattened_in_shape0 / MXFP8_BLOCK_SCALING_SIZE) if t.size(0) != expected: t = t[:expected] buffers.append(t) diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 8bc4e6a15d2..b2ce7658736 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -248,6 +248,13 @@ def quantize_master_weights( fsdp_shard_model_weights = [None] * len(model_weights) else: use_fsdp_shard_model_weights = True + # Validate the entire batch before clearing initialization state or + # populating any per-format work buckets. + for model_weight, fsdp_shard_model_weight in zip(model_weights, fsdp_shard_model_weights): + _validate_fp8_current_scaling_fsdp_hopper_policy( + model_weight, + fsdp_shard_model_weight, + ) for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( model_weights, master_weights, start_offsets, fsdp_shard_model_weights @@ -1203,6 +1210,42 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # --------------------------------------------------------------------------------------------- +def _validate_fp8_current_scaling_fsdp_hopper_policy( + model_weight, + fsdp_shard_model_weight, +): + """Reject unsupported transpose-only current-scaling FSDP updates on Hopper. + + The FSDP shard path can receive a ``Float8Tensor`` and pass it directly to + ``Float8Quantizer.update_quantized``. On Hopper, doing that for a + columnwise-only current-scaling tensor materializes rowwise ``_data`` and + violates the direction pinned by ``HybridQuantizer``. Reject the whole + batch before any scale, cache, or storage mutation instead. + """ + if fsdp_shard_model_weight is None or is_non_tn_fp8_gemm_supported(): + return + + if isinstance(model_weight, HybridQuantizedTensor): + quantizer = model_weight._get_quantizer() + candidates = ( + (model_weight._rowwise_storage, quantizer.rowwise_quantizer), + (model_weight._columnwise_storage, quantizer.columnwise_quantizer), + ) + else: + candidates = ((model_weight, model_weight._get_quantizer()),) + + for storage, quantizer in candidates: + if ( + storage is not None + and isinstance(quantizer, Float8CurrentScalingQuantizer) + and _is_float8_transpose_only(storage) + ): + raise NotImplementedError( + "Columnwise-only per-tensor FP8 quantization is not implemented for " + "FSDP current-scaling updates on this architecture." + ) + + def _route_hybrid_to_buckets( model_weight, master_weight, From 77e0e1c8ebe71bdfb7f8a574e0d957f2f25ef768 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 21 Jul 2026 13:09:34 +0000 Subject: [PATCH 53/97] Guard hybrid partial-master distopt updates Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 128 ++++++++++----------- transformer_engine/pytorch/tensor/utils.py | 50 +++++++- 2 files changed, 107 insertions(+), 71 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index c67359a6baf..2ffd8587b69 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -4802,6 +4802,64 @@ def _reference_fp8_bytes(master, scale, dtype=torch.bfloat16): quantizer.update_quantized(master.reshape(1, -1), temp) return temp._data.reshape(-1) + def test_fp8_current_original_partial_master_raises(self): + """Reject an independently sourced Hybrid column from a partial master shard. + + A one-payload distributed optimizer only communicates the rowwise payload. + Re-creating an independently quantized column from a partial master shard can + therefore produce a different value after checkpoint resume. Construct the + Hopper transpose-only column explicitly so this safety regression does not + depend on columnwise-only kernel support. + """ + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + shape = (4, 8) + row_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + col_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + quantizer = HybridQuantizer( + rowwise_quantizer=row_quantizer, + columnwise_quantizer=col_quantizer, + columnwise_source="original", + ) + + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + rowwise = row_quantizer(source) + columnwise = self._make_transpose_only_float8_weight(shape, col_quantizer) + weight = HybridQuantizedTensor( + shape=shape, + dtype=source.dtype, + rowwise_storage=rowwise, + columnwise_storage=columnwise, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + master_shard = source.reshape(-1)[shape[-1] :].float().contiguous() + + with pytest.raises( + ValueError, + match=( + "partial master shard.*columnwise_source='original'.*" + "columnwise_source='rowwise_dequantized'" + ), + ): + quantize_master_weights( + [weight], + [master_shard], + [shape[-1]], + group=None, + ) + def test_fp8_current_fsdp_transpose_only_raises_before_mutation(self): """Reject Hopper FSDP updates that would materialize rowwise column storage.""" from transformer_engine.pytorch.tensor.utils import quantize_master_weights @@ -4980,76 +5038,6 @@ def test_fp8_current_same_format_full_master(self): torch.testing.assert_close(dq_row.reshape(-1), master_flat, rtol=0.125, atol=0.1) torch.testing.assert_close(dq_col.reshape(-1), master_flat, rtol=0.125, atol=0.1) - @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 - def test_fp8_current_nonzero_start_offset(self): - """Mimic DP-sharded master: master covers logical elements - [start_offset, start_offset + master.numel()) of the full model weight. - Verifies that the shared logical start_offset is honored by both - sub-storages' per-format routings. - """ - from transformer_engine.pytorch.tensor.utils import ( - quantize_master_weights, - post_all_gather_processing, - ) - - group = _ensure_single_rank_dp_group() - hybrid_recipe = _hybrid_recipe_fp8_current() - weight, hp_master_full = _build_hybrid_linear_weight(64, 64, hybrid_recipe) - - half = hp_master_full.numel() // 2 - # Negation preserves the shard amax while guaranteeing that this update - # is observable. The fixed seed puts the full-tensor amax in this shard, - # so the current-scaling factor must remain bitwise stable. - hp_master_shard = -hp_master_full.view(-1)[half:].contiguous() - start_offset = half - - before = {} - for direction, storage in ( - ("rowwise", weight._rowwise_storage), - ("columnwise", weight._columnwise_storage), - ): - assert storage is not None - before[direction] = { - "bytes": self._logical_float8_bytes(storage).clone(), - "scale_inv": storage._scale_inv.clone(), - "dequantized": storage.dequantize(dtype=torch.float32).reshape(-1).clone(), - } - - quantize_master_weights([weight], [hp_master_shard], [start_offset], group=group) - post_all_gather_processing([weight]) - - for direction, storage in ( - ("rowwise", weight._rowwise_storage), - ("columnwise", weight._columnwise_storage), - ): - previous = before[direction] - actual_bytes = self._logical_float8_bytes(storage) - torch.testing.assert_close(storage._scale_inv, previous["scale_inv"], rtol=0, atol=0) - torch.testing.assert_close( - actual_bytes[:start_offset], previous["bytes"][:start_offset], rtol=0, atol=0 - ) - expected_shard_bytes = self._reference_fp8_bytes( - hp_master_shard.to(weight.dtype), - torch.reciprocal(storage._scale_inv), - weight.dtype, - ) - torch.testing.assert_close( - actual_bytes[start_offset:], expected_shard_bytes, rtol=0, atol=0 - ) - assert not torch.equal( - actual_bytes[start_offset:], previous["bytes"][start_offset:] - ), f"{direction} updated shard unexpectedly retained every FP8 byte" - dequantized = storage.dequantize(dtype=torch.float32).reshape(-1) - torch.testing.assert_close( - dequantized[:start_offset], - previous["dequantized"][:start_offset], - rtol=0, - atol=0, - ) - torch.testing.assert_close( - dequantized[start_offset:], hp_master_shard, rtol=0.125, atol=0.1 - ) - @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_fp8_delayed_same_format_full_master(self): """Same-format delayed scaling on both directions. Both sub-storages diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index b2ce7658736..6695509327a 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -250,11 +250,20 @@ def quantize_master_weights( use_fsdp_shard_model_weights = True # Validate the entire batch before clearing initialization state or # populating any per-format work buckets. - for model_weight, fsdp_shard_model_weight in zip(model_weights, fsdp_shard_model_weights): + for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( + model_weights, master_weights, start_offsets, fsdp_shard_model_weights + ): _validate_fp8_current_scaling_fsdp_hopper_policy( model_weight, fsdp_shard_model_weight, ) + if isinstance(model_weight, HybridQuantizedTensor): + _validate_hybrid_partial_master_policy( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, + ) for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( model_weights, master_weights, start_offsets, fsdp_shard_model_weights @@ -1246,6 +1255,45 @@ def _validate_fp8_current_scaling_fsdp_hopper_policy( ) +def _validate_hybrid_partial_master_policy( + model_weight, + master_weight, + start_offset, + fsdp_shard_model_weight, +): + """Reject an unsafe independently sourced Hybrid column before distopt mutation.""" + row_sub = model_weight._rowwise_storage + col_sub = model_weight._columnwise_storage + quantizer = model_weight._get_quantizer() + if ( + master_weight is None + or row_sub is None + or col_sub is None + or quantizer.columnwise_source != "original" + ): + return + + shard_has_both_directions = ( + isinstance(fsdp_shard_model_weight, HybridQuantizedTensor) + and fsdp_shard_model_weight._rowwise_storage is not None + and fsdp_shard_model_weight._columnwise_storage is not None + ) + if shard_has_both_directions: + return + + end_offset = _validate_flat_fragment(model_weight, master_weight, start_offset) + if start_offset == 0 and end_offset == model_weight.numel(): + return + + raise ValueError( + "quantize_master_weights cannot update a two-direction " + "HybridQuantizedTensor from a partial master shard when " + "columnwise_source='original': the one-payload distributed-optimizer " + "path cannot preserve an independently quantized columnwise value. " + "Use columnwise_source='rowwise_dequantized' or provide full-master data." + ) + + def _route_hybrid_to_buckets( model_weight, master_weight, From 461daa1a24cd68e5d7cf3d9613b351d37206cd1f Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 21 Jul 2026 13:53:46 +0000 Subject: [PATCH 54/97] Blackwell CI fixes Signed-off-by: Evgeny --- .../fsdp2_tests/run_fsdp2_model.py | 20 ++- tests/pytorch/distributed/run_hybrid_tp_sp.py | 11 +- .../pytorch/distributed/test_hybrid_tp_sp.py | 169 +++++------------- 3 files changed, 66 insertions(+), 134 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 2c97b04d12f..25f1e7cad14 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -470,6 +470,15 @@ def test_distributed_hybrid(hybrid_recipe_name): torch.manual_seed(42) torch.cuda.manual_seed(42) + # Float8 block scaling requires every local FSDP shard's flattened M + # dimension to contain whole 128-row scale tiles. Keep the historical + # 512-wide model on up to four ranks, and scale it for larger worlds so + # the projection, fused-QKV, and MLP output dimensions all stay aligned. + shard_alignment = 128 * world_size + hidden_size = ((512 + shard_alignment - 1) // shard_alignment) * shard_alignment + ffn_hidden_size = 4 * hidden_size + assert hidden_size % shard_alignment == 0 + kwargs = dict( fuse_qkv_params=True, params_dtype=torch.bfloat16, @@ -479,7 +488,7 @@ def test_distributed_hybrid(hybrid_recipe_name): ) with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): model = torch.nn.Sequential( - *[te.TransformerLayer(512, 2048, 8, **kwargs) for _ in range(2)] + *[te.TransformerLayer(hidden_size, ffn_hidden_size, 8, **kwargs) for _ in range(2)] ) custom_attrs = save_custom_attrs(model) @@ -505,8 +514,8 @@ def _hybrid_param_count(): optimizer = optim.Adam(model.parameters(), lr=1e-3) - input_data = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) - target = torch.randn(128, 16, 512, device=device, dtype=torch.bfloat16) + input_data = torch.randn(128, 16, hidden_size, device=device, dtype=torch.bfloat16) + target = torch.randn(128, 16, hidden_size, device=device, dtype=torch.bfloat16) losses = [] for iteration in range(3): @@ -625,8 +634,11 @@ def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): torch.manual_seed(42) torch.cuda.manual_seed(42) - in_features = 512 + # Keep dim-0 FSDP shards aligned to Float8 block scaling's 128-row tiles. + shard_alignment = 128 * world_size + in_features = ((512 + shard_alignment - 1) // shard_alignment) * shard_alignment out_features = in_features * 3 + assert in_features % shard_alignment == 0 with te.quantized_model_init(enabled=True, recipe=hybrid_recipe): model = te.LayerNormLinear( in_features, diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index a8c20ed0276..61227b5dc37 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -871,7 +871,8 @@ def main(argv=None): parser.add_argument( "--test", type=str, - default="all", + nargs="+", + default=["all"], choices=[ "all", "linear", @@ -882,7 +883,7 @@ def main(argv=None): "layernorm_mlp", "transformer_layer", ], - help="Run only the named test (speeds up iterative debugging)", + help="Run one or more named tests in the same distributed process group", ) args = parser.parse_args(argv) QUANTIZATION = args.quantization @@ -896,10 +897,12 @@ def main(argv=None): "layernorm_mlp": test_layernorm_mlp, "transformer_layer": test_transformer_layer, } - if args.test == "all": + if "all" in args.test: + if len(args.test) != 1: + parser.error("--test all cannot be combined with named tests") tests_to_run = list(test_map.values()) else: - tests_to_run = [test_map[args.test]] + tests_to_run = [test_map[name] for name in args.test] for test_fn in tests_to_run: dist_print(f"=== Starting {test_fn.__name__} ===") diff --git a/tests/pytorch/distributed/test_hybrid_tp_sp.py b/tests/pytorch/distributed/test_hybrid_tp_sp.py index 320acd0b7a3..6030e9d32fa 100644 --- a/tests/pytorch/distributed/test_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/test_hybrid_tp_sp.py @@ -37,12 +37,19 @@ ) -def _run_test(quantization: str, test: str = "all"): +def _run_tests(quantization: str, tests: tuple[str, ...] = ("all",)): + """Run related cases under one torchrun/process-group startup.""" script = TEST_ROOT / "run_hybrid_tp_sp.py" - cmd = LAUNCH_CMD + [str(script), "--quantization", quantization, "--test", test] + cmd = LAUNCH_CMD + [ + str(script), + "--quantization", + quantization, + "--test", + *tests, + ] result = subprocess.run(cmd, env=os.environ, check=False) assert result.returncode == 0, ( - f"run_hybrid_tp_sp.py (quantization={quantization}, test={test})" + f"run_hybrid_tp_sp.py (quantization={quantization}, tests={list(tests)})" f" exited with code {result.returncode}" ) @@ -53,65 +60,34 @@ def _run_test(quantization: str, test: str = "all"): # Exercises TP amax reduction and SP gather paths. -@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") -@xfail_hopper_columnwise_per_tensor_fp8 -def test_hybrid_fp8_linear(): - """Linear TP/SP coverage for hybrid FP8 current scaling.""" - _run_test("hybrid_fp8", "linear") - - -@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") -@xfail_hopper_columnwise_per_tensor_fp8 -def test_hybrid_fp8_linear_vs_vanilla(): - """Same-topology Linear parity against Float8CurrentScaling.""" - _run_test("hybrid_fp8", "linear_vs_vanilla") - - -@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") -@xfail_hopper_columnwise_per_tensor_fp8 -def test_hybrid_fp8_layernorm_linear_vs_vanilla(): - """Same-topology LayerNormLinear parity against vanilla FP8.""" - _run_test("hybrid_fp8", "layernorm_linear_vs_vanilla") - - -@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") -@xfail_hopper_columnwise_per_tensor_fp8 -def test_hybrid_fp8_layernorm_mlp_vs_vanilla(): - """Same-topology LayerNormMLP parity against vanilla FP8.""" - _run_test("hybrid_fp8", "layernorm_mlp_vs_vanilla") - - -@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") -@xfail_hopper_columnwise_per_tensor_fp8 -def test_hybrid_fp8_layernorm_linear(): - """LayerNormLinear TP/SP coverage for hybrid FP8.""" - _run_test("hybrid_fp8", "layernorm_linear") - - -@pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") -@xfail_hopper_columnwise_per_tensor_fp8 -def test_hybrid_fp8_layernorm_mlp(): - """LayerNormMLP TP/SP coverage for hybrid FP8.""" - _run_test("hybrid_fp8", "layernorm_mlp") +_SAME_FORMAT_TESTS = ( + "linear", + "linear_vs_vanilla", + "layernorm_linear_vs_vanilla", + "layernorm_mlp_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", +) @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") @xfail_hopper_columnwise_per_tensor_fp8 -def test_hybrid_fp8_transformer_layer(): - """TransformerLayer TP/SP coverage for hybrid FP8.""" - _run_test("hybrid_fp8", "transformer_layer") +def test_hybrid_fp8(): + """Hybrid FP8 TP/SP coverage and same-topology vanilla parity.""" + _run_tests("hybrid_fp8", _SAME_FORMAT_TESTS) @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_hybrid_fp8_identity_linear(): """Linear TP/SP coverage for FP8 forward plus Identity backward.""" - _run_test("hybrid_fp8_identity", "linear") + _run_tests("hybrid_fp8_identity", ("linear",)) @pytest.mark.skipif(not fp8_available, reason=f"FP8: {reason_for_no_fp8}") def test_identity_all_modules(): """All-Identity TP/SP end-to-end coverage for every supported TE module.""" - _run_test("identity", "all") + _run_tests("identity") # ────────────────────────────────────────────────────────────────────── @@ -121,44 +97,14 @@ def test_identity_all_modules(): @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") -def test_hybrid_mxfp8_linear(): - _run_test("hybrid_mxfp8", "linear") - - -@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") -def test_hybrid_mxfp8_linear_vs_vanilla(): - _run_test("hybrid_mxfp8", "linear_vs_vanilla") - - -@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") -def test_hybrid_mxfp8_layernorm_linear_vs_vanilla(): - _run_test("hybrid_mxfp8", "layernorm_linear_vs_vanilla") - - -@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") -def test_hybrid_mxfp8_layernorm_mlp_vs_vanilla(): - _run_test("hybrid_mxfp8", "layernorm_mlp_vs_vanilla") - - -@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") -def test_hybrid_mxfp8_layernorm_linear(): - _run_test("hybrid_mxfp8", "layernorm_linear") - - -@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") -def test_hybrid_mxfp8_layernorm_mlp(): - _run_test("hybrid_mxfp8", "layernorm_mlp") - - -@pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") -def test_hybrid_mxfp8_transformer_layer(): - _run_test("hybrid_mxfp8", "transformer_layer") +def test_hybrid_mxfp8(): + _run_tests("hybrid_mxfp8", _SAME_FORMAT_TESTS) @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") def test_hybrid_mxfp8_identity_linear(): """Linear TP/SP coverage for MXFP8 forward plus Identity backward.""" - _run_test("hybrid_mxfp8_identity", "linear") + _run_tests("hybrid_mxfp8_identity", ("linear",)) # ────────────────────────────────────────────────────────────────────── @@ -168,28 +114,17 @@ def test_hybrid_mxfp8_identity_linear(): @pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") -def test_hybrid_nvfp4_linear(): - _run_test("hybrid_nvfp4", "linear") - - -@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") -def test_hybrid_nvfp4_linear_vs_vanilla(): - _run_test("hybrid_nvfp4", "linear_vs_vanilla") - - -@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") -def test_hybrid_nvfp4_layernorm_linear(): - _run_test("hybrid_nvfp4", "layernorm_linear") - - -@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") -def test_hybrid_nvfp4_layernorm_mlp(): - _run_test("hybrid_nvfp4", "layernorm_mlp") - - -@pytest.mark.skipif(not nvfp4_available, reason=f"NVFP4: {reason_for_no_nvfp4}") -def test_hybrid_nvfp4_transformer_layer(): - _run_test("hybrid_nvfp4", "transformer_layer") +def test_hybrid_nvfp4(): + _run_tests( + "hybrid_nvfp4", + ( + "linear", + "linear_vs_vanilla", + "layernorm_linear", + "layernorm_mlp", + "transformer_layer", + ), + ) # ────────────────────────────────────────────────────────────────────── @@ -204,26 +139,8 @@ def test_hybrid_nvfp4_transformer_layer(): @pytest.mark.skipif( not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" ) -def test_hybrid_mxfp8_nvfp4_linear(): - _run_test("hybrid_mxfp8_nvfp4", "linear") - - -@pytest.mark.skipif( - not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" -) -def test_hybrid_mxfp8_nvfp4_layernorm_linear(): - _run_test("hybrid_mxfp8_nvfp4", "layernorm_linear") - - -@pytest.mark.skipif( - not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" -) -def test_hybrid_mxfp8_nvfp4_layernorm_mlp(): - _run_test("hybrid_mxfp8_nvfp4", "layernorm_mlp") - - -@pytest.mark.skipif( - not _cross_format_available, reason=f"MXFP8+NVFP4: {_reason_for_no_cross_format}" -) -def test_hybrid_mxfp8_nvfp4_transformer_layer(): - _run_test("hybrid_mxfp8_nvfp4", "transformer_layer") +def test_hybrid_mxfp8_nvfp4(): + _run_tests( + "hybrid_mxfp8_nvfp4", + ("linear", "layernorm_linear", "layernorm_mlp", "transformer_layer"), + ) From 67f0f4824d5b4f832e2500e37fa0cd71d36fc903 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 21 Jul 2026 19:05:35 +0000 Subject: [PATCH 55/97] Docs minor update Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 4 + .../quantization_factory_base.py | 16 +-- .../quantization_factory_zoo.py | 134 +++++++++--------- .../pytorch/tensor/hybrid_tensor.py | 43 +++--- .../pytorch/tensor/identity_tensor.py | 22 +-- .../tensor/storage/hybrid_tensor_storage.py | 11 +- .../tensor/storage/identity_tensor_storage.py | 11 +- 7 files changed, 129 insertions(+), 112 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 2ffd8587b69..27dbf4fb930 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -4860,6 +4860,10 @@ def test_fp8_current_original_partial_master_raises(self): group=None, ) + @pytest.mark.skipif( + is_non_tn_fp8_gemm_supported(), + reason="Hopper-only: Blackwell supports columnwise per-tensor FP8 FSDP updates", + ) def test_fp8_current_fsdp_transpose_only_raises_before_mutation(self): """Reject Hopper FSDP updates that would materialize rowwise column storage.""" from transformer_engine.pytorch.tensor.utils import quantize_master_weights diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py index 7dae0891777..e5125fd92fd 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py @@ -3,12 +3,12 @@ # See LICENSE for license information. """ -Quantizer factory examples using real silicon quantizers. +Quantizer factory examples using native Transformer Engine quantizers. -Each factory below replicates the behaviour of built-in TE recipe but via the -``CustomRecipe`` + ``qfactory`` interface. This is useful when you want to -start from a known-good recipe and then selectively override quantizer settings -for specific layers / tensor types. +For TE ``Linear`` and ``GroupedLinear`` roles, each factory below mirrors the +nominal defaults of a built-in recipe through the ``CustomRecipe`` + +``qfactory`` interface. This provides a built-in-equivalent starting point for +selectively overriding quantizer settings for specific layers or tensor types. Usage (any factory):: @@ -87,8 +87,8 @@ def current_scaling_quantizer_factory( return Float8CurrentScalingQuantizer( fp8_dtype=fp8_dtype, device=torch.device("cuda"), - force_pow_2_scales=False, # constrain scale to powers of 2 - amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero + force_pow_2_scales=False, + amax_epsilon=0.0, ) @@ -131,7 +131,7 @@ def float8_block_scaling_quantizer_factory( fp8_dtype=DType.kFloat8E4M3, rowwise=True, columnwise=True, - amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero + amax_epsilon=0.0, force_pow_2_scales=True, block_scaling_dim=block_scaling_dim, # 1 = 1D (1×128), 2 = 2D (128×128) ) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index aca031d39c0..8b238970507 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -10,9 +10,20 @@ interface to apply *different* quantization recipes to different module/tensor types/instances within the same model, and how to use ``HybridQuantizer`` when rowwise and columnwise tensor directions should use -different formats or sources. +different representations or sources. -Factories are ordered from conservative to more aggressive quantization. +Within the Linear/GroupedLinear and RL-oriented families, examples are roughly +ordered by increasingly aggressive forward quantization. This is an +organizational convention, not an expected accuracy or performance ranking. +These factories demonstrate what can be composed; they are not necessarily +tuned for end-to-end performance. + +When the columnwise/backward representation is higher precision than +the rowwise/forward representation, consider +``columnwise_source="rowwise_dequantized"`` as a starting point. It constructs +backward operands from the value obtained during forward quantization, +improving forward/backward representation consistency. +Note: dequantization does not recover information discarded during forward quantization. Organization: * Linear / grouped-linear recipes (pre-training). Favor more precision @@ -25,12 +36,11 @@ Use these with caution. These are **not** official, supported recipes provided by Transformer Engine -- they are illustrative examples meant to - inspire your own experiments, not drop-in production defaults. While most - of the factories here are grounded in some evidence or rationale (see the - per-factory docstrings), they have not been broadly validated for accuracy, - convergence, or performance across models and hardware. Treat them as - starting points: benchmark and verify on your own workload before relying on - any of them. + inspire your own experiments, not drop-in production defaults. Most include + a motivating rationale in their per-factory docstrings, but they have not + been broadly validated for accuracy, convergence, or performance across + models and hardware. Treat them as starting points: benchmark and verify on + your own workload before relying on any of them. Usage:: @@ -64,7 +74,6 @@ from ..constants import DType from .quantization_factory_base import mxfp8_quantizer_factory, nvfp4_quantizer_factory - # ----------------------------------------------------------------------------- # Linear / GroupedLinear Recipes (pre-training) # ----------------------------------------------------------------------------- @@ -209,52 +218,6 @@ def mxfp8_fwd_dequantized_bwd_quantizer_factory( return mxfp8_quantizer_factory(role) -def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: row-scaled NVFP4 forward, dequantized backward. - - This expresses the linear/grouped-linear equivalent of - ``NVFP4BlockScaling(row_scaled_activation=True, - backward_override="dequantized")`` through per-direction quantizers: - - * ``input`` -> - ``Hybrid(rowwise=row-scaled NVFP4, columnwise=Identity, - columnwise_source="rowwise_dequantized")`` - * ``weight`` -> - ``Hybrid(rowwise=plain NVFP4, columnwise=Identity, - columnwise_source="rowwise_dequantized")`` - * ``grad_output`` -> ``IdentityQuantizer`` - * everything else -> plain NVFP4 - - Row-scaled NVFP4 is fprop-only, so the forward quantizers avoid RHT, - stochastic rounding, and 2D scaling. - - This recipe targets RL-style training use cases and builds on - NVIDIA/TransformerEngine#2931, which introduced row-scaled NVFP4: - https://github.com/NVIDIA/TransformerEngine/pull/2931 - """ - from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer - from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer - - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type == "input": - return HybridQuantizer( - rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), - columnwise_quantizer=IdentityQuantizer(), - columnwise_source="rowwise_dequantized", - ) - if is_linear and role.tensor_type == "weight": - return HybridQuantizer( - rowwise_quantizer=_plain_nvfp4_quantizer(), - columnwise_quantizer=IdentityQuantizer(), - columnwise_source="rowwise_dequantized", - ) - if is_linear and role.tensor_type == "grad_output": - return IdentityQuantizer() - return _plain_nvfp4_quantizer() - - def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -316,6 +279,52 @@ def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( return mxfp8_quantizer_factory(role) +def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: row-scaled NVFP4 forward, dequantized backward. + + This expresses a linear/grouped-linear variant of + ``NVFP4BlockScaling(row_scaled_activation=True, + backward_override="dequantized")`` through per-direction quantizers: + + * ``input`` -> + ``Hybrid(rowwise=row-scaled NVFP4, columnwise=Identity, + columnwise_source="rowwise_dequantized")`` + * ``weight`` -> + ``Hybrid(rowwise=plain NVFP4, columnwise=Identity, + columnwise_source="rowwise_dequantized")`` + * ``grad_output`` -> ``IdentityQuantizer`` + * everything else -> plain NVFP4 + + Row-scaled NVFP4 is fprop-only, so the forward quantizers avoid RHT, + stochastic rounding, and 2D scaling. + + This recipe targets RL-style training use cases and builds on + NVIDIA/TransformerEngine#2931, which introduced row-scaled NVFP4: + https://github.com/NVIDIA/TransformerEngine/pull/2931 + """ + from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer + from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + if is_linear and role.tensor_type == "input": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + if is_linear and role.tensor_type == "grad_output": + return IdentityQuantizer() + return _plain_nvfp4_quantizer() + + # ----------------------------------------------------------------------------- # Linear + Attention Recipes # ----------------------------------------------------------------------------- @@ -329,24 +338,22 @@ def nvfp4_linear_fp8_dpa_factory( This factory demonstrates how to use ``CustomRecipe`` with ``fp8_dpa=True`` to combine NVFP4 quantization for linear layers with FP8 attention. - DPA tensor types (``role.module_type == "dpa"``): + DPA-owned tensor types (``role.module_type == "dpa"``): =========== ============================================================ tensor_type Description =========== ============================================================ ``"qkv"`` Query, Key, Value inputs to the first attention GEMM ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM - ``"o"`` Attention output (O = S·V) ``"do"`` Gradient of the attention output (dO), backward input ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward - ``"dqkv"`` Gradient flowing back to Q, K, V =========== ============================================================ Dispatch logic: * ``role.module_type == "dpa"`` with ``tensor_type in ("s", "dp")`` - -> FP8 delayed scaling (HYBRID, most_recent, history length 1) + -> FP8 delayed scaling (``Format.HYBRID``, most_recent, history length 1) * other DPA roles - -> FP8 current scaling (HYBRID: E4M3 fwd, E5M2 bwd) + -> FP8 current scaling (``Format.HYBRID``: E4M3 fwd, E5M2 bwd) * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in ``role.name``) -> FP8 current scaling placeholder. The fused attention kernel requires FP8-compatible quantizers in all DPA slots, even when the output is @@ -413,21 +420,20 @@ def nvfp4_linear_mxfp8_dpa_factory( ``NVTE_DPA_FP8_RECIPE="MXFP8BlockScaling"`` env override that the built-in recipes would otherwise need is unnecessary. - DPA tensor types (``role.module_type == "dpa"``): + DPA-owned tensor types (``role.module_type == "dpa"``): =========== ============================================================ tensor_type Description =========== ============================================================ ``"qkv"`` Query, Key, Value inputs to the first attention GEMM ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM - ``"o"`` Attention output (O = S·V) ``"do"`` Gradient of the attention output (dO), backward input ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward - ``"dqkv"`` Gradient flowing back to Q, K, V =========== ============================================================ Dispatch logic: - * ``role.module_type == "dpa"`` -> MXFP8 (HYBRID: E4M3 fwd, E5M2 bwd) + * ``role.module_type == "dpa"`` -> MXFP8 + (``Format.HYBRID``: E4M3 fwd, E5M2 bwd) The MXFP8 fused-attention kernel handles the S/dP slots internally, so any quantizer returned for those roles is later nulled out by ``get_attention_quantizers``. Returning MXFP8 is diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 7ab84f32d44..2f94ec24c3b 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Tensor class with hybrid quantized data (different formats for rowwise vs columnwise)""" +"""Tensor and quantizer classes for composed rowwise and columnwise representations.""" from __future__ import annotations from typing import Any, Dict, Iterable, Literal, Optional, Tuple @@ -17,11 +17,13 @@ class HybridQuantizer(Quantizer): - """Quantizer that composes two existing quantizers for different directions. + """Quantizer that composes rowwise and columnwise representations. - Performs two-pass quantization: the rowwise_quantizer produces rowwise - quantized data and the columnwise_quantizer produces columnwise quantized - data. The results are wrapped in a HybridQuantizedTensor. + When both representations are requested, applies ``rowwise_quantizer`` to + produce the rowwise representation and ``columnwise_quantizer`` to produce + the columnwise representation. The results are wrapped in a + ``HybridQuantizedTensor``. The children may use the same or different + Transformer Engine formats, or custom ``Quantizer`` implementations. Parameters ---------- @@ -37,21 +39,23 @@ class HybridQuantizer(Quantizer): Notes ----- + Rowwise and columnwise describe storage and GEMM orientations. Whether a + representation is consumed in the forward or backward pass depends on the + tensor's role in the operation. + ``HybridQuantizer`` pins each sub-quantizer to its designated direction by mutating its usage flags, so it takes ownership of the supplied quantizer instances. The rowwise and columnwise quantizers must be distinct objects. If both directions need shared state, construct two quantizer instances that reference the same external state object. - Reusing a sub-quantizer instance across multiple ``HybridQuantizer`` objects - is unsupported by contract. Catching that robustly would require copying or - ownership tracking, both of which are more intrusive, so only the direct - rowwise/columnwise aliasing case is enforced. + Each ``HybridQuantizer`` must receive its own sub-quantizer instances; do not + reuse a sub-quantizer instance across multiple ``HybridQuantizer`` objects. Example ------- - MXFP8 forward plus high-precision backward from the rowwise-dequantized - forward value can be expressed as:: + MXFP8 rowwise data plus a high-precision columnwise representation derived + from the rowwise value can be expressed as:: HybridQuantizer( rowwise_quantizer=mxfp8_quantizer, @@ -405,11 +409,12 @@ def _get_compatible_recipe(self): class HybridQuantizedTensor(HybridQuantizedTensorStorage, QuantizedTensor): - """Quantized tensor holding data in two different formats per direction. + """Tensor holding independently produced rowwise and columnwise representations. - The tensor presents as having a standard, higher-precision dtype, but - internally stores rowwise data in one quantized format and columnwise - data in another. + The tensor presents as having a standard logical dtype, but + internally stores representations produced by the two sub-quantizers. + These may use the same format, different formats, or a high-precision + representation such as ``IdentityTensorStorage``. Parameters ---------- @@ -417,10 +422,10 @@ class HybridQuantizedTensor(HybridQuantizedTensorStorage, QuantizedTensor): Tensor dimensions. dtype : torch.dtype Nominal tensor datatype. - rowwise_storage : QuantizedTensorStorage - Sub-storage for rowwise quantized data. - columnwise_storage : QuantizedTensorStorage - Sub-storage for columnwise quantized data. + rowwise_storage : QuantizedTensorStorage, optional + Sub-storage for the rowwise representation. + columnwise_storage : QuantizedTensorStorage, optional + Sub-storage for the columnwise representation. quantizer : HybridQuantizer Parent hybrid quantizer that owns the rowwise and columnwise sub-quantizers. requires_grad : bool, default = False diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index cae581fe98e..e418f9e7216 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -4,12 +4,13 @@ """High-precision passthrough quantizer and tensor. -``IdentityQuantizer`` is a no-op "quantizer": it keeps the input tensor in its -original high precision instead of casting to a low-precision format. It exists -so the ``CustomRecipe`` + ``qfactory`` machinery can express *unquantized* -tensors and, composed inside a :class:`HybridQuantizer`, *unquantized -directions* (e.g. high-precision forward + quantized backward, or vice versa) -without scattering ``None``/``isinstance`` special-cases across the modules. +``IdentityQuantizer`` stores a tensor directly without a low-precision encoding, +preserving the input dtype by default. It exists so the ``CustomRecipe`` + +``qfactory`` machinery can express *high-precision* tensors and, composed inside +a :class:`HybridQuantizer`, *high-precision directions* without scattering +``None``/``isinstance`` special-cases across the modules. Here high precision +means the held compute dtype, typically BF16, FP16, or FP32; it does not mean +FP32 specifically. """ from __future__ import annotations @@ -23,11 +24,12 @@ class IdentityQuantizer(Quantizer): - """Quantizer that performs no quantization (high-precision passthrough). + """Quantizer that produces a high-precision passthrough representation. Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`) - wrapping the original high-precision tensor. ``general_gemm`` materializes - it back to a plain tensor, so any GEMM consuming it runs in high precision. + holding the tensor directly, without a low-precision encoding. + ``general_gemm`` materializes it as a plain tensor, so a GEMM consumes it + in the held dtype. Parameters ---------- @@ -161,7 +163,7 @@ class IdentityTensor(IdentityTensorStorage, QuantizedTensor): """High-precision passthrough tensor produced by :class:`IdentityQuantizer`. Presents as a standard tensor of its nominal dtype; internally it just - holds the original high-precision data (no quantization). + holds data directly in that dtype, without a low-precision encoding. """ def __repr__(self, *, tensor_contents=None): diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index 86a9a7534c9..fb08dfabcea 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Mixin class holding data specific for HybridQuantizedTensor""" +"""Storage class for composed rowwise and columnwise representations.""" from __future__ import annotations from collections.abc import Iterable @@ -14,11 +14,12 @@ class HybridQuantizedTensorStorage(QuantizedTensorStorage): - """Storage that composes two QuantizedTensorStorage instances. + """Storage that composes rowwise and columnwise sub-storages. - One sub-storage provides rowwise quantized data and the other provides - columnwise quantized data. This enables mixed-precision quantization - where, for example, rowwise data is FP8 and columnwise data is FP4. + One sub-storage provides the rowwise representation and the other provides + the columnwise representation. Either may be absent when that direction is + not needed. The sub-storages may use different formats or include an + high-precision representation such as ``IdentityTensorStorage``. """ diff --git a/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py index bcf606b7e66..c1dea51097f 100644 --- a/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/identity_tensor_storage.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Mixin class holding data for IdentityTensor (high-precision passthrough).""" +"""Storage class for a high-precision IdentityTensor representation.""" from __future__ import annotations from typing import Any, Dict, Optional, Tuple @@ -14,14 +14,13 @@ class IdentityTensorStorage(QuantizedTensorStorage): - """Passthrough storage that holds a high-precision (unquantized) tensor. + """Passthrough storage that holds a high-precision tensor directly. Produced by :class:`IdentityQuantizer`. It implements the ``QuantizedTensorStorage`` interface so it can flow through the same module / GEMM / save-for-backward / FSDP machinery as the real quantized - storages, but it performs no quantization: it simply carries the original - high-precision tensor. ``general_gemm`` materializes it back to that plain - tensor (so the matmul runs in high precision). + storages, but it uses no low-precision encoding. ``general_gemm`` + materializes it as a plain tensor in the held dtype. The data is direction-agnostic -- the same tensor serves both the rowwise and columnwise directions (the GEMM transposes via its layout flags), so a @@ -99,7 +98,7 @@ def get_data_tensors(self, rowwise_data: bool = True, columnwise_data: bool = Tr raise ValueError("No data to get, both rowwise_data and columnwise_data are False") def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: - """Return the held high-precision tensor (no-op dequantization).""" + """Return the held high-precision tensor, casting when requested.""" if self._hp_data is None: raise RuntimeError("IdentityTensorStorage has no data to dequantize") if dtype is None: From a6fe3447c6fa91c7ad6c16a68609f9741d468eda Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 22 Jul 2026 12:32:36 +0000 Subject: [PATCH 56/97] Improve fsdp2 and other tests Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/fsdp2_utils.py | 88 +- .../fsdp2_tests/run_fsdp2_model.py | 95 +- tests/pytorch/distributed/run_hybrid_tp_sp.py | 126 +- tests/pytorch/distributed/test_torch_fsdp2.py | 2 + tests/pytorch/hybrid_quantization_utils.py | 362 ++++ tests/pytorch/test_cpu_offloading.py | 46 +- tests/pytorch/test_cpu_offloading_v1.py | 49 +- tests/pytorch/test_hybrid_quantization.py | 1699 +---------------- .../pytorch/test_hybrid_quantization_fsdp2.py | 1570 +++++++++++++++ tests/pytorch/test_identity_quantizer.py | 73 +- 10 files changed, 2082 insertions(+), 2028 deletions(-) create mode 100644 tests/pytorch/hybrid_quantization_utils.py create mode 100644 tests/pytorch/test_hybrid_quantization_fsdp2.py diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index fff55471025..042333843b5 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -5,11 +5,15 @@ """Shared utility functions for FSDP2 distributed tests.""" import transformer_engine.common.recipe -from transformer_engine.pytorch import HybridQuantizer, IdentityQuantizer, QuantizedTensor -from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - current_scaling_quantizer_factory, - float8_block_scaling_quantizer_factory, - mxfp8_quantizer_factory, +from transformer_engine.pytorch import QuantizedTensor + +from hybrid_quantization_utils import ( + hybrid_float8_block_qfactory, + hybrid_fp8_current_identity_qfactory, + hybrid_fp8_current_qfactory, + hybrid_mixed_mxfp8_fp8_qfactory, + hybrid_mxfp8_qfactory, + identity_qfactory, ) @@ -17,80 +21,18 @@ def get_recipe_from_string(recipe): return getattr(transformer_engine.common.recipe, recipe)() -def _hybrid_fp8_current_qfactory(role): - """FP8 current-scaling rowwise + FP8 current-scaling columnwise.""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=current_scaling_quantizer_factory(role), - columnwise_quantizer=current_scaling_quantizer_factory(role), - ) - return current_scaling_quantizer_factory(role) - - -def _hybrid_mxfp8_qfactory(role): - """MXFP8 rowwise + MXFP8 columnwise.""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=mxfp8_quantizer_factory(role), - columnwise_quantizer=mxfp8_quantizer_factory(role), - ) - return mxfp8_quantizer_factory(role) - - -def _hybrid_float8_block_qfactory(role): - """Float8 block-scaling rowwise + Float8 block-scaling columnwise.""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=float8_block_scaling_quantizer_factory(role), - columnwise_quantizer=float8_block_scaling_quantizer_factory(role), - ) - return float8_block_scaling_quantizer_factory(role) - - -def _hybrid_mixed_mxfp8_fp8_qfactory(role): - """MXFP8 rowwise + FP8 current columnwise (cross-format hybrid).""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=mxfp8_quantizer_factory(role), - columnwise_quantizer=current_scaling_quantizer_factory(role), - ) - return current_scaling_quantizer_factory(role) - - -def _hybrid_fp8_current_identity_qfactory(role): - """FP8 current forward + high-precision backward via Identity.""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=current_scaling_quantizer_factory(role), - columnwise_quantizer=IdentityQuantizer(), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return IdentityQuantizer() - return current_scaling_quantizer_factory(role) - - -def _identity_qfactory(role): # pylint: disable=unused-argument - """High-precision passthrough for every quantizer slot.""" - return IdentityQuantizer() - - # CustomRecipe has dynamic TE extra-state handling. Once FP8 state is # initialized, TE's get_extra_state() pickles the recipe on save, so # checkpoint-test qfactories must be module-level and picklable. On load, # payloads without delayed-scaling state are identified and ignored without # unpickling. See ``run_fsdp2_fused_adam.py::test_hybrid_dcp_output_parity``. _HYBRID_QFACTORIES = { - "HybridFP8CurrentScaling": _hybrid_fp8_current_qfactory, - "HybridMXFP8": _hybrid_mxfp8_qfactory, - "HybridFloat8BlockScaling": _hybrid_float8_block_qfactory, - "HybridMixed_MXFP8_FP8": _hybrid_mixed_mxfp8_fp8_qfactory, - "HybridFP8CurrentScalingIdentity": _hybrid_fp8_current_identity_qfactory, - "Identity": _identity_qfactory, + "HybridFP8CurrentScaling": hybrid_fp8_current_qfactory, + "HybridMXFP8": hybrid_mxfp8_qfactory, + "HybridFloat8BlockScaling": hybrid_float8_block_qfactory, + "HybridMixed_MXFP8_FP8": hybrid_mixed_mxfp8_fp8_qfactory, + "HybridFP8CurrentScalingIdentity": hybrid_fp8_current_identity_qfactory, + "Identity": identity_qfactory, } diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 25f1e7cad14..77dad90b183 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -212,61 +212,14 @@ def shard_model_with_fsdp2(model, mesh): @torch.no_grad() -def _check_fp8_fsdp2_allgather(model): - # Do manual allgather in fp32 and match against fp8 allgather done - # with fsdp2 - # FP32 manual weight allgather - fp32_allgathered_params = {} - for name, param in model.named_parameters(): - assert isinstance(param, DTensor) - local_tensor = param._local_tensor - device_mesh = param.device_mesh - dist_group = ( - device_mesh.get_group(mesh_dim="shard") - if device_mesh.ndim > 1 - else device_mesh.get_group() - ) - # Perform manual allgather on local_tensor. zeros_like will create hp tensor since - # torch_dispatch for local_tensor will go down the dequantization route. - gathered_tensor = [ - torch.zeros_like(local_tensor) for _ in range(dist.get_world_size(group=dist_group)) - ] - dist.all_gather(gathered_tensor, local_tensor.dequantize(), group=dist_group) - full_tensor = torch.cat(gathered_tensor, dim=0) - fp32_allgathered_params[name] = full_tensor - # FP8 allgather using FSDP2 - for module in model.modules(): - # Not all modules are wrapped/sharded with FSDP2. - if hasattr(module, "unshard"): - module.unshard() - # Make sure allgathered parameters match exactly - for name, param in model.named_parameters(): - # NVFP4 scale unpad/repad through FSDP2 introduces small numerical - # differences vs the manual dequantize-then-allgather path. - if isinstance(param, NVFP4Tensor): - tols = dict(atol=5e-4, rtol=5e-3) - else: - tols = {} - torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **tols) - # Revert model to original sharded state - for module in model.modules(): - # Not all modules are wrapped/sharded with FSDP2. - if hasattr(module, "reshard"): - module.reshard() - +def _check_fp8_fsdp2_allgather(model, *, tols=None): + """Compare FSDP2's quantized all-gather against a manual HP all-gather. -@torch.no_grad() -def _check_hybrid_fsdp2_allgather(model): - """All-gather correctness for hybrid quantized params under FSDP2. - - Mirrors :func:`_check_fp8_fsdp2_allgather` but for hybrid params: compares - FSDP2's quantized all-gather (``unshard`` -> ``dequantize``) against a manual - fp32 dequantize-then-allgather of the local shards. This is self-referential - (no vanilla reference model), so it is robust to the norm-fusion difference - between hybrid and vanilla recipes -- hybrid auto-disables ``with_quantized_norm`` - while vanilla fuses, so a hybrid-vs-vanilla comparison would not match. + ``tols=None`` preserves the format-specific defaults used by the original + FP8 checker. Callers for other quantized tensor compositions can provide a + single tolerance dictionary that applies to every parameter. """ - # Manual fp32 weight allgather of the (possibly hybrid) local shards. + # Manual high-precision weight all-gather. fp32_allgathered_params = {} for name, param in model.named_parameters(): assert isinstance(param, DTensor) @@ -277,25 +230,36 @@ def _check_hybrid_fsdp2_allgather(model): if device_mesh.ndim > 1 else device_mesh.get_group() ) - # ``dequantize`` is a no-op on plain (non-quantized) bf16 params such as - # LayerNorm weights/biases, and dequantizes hybrid sub-storages otherwise. + # Materialize high precision explicitly so this also works for composite + # quantized tensors such as HybridQuantizedTensor. local_hp = local_tensor.dequantize() gathered_tensor = [ torch.zeros_like(local_hp) for _ in range(dist.get_world_size(group=dist_group)) ] dist.all_gather(gathered_tensor, local_hp, group=dist_group) - fp32_allgathered_params[name] = torch.cat(gathered_tensor, dim=0) - # Quantized allgather via FSDP2. + full_tensor = torch.cat(gathered_tensor, dim=0) + fp32_allgathered_params[name] = full_tensor + # Quantized all-gather using FSDP2. for module in model.modules(): + # Not all modules are wrapped/sharded with FSDP2. if hasattr(module, "unshard"): module.unshard() - # Hybrid sub-storages (e.g. MXFP8 scale unpad/repad through FSDP2) can introduce - # small numerical differences vs the manual dequantize-then-allgather path. - tols = dict(atol=5e-4, rtol=5e-3) + # Make sure all-gathered parameters match. for name, param in model.named_parameters(): - torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **tols) + if tols is not None: + param_tols = tols + elif isinstance(param, NVFP4Tensor): + # NVFP4 scale unpad/repad through FSDP2 introduces small numerical + # differences vs the manual dequantize-then-allgather path. + param_tols = dict(atol=5e-4, rtol=5e-3) + else: + param_tols = {} + torch.testing.assert_close( + param.dequantize(), fp32_allgathered_params[name], **param_tols + ) # Revert model to original sharded state. for module in model.modules(): + # Not all modules are wrapped/sharded with FSDP2. if hasattr(module, "reshard"): module.reshard() @@ -512,7 +476,8 @@ def _hybrid_param_count(): hybrid_count = _hybrid_param_count() assert hybrid_count > 0, "No HybridQuantizedTensor local tensors after sharding" - optimizer = optim.Adam(model.parameters(), lr=1e-3) + optimizer_lr = 1e-3 if hidden_size <= 512 else 1e-4 + optimizer = optim.Adam(model.parameters(), lr=optimizer_lr) input_data = torch.randn(128, 16, hidden_size, device=device, dtype=torch.bfloat16) target = torch.randn(128, 16, hidden_size, device=device, dtype=torch.bfloat16) @@ -542,7 +507,9 @@ def _hybrid_param_count(): ), "HybridQuantizedTensor params lost their quantized type after optimizer.step()" # FSDP2 quantized all-gather must match a manual fp32 dequant-then-allgather. - _check_hybrid_fsdp2_allgather(model) + # Hybrid sub-storages (e.g. MXFP8 scale unpad/repad through FSDP2) can + # introduce small differences vs the manual dequantize-then-allgather path. + _check_fp8_fsdp2_allgather(model, tols=dict(atol=5e-4, rtol=5e-3)) def test_distributed_hybrid_identity_all(): @@ -615,7 +582,7 @@ def _identity_param_count(): _identity_param_count() == identity_count ), "IdentityTensor params lost their quantized type after optimizer.step()" - _check_hybrid_fsdp2_allgather(model) + _check_fp8_fsdp2_allgather(model, tols={}) def test_distributed_hybrid_reshard_after_forward(hybrid_recipe_name): diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index 61227b5dc37..6094b6cc519 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -10,7 +10,6 @@ import datetime import os import sys -from pathlib import Path import torch import torch.distributed as dist @@ -18,20 +17,17 @@ import transformer_engine.pytorch as te from transformer_engine.common import recipe as te_recipe -from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - nvfp4_quantizer_factory, -) -from transformer_engine.pytorch import ( - Float8CurrentScalingQuantizer, - HybridQuantizer, - IdentityQuantizer, - MXFP8Quantizer, -) -# Sibling helper shared with distributed numerics tests. -TEST_ROOT = Path(__file__).parent.resolve() -sys.path.insert(0, str(TEST_ROOT)) -from run_layer_with_overlap import _compare_tensors # noqa: E402 +from hybrid_quantization_utils import ( + hybrid_fp8_current_e5m2_grads_qfactory, + hybrid_fp8_current_identity_qfactory, + hybrid_mxfp8_identity_qfactory, + hybrid_mxfp8_qfactory, + hybrid_nvfp4_qfactory, + hybrid_tp_mxfp8_nvfp4_qfactory, + identity_qfactory, +) +from distributed.run_layer_with_overlap import _compare_tensors # ── Global state ───────────────────────────────────────────────────── @@ -49,110 +45,22 @@ LOSS_FN = nn.MSELoss() -# ── Hybrid recipe factories ────────────────────────────────────────── -# -# Factories stay small; these tests target TP/SP plumbing. - - -def _make_fp8_current_quantizer(*, fp8_dtype=te.DType.kFloat8E4M3): - return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") - - -def _make_mxfp8_quantizer(*, fp8_dtype=te.DType.kFloat8E4M3): - return MXFP8Quantizer(fp8_dtype=fp8_dtype) - - -def _hybrid_fp8_qfactory(role): - """FP8 current scaling; backward operands use E5M2.""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=_make_fp8_current_quantizer(), - columnwise_quantizer=_make_fp8_current_quantizer(), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_fp8_current_quantizer(fp8_dtype=te.DType.kFloat8E5M2) - return _make_fp8_current_quantizer() - - -def _hybrid_mxfp8_qfactory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=_make_mxfp8_quantizer(), - columnwise_quantizer=_make_mxfp8_quantizer(), - ) - return _make_mxfp8_quantizer() - - -def _hybrid_fp8_identity_qfactory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=_make_fp8_current_quantizer(), - columnwise_quantizer=IdentityQuantizer(), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return IdentityQuantizer() - return _make_fp8_current_quantizer() - - -def _hybrid_mxfp8_identity_qfactory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=_make_mxfp8_quantizer(), - columnwise_quantizer=IdentityQuantizer(), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return IdentityQuantizer() - return _make_mxfp8_quantizer() - - -def _identity_qfactory(role): # pylint: disable=unused-argument - return IdentityQuantizer() - - -def _hybrid_nvfp4_qfactory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=nvfp4_quantizer_factory(role), - columnwise_quantizer=nvfp4_quantizer_factory(role), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return nvfp4_quantizer_factory(role) - return nvfp4_quantizer_factory(role) - - -def _hybrid_mxfp8_nvfp4_qfactory(role): - """Cross-format recipe: MXFP8 rowwise, NVFP4 columnwise.""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return nvfp4_quantizer_factory(role) - # Forward/boundary roles keep both formats available. - return HybridQuantizer( - rowwise_quantizer=_make_mxfp8_quantizer(), - columnwise_quantizer=nvfp4_quantizer_factory(role), - ) - - def hybrid_recipe(): """Return a fresh CustomRecipe for the selected test recipe.""" if QUANTIZATION == "hybrid_fp8": - return te_recipe.CustomRecipe(qfactory=_hybrid_fp8_qfactory) + return te_recipe.CustomRecipe(qfactory=hybrid_fp8_current_e5m2_grads_qfactory) if QUANTIZATION == "hybrid_mxfp8": - return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_qfactory) + return te_recipe.CustomRecipe(qfactory=hybrid_mxfp8_qfactory) if QUANTIZATION == "hybrid_fp8_identity": - return te_recipe.CustomRecipe(qfactory=_hybrid_fp8_identity_qfactory) + return te_recipe.CustomRecipe(qfactory=hybrid_fp8_current_identity_qfactory) if QUANTIZATION == "hybrid_mxfp8_identity": - return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_identity_qfactory) + return te_recipe.CustomRecipe(qfactory=hybrid_mxfp8_identity_qfactory) if QUANTIZATION == "identity": - return te_recipe.CustomRecipe(qfactory=_identity_qfactory) + return te_recipe.CustomRecipe(qfactory=identity_qfactory) if QUANTIZATION == "hybrid_nvfp4": - return te_recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) + return te_recipe.CustomRecipe(qfactory=hybrid_nvfp4_qfactory) if QUANTIZATION == "hybrid_mxfp8_nvfp4": - return te_recipe.CustomRecipe(qfactory=_hybrid_mxfp8_nvfp4_qfactory) + return te_recipe.CustomRecipe(qfactory=hybrid_tp_mxfp8_nvfp4_qfactory) raise ValueError(f"Unknown hybrid QUANTIZATION={QUANTIZATION!r}") diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index e3c33b0b36d..02b1cb037ee 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -42,6 +42,8 @@ def test_fsdp2_model_tests(): "-v", "-s", "--tb=short", + "-k", + "not hybrid", ], valid_returncodes=(0, 5), env=os.environ, diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py new file mode 100644 index 00000000000..6ad14f4471a --- /dev/null +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -0,0 +1,362 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared explicit quantizer factories for hybrid quantization tests. + +These are intentionally ordinary module-level functions rather than a public +factory abstraction. Keeping them module-level makes ``CustomRecipe`` objects +that reference them picklable in distributed checkpoint tests. +""" + +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe +from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + current_scaling_quantizer_factory, + float8_block_scaling_quantizer_factory, + mxfp8_quantizer_factory, + nvfp4_quantizer_factory, +) + +_LINEAR_MODULE_TYPES = ("linear", "grouped_linear") +_FORWARD_TENSOR_TYPES = ("input", "weight", "output") +_GRAD_TENSOR_TYPES = ("grad_output", "grad_input") + + +def _is_linear_role(role): + return role is not None and role.module_type in _LINEAR_MODULE_TYPES + + +def _make_fp8_current(*, fp8_dtype=te.DType.kFloat8E4M3): + return te.Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") + + +def _make_mxfp8(*, fp8_dtype=te.DType.kFloat8E4M3): + return te.MXFP8Quantizer(fp8_dtype=fp8_dtype) + + +def fp8_e4m3_factory(): + """Construct the default E4M3 current-scaling test quantizer.""" + return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") + + +def fp8_e5m2_factory(): + """Construct the default E5M2 current-scaling test quantizer.""" + return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E5M2, device="cuda") + + +def mxfp8_e4m3_factory(): + """Construct the default E4M3 MXFP8 test quantizer.""" + return te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + + +def make_fp8_quantizer(*, rowwise=True, columnwise=True): + """Construct the standard current-scaling FP8 test quantizer.""" + return te.Float8CurrentScalingQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + device="cuda", + rowwise=rowwise, + columnwise=columnwise, + ) + + +def make_nvfp4_quantizer(*, rowwise=True, columnwise=True): + """Construct the standard NVFP4 test quantizer.""" + return te.NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=rowwise, + columnwise=columnwise, + ) + + +def make_hybrid_quantizer_fp8_row_fp4_col(): + """Construct a hybrid quantizer with FP8 rowwise and NVFP4 columnwise.""" + return te.HybridQuantizer( + rowwise_quantizer=make_fp8_quantizer(), + columnwise_quantizer=make_nvfp4_quantizer(), + ) + + +def as_data_tensor_tuple(storage): + """Return a storage's raw buffers as a tuple without copying them.""" + if storage is None: + return () + tensors = storage.get_data_tensors() + return tensors if isinstance(tensors, tuple) else (tensors,) + + +def snapshot_storage_tensor_metadata(storage, *, clone=False): + """Capture every tensor-valued concrete-storage metadata field.""" + if storage is None: + return None + tensor_metadata = {} + for name, value in storage.get_metadata().items(): + if isinstance(value, torch.Tensor) or value is None: + snapshot_value = ( + value.detach().clone() if clone and value is not None else value + ) + if ( + value is not None + and getattr(storage, "_is_2D_scaled", False) + and name in ("rowwise_scale_inv", "columnwise_scale_inv") + ): + # Float8Block 2D scales pad one tile dimension. Kernels do not + # initialize or consume that padding, so canonicalize it. + snapshot_value = value.detach().clone() + m, n = storage._fsdp_logical_mn() + block_len = storage._FSDP_BLOCK_LEN + m_tiles = (m + block_len - 1) // block_len + n_tiles = (n + block_len - 1) // block_len + if name == "rowwise_scale_inv": + snapshot_value[:, n_tiles:] = 0 + else: + snapshot_value[:, m_tiles:] = 0 + tensor_metadata[name] = snapshot_value + return { + "storage_type": type(storage).__name__, + "tensor_metadata": tensor_metadata, + } + + +def assert_nested_state_exact(actual, expected, *, path="state"): + """Recursively compare nested state, using zero tolerance for tensors.""" + if isinstance(expected, torch.Tensor): + assert isinstance( + actual, torch.Tensor + ), f"{path}: expected Tensor, got {type(actual)}" + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0, msg=path) + return + if isinstance(expected, dict): + assert isinstance(actual, dict), f"{path}: expected dict, got {type(actual)}" + assert actual.keys() == expected.keys(), f"{path}: dictionary keys differ" + for key in expected: + assert_nested_state_exact(actual[key], expected[key], path=f"{path}.{key}") + return + if isinstance(expected, (list, tuple)): + assert isinstance( + actual, type(expected) + ), f"{path}: expected {type(expected).__name__}, got {type(actual).__name__}" + assert len(actual) == len(expected), f"{path}: sequence lengths differ" + for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): + assert_nested_state_exact( + actual_item, + expected_item, + path=f"{path}[{index}]", + ) + return + assert actual == expected, f"{path}: {actual!r} != {expected!r}" + + +def assert_storage_data_exact(actual, expected, *, context): + """Assert every tensor-valued data, scale, and amax metadata field exactly.""" + assert_nested_state_exact( + snapshot_storage_tensor_metadata(actual), + snapshot_storage_tensor_metadata(expected), + path=context, + ) + + +def assert_hybrid_tensor_exact(actual, expected, *, context): + """Compare both hybrid directions, including metadata and dequantization.""" + for direction in ("rowwise", "columnwise"): + actual_storage = getattr(actual, f"{direction}_sub_storage") + expected_storage = getattr(expected, f"{direction}_sub_storage") + assert_storage_data_exact( + actual_storage, + expected_storage, + context=f"{context} {direction}", + ) + if expected_storage is None: + continue + try: + expected_dequantized = expected_storage.dequantize() + except NotImplementedError: + continue + torch.testing.assert_close( + actual_storage.dequantize(), + expected_dequantized, + rtol=0.0, + atol=0.0, + msg=f"{context} {direction} dequantized value differs", + ) + + +def make_role_aware_quantizer(factory, role): + """Construct a quantizer with the standard Block-FP8 GEMM geometry.""" + quantizer = factory() + if isinstance(quantizer, te.Float8BlockQuantizer): + is_weight = ( + role is not None + and role.module_type in _LINEAR_MODULE_TYPES + and role.tensor_type == "weight" + ) + quantizer.block_scaling_dim = 2 if is_weight else 1 + return quantizer + + +def hybrid_custom_recipe(row_factory, col_factory, grad_factory=None): + """Build a CustomRecipe with hybrid forward and configurable grad quantizers.""" + if grad_factory is None: + grad_factory = col_factory + + def qfactory(role): + is_linear = _is_linear_role(role) + if is_linear and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=make_role_aware_quantizer(row_factory, role), + columnwise_quantizer=make_role_aware_quantizer(col_factory, role), + ) + if is_linear and role.tensor_type in _GRAD_TENSOR_TYPES: + return make_role_aware_quantizer(grad_factory, role) + return make_role_aware_quantizer(row_factory, role) + + return recipe.CustomRecipe(qfactory=qfactory) + + +def hybrid_fp8_current_qfactory(role): + """FP8 current scaling in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=current_scaling_quantizer_factory(role), + columnwise_quantizer=current_scaling_quantizer_factory(role), + ) + return current_scaling_quantizer_factory(role) + + +def hybrid_fp8_current_e5m2_grads_qfactory(role): + """FP8 current-scaling hybrid with E5M2 for both grad boundary roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_fp8_current(), + columnwise_quantizer=_make_fp8_current(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return _make_fp8_current(fp8_dtype=te.DType.kFloat8E5M2) + return _make_fp8_current() + + +def hybrid_mxfp8_qfactory(role): + """MXFP8 in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=mxfp8_quantizer_factory(role), + ) + return mxfp8_quantizer_factory(role) + + +def hybrid_float8_block_qfactory(role): + """Float8 block scaling in both hybrid directions for forward roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=float8_block_scaling_quantizer_factory(role), + columnwise_quantizer=float8_block_scaling_quantizer_factory(role), + ) + return float8_block_scaling_quantizer_factory(role) + + +def hybrid_block_fp8_e4m3_qfactory(role): + """Hybrid E4M3 Block-FP8 with role-aware 1D/2D block geometry.""" + is_linear = _is_linear_role(role) + is_weight = is_linear and role.tensor_type == "weight" + block_scaling_dim = 2 if is_weight else 1 + + def make_quantizer(): + return te.Float8BlockQuantizer( + fp8_dtype=te.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=block_scaling_dim, + ) + + if is_linear and role.tensor_type in _GRAD_TENSOR_TYPES: + return make_quantizer() + return te.HybridQuantizer( + rowwise_quantizer=make_quantizer(), + columnwise_quantizer=make_quantizer(), + ) + + +def hybrid_mixed_mxfp8_fp8_qfactory(role): + """MXFP8 rowwise plus FP8 current-scaling columnwise.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=current_scaling_quantizer_factory(role), + ) + return current_scaling_quantizer_factory(role) + + +def hybrid_fp8_current_identity_qfactory(role): + """FP8 current-scaling forward plus Identity backward.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=current_scaling_quantizer_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return te.IdentityQuantizer() + return current_scaling_quantizer_factory(role) + + +def hybrid_mxfp8_identity_qfactory(role): + """MXFP8 forward plus Identity backward.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=te.IdentityQuantizer(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return te.IdentityQuantizer() + return mxfp8_quantizer_factory(role) + + +def identity_qfactory(role): # pylint: disable=unused-argument + """High-precision passthrough for every quantizer slot.""" + return te.IdentityQuantizer() + + +def hybrid_nvfp4_qfactory(role): + """NVFP4 in both hybrid directions for forward tensor roles.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + return nvfp4_quantizer_factory(role) + + +def hybrid_tp_mxfp8_nvfp4_qfactory(role): + """TP/SP MXFP8 rowwise plus NVFP4 columnwise, including boundary roles.""" + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return nvfp4_quantizer_factory(role) + return te.HybridQuantizer( + rowwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + + +def hybrid_fp8_mxfp8_qfactory(role): + """FP8 current-scaling rowwise plus MXFP8 columnwise for CPU-offload tests.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_fp8_current(), + columnwise_quantizer=_make_mxfp8(), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return _make_mxfp8(fp8_dtype=te.DType.kFloat8E5M2) + return _make_fp8_current() + + +def hybrid_mxfp8_nvfp4_qfactory(role): + """MXFP8 rowwise plus NVFP4 columnwise for CPU-offload tests.""" + if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: + return te.HybridQuantizer( + rowwise_quantizer=_make_mxfp8(), + columnwise_quantizer=nvfp4_quantizer_factory(role), + ) + if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: + return nvfp4_quantizer_factory(role) + return _make_mxfp8() diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 773ced21e7a..3de06f38475 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -19,8 +19,9 @@ from transformer_engine.pytorch.fp8 import FP8GlobalStateManager import transformer_engine.pytorch as te from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - nvfp4_quantizer_factory, +from hybrid_quantization_utils import ( + hybrid_fp8_mxfp8_qfactory, + hybrid_mxfp8_nvfp4_qfactory, ) from utils import ModelConfig, recipe_id, skip_unsupported_backward_override @@ -31,43 +32,6 @@ nvfp4_available, _ = FP8GlobalStateManager.is_nvfp4_available() -def _hybrid_fp8_mxfp8_qfactory(role): - """Hybrid CustomRecipe factory: FP8 current-scaling rowwise + MXFP8 columnwise. - - Forward roles -> HybridQuantizer; backward roles -> plain MXFP8 so - dgrad/wgrad operand pairs share a single scaling mode. Catch-all - returns plain FP8 for non-linear roles used by layernorm_linear, - layernorm_mlp, multihead_attention, and transformer_layer. - """ - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return te.HybridQuantizer( - rowwise_quantizer=te.Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda"), - columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E5M2) - return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") - - -def _hybrid_mxfp8_nvfp4_qfactory(role): - """Hybrid CustomRecipe factory: MXFP8 rowwise + NVFP4 columnwise. - - Mirrors ``mxfp8_fwd_nvfp4_bwd_quantizer_factory`` from - ``custom_recipes/quantization_factory_zoo.py``. grad_output uses plain NVFP4 - (both directions) so wgrad's columnwise operand matches. - """ - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return te.HybridQuantizer( - rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3), - columnwise_quantizer=nvfp4_quantizer_factory(role), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return nvfp4_quantizer_factory(role) - return te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) - - def nvfp4_row_scaled(): nvfp4_recipe = recipe.NVFP4BlockScaling( disable_rht=True, @@ -106,9 +70,9 @@ def nvfp4_4over6(): quantization_recipes.append(nvfp4_4over6()) quantization_recipes.append(nvfp4_row_scaled()) if fp8_available and mxfp8_available: - quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_fp8_mxfp8_qfactory)) + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_fp8_mxfp8_qfactory)) if mxfp8_available and nvfp4_available: - quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_mxfp8_nvfp4_qfactory)) + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_mxfp8_nvfp4_qfactory)) model_config = { diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index b942fc8ab09..19271a5f42a 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -11,10 +11,10 @@ import torch import transformer_engine.pytorch as te -import transformer_engine_torch as tex from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - nvfp4_quantizer_factory, +from hybrid_quantization_utils import ( + hybrid_fp8_mxfp8_qfactory, + hybrid_mxfp8_nvfp4_qfactory, ) from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported @@ -26,52 +26,13 @@ nvfp4_available = te.is_nvfp4_available() -def _hybrid_fp8_mxfp8_qfactory(role): - """Hybrid CustomRecipe factory: FP8 current-scaling rowwise + MXFP8 columnwise. - - Forward roles get a HybridQuantizer; backward/grad roles get a plain - MXFP8 quantizer so dgrad/wgrad GEMMs see a single scaling mode per - operand pair. Catch-all returns plain FP8 for non-linear roles - (layernorm_linear, layernorm_mlp, multihead_attention, transformer_layer). - """ - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return te.HybridQuantizer( - rowwise_quantizer=te.Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - columnwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2) - return te.Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - - -def _hybrid_mxfp8_nvfp4_qfactory(role): - """Hybrid CustomRecipe factory: MXFP8 rowwise + NVFP4 columnwise. - - Mirrors the ``mxfp8_fwd_nvfp4_bwd_quantizer_factory`` headline recipe - from ``custom_recipes/quantization_factory_zoo.py``. grad_output uses plain - NVFP4 (both directions) so wgrad's columnwise operand matches. - """ - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return te.HybridQuantizer( - rowwise_quantizer=te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - columnwise_quantizer=nvfp4_quantizer_factory(role), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return nvfp4_quantizer_factory(role) - return te.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) - - quantization_recipes: Optional[recipe.Recipe] = [None] if fp8_available: quantization_recipes.extend((recipe.Float8CurrentScaling(), recipe.DelayedScaling())) if fp8_available and mxfp8_available: - quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_fp8_mxfp8_qfactory)) + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_fp8_mxfp8_qfactory)) if mxfp8_available and nvfp4_available: - quantization_recipes.append(recipe.CustomRecipe(qfactory=_hybrid_mxfp8_nvfp4_qfactory)) + quantization_recipes.append(recipe.CustomRecipe(qfactory=hybrid_mxfp8_nvfp4_qfactory)) hybrid_quantization_recipes = [ item for item in quantization_recipes if item is not None and item.custom() diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 27dbf4fb930..c03101a9857 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -11,6 +11,23 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex +from hybrid_quantization_utils import ( + as_data_tensor_tuple as _as_data_tensor_tuple, + assert_hybrid_tensor_exact as _assert_hybrid_tensor_exact, + assert_nested_state_exact as _assert_nested_state_exact, + assert_storage_data_exact as _assert_storage_data_exact, + fp8_e4m3_factory as _fp8_row_factory, + fp8_e5m2_factory as _fp8_grad_factory, + hybrid_block_fp8_e4m3_qfactory as _hybrid_block_fp8_qfactory, + hybrid_custom_recipe as _hybrid_custom_recipe, + hybrid_fp8_current_e5m2_grads_qfactory as _hybrid_fp8_current_qfactory, + make_fp8_quantizer as _make_fp8_quantizer, + make_hybrid_quantizer_fp8_row_fp4_col as _make_hybrid_quantizer_fp8_row_fp4_col, + make_nvfp4_quantizer as _make_nvfp4_quantizer, + make_role_aware_quantizer as _make_role_aware_quantizer, + mxfp8_e4m3_factory as _mxfp8_factory, + snapshot_storage_tensor_metadata as _snapshot_storage_tensor_metadata, +) from transformer_engine.common import recipe from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( nvfp4_quantizer_factory, @@ -47,6 +64,8 @@ ) from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported +_fp8_col_factory = _fp8_row_factory + fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) @@ -89,31 +108,6 @@ ) -def _make_fp8_quantizer(*, rowwise=True, columnwise=True): - return Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - device="cuda", - rowwise=rowwise, - columnwise=columnwise, - ) - - -def _make_nvfp4_quantizer(*, rowwise=True, columnwise=True): - return NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, - rowwise=rowwise, - columnwise=columnwise, - ) - - -def _make_hybrid_quantizer_fp8_row_fp4_col(): - """FP8 rowwise + NVFP4 columnwise.""" - return HybridQuantizer( - rowwise_quantizer=_make_fp8_quantizer(), - columnwise_quantizer=_make_nvfp4_quantizer(), - ) - - def _make_hybrid_quantizer_fp4_row_fp8_col(): """NVFP4 rowwise + FP8 columnwise (reversed direction).""" return HybridQuantizer( @@ -122,60 +116,6 @@ def _make_hybrid_quantizer_fp4_row_fp8_col(): ) -def _as_data_tensor_tuple(storage): - """Return a storage's raw buffers as a tuple without copying them.""" - if storage is None: - return () - tensors = storage.get_data_tensors() - return tensors if isinstance(tensors, tuple) else (tensors,) - - -def _snapshot_storage_tensor_metadata(storage, *, clone=False): - """Capture every tensor-valued concrete-storage metadata field without mutation.""" - if storage is None: - return None - metadata = storage.get_metadata() - tensor_metadata = {} - for name, value in metadata.items(): - # Optional tensor fields are represented by ``None``. Including them - # makes missing rowwise/columnwise data, scales, and amax explicit. - if isinstance(value, torch.Tensor) or value is None: - snapshot_value = value.detach().clone() if clone and value is not None else value - if ( - value is not None - and getattr(storage, "_is_2D_scaled", False) - and name in ("rowwise_scale_inv", "columnwise_scale_inv") - ): - # Float8Block 2D scales pad one tile dimension to a multiple - # of four. Kernels never initialize or consume that padding, - # and FSDP deliberately strips then zero-repads it. Canonicalize - # only those non-logical elements while still comparing the - # complete tensor field, shape, presence, and every valid scale. - snapshot_value = value.detach().clone() - m, n = storage._fsdp_logical_mn() - block_len = storage._FSDP_BLOCK_LEN - m_tiles = (m + block_len - 1) // block_len - n_tiles = (n + block_len - 1) // block_len - if name == "rowwise_scale_inv": - snapshot_value[:, n_tiles:] = 0 - else: - snapshot_value[:, m_tiles:] = 0 - tensor_metadata[name] = snapshot_value - return { - "storage_type": type(storage).__name__, - "tensor_metadata": tensor_metadata, - } - - -def _assert_storage_data_exact(actual, expected, *, context): - """Assert every tensor-valued data/scale/amax metadata field exactly.""" - _assert_nested_state_exact( - _snapshot_storage_tensor_metadata(actual), - _snapshot_storage_tensor_metadata(expected), - path=context, - ) - - def _clone_nested_state(value): """Clone tensors in nested checkpoint state so later steps cannot mutate snapshots.""" if isinstance(value, torch.Tensor): @@ -189,61 +129,6 @@ def _clone_nested_state(value): return value -def _assert_nested_state_exact(actual, expected, *, path="state"): - """Recursively compare checkpoint state, using zero tolerance for tensors.""" - if isinstance(expected, torch.Tensor): - assert isinstance(actual, torch.Tensor), f"{path}: expected Tensor, got {type(actual)}" - torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0, msg=path) - return - if isinstance(expected, dict): - assert isinstance(actual, dict), f"{path}: expected dict, got {type(actual)}" - assert actual.keys() == expected.keys(), f"{path}: dictionary keys differ" - for key in expected: - _assert_nested_state_exact(actual[key], expected[key], path=f"{path}.{key}") - return - if isinstance(expected, (list, tuple)): - assert isinstance( - actual, type(expected) - ), f"{path}: expected {type(expected).__name__}, got {type(actual).__name__}" - assert len(actual) == len(expected), f"{path}: sequence lengths differ" - for index, (actual_item, expected_item) in enumerate(zip(actual, expected)): - _assert_nested_state_exact( - actual_item, - expected_item, - path=f"{path}[{index}]", - ) - return - assert actual == expected, f"{path}: {actual!r} != {expected!r}" - - -def _assert_hybrid_tensor_exact(actual, expected, *, context): - """Compare both hybrid directions, including all metadata and native dequantization.""" - for direction in ("rowwise", "columnwise"): - actual_storage = getattr(actual, f"{direction}_sub_storage") - expected_storage = getattr(expected, f"{direction}_sub_storage") - _assert_storage_data_exact( - actual_storage, - expected_storage, - context=f"{context} {direction}", - ) - if expected_storage is None: - continue - try: - expected_dequantized = expected_storage.dequantize() - except NotImplementedError: - # Some direction-only formats (notably NVFP4 columnwise) expose - # bytes/scales for GEMM but intentionally do not dequantize. - continue - actual_dequantized = actual_storage.dequantize() - torch.testing.assert_close( - actual_dequantized, - expected_dequantized, - rtol=0.0, - atol=0.0, - msg=f"{context} {direction} dequantized value differs", - ) - - def _snapshot_model_parameters(model): """Capture normal values and all tensor metadata in both hybrid directions.""" snapshot = {} @@ -3194,22 +3079,6 @@ def _make_block_quantizer_e5m2(*, rowwise=True, columnwise=True): } -def _make_role_aware_quantizer(factory, role): - """Construct a quantizer with the standard Block-FP8 GEMM geometry.""" - quantizer = factory() - if isinstance(quantizer, Float8BlockQuantizer): - is_weight = ( - role is not None - and role.module_type in ("linear", "grouped_linear") - and role.tensor_type == "weight" - ) - # Match Float8BlockScaling: 2D blocks for weights, 1D blocks for - # activations and gradients. A GEMM cannot consume 2D-scaled inputs - # on both sides. - quantizer.block_scaling_dim = 2 if is_weight else 1 - return quantizer - - def _build_cross_format_params(): """Build parametrize list for all stateless cross-format hybrid combos.""" combos = [ @@ -4275,38 +4144,6 @@ def make_quantizer(): # Quantized Parameters (quantized_model_init) tests for hybrid quantization # =========================================================================== - -def _hybrid_custom_recipe(row_factory, col_factory, grad_factory=None): - """Build a CustomRecipe where forward roles use HybridQuantizer and - backward roles use a plain quantizer (or hybrid if grad_factory builds one). - - Parameters - ---------- - row_factory : callable() -> Quantizer - Creates the rowwise sub-quantizer for forward roles. - col_factory : callable() -> Quantizer - Creates the columnwise sub-quantizer for forward roles. - grad_factory : callable() -> Quantizer, optional - Creates the quantizer for grad_output/grad_input roles. - If None, uses col_factory (matching columnwise format for wgrad compatibility). - """ - if grad_factory is None: - grad_factory = col_factory - - def qfactory(role): - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=_make_role_aware_quantizer(row_factory, role), - columnwise_quantizer=_make_role_aware_quantizer(col_factory, role), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return _make_role_aware_quantizer(grad_factory, role) - return _make_role_aware_quantizer(row_factory, role) - - return recipe.CustomRecipe(qfactory=qfactory) - - # --------------------------------------------------------------------------- # 1. quantized_model_init: model creation and parameter type verification # --------------------------------------------------------------------------- @@ -5862,21 +5699,6 @@ def test_mixed_format_sub_storage_types(self): # --------------------------------------------------------------------------- -def _hybrid_fp8_current_qfactory(role): - """Hybrid FP8 current scaling (E4M3 both dirs, E5M2 for grad).""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight", "output"): - return HybridQuantizer( - rowwise_quantizer=Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), - columnwise_quantizer=Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, device="cuda" - ), - ) - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - - def _hybrid_mxfp8_qfactory(role): """Hybrid MXFP8 (E4M3 both dirs).""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -5888,34 +5710,6 @@ def _hybrid_mxfp8_qfactory(role): ) -def _hybrid_block_fp8_qfactory(role): - """Hybrid block FP8 (E4M3 both dirs).""" - is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - is_weight = is_linear and role.tensor_type == "weight" - dim = 2 if is_weight else 1 - if is_linear and role.tensor_type in ("grad_output", "grad_input"): - return Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, - columnwise=True, - block_scaling_dim=dim, - ) - return HybridQuantizer( - rowwise_quantizer=Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, - columnwise=True, - block_scaling_dim=dim, - ), - columnwise_quantizer=Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, - columnwise=True, - block_scaling_dim=dim, - ), - ) - - def _hybrid_nvfp4_qfactory(role): """Hybrid NVFP4 (E2M1 both dirs, base role behavior).""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -6355,1460 +6149,7 @@ def test_checkpoint_resume_training(self): # --------------------------------------------------------------------------- -# 11. FSDP2 prerequisites: __torch_dispatch__ ops that FSDP2 relies on -# --------------------------------------------------------------------------- - -aten = torch.ops.aten - - -def _make_hybrid_param_for_dispatch( - row_factory, col_factory, grad_factory=None, in_features=256, out_features=256 -): - """Create a HybridQuantizedTensor weight via quantized_model_init for dispatch tests.""" - hybrid_recipe = _hybrid_custom_recipe(row_factory, col_factory, grad_factory) - with quantized_model_init(enabled=True, recipe=hybrid_recipe): - model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() - return model.weight - - -def _fp8_row_factory(): - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - - -def _fp8_col_factory(): - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - - -def _fp8_grad_factory(): - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") - - -def _mxfp8_factory(): - return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) - - -_dispatch_configs = [ - pytest.param( - "fp8_fp8", - id="same-format-fp8", - marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, - ), -] -if mxfp8_available: - _dispatch_configs.append(pytest.param("mxfp8_mxfp8", id="same-format-mxfp8")) - - -def _get_dispatch_hybrid_param(config_name): - """Return a HybridQuantizedTensor weight for the given config.""" - if config_name == "fp8_fp8": - return _make_hybrid_param_for_dispatch( - _fp8_row_factory, - _fp8_col_factory, - _fp8_grad_factory, - ) - elif config_name == "mxfp8_mxfp8": - return _make_hybrid_param_for_dispatch( - _mxfp8_factory, - _mxfp8_factory, - grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), - ) - else: - raise ValueError(f"Unknown config: {config_name}") - - -@requires_fp8 -class TestFloat8TransposeOnlySplit: - """Regression coverage for columnwise-only Float8 split metadata. - - A columnwise-only per-tensor Float8 sub-storage may have ``_data=None`` and - store its bytes in ``_transpose`` with physical shape ``[K, M]``. Splitting - that tensor must still produce pieces whose wrapper shape is the logical - row-major shape ``[M_i, K]``; otherwise HybridQuantizedTensor uses the - transposed shape when rowwise storage is absent. - """ - - @staticmethod - def _make_transpose_only_float8_tensor(shape=(12, 16)): - m, k = shape - data_transpose = torch.empty((k, m), dtype=torch.uint8, device="cuda") - return Float8Tensor( - shape=shape, - dtype=torch.bfloat16, - data=None, - data_transpose=data_transpose, - fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), - fp8_dtype=tex.DType.kFloat8E4M3, - requires_grad=False, - device="cuda", - ) - - @pytest.mark.parametrize( - "split_size,dim,expected_shapes,expected_transpose_shapes", - [ - (5, 0, [(5, 16), (5, 16), (2, 16)], [(16, 5), (16, 5), (16, 2)]), - (6, 1, [(12, 6), (12, 6), (12, 4)], [(6, 12), (6, 12), (4, 12)]), - ], - ) - def test_float8_split_uses_logical_shape_for_transpose_only_storage( - self, split_size, dim, expected_shapes, expected_transpose_shapes - ): - tensor = self._make_transpose_only_float8_tensor() - - pieces = torch.split(tensor, split_size, dim=dim) - - assert [tuple(piece.shape) for piece in pieces] == expected_shapes - assert [tuple(piece._transpose.shape) for piece in pieces] == expected_transpose_shapes - assert all(piece._data is None for piece in pieces) - assert all(piece._transpose_invalid is False for piece in pieces) - - def test_float8_view_preserves_transpose_only_storage(self): - tensor = self._make_transpose_only_float8_tensor() - - viewed = tensor.view(12, 16) - - assert tuple(viewed.shape) == (12, 16) - assert viewed._data is None - assert viewed._transpose is not None - assert tuple(viewed._transpose.shape) == (16, 12) - assert viewed._transpose_invalid is False - - def test_float8_aten_view_preserves_transpose_only_storage(self): - tensor = self._make_transpose_only_float8_tensor() - - viewed = torch.ops.aten.view.default(tensor, [12, 16]) - - assert tuple(viewed.shape) == (12, 16) - assert viewed._data is None - assert viewed._transpose is not None - assert tuple(viewed._transpose.shape) == (16, 12) - assert viewed._transpose_invalid is False - - def test_float8_storage_view_preserves_transpose_only_storage(self): - data_transpose = torch.empty((16, 12), dtype=torch.uint8, device="cuda") - storage = Float8TensorStorage( - data=None, - data_transpose=data_transpose, - fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), - fp8_dtype=tex.DType.kFloat8E4M3, - fake_dtype=torch.bfloat16, - ) - - viewed = storage.view(torch.Size((12, 16))) - - assert viewed._data is None - assert viewed._transpose is not None - assert tuple(viewed._transpose.shape) == (16, 12) - assert viewed._transpose_invalid is False - - def test_float8_shape_changing_view_raises_for_transpose_only_storage(self): - tensor = self._make_transpose_only_float8_tensor() - - with pytest.raises(NotImplementedError, match="columnwise-only data"): - tensor.view(6, 32) - - def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self): - columnwise = self._make_transpose_only_float8_tensor() - quantizer = HybridQuantizer( - rowwise_quantizer=_make_fp8_quantizer(), - columnwise_quantizer=_make_fp8_quantizer(), - ) - quantizer.set_usage(rowwise=False, columnwise=True) - hybrid = HybridQuantizedTensor( - shape=columnwise.shape, - dtype=columnwise.dtype, - rowwise_storage=None, - columnwise_storage=columnwise, - quantizer=quantizer, - device="cuda", - ) - - pieces = torch.split(hybrid, 5, dim=0) - - assert [tuple(piece.shape) for piece in pieces] == [ - (5, 16), - (5, 16), - (2, 16), - ] - assert all(piece.rowwise_sub_storage is None for piece in pieces) - assert [tuple(piece.columnwise_sub_storage.shape) for piece in pieces] == [ - (5, 16), - (5, 16), - (2, 16), - ] - assert [tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces] == [ - (16, 5), - (16, 5), - (16, 2), - ] - - @staticmethod - def _make_valid_transpose_only_float8_tensor(shape=(12, 16)): - """Build transpose-only storage with known, numerically valid FP8 bytes.""" - source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") - rowwise_quantizer = Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, - device="cuda", - rowwise=True, - columnwise=False, - ) - rowwise = rowwise_quantizer(source) - columnwise_quantizer = Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, - device="cuda", - rowwise=False, - columnwise=True, - ) - tensor = Float8Tensor( - shape=shape, - dtype=torch.bfloat16, - data=None, - data_transpose=rowwise._data.movedim(-1, 0).contiguous(), - fp8_scale_inv=rowwise._scale_inv.detach().clone(), - fp8_dtype=tex.DType.kFloat8E4M3, - quantizer=columnwise_quantizer, - requires_grad=False, - device="cuda", - ) - return tensor, rowwise.dequantize() - - @pytest.mark.parametrize("weights_only", (False, True)) - def test_serialization_preserves_transpose_only_payload(self, weights_only): - tensor, expected = self._make_valid_transpose_only_float8_tensor() - buffer = io.BytesIO() - - torch.save(tensor, buffer) - buffer.seek(0) - loaded = torch.load(buffer, weights_only=weights_only) - - assert loaded._data is None - assert loaded._transpose_invalid is False - assert torch.equal(loaded._transpose, tensor._transpose) - torch.testing.assert_close(loaded.dequantize(), expected, rtol=0.0, atol=0.0) - - def test_dequantize_from_transpose_only_payload(self): - tensor, expected = self._make_valid_transpose_only_float8_tensor() - - torch.testing.assert_close(tensor.dequantize(), expected, rtol=0.0, atol=0.0) - - def test_slice_and_select_preserve_transpose_only_payload(self): - tensor, expected = self._make_valid_transpose_only_float8_tensor() - - sliced = tensor[2:8] - selected = torch.select(tensor, 1, 3) - - assert isinstance(sliced, Float8Tensor) - assert sliced._data is None - assert sliced._transpose.shape == torch.Size((16, 6)) - torch.testing.assert_close(sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0) - assert isinstance(selected, Float8Tensor) - assert selected._data is None - assert selected._transpose.shape == torch.Size((12,)) - torch.testing.assert_close(selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0) - - def test_as_strided_row_shard_preserves_transpose_only_payload(self): - tensor, expected = self._make_valid_transpose_only_float8_tensor() - - shard = torch.as_strided(tensor, (5, 16), (16, 1), 16) - - assert isinstance(shard, Float8Tensor) - assert shard._data is None - assert shard._transpose.shape == torch.Size((16, 5)) - torch.testing.assert_close(shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0) - - def test_as_strided_falls_back_for_nonrepresentable_layout(self): - tensor, expected = self._make_valid_transpose_only_float8_tensor() - - output = torch.as_strided(tensor, (16, 12), (1, 16), 0) - - assert type(output) is torch.Tensor - reference = torch.as_strided(expected, (16, 12), (1, 16), 0) - torch.testing.assert_close(output, reference, rtol=0.0, atol=0.0) - - def test_new_zeros_preserves_transpose_only_layout(self): - tensor, _ = self._make_valid_transpose_only_float8_tensor() - - output = tensor.new_zeros((5, 16)) - - assert isinstance(output, Float8Tensor) - assert output.shape == torch.Size((5, 16)) - assert output._data is None - assert output._transpose.shape == torch.Size((16, 5)) - assert output._transpose.dtype == torch.uint8 - assert torch.count_nonzero(output._transpose).item() == 0 - torch.testing.assert_close( - output.dequantize(), - torch.zeros((5, 16), dtype=output.dtype, device=output.device), - rtol=0.0, - atol=0.0, - ) - - -class TestHybridNewZeros: - """Public new_zeros semantics independent of FSDP-specific allocation.""" - - @staticmethod - def _make_identity_hybrid(): - quantizer = HybridQuantizer( - rowwise_quantizer=IdentityQuantizer(), - columnwise_quantizer=IdentityQuantizer(), - ) - source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) - return quantizer, source - - @pytest.mark.parametrize( - "keep_rowwise,keep_columnwise", - [(True, True), (True, False), (False, True)], - ) - def test_initializes_every_present_direction_and_preserves_kwargs( - self, - keep_rowwise, - keep_columnwise, - ): - quantizer, source = self._make_identity_hybrid() - if not keep_rowwise: - source.update_usage(rowwise_usage=False) - if not keep_columnwise: - source.update_usage(columnwise_usage=False) - - # Direction selection must come from source storage, not mutable parent - # usage flags, and new_zeros must not change those flags. - quantizer.set_usage(rowwise=False, columnwise=False) - result = torch.ops.aten.new_zeros.default( - source, - [3, 5], - dtype=torch.float32, - device=source.device, - ) - - assert isinstance(result, HybridQuantizedTensor) - assert result.shape == torch.Size((3, 5)) - assert result.dtype == torch.float32 - assert result.device == source.device - assert (result.rowwise_sub_storage is not None) is keep_rowwise - assert (result.columnwise_sub_storage is not None) is keep_columnwise - assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} - assert result._quantizer is not quantizer - assert result._quantizer.get_usages() == { - "rowwise": keep_rowwise, - "columnwise": keep_columnwise, - } - - for sub_storage in ( - result.rowwise_sub_storage, - result.columnwise_sub_storage, - ): - if sub_storage is not None: - torch.testing.assert_close( - sub_storage.dequantize(), - torch.zeros((3, 5), dtype=torch.float32, device=source.device), - rtol=0.0, - atol=0.0, - ) - - # A plain-source copy routes through the result's parent quantizer. It - # must update every allocated direction even though the source parent - # had both of its mutable usage flags disabled before new_zeros. - plain_source = torch.full((3, 5), 7.0, dtype=torch.float32, device=source.device) - result.copy_(plain_source) - for sub_storage in ( - result.rowwise_sub_storage, - result.columnwise_sub_storage, - ): - if sub_storage is not None: - torch.testing.assert_close( - sub_storage.dequantize(), plain_source, rtol=0.0, atol=0.0 - ) - assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} - - def test_identity_substorages_allow_integer_dtype(self): - _, source = self._make_identity_hybrid() - result = source.new_zeros((2, 3), dtype=torch.int32) - - assert result.dtype == torch.int32 - for sub_storage in ( - result.rowwise_sub_storage, - result.columnwise_sub_storage, - ): - assert sub_storage.dequantize().dtype == torch.int32 - assert torch.count_nonzero(sub_storage.dequantize()).item() == 0 - - @requires_fp8 - @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 - @pytest.mark.parametrize("unsupported_dtype", (torch.int32, torch.float64, torch.bool)) - def test_non_identity_substorage_rejects_unsupported_dtype(self, unsupported_dtype): - quantizer = HybridQuantizer( - rowwise_quantizer=_make_fp8_quantizer(), - columnwise_quantizer=_make_fp8_quantizer(), - ) - source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) - - with pytest.raises(TypeError, match="new_zeros only supports"): - source.new_zeros((2, 3), dtype=unsupported_dtype) - - def test_does_not_invoke_live_quantizers_or_consume_rng(self, monkeypatch): - quantizer, source = self._make_identity_hybrid() - cpu_rng_before = torch.get_rng_state().clone() - cuda_rng_before = torch.cuda.get_rng_state(source.device).clone() - - def fail_live_make_empty(*args, **kwargs): - raise AssertionError("new_zeros invoked a live quantizer") - - monkeypatch.setattr(quantizer, "make_empty", fail_live_make_empty) - monkeypatch.setattr( - source.rowwise_sub_storage._quantizer, - "make_empty", - fail_live_make_empty, - ) - monkeypatch.setattr( - source.columnwise_sub_storage._quantizer, - "make_empty", - fail_live_make_empty, - ) - - result = source.new_zeros((2, 6)) - - assert isinstance(result, HybridQuantizedTensor) - assert torch.equal(torch.get_rng_state(), cpu_rng_before) - assert torch.equal(torch.cuda.get_rng_state(source.device), cuda_rng_before) - - def test_rejects_empty_hybrid(self): - _, source = self._make_identity_hybrid() - source.update_usage(rowwise_usage=False, columnwise_usage=False) - - with pytest.raises(RuntimeError, match="at least one present sub-storage"): - source.new_zeros((2, 6)) - - @requires_fp8_and_nvfp4 - def test_initializes_nvfp4_data_scale_and_amax_buffers(self): - quantizer = _make_hybrid_quantizer_fp8_row_fp4_col() - source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) - - result = source.new_zeros(source.shape) - - assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) - assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) - torch.testing.assert_close( - result.dequantize(), - torch.zeros_like(result.dequantize()), - rtol=0.0, - atol=0.0, - ) - for sub_storage in ( - result.rowwise_sub_storage, - result.columnwise_sub_storage, - ): - buffers, storage = sub_storage.prepare_for_saving() - try: - assert all( - buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers - ) - finally: - assert storage.restore_from_saved(buffers) == [] - - @pytest.mark.skipif( - not fp8_block_scaling_available, - reason=f"Float8Blockwise: {reason_for_no_fp8_block_scaling}", - ) - def test_initializes_float8_block_data_and_scale_buffers(self): - quantizer = HybridQuantizer( - rowwise_quantizer=Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, - columnwise=True, - block_scaling_dim=2, - ), - columnwise_quantizer=Float8BlockQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, - columnwise=True, - block_scaling_dim=2, - ), - ) - source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) - - result = source.new_zeros(source.shape) - - assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) - assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) - torch.testing.assert_close( - result.dequantize(), - torch.zeros_like(result.dequantize()), - rtol=0.0, - atol=0.0, - ) - for sub_storage in ( - result.rowwise_sub_storage, - result.columnwise_sub_storage, - ): - buffers, storage = sub_storage.prepare_for_saving() - try: - assert all( - buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers - ) - finally: - assert storage.restore_from_saved(buffers) == [] - - -@requires_fp8 -class TestHybridTorchDispatchFSDP2Ops: - """Test aten ops that FSDP2 relies on to preserve the HybridQuantizedTensor type. - - Each op is called directly via torch.ops.aten and the result is verified to - still be HybridQuantizedTensor with valid sub-storages. - """ - - @pytest.fixture(params=_dispatch_configs) - def hybrid_param(self, request): - torch.manual_seed(42) - return _get_dispatch_hybrid_param(request.param) - - def test_split_preserves_hybrid_type(self, hybrid_param): - """torch.split must return a list of HybridQuantizedTensor pieces.""" - dim0 = hybrid_param.shape[0] - chunk_size = dim0 // 2 - pieces = torch.split(hybrid_param, chunk_size, dim=0) - expected_rowwise = torch.split(hybrid_param.rowwise_sub_storage, chunk_size, dim=0) - expected_columnwise = torch.split(hybrid_param.columnwise_sub_storage, chunk_size, dim=0) - assert len(pieces) == len(expected_rowwise) == len(expected_columnwise) - - assert len(pieces) >= 2 - for piece in pieces: - assert isinstance( - piece, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(piece).__name__}" - assert piece.rowwise_sub_storage is not None - assert piece.columnwise_sub_storage is not None - for index, (piece, expected_row, expected_column) in enumerate( - zip(pieces, expected_rowwise, expected_columnwise) - ): - _assert_storage_data_exact( - piece.rowwise_sub_storage, - expected_row, - context=f"split {index} rowwise", - ) - _assert_storage_data_exact( - piece.columnwise_sub_storage, - expected_column, - context=f"split {index} columnwise", - ) - - total_rows = sum(p.shape[0] for p in pieces) - assert total_rows == dim0 - - orig_deq = hybrid_param.dequantize() - reassembled = torch.cat([p.dequantize() for p in pieces], dim=0) - torch.testing.assert_close(orig_deq, reassembled, rtol=0.0, atol=0.0) - - def test_split_sub_storage_types_preserved(self, hybrid_param): - """After split, sub-storage types must match the original.""" - orig_row_type = type(hybrid_param.rowwise_sub_storage) - orig_col_type = type(hybrid_param.columnwise_sub_storage) - - chunk_size = hybrid_param.shape[0] // 2 - pieces = torch.split(hybrid_param, chunk_size, dim=0) - for piece in pieces: - assert type(piece.rowwise_sub_storage) is orig_row_type - assert type(piece.columnwise_sub_storage) is orig_col_type - - @requires_mxfp8 - def test_split_rejects_mxfp8_high_precision_fallback(self): - """Fail before wrapping unquantizable MXFP8 shards for Hybrid FSDP2.""" - quantizer = HybridQuantizer( - rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - ) - tensor = quantizer(torch.randn(64, 64, dtype=torch.bfloat16, device="cuda")) - - with pytest.raises(NotImplementedError, match="local shape.*divisible by 32"): - torch.split(tensor, 16, dim=0) - - def test_view_preserves_hybrid_type(self, hybrid_param): - """view must return a HybridQuantizedTensor (used by FSDP2 reset_sharded_param).""" - shape_2d = hybrid_param.shape - result = aten.view.default(hybrid_param, list(shape_2d)) - assert isinstance( - result, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" - assert result.rowwise_sub_storage is not None - assert result.columnwise_sub_storage is not None - - def test_view_same_shape_preserves_hybrid(self, hybrid_param): - """view with same shape must return HybridQuantizedTensor.""" - shape_2d = list(hybrid_param.shape) - result = aten.view.default(hybrid_param, shape_2d) - assert isinstance( - result, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" - - def test_as_strided_noop_preserves_hybrid(self, hybrid_param): - """as_strided with matching shape/strides is a no-op that preserves type.""" - shape = tuple(hybrid_param.size()) - strides = (shape[-1], 1) - result = aten.as_strided.default(hybrid_param, list(shape), list(strides)) - assert isinstance( - result, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" - assert result.rowwise_sub_storage is not None - assert result.columnwise_sub_storage is not None - - def test_slice_noop_preserves_hybrid(self, hybrid_param): - """slice with full range is a no-op that preserves type.""" - result = aten.slice.Tensor(hybrid_param, 0, 0, hybrid_param.size(0)) - assert isinstance( - result, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" - assert result.rowwise_sub_storage is not None - - def test_copy_between_hybrid_tensors(self, hybrid_param): - """copy_ between compatible HybridQuantizedTensors copies quantized data directly.""" - src_deq = hybrid_param.dequantize().clone() - dst = hybrid_param._quantizer.make_empty( - shape=hybrid_param.shape, - dtype=hybrid_param.dtype, - device=hybrid_param.device, - ) - assert isinstance(dst, HybridQuantizedTensor) - - aten.copy_.default(dst, hybrid_param) - dst_deq = dst.dequantize() - torch.testing.assert_close(src_deq, dst_deq, rtol=0.0, atol=0.0) - _assert_hybrid_tensor_exact(dst, hybrid_param, context="hybrid copy_") - - def test_copy_between_mismatched_usages_raises_atomically(self, hybrid_param): - quantizer = hybrid_param._quantizer.copy() - src = quantizer.quantize(torch.ones_like(hybrid_param.dequantize())) - dst = quantizer.quantize(torch.zeros_like(hybrid_param.dequantize())) - src.update_usage(columnwise_usage=False) - dst_before = tuple( - None if tensor is None else tensor.clone() for tensor in _as_data_tensor_tuple(dst) - ) - - with pytest.raises( - NotImplementedError, - match="requires matching rowwise/columnwise usages", - ): - aten.copy_.default(dst, src) - - dst_after = _as_data_tensor_tuple(dst) - assert len(dst_after) == len(dst_before) - for before, after in zip(dst_before, dst_after): - if before is None: - assert after is None - else: - torch.testing.assert_close(after, before, rtol=0.0, atol=0.0) - - def test_copy_from_bf16_to_hybrid(self, hybrid_param): - """copy_ from BF16 into HybridQuantizedTensor triggers quantize_.""" - param = hybrid_param.detach() - bf16_data = torch.randn_like(param.dequantize()) - expected = param._quantizer.quantize(bf16_data) - aten.copy_.default(param, bf16_data) - assert isinstance(param, HybridQuantizedTensor) - _assert_hybrid_tensor_exact(param, expected, context="BF16 copy_") - torch.testing.assert_close(param.dequantize(), expected.dequantize(), rtol=0.0, atol=0.0) - - def test_new_zeros_returns_hybrid(self, hybrid_param): - """new_zeros returns initialized storage that remains FSDP-copyable.""" - new_shape = list(hybrid_param.shape) - result = aten.new_zeros.default(hybrid_param, new_shape) - - assert isinstance( - result, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" - assert result.shape == hybrid_param.shape - assert result.rowwise_sub_storage is not None - assert result.columnwise_sub_storage is not None - assert type(result.rowwise_sub_storage) is type(hybrid_param.rowwise_sub_storage) - assert type(result.columnwise_sub_storage) is type(hybrid_param.columnwise_sub_storage) - torch.testing.assert_close( - result.dequantize(), - torch.zeros_like(result.dequantize()), - rtol=0.0, - atol=0.0, - ) - - # FSDP2 overwrites the initialized destination via copy_ after gather. - aten.copy_.default(result, hybrid_param) - torch.testing.assert_close( - result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 - ) - _assert_hybrid_tensor_exact(result, hybrid_param, context="new_zeros then copy_") - - def test_empty_like_returns_hybrid(self, hybrid_param): - """empty_like must return a HybridQuantizedTensor.""" - result = aten.empty_like.default(hybrid_param) - assert isinstance( - result, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" - assert result.shape == hybrid_param.shape - assert result.rowwise_sub_storage is not None - - def test_clone_returns_hybrid(self, hybrid_param): - """clone must return an independent HybridQuantizedTensor with same data.""" - result = aten.clone.default(hybrid_param) - assert isinstance( - result, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" - assert result is not hybrid_param - torch.testing.assert_close( - result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 - ) - _assert_hybrid_tensor_exact(result, hybrid_param, context="clone") - - -# --------------------------------------------------------------------------- -# 12. FSDP2 prerequisites: fsdp_pre_all_gather protocol -# --------------------------------------------------------------------------- - - -def _make_fsdp_protocol_param(config_name): - """Create a HybridQuantizedTensor weight for FSDP protocol tests.""" - if config_name == "fp8_fp8": - r = _hybrid_custom_recipe(_fp8_row_factory, _fp8_col_factory, _fp8_grad_factory) - elif config_name == "mxfp8_fp8": - r = _hybrid_custom_recipe(_mxfp8_factory, _fp8_col_factory, _fp8_grad_factory) - elif config_name == "block_fp8": - r = recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) - else: - raise ValueError(f"Unknown config: {config_name}") - with quantized_model_init(enabled=True, recipe=r): - model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() - return model.weight - - -_fsdp_protocol_configs = [ - pytest.param( - "fp8_fp8", - id="same-format", - marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, - ) -] -if mxfp8_available: - _fsdp_protocol_configs.append(pytest.param("mxfp8_fp8", id="mixed-mxfp8-fp8")) -if fp8_block_scaling_available: - _fsdp_protocol_configs.append(pytest.param("block_fp8", id="same-format-block-fp8")) - - -@requires_fp8 -class TestHybridFsdpPreAllGatherProtocol: - """Test the fsdp_pre_all_gather method on HybridQuantizedTensor. - - These tests call the method directly (no actual all-gather communication) - to verify the protocol contract: returns (sharded_tensors, metadata) where - sharded_tensors is a tuple of plain torch.Tensor. - """ - - @pytest.fixture(params=_fsdp_protocol_configs) - def hybrid_param(self, request): - torch.manual_seed(42) - return _make_fsdp_protocol_param(request.param) - - def test_pre_all_gather_returns_tuple_pair(self, hybrid_param): - """fsdp_pre_all_gather returns (sharded_tensors, metadata).""" - sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - assert isinstance( - sharded_tensors, tuple - ), f"sharded_tensors should be tuple, got {type(sharded_tensors).__name__}" - assert len(sharded_tensors) > 0, "sharded_tensors should not be empty" - assert isinstance( - metadata, tuple - ), f"metadata should be tuple, got {type(metadata).__name__}" - - def test_pre_all_gather_buffers_are_plain_tensors(self, hybrid_param): - """Every element in sharded_tensors must be a plain torch.Tensor.""" - sharded_tensors, _ = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - for i, t in enumerate(sharded_tensors): - assert isinstance( - t, torch.Tensor - ), f"sharded_tensors[{i}] should be torch.Tensor, got {type(t).__name__}" - assert not isinstance( - t, QuantizedTensor - ), f"sharded_tensors[{i}] should NOT be QuantizedTensor subclass" - - def test_pre_all_gather_buffer_count_consistent(self, hybrid_param): - """Buffer count must be the same across repeated calls (FSDP2 buffer reuse).""" - sharded_1, _ = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - sharded_2, _ = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - assert len(sharded_1) == len( - sharded_2 - ), f"Buffer count changed: {len(sharded_1)} vs {len(sharded_2)}" - - def test_pre_all_gather_metadata_sufficient_for_reconstruction(self, hybrid_param): - """Metadata must contain enough info to reconstruct the tensor.""" - _, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - assert metadata is not None - assert len(metadata) > 0, "metadata should not be empty" - - -# --------------------------------------------------------------------------- -# 13. FSDP2 prerequisites: fsdp_post_all_gather protocol -# --------------------------------------------------------------------------- - - -@requires_fp8 -class TestHybridFsdpPostAllGatherProtocol: - """Test the fsdp_post_all_gather method on HybridQuantizedTensor. - - Simulates the post-all-gather phase by passing the sharded_tensors - from pre_all_gather directly (mimicking a single-rank all-gather). - """ - - @pytest.fixture(params=_fsdp_protocol_configs) - def hybrid_param(self, request): - torch.manual_seed(42) - return _make_fsdp_protocol_param(request.param) - - def test_post_all_gather_first_call_returns_hybrid_tensor(self, hybrid_param): - """With out=None, post_all_gather returns (HybridQuantizedTensor, outputs).""" - sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - result, ag_outputs = hybrid_param.fsdp_post_all_gather( - sharded_tensors, - metadata, - hybrid_param.dtype, - out=None, - ) - assert isinstance( - result, HybridQuantizedTensor - ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" - assert result.shape == hybrid_param.shape - assert result.rowwise_sub_storage is not None - assert result.columnwise_sub_storage is not None - - def test_post_all_gather_buffer_reuse(self, hybrid_param): - """On second call with out=previous, the same object is returned (buffer reuse).""" - sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - first_result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, - metadata, - hybrid_param.dtype, - out=None, - ) - - second_result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, - metadata, - hybrid_param.dtype, - out=first_result, - ) - assert ( - second_result is first_result - ), "Buffer reuse: post_all_gather(out=prev) should return the same object" - _assert_hybrid_tensor_exact(second_result, hybrid_param, context="post-all-gather reuse") - - def test_post_all_gather_dequantize_matches_original(self, hybrid_param): - """Reconstructed tensor should dequantize close to the original.""" - orig_deq = hybrid_param.dequantize() - - sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, - metadata, - hybrid_param.dtype, - out=None, - ) - result_deq = result.dequantize() - _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather") - torch.testing.assert_close(orig_deq, result_deq, rtol=0.0, atol=0.0) - - def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): - """Reconstructed tensor's sub-storages match the original types.""" - orig_row_type = type(hybrid_param.rowwise_sub_storage) - orig_col_type = type(hybrid_param.columnwise_sub_storage) - - sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, - metadata, - hybrid_param.dtype, - out=None, - ) - assert type(result.rowwise_sub_storage) is orig_row_type - _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather storage types") - assert type(result.columnwise_sub_storage) is orig_col_type - - -# --------------------------------------------------------------------------- -# 14. FSDP2 prerequisites: pre/post roundtrip -# --------------------------------------------------------------------------- - - -@requires_fp8 -class TestHybridFsdpRoundtrip: - """End-to-end single-process roundtrip (pre -> post) without communication.""" - - @pytest.fixture(params=_fsdp_protocol_configs) - def hybrid_param(self, request): - torch.manual_seed(42) - return _make_fsdp_protocol_param(request.param) - - def test_pre_post_roundtrip_preserves_data(self, hybrid_param): - """pre_all_gather -> post_all_gather(out=None) -> dequantize matches original.""" - orig_deq = hybrid_param.dequantize() - - sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, - metadata, - hybrid_param.dtype, - out=None, - ) - _assert_hybrid_tensor_exact(result, hybrid_param, context="FSDP roundtrip") - torch.testing.assert_close(orig_deq, result.dequantize(), rtol=0.0, atol=0.0) - - def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): - """Second roundtrip with out=previous preserves data (iteration 2+ simulation).""" - sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - first_result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors, - metadata, - hybrid_param.dtype, - out=None, - ) - - sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - second_result, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors_2, - metadata_2, - hybrid_param.dtype, - out=first_result, - ) - assert second_result is first_result - torch.testing.assert_close( - hybrid_param.dequantize(), second_result.dequantize(), rtol=0.0, atol=0.0 - ) - _assert_hybrid_tensor_exact(second_result, hybrid_param, context="FSDP roundtrip reuse") - - @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 - def test_scale_refresh_across_iterations(self): - """After a sharded optimizer-style requantize, iter-2+ gathers see the new scale. - - Per-tensor FP8 does NOT include ``_scale_inv`` in ``fsdp_buffer_fields`` - (only ``_data`` is gathered; the scalar scale travels via iter-1 - metadata). This relies on the invariant that the sharded and gathered - ``Float8Tensor`` s share the same ``_scale_inv`` tensor object, and - that ``Float8CurrentScalingQuantizer.update_quantized`` writes the new - scale in place rather than replacing the tensor reference. If either - invariant broke, the gathered copy would carry a stale scale on - iter-2+ and silently apply the wrong dequantization. - - This test locks the invariant down by forcing a radically different - scale between iterations and asserting the gathered tensor's - dequantization tracks the sharded one. - """ - torch.manual_seed(42) - hybrid_recipe = _hybrid_custom_recipe( - _fp8_row_factory, - _fp8_col_factory, - _fp8_grad_factory, - ) - with quantized_model_init(enabled=True, recipe=hybrid_recipe): - model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() - hybrid_param = model.weight - - # Iter-1 gather with the initial (small-magnitude) weights - sharded_tensors_1, metadata_1 = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - gathered, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors_1, - metadata_1, - hybrid_param.dtype, - out=None, - ) - - # Simulate an optimizer writeback that produces a much larger weight; - # Float8CurrentScalingQuantizer.update_quantized must recompute - # _scale_inv for this range. If the gathered copy didn't see the new - # scale, the dequantize below would disagree with the sharded copy. - huge_master = torch.randn_like(hybrid_param.dequantize()) * 100.0 - hybrid_param._quantizer.update_quantized(huge_master, hybrid_param) - - # Iter-2+ path: reuse the gathered buffer - sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( - mesh=None, - orig_size=hybrid_param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - gathered_refreshed, _ = hybrid_param.fsdp_post_all_gather( - sharded_tensors_2, - metadata_2, - hybrid_param.dtype, - out=gathered, - ) - assert gathered_refreshed is gathered - - # The gathered copy must now reflect the new sharded scale, not the - # tiny original scale. - torch.testing.assert_close( - hybrid_param.dequantize(), - gathered_refreshed.dequantize(), - rtol=0.0, - atol=0.0, - ) - _assert_hybrid_tensor_exact(gathered_refreshed, hybrid_param, context="FSDP scale refresh") - # And the magnitude really did change (sanity: this test would pass - # vacuously if update_quantized didn't actually change anything). - assert gathered_refreshed.dequantize().abs().max() > 10.0, ( - "update_quantized did not produce a sufficiently different " - "weight; the scale-refresh invariant is not being exercised" - ) - - def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): - """Hybrid FSDP2 with an NVFP4 sub-storage must raise a clear error. - - NVIDIA/TransformerEngine#3158 tracks the missing NVFP4 FSDP2 support: - packed FP4 dim-0 alignment, columnwise dequantization, and RHT-cache - handling are not implemented. Until that lands, hybrid pre-all-gather must - refuse an NVFP4 sub-storage cleanly via the ``fsdp_buffer_fields`` - protocol rather than silently producing wrong data. - - This test pins that contract: any hybrid whose sub-storage does not - implement ``fsdp_buffer_fields`` raises ``NotImplementedError`` at - ``fsdp_pre_all_gather`` time. The prior version of this test - inadvertently asserted the opposite when buffer extraction used - implicit ``get_metadata()``-based tensor scanning. - """ - if not (fp8_available and nvfp4_available): - pytest.skip("Requires FP8 + NVFP4 support") - - hybrid_recipe = _hybrid_custom_recipe( - row_factory=lambda: NVFP4Quantizer(), - col_factory=lambda: NVFP4Quantizer(), - ) - with quantized_model_init(enabled=True, recipe=hybrid_recipe): - model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() - param = model.weight - - # Clean refusal: hybrid's pre_all_gather raises an NVFP4-specific - # message identifying the unsupported sub-storage protocol, not a generic - # "NVFP4Tensor does not implement fsdp_buffer_fields" from deep inside - # the base class. - with pytest.raises(NotImplementedError) as exc_info: - param.fsdp_pre_all_gather( - mesh=None, - orig_size=param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - msg = str(exc_info.value) - assert "NVFP4Tensor" in msg - assert "rowwise sub-storage" in msg - assert "Use a supported sub-quantizer" in msg - assert "fsdp_buffer_fields" in msg - - -# --------------------------------------------------------------------------- -# 15. FSDP2 prerequisites: make_like correctness -# --------------------------------------------------------------------------- - - -@requires_fp8 -@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 -class TestHybridMakeLike: - """Test that make_like produces correct copies for __torch_dispatch__ usage.""" - - def _make_hybrid_param(self): - hybrid_recipe = _hybrid_custom_recipe( - _fp8_row_factory, - _fp8_col_factory, - _fp8_grad_factory, - ) - with quantized_model_init(enabled=True, recipe=hybrid_recipe): - model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() - return model.weight - - def test_make_like_preserves_sub_storages(self): - """make_like result has the same sub-storage types, quantizers, and dtype.""" - param = self._make_hybrid_param() - copy = HybridQuantizedTensor.make_like(param) - - assert isinstance(copy, HybridQuantizedTensor) - assert copy.dtype == param.dtype - assert copy.shape == param.shape - assert type(copy.rowwise_sub_storage) is type(param.rowwise_sub_storage) - assert type(copy.columnwise_sub_storage) is type(param.columnwise_sub_storage) - _assert_hybrid_tensor_exact(copy, param, context="make_like") - torch.testing.assert_close(copy.dequantize(), param.dequantize(), rtol=0.0, atol=0.0) - - def test_make_like_is_independent(self): - """make_like result should not share the same tensor identity.""" - param = self._make_hybrid_param() - copy = HybridQuantizedTensor.make_like(param) - assert copy is not param - - -# --------------------------------------------------------------------------- -# 15b. Hopper-only paths: columnwise-only Float8 sub-storage -# --------------------------------------------------------------------------- -# -# On architectures where ``is_non_tn_fp8_gemm_supported()`` returns False -# (Hopper sm_90, L40 sm_89), per-tensor FP8 GEMM only supports the TN -# layout — non-TN layouts are simulated by feeding pre-transposed data. -# So a columnwise-only ``Float8TensorStorage`` (used as a hybrid sub- -# storage) holds its quantized data in ``_transpose`` instead of -# ``_data``, with ``_data = None``. -# -# This is the exact layout the FSDP2 buffer protocol must recognize -# when the sub-storage is part of a ``HybridQuantizedTensor`` parameter. -# These tests pin the contracts that would break if the buffer -# protocol regressed to the unconditional ``("_data",)`` field name -# (which would all-gather a ``None`` tensor on Hopper). -# -# Skip on Blackwell where the C++ kernel always populates ``_data`` and -# the columnwise-only Float8 path doesn't exercise ``_transpose``. - -requires_hopper_fp8 = pytest.mark.skipif( - is_non_tn_fp8_gemm_supported() or not fp8_available, - reason=( - "Hopper-only: requires per-tensor FP8 with non-TN GEMM unsupported " - "(Hopper sm_90 / L40 sm_89). On Blackwell the C++ kernel populates " - "_data even for columnwise-only mode, so the _transpose-only path " - "is not exercised." - ), -) - - -@requires_hopper_fp8 -class TestHybridFloat8ColumnwiseOnlyHopperPath: - """Hopper transpose-only Float8 uses an M-major FSDP transport layout.""" - - @staticmethod - def _make_columnwise_only_float8_storage(shape=(32, 64), value=1.0, quantizer=None): - source = torch.full(shape, value, device="cuda", dtype=torch.bfloat16) - byte_quantizer = Float8Quantizer( - scale=torch.ones(1, device="cuda", dtype=torch.float32), - amax=torch.zeros(1, device="cuda", dtype=torch.float32), - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=True, - columnwise=False, - ) - rowwise = byte_quantizer(source) - if quantizer is None: - quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - device="cuda", - rowwise=False, - columnwise=True, - ) - storage = Float8Tensor( - shape=shape, - dtype=source.dtype, - data=None, - data_transpose=rowwise._data.movedim(-1, 0).contiguous(), - fp8_scale_inv=rowwise._scale_inv.clone(), - fp8_dtype=tex.DType.kFloat8E4M3, - quantizer=quantizer, - requires_grad=False, - device="cuda", - ) - return storage, rowwise.dequantize() - - @classmethod - def _make_hybrid_shard(cls, value): - quantizer = HybridQuantizer( - rowwise_quantizer=IdentityQuantizer(), - columnwise_quantizer=Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - device="cuda", - ), - ) - source = torch.full((32, 64), value, device="cuda", dtype=torch.bfloat16) - rowwise = quantizer.rowwise_quantizer(source) - columnwise, expected = cls._make_columnwise_only_float8_storage( - value=value, quantizer=quantizer.columnwise_quantizer - ) - hybrid = HybridQuantizedTensor( - shape=source.shape, - dtype=source.dtype, - rowwise_storage=rowwise, - columnwise_storage=columnwise, - quantizer=quantizer, - requires_grad=False, - device="cuda", - ) - return hybrid, expected - - def test_columnwise_only_float8_fsdp_buffer_fields_returns_transpose(self): - out, _ = self._make_columnwise_only_float8_storage() - - assert out._data is None - assert out._transpose is not None - assert not out._transpose_invalid - assert out.fsdp_buffer_fields() == ("_transpose",) - - def test_columnwise_only_float8_fsdp_extract_uses_m_major_transport(self): - out, _ = self._make_columnwise_only_float8_storage() - - buffers, metadata = out.fsdp_extract_buffers() - - assert len(buffers) == 1 - assert buffers[0].shape == out.shape - assert buffers[0].is_contiguous() - torch.testing.assert_close(buffers[0], out._transpose.movedim(0, -1).contiguous()) - assert metadata == { - "field_names": ("_transpose",), - "transport_layout": "columnwise_m_major", - } - - def test_columnwise_only_float8_fsdp_assign_restores_columnwise_layout(self): - out, expected = self._make_columnwise_only_float8_storage() - buffers, metadata = out.fsdp_extract_buffers() - gathered = torch.cat((buffers[0], buffers[0]), dim=0) - rebuilt = Float8Tensor.make_like(out, shape=gathered.shape) - - rebuilt.fsdp_assign_gathered((gathered,), metadata) - - assert rebuilt.shape == torch.Size((64, 64)) - assert rebuilt._data is None - assert rebuilt._transpose.shape == torch.Size((64, 64)) - assert not rebuilt._transpose_invalid - torch.testing.assert_close(rebuilt._transpose, gathered.movedim(-1, 0).contiguous()) - torch.testing.assert_close(rebuilt.dequantize(), torch.cat((expected, expected), dim=0)) - - def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): - shards = [self._make_hybrid_shard(value) for value in (1.0, 2.0)] - extracted = [ - shard.fsdp_pre_all_gather( - mesh=None, - orig_size=shard.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - for shard, _ in shards - ] - gathered = tuple( - torch.cat((extracted[0][0][i], extracted[1][0][i]), dim=0) - for i in range(len(extracted[0][0])) - ) - expected = torch.cat((shards[0][1], shards[1][1]), dim=0) - - rebuilt, _ = shards[0][0].fsdp_post_all_gather(gathered, extracted[0][1], torch.bfloat16) - - assert rebuilt.shape == torch.Size((64, 64)) - assert rebuilt._columnwise_storage.shape == torch.Size((64, 64)) - assert rebuilt._columnwise_storage._data is None - assert rebuilt._columnwise_storage._transpose.shape == torch.Size((64, 64)) - torch.testing.assert_close(rebuilt._columnwise_storage.dequantize(), expected) - - reused, _ = shards[0][0].fsdp_post_all_gather( - gathered, extracted[0][1], torch.bfloat16, out=rebuilt - ) - assert reused is rebuilt - assert reused._columnwise_storage._transpose.shape == torch.Size((64, 64)) - torch.testing.assert_close(reused._columnwise_storage.dequantize(), expected) - - def test_fsdp_buffer_fields_falls_back_to_data_when_both_present(self): - quantizer = Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - device="cuda", - ) - out = quantizer(torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)) - - assert out._data is not None - assert out._transpose is not None - assert out.fsdp_buffer_fields() == ("_data",) - buffers, metadata = out.fsdp_extract_buffers() - assert len(buffers) == 1 - assert buffers[0] is out._data - assert metadata == { - "field_names": ("_data",), - "transport_layout": "native", - } - - -@requires_hopper_fp8 -class TestPerTensorFloat8ColumnwiseOnlyGuard: - """Per-tensor FP8 cannot yet emit only the physical transpose.""" - - @staticmethod - def _make_quantizer(scaling): - if scaling == "current": - return Float8CurrentScalingQuantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - device="cuda", - rowwise=False, - columnwise=True, - ) - return Float8Quantizer( - scale=torch.ones(1, device="cuda", dtype=torch.float32), - amax=torch.zeros(1, device="cuda", dtype=torch.float32), - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=False, - columnwise=True, - ) - - @pytest.mark.parametrize("scaling", ("current", "delayed")) - def test_initial_quantize_raises_not_implemented(self, scaling): - quantizer = self._make_quantizer(scaling) - source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) - - with pytest.raises( - NotImplementedError, - match="Columnwise-only per-tensor FP8 quantization is not implemented", - ): - quantizer(source) - - @pytest.mark.parametrize("scaling", ("current", "delayed")) - def test_update_quantized_raises_not_implemented(self, scaling): - quantizer = self._make_quantizer(scaling) - source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) - output = quantizer.make_empty( - source.shape, - dtype=source.dtype, - device=source.device, - ) - assert output._data is None - assert output._transpose is not None - - with pytest.raises( - NotImplementedError, - match="Columnwise-only per-tensor FP8 quantization is not implemented", - ): - quantizer.update_quantized(source, output) - - -@requires_hopper_fp8 -@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 -class TestHybridFsdpPostAllGatherUpdateUsage: - """``HybridQuantizedTensor.fsdp_post_all_gather`` must call - ``update_usage`` on each sub-storage after writing gathered data - (mirroring vanilla ``Float8Tensor.fsdp_post_all_gather:888``). - Without it, on Hopper a previously-cached ``_transpose`` from the - prior iteration is silently reused with the new ``_data``, producing - incorrect dgrad / wgrad GEMMs. - """ - - def _make_param(self): - hybrid_recipe = _hybrid_custom_recipe( - _fp8_row_factory, - _fp8_col_factory, - _fp8_grad_factory, - ) - with quantized_model_init(enabled=True, recipe=hybrid_recipe): - model = Linear(64, 64, params_dtype=torch.bfloat16).cuda() - return model.weight - - def test_iter2_invalidates_stale_transpose_on_rowwise_substorage(self): - """Simulates iter-2+ buffer reuse: pre-existing ``out`` with a - possibly-stale ``_transpose`` cache; after ``fsdp_post_all_gather`` - the rowwise sub-storage's ``_transpose`` must be invalidated / - regenerated to match the freshly gathered ``_data``. - """ - param = self._make_param() - # Build a plausible iter-2+ "out" with stale state. - out = HybridQuantizedTensor.make_like(param) - # Rowwise sub-storage on Hopper has _data populated. Force a stale - # _transpose and invalidate flag to mimic the regression scenario. - if out._rowwise_storage._transpose is None: - # Set up a fake stale transpose (non-None, marked invalid by - # the mismatching shape would catch nothing, so just plant - # a tensor and clear the invalid flag to "valid"). - out._rowwise_storage._transpose = torch.zeros_like(out._rowwise_storage._data).t() - out._rowwise_storage._transpose_invalid = False - stale_transpose_id = id(out._rowwise_storage._transpose) - - # Drive a real all-gather round trip via the protocol - sharded_tensors, metadata = param.fsdp_pre_all_gather( - mesh=None, - orig_size=param.shape, - contiguous_orig_stride=None, - module=None, - mp_policy=None, - ) - out2, _ = param.fsdp_post_all_gather(sharded_tensors, metadata, param.dtype, out=out) - - # After fsdp_post_all_gather, the rowwise sub-quantizer is pinned - # columnwise=False, so update_usage(rowwise=True, columnwise=False) - # must clear the stale _transpose (preventing the silent - # stale-cache regression on Hopper). - assert out2._rowwise_storage._transpose is None or ( - out2._rowwise_storage._transpose_invalid - and id(out2._rowwise_storage._transpose) != stale_transpose_id - ), "Stale _transpose was not invalidated after fsdp_post_all_gather" - - -# --------------------------------------------------------------------------- -# 16. Activation recomputation (torch.utils.checkpoint / te.checkpoint) +# 11. Activation recomputation (torch.utils.checkpoint / te.checkpoint) # --------------------------------------------------------------------------- diff --git a/tests/pytorch/test_hybrid_quantization_fsdp2.py b/tests/pytorch/test_hybrid_quantization_fsdp2.py new file mode 100644 index 00000000000..055675bb465 --- /dev/null +++ b/tests/pytorch/test_hybrid_quantization_fsdp2.py @@ -0,0 +1,1570 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for HybridQuantizedTensor behavior required by PyTorch FSDP2.""" + +import io + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex + +from hybrid_quantization_utils import ( + as_data_tensor_tuple as _as_data_tensor_tuple, + assert_hybrid_tensor_exact as _assert_hybrid_tensor_exact, + assert_storage_data_exact as _assert_storage_data_exact, + fp8_e4m3_factory as _fp8_row_factory, + fp8_e5m2_factory as _fp8_grad_factory, + hybrid_block_fp8_e4m3_qfactory as _hybrid_block_fp8_qfactory, + hybrid_custom_recipe as _hybrid_custom_recipe, + make_fp8_quantizer as _make_fp8_quantizer, + make_hybrid_quantizer_fp8_row_fp4_col as _make_hybrid_quantizer_fp8_row_fp4_col, + mxfp8_e4m3_factory as _mxfp8_factory, +) +from transformer_engine.common import recipe +from transformer_engine.pytorch import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + Float8Tensor, + Float8TensorStorage, + HybridQuantizedTensor, + HybridQuantizer, + IdentityQuantizer, + Linear, + MXFP8Quantizer, + NVFP4Quantizer, + QuantizedTensor, + quantized_model_init, +) +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported + +_fp8_col_factory = _fp8_row_factory + + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = ( + te.is_fp8_block_scaling_available(return_reason=True) +) + +_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( + condition=not is_non_tn_fp8_gemm_supported(), + raises=NotImplementedError, + strict=True, + reason=( + "Hopper does not yet support columnwise-only per-tensor FP8 quantization; " + "tracked by NVIDIA/TransformerEngine#3158" + ), +) + +requires_fp8 = pytest.mark.skipif( + not fp8_available, + reason=f"FP8: {reason_for_no_fp8}", +) + +requires_fp8_and_nvfp4 = pytest.mark.skipif( + not (fp8_available and nvfp4_available), + reason=f"FP8: {reason_for_no_fp8}; NVFP4: {reason_for_no_nvfp4}", +) + +requires_mxfp8 = pytest.mark.skipif( + not mxfp8_available, + reason=f"MXFP8: {reason_for_no_mxfp8}", +) + + +# --------------------------------------------------------------------------- +# 1. __torch_dispatch__ operations that FSDP2 relies on +# --------------------------------------------------------------------------- + +aten = torch.ops.aten + + +def _make_hybrid_param_for_dispatch( + row_factory, col_factory, grad_factory=None, in_features=256, out_features=256 +): + """Create a HybridQuantizedTensor weight via quantized_model_init for dispatch tests.""" + hybrid_recipe = _hybrid_custom_recipe(row_factory, col_factory, grad_factory) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() + return model.weight + + +_dispatch_configs = [ + pytest.param( + "fp8_fp8", + id="same-format-fp8", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ), +] +if mxfp8_available: + _dispatch_configs.append(pytest.param("mxfp8_mxfp8", id="same-format-mxfp8")) + + +def _get_dispatch_hybrid_param(config_name): + """Return a HybridQuantizedTensor weight for the given config.""" + if config_name == "fp8_fp8": + return _make_hybrid_param_for_dispatch( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + elif config_name == "mxfp8_mxfp8": + return _make_hybrid_param_for_dispatch( + _mxfp8_factory, + _mxfp8_factory, + grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + ) + else: + raise ValueError(f"Unknown config: {config_name}") + + +@requires_fp8 +class TestFloat8TransposeOnlySplit: + """Regression coverage for columnwise-only Float8 split metadata. + + A columnwise-only per-tensor Float8 sub-storage may have ``_data=None`` and + store its bytes in ``_transpose`` with physical shape ``[K, M]``. Splitting + that tensor must still produce pieces whose wrapper shape is the logical + row-major shape ``[M_i, K]``; otherwise HybridQuantizedTensor uses the + transposed shape when rowwise storage is absent. + """ + + @staticmethod + def _make_transpose_only_float8_tensor(shape=(12, 16)): + m, k = shape + data_transpose = torch.empty((k, m), dtype=torch.uint8, device="cuda") + return Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=data_transpose, + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + requires_grad=False, + device="cuda", + ) + + @pytest.mark.parametrize( + "split_size,dim,expected_shapes,expected_transpose_shapes", + [ + (5, 0, [(5, 16), (5, 16), (2, 16)], [(16, 5), (16, 5), (16, 2)]), + (6, 1, [(12, 6), (12, 6), (12, 4)], [(6, 12), (6, 12), (4, 12)]), + ], + ) + def test_float8_split_uses_logical_shape_for_transpose_only_storage( + self, split_size, dim, expected_shapes, expected_transpose_shapes + ): + tensor = self._make_transpose_only_float8_tensor() + + pieces = torch.split(tensor, split_size, dim=dim) + + assert [tuple(piece.shape) for piece in pieces] == expected_shapes + assert [ + tuple(piece._transpose.shape) for piece in pieces + ] == expected_transpose_shapes + assert all(piece._data is None for piece in pieces) + assert all(piece._transpose_invalid is False for piece in pieces) + + def test_float8_view_preserves_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + viewed = tensor.view(12, 16) + + assert tuple(viewed.shape) == (12, 16) + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_aten_view_preserves_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + viewed = torch.ops.aten.view.default(tensor, [12, 16]) + + assert tuple(viewed.shape) == (12, 16) + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_storage_view_preserves_transpose_only_storage(self): + data_transpose = torch.empty((16, 12), dtype=torch.uint8, device="cuda") + storage = Float8TensorStorage( + data=None, + data_transpose=data_transpose, + fp8_scale_inv=torch.ones(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + fake_dtype=torch.bfloat16, + ) + + viewed = storage.view(torch.Size((12, 16))) + + assert viewed._data is None + assert viewed._transpose is not None + assert tuple(viewed._transpose.shape) == (16, 12) + assert viewed._transpose_invalid is False + + def test_float8_shape_changing_view_raises_for_transpose_only_storage(self): + tensor = self._make_transpose_only_float8_tensor() + + with pytest.raises(NotImplementedError, match="columnwise-only data"): + tensor.view(6, 32) + + def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self): + columnwise = self._make_transpose_only_float8_tensor() + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + quantizer.set_usage(rowwise=False, columnwise=True) + hybrid = HybridQuantizedTensor( + shape=columnwise.shape, + dtype=columnwise.dtype, + rowwise_storage=None, + columnwise_storage=columnwise, + quantizer=quantizer, + device="cuda", + ) + + pieces = torch.split(hybrid, 5, dim=0) + + assert [tuple(piece.shape) for piece in pieces] == [ + (5, 16), + (5, 16), + (2, 16), + ] + assert all(piece.rowwise_sub_storage is None for piece in pieces) + assert [tuple(piece.columnwise_sub_storage.shape) for piece in pieces] == [ + (5, 16), + (5, 16), + (2, 16), + ] + assert [ + tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces + ] == [ + (16, 5), + (16, 5), + (16, 2), + ] + + @staticmethod + def _make_valid_transpose_only_float8_tensor(shape=(12, 16)): + """Build transpose-only storage with known, numerically valid FP8 bytes.""" + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + rowwise_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + rowwise = rowwise_quantizer(source) + columnwise_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + tensor = Float8Tensor( + shape=shape, + dtype=torch.bfloat16, + data=None, + data_transpose=rowwise._data.movedim(-1, 0).contiguous(), + fp8_scale_inv=rowwise._scale_inv.detach().clone(), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=columnwise_quantizer, + requires_grad=False, + device="cuda", + ) + return tensor, rowwise.dequantize() + + @pytest.mark.parametrize("weights_only", (False, True)) + def test_serialization_preserves_transpose_only_payload(self, weights_only): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + buffer = io.BytesIO() + + torch.save(tensor, buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=weights_only) + + assert loaded._data is None + assert loaded._transpose_invalid is False + assert torch.equal(loaded._transpose, tensor._transpose) + torch.testing.assert_close(loaded.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_dequantize_from_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + torch.testing.assert_close(tensor.dequantize(), expected, rtol=0.0, atol=0.0) + + def test_slice_and_select_preserve_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + sliced = tensor[2:8] + selected = torch.select(tensor, 1, 3) + + assert isinstance(sliced, Float8Tensor) + assert sliced._data is None + assert sliced._transpose.shape == torch.Size((16, 6)) + torch.testing.assert_close( + sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0 + ) + assert isinstance(selected, Float8Tensor) + assert selected._data is None + assert selected._transpose.shape == torch.Size((12,)) + torch.testing.assert_close( + selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0 + ) + + def test_as_strided_row_shard_preserves_transpose_only_payload(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + shard = torch.as_strided(tensor, (5, 16), (16, 1), 16) + + assert isinstance(shard, Float8Tensor) + assert shard._data is None + assert shard._transpose.shape == torch.Size((16, 5)) + torch.testing.assert_close( + shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0 + ) + + def test_as_strided_falls_back_for_nonrepresentable_layout(self): + tensor, expected = self._make_valid_transpose_only_float8_tensor() + + output = torch.as_strided(tensor, (16, 12), (1, 16), 0) + + assert type(output) is torch.Tensor + reference = torch.as_strided(expected, (16, 12), (1, 16), 0) + torch.testing.assert_close(output, reference, rtol=0.0, atol=0.0) + + def test_new_zeros_preserves_transpose_only_layout(self): + tensor, _ = self._make_valid_transpose_only_float8_tensor() + + output = tensor.new_zeros((5, 16)) + + assert isinstance(output, Float8Tensor) + assert output.shape == torch.Size((5, 16)) + assert output._data is None + assert output._transpose.shape == torch.Size((16, 5)) + assert output._transpose.dtype == torch.uint8 + assert torch.count_nonzero(output._transpose).item() == 0 + torch.testing.assert_close( + output.dequantize(), + torch.zeros((5, 16), dtype=output.dtype, device=output.device), + rtol=0.0, + atol=0.0, + ) + + +class TestHybridNewZeros: + """Public new_zeros semantics independent of FSDP-specific allocation.""" + + @staticmethod + def _make_identity_hybrid(): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) + return quantizer, source + + @pytest.mark.parametrize( + "keep_rowwise,keep_columnwise", + [(True, True), (True, False), (False, True)], + ) + def test_initializes_every_present_direction_and_preserves_kwargs( + self, + keep_rowwise, + keep_columnwise, + ): + quantizer, source = self._make_identity_hybrid() + if not keep_rowwise: + source.update_usage(rowwise_usage=False) + if not keep_columnwise: + source.update_usage(columnwise_usage=False) + + # Direction selection must come from source storage, not mutable parent + # usage flags, and new_zeros must not change those flags. + quantizer.set_usage(rowwise=False, columnwise=False) + result = torch.ops.aten.new_zeros.default( + source, + [3, 5], + dtype=torch.float32, + device=source.device, + ) + + assert isinstance(result, HybridQuantizedTensor) + assert result.shape == torch.Size((3, 5)) + assert result.dtype == torch.float32 + assert result.device == source.device + assert (result.rowwise_sub_storage is not None) is keep_rowwise + assert (result.columnwise_sub_storage is not None) is keep_columnwise + assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} + assert result._quantizer is not quantizer + assert result._quantizer.get_usages() == { + "rowwise": keep_rowwise, + "columnwise": keep_columnwise, + } + + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + if sub_storage is not None: + torch.testing.assert_close( + sub_storage.dequantize(), + torch.zeros((3, 5), dtype=torch.float32, device=source.device), + rtol=0.0, + atol=0.0, + ) + + # A plain-source copy routes through the result's parent quantizer. It + # must update every allocated direction even though the source parent + # had both of its mutable usage flags disabled before new_zeros. + plain_source = torch.full( + (3, 5), 7.0, dtype=torch.float32, device=source.device + ) + result.copy_(plain_source) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + if sub_storage is not None: + torch.testing.assert_close( + sub_storage.dequantize(), plain_source, rtol=0.0, atol=0.0 + ) + assert quantizer.get_usages() == {"rowwise": False, "columnwise": False} + + def test_identity_substorages_allow_integer_dtype(self): + _, source = self._make_identity_hybrid() + result = source.new_zeros((2, 3), dtype=torch.int32) + + assert result.dtype == torch.int32 + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + assert sub_storage.dequantize().dtype == torch.int32 + assert torch.count_nonzero(sub_storage.dequantize()).item() == 0 + + @requires_fp8 + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + @pytest.mark.parametrize( + "unsupported_dtype", (torch.int32, torch.float64, torch.bool) + ) + def test_non_identity_substorage_rejects_unsupported_dtype(self, unsupported_dtype): + quantizer = HybridQuantizer( + rowwise_quantizer=_make_fp8_quantizer(), + columnwise_quantizer=_make_fp8_quantizer(), + ) + source = quantizer(torch.ones((4, 8), dtype=torch.bfloat16, device="cuda")) + + with pytest.raises(TypeError, match="new_zeros only supports"): + source.new_zeros((2, 3), dtype=unsupported_dtype) + + def test_does_not_invoke_live_quantizers_or_consume_rng(self, monkeypatch): + quantizer, source = self._make_identity_hybrid() + cpu_rng_before = torch.get_rng_state().clone() + cuda_rng_before = torch.cuda.get_rng_state(source.device).clone() + + def fail_live_make_empty(*args, **kwargs): + raise AssertionError("new_zeros invoked a live quantizer") + + monkeypatch.setattr(quantizer, "make_empty", fail_live_make_empty) + monkeypatch.setattr( + source.rowwise_sub_storage._quantizer, + "make_empty", + fail_live_make_empty, + ) + monkeypatch.setattr( + source.columnwise_sub_storage._quantizer, + "make_empty", + fail_live_make_empty, + ) + + result = source.new_zeros((2, 6)) + + assert isinstance(result, HybridQuantizedTensor) + assert torch.equal(torch.get_rng_state(), cpu_rng_before) + assert torch.equal(torch.cuda.get_rng_state(source.device), cuda_rng_before) + + def test_rejects_empty_hybrid(self): + _, source = self._make_identity_hybrid() + source.update_usage(rowwise_usage=False, columnwise_usage=False) + + with pytest.raises(RuntimeError, match="at least one present sub-storage"): + source.new_zeros((2, 6)) + + @requires_fp8_and_nvfp4 + def test_initializes_nvfp4_data_scale_and_amax_buffers(self): + quantizer = _make_hybrid_quantizer_fp8_row_fp4_col() + source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) + + result = source.new_zeros(source.shape) + + assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type( + source.columnwise_sub_storage + ) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + buffers, storage = sub_storage.prepare_for_saving() + try: + assert all( + buffer is None or torch.count_nonzero(buffer).item() == 0 + for buffer in buffers + ) + finally: + assert storage.restore_from_saved(buffers) == [] + + @pytest.mark.skipif( + not fp8_block_scaling_available, + reason=f"Float8Blockwise: {reason_for_no_fp8_block_scaling}", + ) + def test_initializes_float8_block_data_and_scale_buffers(self): + quantizer = HybridQuantizer( + rowwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + columnwise_quantizer=Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + block_scaling_dim=2, + ), + ) + source = quantizer(torch.randn((128, 256), dtype=torch.bfloat16, device="cuda")) + + result = source.new_zeros(source.shape) + + assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type( + source.columnwise_sub_storage + ) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + for sub_storage in ( + result.rowwise_sub_storage, + result.columnwise_sub_storage, + ): + buffers, storage = sub_storage.prepare_for_saving() + try: + assert all( + buffer is None or torch.count_nonzero(buffer).item() == 0 + for buffer in buffers + ) + finally: + assert storage.restore_from_saved(buffers) == [] + + +@requires_fp8 +class TestHybridTorchDispatchFSDP2Ops: + """Test aten ops that FSDP2 relies on to preserve the HybridQuantizedTensor type. + + Each op is called directly via torch.ops.aten and the result is verified to + still be HybridQuantizedTensor with valid sub-storages. + """ + + @pytest.fixture(params=_dispatch_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _get_dispatch_hybrid_param(request.param) + + def test_split_preserves_hybrid_type(self, hybrid_param): + """torch.split must return a list of HybridQuantizedTensor pieces.""" + dim0 = hybrid_param.shape[0] + chunk_size = dim0 // 2 + pieces = torch.split(hybrid_param, chunk_size, dim=0) + expected_rowwise = torch.split( + hybrid_param.rowwise_sub_storage, chunk_size, dim=0 + ) + expected_columnwise = torch.split( + hybrid_param.columnwise_sub_storage, chunk_size, dim=0 + ) + assert len(pieces) == len(expected_rowwise) == len(expected_columnwise) + + assert len(pieces) >= 2 + for piece in pieces: + assert isinstance( + piece, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(piece).__name__}" + assert piece.rowwise_sub_storage is not None + assert piece.columnwise_sub_storage is not None + for index, (piece, expected_row, expected_column) in enumerate( + zip(pieces, expected_rowwise, expected_columnwise) + ): + _assert_storage_data_exact( + piece.rowwise_sub_storage, + expected_row, + context=f"split {index} rowwise", + ) + _assert_storage_data_exact( + piece.columnwise_sub_storage, + expected_column, + context=f"split {index} columnwise", + ) + + total_rows = sum(p.shape[0] for p in pieces) + assert total_rows == dim0 + + orig_deq = hybrid_param.dequantize() + reassembled = torch.cat([p.dequantize() for p in pieces], dim=0) + torch.testing.assert_close(orig_deq, reassembled, rtol=0.0, atol=0.0) + + def test_split_sub_storage_types_preserved(self, hybrid_param): + """After split, sub-storage types must match the original.""" + orig_row_type = type(hybrid_param.rowwise_sub_storage) + orig_col_type = type(hybrid_param.columnwise_sub_storage) + + chunk_size = hybrid_param.shape[0] // 2 + pieces = torch.split(hybrid_param, chunk_size, dim=0) + for piece in pieces: + assert type(piece.rowwise_sub_storage) is orig_row_type + assert type(piece.columnwise_sub_storage) is orig_col_type + + @requires_mxfp8 + def test_split_rejects_mxfp8_high_precision_fallback(self): + """Fail before wrapping unquantizable MXFP8 shards for Hybrid FSDP2.""" + quantizer = HybridQuantizer( + rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), + ) + tensor = quantizer(torch.randn(64, 64, dtype=torch.bfloat16, device="cuda")) + + with pytest.raises(NotImplementedError, match="local shape.*divisible by 32"): + torch.split(tensor, 16, dim=0) + + def test_view_preserves_hybrid_type(self, hybrid_param): + """view must return a HybridQuantizedTensor (used by FSDP2 reset_sharded_param).""" + shape_2d = hybrid_param.shape + result = aten.view.default(hybrid_param, list(shape_2d)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_view_same_shape_preserves_hybrid(self, hybrid_param): + """view with same shape must return HybridQuantizedTensor.""" + shape_2d = list(hybrid_param.shape) + result = aten.view.default(hybrid_param, shape_2d) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + + def test_as_strided_noop_preserves_hybrid(self, hybrid_param): + """as_strided with matching shape/strides is a no-op that preserves type.""" + shape = tuple(hybrid_param.size()) + strides = (shape[-1], 1) + result = aten.as_strided.default(hybrid_param, list(shape), list(strides)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_slice_noop_preserves_hybrid(self, hybrid_param): + """slice with full range is a no-op that preserves type.""" + result = aten.slice.Tensor(hybrid_param, 0, 0, hybrid_param.size(0)) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.rowwise_sub_storage is not None + + def test_copy_between_hybrid_tensors(self, hybrid_param): + """copy_ between compatible HybridQuantizedTensors copies quantized data directly.""" + src_deq = hybrid_param.dequantize().clone() + dst = hybrid_param._quantizer.make_empty( + shape=hybrid_param.shape, + dtype=hybrid_param.dtype, + device=hybrid_param.device, + ) + assert isinstance(dst, HybridQuantizedTensor) + + aten.copy_.default(dst, hybrid_param) + dst_deq = dst.dequantize() + torch.testing.assert_close(src_deq, dst_deq, rtol=0.0, atol=0.0) + _assert_hybrid_tensor_exact(dst, hybrid_param, context="hybrid copy_") + + def test_copy_between_mismatched_usages_raises_atomically(self, hybrid_param): + quantizer = hybrid_param._quantizer.copy() + src = quantizer.quantize(torch.ones_like(hybrid_param.dequantize())) + dst = quantizer.quantize(torch.zeros_like(hybrid_param.dequantize())) + src.update_usage(columnwise_usage=False) + dst_before = tuple( + None if tensor is None else tensor.clone() + for tensor in _as_data_tensor_tuple(dst) + ) + + with pytest.raises( + NotImplementedError, + match="requires matching rowwise/columnwise usages", + ): + aten.copy_.default(dst, src) + + dst_after = _as_data_tensor_tuple(dst) + assert len(dst_after) == len(dst_before) + for before, after in zip(dst_before, dst_after): + if before is None: + assert after is None + else: + torch.testing.assert_close(after, before, rtol=0.0, atol=0.0) + + def test_copy_from_bf16_to_hybrid(self, hybrid_param): + """copy_ from BF16 into HybridQuantizedTensor triggers quantize_.""" + param = hybrid_param.detach() + bf16_data = torch.randn_like(param.dequantize()) + expected = param._quantizer.quantize(bf16_data) + aten.copy_.default(param, bf16_data) + assert isinstance(param, HybridQuantizedTensor) + _assert_hybrid_tensor_exact(param, expected, context="BF16 copy_") + torch.testing.assert_close( + param.dequantize(), expected.dequantize(), rtol=0.0, atol=0.0 + ) + + def test_new_zeros_returns_hybrid(self, hybrid_param): + """new_zeros returns initialized storage that remains FSDP-copyable.""" + new_shape = list(hybrid_param.shape) + result = aten.new_zeros.default(hybrid_param, new_shape) + + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + assert type(result.rowwise_sub_storage) is type( + hybrid_param.rowwise_sub_storage + ) + assert type(result.columnwise_sub_storage) is type( + hybrid_param.columnwise_sub_storage + ) + torch.testing.assert_close( + result.dequantize(), + torch.zeros_like(result.dequantize()), + rtol=0.0, + atol=0.0, + ) + + # FSDP2 overwrites the initialized destination via copy_ after gather. + aten.copy_.default(result, hybrid_param) + torch.testing.assert_close( + result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact( + result, hybrid_param, context="new_zeros then copy_" + ) + + def test_empty_like_returns_hybrid(self, hybrid_param): + """empty_like must return a HybridQuantizedTensor.""" + result = aten.empty_like.default(hybrid_param) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + + def test_clone_returns_hybrid(self, hybrid_param): + """clone must return an independent HybridQuantizedTensor with same data.""" + result = aten.clone.default(hybrid_param) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result is not hybrid_param + torch.testing.assert_close( + result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="clone") + + +# --------------------------------------------------------------------------- +# 2. fsdp_pre_all_gather protocol +# --------------------------------------------------------------------------- + + +def _make_fsdp_protocol_param(config_name): + """Create a HybridQuantizedTensor weight for FSDP protocol tests.""" + if config_name == "fp8_fp8": + r = _hybrid_custom_recipe(_fp8_row_factory, _fp8_col_factory, _fp8_grad_factory) + elif config_name == "mxfp8_fp8": + r = _hybrid_custom_recipe(_mxfp8_factory, _fp8_col_factory, _fp8_grad_factory) + elif config_name == "block_fp8": + r = recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) + else: + raise ValueError(f"Unknown config: {config_name}") + with quantized_model_init(enabled=True, recipe=r): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model.weight + + +_fsdp_protocol_configs = [ + pytest.param( + "fp8_fp8", + id="same-format", + marks=_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8, + ) +] +if mxfp8_available: + _fsdp_protocol_configs.append(pytest.param("mxfp8_fp8", id="mixed-mxfp8-fp8")) +if fp8_block_scaling_available: + _fsdp_protocol_configs.append(pytest.param("block_fp8", id="same-format-block-fp8")) + + +@requires_fp8 +class TestHybridFsdpPreAllGatherProtocol: + """Test the fsdp_pre_all_gather method on HybridQuantizedTensor. + + These tests call the method directly (no actual all-gather communication) + to verify the protocol contract: returns (sharded_tensors, metadata) where + sharded_tensors is a tuple of plain torch.Tensor. + """ + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_pre_all_gather_returns_tuple_pair(self, hybrid_param): + """fsdp_pre_all_gather returns (sharded_tensors, metadata).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert isinstance( + sharded_tensors, tuple + ), f"sharded_tensors should be tuple, got {type(sharded_tensors).__name__}" + assert len(sharded_tensors) > 0, "sharded_tensors should not be empty" + assert isinstance( + metadata, tuple + ), f"metadata should be tuple, got {type(metadata).__name__}" + + def test_pre_all_gather_buffers_are_plain_tensors(self, hybrid_param): + """Every element in sharded_tensors must be a plain torch.Tensor.""" + sharded_tensors, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + for i, t in enumerate(sharded_tensors): + assert isinstance( + t, torch.Tensor + ), f"sharded_tensors[{i}] should be torch.Tensor, got {type(t).__name__}" + assert not isinstance( + t, QuantizedTensor + ), f"sharded_tensors[{i}] should NOT be QuantizedTensor subclass" + + def test_pre_all_gather_buffer_count_consistent(self, hybrid_param): + """Buffer count must be the same across repeated calls (FSDP2 buffer reuse).""" + sharded_1, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + sharded_2, _ = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert len(sharded_1) == len( + sharded_2 + ), f"Buffer count changed: {len(sharded_1)} vs {len(sharded_2)}" + + def test_pre_all_gather_metadata_sufficient_for_reconstruction(self, hybrid_param): + """Metadata must contain enough info to reconstruct the tensor.""" + _, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + assert metadata is not None + assert len(metadata) > 0, "metadata should not be empty" + + +# --------------------------------------------------------------------------- +# 3. fsdp_post_all_gather protocol +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFsdpPostAllGatherProtocol: + """Test the fsdp_post_all_gather method on HybridQuantizedTensor. + + Simulates the post-all-gather phase by passing the sharded_tensors + from pre_all_gather directly (mimicking a single-rank all-gather). + """ + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_post_all_gather_first_call_returns_hybrid_tensor(self, hybrid_param): + """With out=None, post_all_gather returns (HybridQuantizedTensor, outputs).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, ag_outputs = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + assert isinstance( + result, HybridQuantizedTensor + ), f"Expected HybridQuantizedTensor, got {type(result).__name__}" + assert result.shape == hybrid_param.shape + assert result.rowwise_sub_storage is not None + assert result.columnwise_sub_storage is not None + + def test_post_all_gather_buffer_reuse(self, hybrid_param): + """On second call with out=previous, the same object is returned (buffer reuse).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + first_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + + second_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=first_result, + ) + assert ( + second_result is first_result + ), "Buffer reuse: post_all_gather(out=prev) should return the same object" + _assert_hybrid_tensor_exact( + second_result, hybrid_param, context="post-all-gather reuse" + ) + + def test_post_all_gather_dequantize_matches_original(self, hybrid_param): + """Reconstructed tensor should dequantize close to the original.""" + orig_deq = hybrid_param.dequantize() + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + result_deq = result.dequantize() + _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather") + torch.testing.assert_close(orig_deq, result_deq, rtol=0.0, atol=0.0) + + def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): + """Reconstructed tensor's sub-storages match the original types.""" + orig_row_type = type(hybrid_param.rowwise_sub_storage) + orig_col_type = type(hybrid_param.columnwise_sub_storage) + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + assert type(result.rowwise_sub_storage) is orig_row_type + _assert_hybrid_tensor_exact( + result, hybrid_param, context="post-all-gather storage types" + ) + assert type(result.columnwise_sub_storage) is orig_col_type + + +# --------------------------------------------------------------------------- +# 4. Pre/post all-gather roundtrip +# --------------------------------------------------------------------------- + + +@requires_fp8 +class TestHybridFsdpRoundtrip: + """End-to-end single-process roundtrip (pre -> post) without communication.""" + + @pytest.fixture(params=_fsdp_protocol_configs) + def hybrid_param(self, request): + torch.manual_seed(42) + return _make_fsdp_protocol_param(request.param) + + def test_pre_post_roundtrip_preserves_data(self, hybrid_param): + """pre_all_gather -> post_all_gather(out=None) -> dequantize matches original.""" + orig_deq = hybrid_param.dequantize() + + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="FSDP roundtrip") + torch.testing.assert_close(orig_deq, result.dequantize(), rtol=0.0, atol=0.0) + + def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): + """Second roundtrip with out=previous preserves data (iteration 2+ simulation).""" + sharded_tensors, metadata = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + first_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors, + metadata, + hybrid_param.dtype, + out=None, + ) + + sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + second_result, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_2, + metadata_2, + hybrid_param.dtype, + out=first_result, + ) + assert second_result is first_result + torch.testing.assert_close( + hybrid_param.dequantize(), second_result.dequantize(), rtol=0.0, atol=0.0 + ) + _assert_hybrid_tensor_exact( + second_result, hybrid_param, context="FSDP roundtrip reuse" + ) + + @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + def test_scale_refresh_across_iterations(self): + """After a sharded optimizer-style requantize, iter-2+ gathers see the new scale. + + Per-tensor FP8 does NOT include ``_scale_inv`` in ``fsdp_buffer_fields`` + (only ``_data`` is gathered; the scalar scale travels via iter-1 + metadata). This relies on the invariant that the sharded and gathered + ``Float8Tensor`` s share the same ``_scale_inv`` tensor object, and + that ``Float8CurrentScalingQuantizer.update_quantized`` writes the new + scale in place rather than replacing the tensor reference. If either + invariant broke, the gathered copy would carry a stale scale on + iter-2+ and silently apply the wrong dequantization. + + This test locks the invariant down by forcing a radically different + scale between iterations and asserting the gathered tensor's + dequantization tracks the sharded one. + """ + torch.manual_seed(42) + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + hybrid_param = model.weight + + # Iter-1 gather with the initial (small-magnitude) weights + sharded_tensors_1, metadata_1 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + gathered, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_1, + metadata_1, + hybrid_param.dtype, + out=None, + ) + + # Simulate an optimizer writeback that produces a much larger weight; + # Float8CurrentScalingQuantizer.update_quantized must recompute + # _scale_inv for this range. If the gathered copy didn't see the new + # scale, the dequantize below would disagree with the sharded copy. + huge_master = torch.randn_like(hybrid_param.dequantize()) * 100.0 + hybrid_param._quantizer.update_quantized(huge_master, hybrid_param) + + # Iter-2+ path: reuse the gathered buffer + sharded_tensors_2, metadata_2 = hybrid_param.fsdp_pre_all_gather( + mesh=None, + orig_size=hybrid_param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + gathered_refreshed, _ = hybrid_param.fsdp_post_all_gather( + sharded_tensors_2, + metadata_2, + hybrid_param.dtype, + out=gathered, + ) + assert gathered_refreshed is gathered + + # The gathered copy must now reflect the new sharded scale, not the + # tiny original scale. + torch.testing.assert_close( + hybrid_param.dequantize(), + gathered_refreshed.dequantize(), + rtol=0.0, + atol=0.0, + ) + _assert_hybrid_tensor_exact( + gathered_refreshed, hybrid_param, context="FSDP scale refresh" + ) + # And the magnitude really did change (sanity: this test would pass + # vacuously if update_quantized didn't actually change anything). + assert gathered_refreshed.dequantize().abs().max() > 10.0, ( + "update_quantized did not produce a sufficiently different " + "weight; the scale-refresh invariant is not being exercised" + ) + + def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): + """Hybrid FSDP2 with an NVFP4 sub-storage must raise a clear error. + + NVIDIA/TransformerEngine#3158 tracks the missing NVFP4 FSDP2 support: + packed FP4 dim-0 alignment, columnwise dequantization, and RHT-cache + handling are not implemented. Until that lands, hybrid pre-all-gather must + refuse an NVFP4 sub-storage cleanly via the ``fsdp_buffer_fields`` + protocol rather than silently producing wrong data. + + This test pins that contract: any hybrid whose sub-storage does not + implement ``fsdp_buffer_fields`` raises ``NotImplementedError`` at + ``fsdp_pre_all_gather`` time. The prior version of this test + inadvertently asserted the opposite when buffer extraction used + implicit ``get_metadata()``-based tensor scanning. + """ + if not (fp8_available and nvfp4_available): + pytest.skip("Requires FP8 + NVFP4 support") + + hybrid_recipe = _hybrid_custom_recipe( + row_factory=lambda: NVFP4Quantizer(), + col_factory=lambda: NVFP4Quantizer(), + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + param = model.weight + + # Clean refusal: hybrid's pre_all_gather raises an NVFP4-specific + # message identifying the unsupported sub-storage protocol, not a generic + # "NVFP4Tensor does not implement fsdp_buffer_fields" from deep inside + # the base class. + with pytest.raises(NotImplementedError) as exc_info: + param.fsdp_pre_all_gather( + mesh=None, + orig_size=param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + msg = str(exc_info.value) + assert "NVFP4Tensor" in msg + assert "rowwise sub-storage" in msg + assert "Use a supported sub-quantizer" in msg + assert "fsdp_buffer_fields" in msg + + +# --------------------------------------------------------------------------- +# 5. make_like correctness +# --------------------------------------------------------------------------- + + +@requires_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridMakeLike: + """Test that make_like produces correct copies for __torch_dispatch__ usage.""" + + def _make_hybrid_param(self): + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model.weight + + def test_make_like_preserves_sub_storages(self): + """make_like result has the same sub-storage types, quantizers, and dtype.""" + param = self._make_hybrid_param() + copy = HybridQuantizedTensor.make_like(param) + + assert isinstance(copy, HybridQuantizedTensor) + assert copy.dtype == param.dtype + assert copy.shape == param.shape + assert type(copy.rowwise_sub_storage) is type(param.rowwise_sub_storage) + assert type(copy.columnwise_sub_storage) is type(param.columnwise_sub_storage) + _assert_hybrid_tensor_exact(copy, param, context="make_like") + torch.testing.assert_close( + copy.dequantize(), param.dequantize(), rtol=0.0, atol=0.0 + ) + + def test_make_like_is_independent(self): + """make_like result should not share the same tensor identity.""" + param = self._make_hybrid_param() + copy = HybridQuantizedTensor.make_like(param) + assert copy is not param + + +# --------------------------------------------------------------------------- +# 5b. Hopper-only paths: columnwise-only Float8 sub-storage +# --------------------------------------------------------------------------- +# +# On architectures where ``is_non_tn_fp8_gemm_supported()`` returns False +# (Hopper sm_90, L40 sm_89), per-tensor FP8 GEMM only supports the TN +# layout — non-TN layouts are simulated by feeding pre-transposed data. +# So a columnwise-only ``Float8TensorStorage`` (used as a hybrid sub- +# storage) holds its quantized data in ``_transpose`` instead of +# ``_data``, with ``_data = None``. +# +# This is the exact layout the FSDP2 buffer protocol must recognize +# when the sub-storage is part of a ``HybridQuantizedTensor`` parameter. +# These tests pin the contracts that would break if the buffer +# protocol regressed to the unconditional ``("_data",)`` field name +# (which would all-gather a ``None`` tensor on Hopper). +# +# Skip on Blackwell where the C++ kernel always populates ``_data`` and +# the columnwise-only Float8 path doesn't exercise ``_transpose``. + +requires_hopper_fp8 = pytest.mark.skipif( + is_non_tn_fp8_gemm_supported() or not fp8_available, + reason=( + "Hopper-only: requires per-tensor FP8 with non-TN GEMM unsupported " + "(Hopper sm_90 / L40 sm_89). On Blackwell the C++ kernel populates " + "_data even for columnwise-only mode, so the _transpose-only path " + "is not exercised." + ), +) + + +@requires_hopper_fp8 +class TestHybridFloat8ColumnwiseOnlyHopperPath: + """Hopper transpose-only Float8 uses an M-major FSDP transport layout.""" + + @staticmethod + def _make_columnwise_only_float8_storage(shape=(32, 64), value=1.0, quantizer=None): + source = torch.full(shape, value, device="cuda", dtype=torch.bfloat16) + byte_quantizer = Float8Quantizer( + scale=torch.ones(1, device="cuda", dtype=torch.float32), + amax=torch.zeros(1, device="cuda", dtype=torch.float32), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + rowwise = byte_quantizer(source) + if quantizer is None: + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + storage = Float8Tensor( + shape=shape, + dtype=source.dtype, + data=None, + data_transpose=rowwise._data.movedim(-1, 0).contiguous(), + fp8_scale_inv=rowwise._scale_inv.clone(), + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + return storage, rowwise.dequantize() + + @classmethod + def _make_hybrid_shard(cls, value): + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ), + ) + source = torch.full((32, 64), value, device="cuda", dtype=torch.bfloat16) + rowwise = quantizer.rowwise_quantizer(source) + columnwise, expected = cls._make_columnwise_only_float8_storage( + value=value, quantizer=quantizer.columnwise_quantizer + ) + hybrid = HybridQuantizedTensor( + shape=source.shape, + dtype=source.dtype, + rowwise_storage=rowwise, + columnwise_storage=columnwise, + quantizer=quantizer, + requires_grad=False, + device="cuda", + ) + return hybrid, expected + + def test_columnwise_only_float8_fsdp_buffer_fields_returns_transpose(self): + out, _ = self._make_columnwise_only_float8_storage() + + assert out._data is None + assert out._transpose is not None + assert not out._transpose_invalid + assert out.fsdp_buffer_fields() == ("_transpose",) + + def test_columnwise_only_float8_fsdp_extract_uses_m_major_transport(self): + out, _ = self._make_columnwise_only_float8_storage() + + buffers, metadata = out.fsdp_extract_buffers() + + assert len(buffers) == 1 + assert buffers[0].shape == out.shape + assert buffers[0].is_contiguous() + torch.testing.assert_close( + buffers[0], out._transpose.movedim(0, -1).contiguous() + ) + assert metadata == { + "field_names": ("_transpose",), + "transport_layout": "columnwise_m_major", + } + + def test_columnwise_only_float8_fsdp_assign_restores_columnwise_layout(self): + out, expected = self._make_columnwise_only_float8_storage() + buffers, metadata = out.fsdp_extract_buffers() + gathered = torch.cat((buffers[0], buffers[0]), dim=0) + rebuilt = Float8Tensor.make_like(out, shape=gathered.shape) + + rebuilt.fsdp_assign_gathered((gathered,), metadata) + + assert rebuilt.shape == torch.Size((64, 64)) + assert rebuilt._data is None + assert rebuilt._transpose.shape == torch.Size((64, 64)) + assert not rebuilt._transpose_invalid + torch.testing.assert_close( + rebuilt._transpose, gathered.movedim(-1, 0).contiguous() + ) + torch.testing.assert_close( + rebuilt.dequantize(), torch.cat((expected, expected), dim=0) + ) + + def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): + shards = [self._make_hybrid_shard(value) for value in (1.0, 2.0)] + extracted = [ + shard.fsdp_pre_all_gather( + mesh=None, + orig_size=shard.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + for shard, _ in shards + ] + gathered = tuple( + torch.cat((extracted[0][0][i], extracted[1][0][i]), dim=0) + for i in range(len(extracted[0][0])) + ) + expected = torch.cat((shards[0][1], shards[1][1]), dim=0) + + rebuilt, _ = shards[0][0].fsdp_post_all_gather( + gathered, extracted[0][1], torch.bfloat16 + ) + + assert rebuilt.shape == torch.Size((64, 64)) + assert rebuilt._columnwise_storage.shape == torch.Size((64, 64)) + assert rebuilt._columnwise_storage._data is None + assert rebuilt._columnwise_storage._transpose.shape == torch.Size((64, 64)) + torch.testing.assert_close(rebuilt._columnwise_storage.dequantize(), expected) + + reused, _ = shards[0][0].fsdp_post_all_gather( + gathered, extracted[0][1], torch.bfloat16, out=rebuilt + ) + assert reused is rebuilt + assert reused._columnwise_storage._transpose.shape == torch.Size((64, 64)) + torch.testing.assert_close(reused._columnwise_storage.dequantize(), expected) + + def test_fsdp_buffer_fields_falls_back_to_data_when_both_present(self): + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ) + out = quantizer(torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)) + + assert out._data is not None + assert out._transpose is not None + assert out.fsdp_buffer_fields() == ("_data",) + buffers, metadata = out.fsdp_extract_buffers() + assert len(buffers) == 1 + assert buffers[0] is out._data + assert metadata == { + "field_names": ("_data",), + "transport_layout": "native", + } + + +@requires_hopper_fp8 +class TestPerTensorFloat8ColumnwiseOnlyGuard: + """Per-tensor FP8 cannot yet emit only the physical transpose.""" + + @staticmethod + def _make_quantizer(scaling): + if scaling == "current": + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + return Float8Quantizer( + scale=torch.ones(1, device="cuda", dtype=torch.float32), + amax=torch.zeros(1, device="cuda", dtype=torch.float32), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=False, + columnwise=True, + ) + + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_initial_quantize_raises_not_implemented(self, scaling): + quantizer = self._make_quantizer(scaling) + source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantizer(source) + + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_update_quantized_raises_not_implemented(self, scaling): + quantizer = self._make_quantizer(scaling) + source = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + output = quantizer.make_empty( + source.shape, + dtype=source.dtype, + device=source.device, + ) + assert output._data is None + assert output._transpose is not None + + with pytest.raises( + NotImplementedError, + match="Columnwise-only per-tensor FP8 quantization is not implemented", + ): + quantizer.update_quantized(source, output) + + +@requires_hopper_fp8 +@_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 +class TestHybridFsdpPostAllGatherUpdateUsage: + """``HybridQuantizedTensor.fsdp_post_all_gather`` must call + ``update_usage`` on each sub-storage after writing gathered data + (mirroring vanilla ``Float8Tensor.fsdp_post_all_gather:888``). + Without it, on Hopper a previously-cached ``_transpose`` from the + prior iteration is silently reused with the new ``_data``, producing + incorrect dgrad / wgrad GEMMs. + """ + + def _make_param(self): + hybrid_recipe = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + ) + with quantized_model_init(enabled=True, recipe=hybrid_recipe): + model = Linear(64, 64, params_dtype=torch.bfloat16).cuda() + return model.weight + + def test_iter2_invalidates_stale_transpose_on_rowwise_substorage(self): + """Simulates iter-2+ buffer reuse: pre-existing ``out`` with a + possibly-stale ``_transpose`` cache; after ``fsdp_post_all_gather`` + the rowwise sub-storage's ``_transpose`` must be invalidated / + regenerated to match the freshly gathered ``_data``. + """ + param = self._make_param() + # Build a plausible iter-2+ "out" with stale state. + out = HybridQuantizedTensor.make_like(param) + # Rowwise sub-storage on Hopper has _data populated. Force a stale + # _transpose and invalidate flag to mimic the regression scenario. + if out._rowwise_storage._transpose is None: + # Set up a fake stale transpose (non-None, marked invalid by + # the mismatching shape would catch nothing, so just plant + # a tensor and clear the invalid flag to "valid"). + out._rowwise_storage._transpose = torch.zeros_like( + out._rowwise_storage._data + ).t() + out._rowwise_storage._transpose_invalid = False + stale_transpose_id = id(out._rowwise_storage._transpose) + + # Drive a real all-gather round trip via the protocol + sharded_tensors, metadata = param.fsdp_pre_all_gather( + mesh=None, + orig_size=param.shape, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + out2, _ = param.fsdp_post_all_gather( + sharded_tensors, metadata, param.dtype, out=out + ) + + # After fsdp_post_all_gather, the rowwise sub-quantizer is pinned + # columnwise=False, so update_usage(rowwise=True, columnwise=False) + # must clear the stale _transpose (preventing the silent + # stale-cache regression on Hopper). + assert out2._rowwise_storage._transpose is None or ( + out2._rowwise_storage._transpose_invalid + and id(out2._rowwise_storage._transpose) != stale_transpose_id + ), "Stale _transpose was not invalidated after fsdp_post_all_gather" diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 48f5f04c71a..c07614d87e7 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -472,6 +472,7 @@ def test_hybrid_quantizer_copy_preserves_parent_flags(self): assert out.columnwise_quantizer.rowwise_usage is False assert out.columnwise_quantizer.columnwise_usage is True + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) def test_te_ops_basic_linear_accepts_hybrid_identity_quantized_weight(self): import transformer_engine.pytorch.ops as te_ops @@ -523,6 +524,7 @@ def qfactory(role): # pylint: disable=unused-argument ), ], ) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) def test_te_ops_quantize_then_gelu_accepts_identity_backed_tensors(self, qfactory): import transformer_engine.pytorch.ops as te_ops @@ -980,13 +982,13 @@ def test_quantize_master_weights_hybrid_identity_fp8_current(self): # ── te.Linear integration ──────────────────────────────────────────── -def _make_linears(in_f, out_f, seed=1234, dtype=torch.bfloat16): +def _make_linears(in_f, out_f, seed=1234, dtype=torch.bfloat16, bias=True): torch.manual_seed(seed) torch.cuda.manual_seed(seed) - ref = te.Linear(in_f, out_f, params_dtype=dtype).cuda() + ref = te.Linear(in_f, out_f, bias=bias, params_dtype=dtype).cuda() torch.manual_seed(seed) torch.cuda.manual_seed(seed) - test = te.Linear(in_f, out_f, params_dtype=dtype).cuda() + test = te.Linear(in_f, out_f, bias=bias, params_dtype=dtype).cuda() with torch.no_grad(): for p_test, p_ref in zip(test.parameters(), ref.parameters()): p_test.copy_(p_ref) @@ -1062,19 +1064,25 @@ def _fwd_bwd_checkpoint(model, x, recipe, use_reentrant): ) -def _make_identity_module(module_name, seed=1234, dtype=torch.bfloat16): +def _make_identity_module(module_name, seed=1234, dtype=torch.bfloat16, bias=True): hidden_size = 64 ffn_hidden_size = 128 torch.manual_seed(seed) torch.cuda.manual_seed(seed) if module_name == "Linear": - return te.Linear(hidden_size, hidden_size, params_dtype=dtype).cuda() + return te.Linear(hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() if module_name == "LayerNormLinear": - return te.LayerNormLinear(hidden_size, hidden_size, params_dtype=dtype).cuda() + return te.LayerNormLinear( + hidden_size, hidden_size, bias=bias, params_dtype=dtype + ).cuda() if module_name == "LayerNormMLP": - return te.LayerNormMLP(hidden_size, ffn_hidden_size, params_dtype=dtype).cuda() + return te.LayerNormMLP( + hidden_size, ffn_hidden_size, bias=bias, params_dtype=dtype + ).cuda() if module_name == "GroupedLinear": - return te.GroupedLinear(2, hidden_size, hidden_size, params_dtype=dtype).cuda() + return te.GroupedLinear( + 2, hidden_size, hidden_size, bias=bias, params_dtype=dtype + ).cuda() if module_name == "TransformerLayer": return te.TransformerLayer( hidden_size, @@ -1082,14 +1090,15 @@ def _make_identity_module(module_name, seed=1234, dtype=torch.bfloat16): 4, hidden_dropout=0.0, attention_dropout=0.0, + bias=bias, params_dtype=dtype, ).cuda() raise ValueError(module_name) -def _make_identity_module_pair(module_name, seed=1234, dtype=torch.bfloat16): - ref = _make_identity_module(module_name, seed=seed, dtype=dtype) - test = _make_identity_module(module_name, seed=seed + 1, dtype=dtype) +def _make_identity_module_pair(module_name, seed=1234, dtype=torch.bfloat16, bias=True): + ref = _make_identity_module(module_name, seed=seed, dtype=dtype, bias=bias) + test = _make_identity_module(module_name, seed=seed + 1, dtype=dtype, bias=bias) with torch.no_grad(): for p_test, p_ref in zip(test.parameters(), ref.parameters()): p_test.copy_(p_ref) @@ -1138,7 +1147,11 @@ class TestIdentityTEModuleCoverage: ], ) def test_identity_recipe_matches_bf16_bitwise(self, module_name, qfactory): - ref, test = _make_identity_module_pair(module_name, seed=7300) + # Keep the GEMM topology identical. Identity-backed grad-output uses an + # unfused bgrad path, while the plain BF16 path may fuse bgrad with + # wgrad. Those mathematically equivalent kernels need not accumulate + # wgrad in the same order on every architecture. + ref, test = _make_identity_module_pair(module_name, seed=7300, bias=False) x = _identity_module_input(module_name) recipe = CustomRecipe(qfactory=qfactory) @@ -1530,10 +1543,34 @@ def _input(self): torch.manual_seed(7) return torch.randn(self.BATCH, self.IN_F, device="cuda", dtype=torch.bfloat16) + @pytest.mark.parametrize( + "qfactory", + [ + pytest.param(identity_all_factory, id="plain_identity"), + pytest.param(hybrid_all_identity_factory, id="hybrid_identity"), + ], + ) + def test_identity_bias_path_matches_bf16_bitwise(self, qfactory): + """The unfused Identity bgrad path must preserve exact bias semantics. + + Weight-gradient bitwise parity is covered with ``bias=False`` so both + sides use the same GEMM topology. This bias-enabled test separately + locks output, input-gradient, and bias-gradient parity to zero tolerance. + """ + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=True) + x = self._input() + + y_ref, dx_ref, _ = _fwd_bwd(ref, x, recipe=None) + y_id, dx_id, _ = _fwd_bwd(test, x, recipe=CustomRecipe(qfactory=qfactory)) + + torch.testing.assert_close(y_id, y_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(dx_id, dx_ref, rtol=0.0, atol=0.0) + torch.testing.assert_close(test.bias.grad, ref.bias.grad, rtol=0.0, atol=0.0) + def test_whole_layer_hp_matches_bf16_bitwise(self): """Identity for every slot => all GEMMs high precision => bitwise-equal to a plain BF16 te.Linear (no autocast).""" - ref, test = _make_linears(self.IN_F, self.OUT_F) + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) x = self._input() y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) @@ -1598,7 +1635,7 @@ def test_hybrid_all_identity_matches_bf16_bitwise(self): plain BF16 te.Linear. Complements the non-hybrid whole-layer-HP test: this exercises HybridQuantizedTensor with Identity sub-storages in both directions and the per-operand unwrap of every GEMM.""" - ref, test = _make_linears(self.IN_F, self.OUT_F) + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) x = self._input() y_ref, dx_ref, wg_ref = _fwd_bwd(ref, x, recipe=None) @@ -1622,7 +1659,7 @@ def test_identity_reproduces_backward_override_high_precision_bitwise(self): grad = Identity); the reference uses the global ``backward_override`` knob. Identical forward FP8 + identical original HP backward operands => bitwise. """ - ref, test = _make_linears(self.IN_F, self.OUT_F) + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) x = self._input() y_bo, dx_bo, wg_bo = _fwd_bwd( @@ -1646,7 +1683,7 @@ def test_identity_reproduces_backward_override_dequantized_bitwise(self): the hybrid path stores Identity columnwise data sourced from the rowwise dequantized value and uses Identity grad tensors. """ - ref, test = _make_linears(self.IN_F, self.OUT_F) + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) x = self._input() y_bo, dx_bo, wg_bo = _fwd_bwd( @@ -1668,7 +1705,7 @@ def test_identity_matches_bf16_multistep_training_bitwise(self): """Multi-step SGD: an all-Identity recipe must track a plain BF16 te.Linear bitwise across optimizer steps (no drift from workspace caching or any hidden state).""" - ref, test = _make_linears(self.IN_F, self.OUT_F) + ref, test = _make_linears(self.IN_F, self.OUT_F, bias=False) opt_ref = torch.optim.SGD(ref.parameters(), lr=0.1) opt_test = torch.optim.SGD(test.parameters(), lr=0.1) recipe = CustomRecipe(qfactory=identity_all_factory) @@ -1770,7 +1807,7 @@ def test_quantized_model_init_identity_training_loss_decreases_bitwise(self): @pytest.mark.parametrize("use_reentrant", [True, False]) def test_identity_activation_recompute_matches_bf16_bitwise(self, use_reentrant): """All-Identity recompute should be exactly the BF16 no-checkpoint path.""" - ref, test = _make_linears(self.IN_F, self.OUT_F, seed=4242) + ref, test = _make_linears(self.IN_F, self.OUT_F, seed=4242, bias=False) recipe = CustomRecipe(qfactory=identity_all_factory) x = self._input() From 04caed2a2f50832b16ec7dca2f17bb1df1d52144 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 22 Jul 2026 12:35:42 +0000 Subject: [PATCH 57/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../fsdp2_tests/run_fsdp2_model.py | 4 +- tests/pytorch/hybrid_quantization_utils.py | 8 +- .../pytorch/test_hybrid_quantization_fsdp2.py | 117 +++++------------- tests/pytorch/test_identity_quantizer.py | 12 +- 4 files changed, 37 insertions(+), 104 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 77dad90b183..39d825e7013 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -254,9 +254,7 @@ def _check_fp8_fsdp2_allgather(model, *, tols=None): param_tols = dict(atol=5e-4, rtol=5e-3) else: param_tols = {} - torch.testing.assert_close( - param.dequantize(), fp32_allgathered_params[name], **param_tols - ) + torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **param_tols) # Revert model to original sharded state. for module in model.modules(): # Not all modules are wrapped/sharded with FSDP2. diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py index 6ad14f4471a..dfbe8d4c344 100644 --- a/tests/pytorch/hybrid_quantization_utils.py +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -94,9 +94,7 @@ def snapshot_storage_tensor_metadata(storage, *, clone=False): tensor_metadata = {} for name, value in storage.get_metadata().items(): if isinstance(value, torch.Tensor) or value is None: - snapshot_value = ( - value.detach().clone() if clone and value is not None else value - ) + snapshot_value = value.detach().clone() if clone and value is not None else value if ( value is not None and getattr(storage, "_is_2D_scaled", False) @@ -123,9 +121,7 @@ def snapshot_storage_tensor_metadata(storage, *, clone=False): def assert_nested_state_exact(actual, expected, *, path="state"): """Recursively compare nested state, using zero tolerance for tensors.""" if isinstance(expected, torch.Tensor): - assert isinstance( - actual, torch.Tensor - ), f"{path}: expected Tensor, got {type(actual)}" + assert isinstance(actual, torch.Tensor), f"{path}: expected Tensor, got {type(actual)}" torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0, msg=path) return if isinstance(expected, dict): diff --git a/tests/pytorch/test_hybrid_quantization_fsdp2.py b/tests/pytorch/test_hybrid_quantization_fsdp2.py index 055675bb465..cda3e7184af 100644 --- a/tests/pytorch/test_hybrid_quantization_fsdp2.py +++ b/tests/pytorch/test_hybrid_quantization_fsdp2.py @@ -48,8 +48,8 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) -fp8_block_scaling_available, reason_for_no_fp8_block_scaling = ( - te.is_fp8_block_scaling_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True ) _XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 = pytest.mark.xfail( @@ -165,9 +165,7 @@ def test_float8_split_uses_logical_shape_for_transpose_only_storage( pieces = torch.split(tensor, split_size, dim=dim) assert [tuple(piece.shape) for piece in pieces] == expected_shapes - assert [ - tuple(piece._transpose.shape) for piece in pieces - ] == expected_transpose_shapes + assert [tuple(piece._transpose.shape) for piece in pieces] == expected_transpose_shapes assert all(piece._data is None for piece in pieces) assert all(piece._transpose_invalid is False for piece in pieces) @@ -245,9 +243,7 @@ def test_hybrid_split_uses_columnwise_logical_shape_when_rowwise_is_absent(self) (5, 16), (2, 16), ] - assert [ - tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces - ] == [ + assert [tuple(piece.columnwise_sub_storage._transpose.shape) for piece in pieces] == [ (16, 5), (16, 5), (16, 2), @@ -311,15 +307,11 @@ def test_slice_and_select_preserve_transpose_only_payload(self): assert isinstance(sliced, Float8Tensor) assert sliced._data is None assert sliced._transpose.shape == torch.Size((16, 6)) - torch.testing.assert_close( - sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(sliced.dequantize(), expected[2:8], rtol=0.0, atol=0.0) assert isinstance(selected, Float8Tensor) assert selected._data is None assert selected._transpose.shape == torch.Size((12,)) - torch.testing.assert_close( - selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(selected.dequantize(), expected[:, 3], rtol=0.0, atol=0.0) def test_as_strided_row_shard_preserves_transpose_only_payload(self): tensor, expected = self._make_valid_transpose_only_float8_tensor() @@ -329,9 +321,7 @@ def test_as_strided_row_shard_preserves_transpose_only_payload(self): assert isinstance(shard, Float8Tensor) assert shard._data is None assert shard._transpose.shape == torch.Size((16, 5)) - torch.testing.assert_close( - shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(shard.dequantize(), expected[1:6], rtol=0.0, atol=0.0) def test_as_strided_falls_back_for_nonrepresentable_layout(self): tensor, expected = self._make_valid_transpose_only_float8_tensor() @@ -426,9 +416,7 @@ def test_initializes_every_present_direction_and_preserves_kwargs( # A plain-source copy routes through the result's parent quantizer. It # must update every allocated direction even though the source parent # had both of its mutable usage flags disabled before new_zeros. - plain_source = torch.full( - (3, 5), 7.0, dtype=torch.float32, device=source.device - ) + plain_source = torch.full((3, 5), 7.0, dtype=torch.float32, device=source.device) result.copy_(plain_source) for sub_storage in ( result.rowwise_sub_storage, @@ -454,9 +442,7 @@ def test_identity_substorages_allow_integer_dtype(self): @requires_fp8 @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 - @pytest.mark.parametrize( - "unsupported_dtype", (torch.int32, torch.float64, torch.bool) - ) + @pytest.mark.parametrize("unsupported_dtype", (torch.int32, torch.float64, torch.bool)) def test_non_identity_substorage_rejects_unsupported_dtype(self, unsupported_dtype): quantizer = HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), @@ -508,9 +494,7 @@ def test_initializes_nvfp4_data_scale_and_amax_buffers(self): result = source.new_zeros(source.shape) assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) - assert type(result.columnwise_sub_storage) is type( - source.columnwise_sub_storage - ) + assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) torch.testing.assert_close( result.dequantize(), torch.zeros_like(result.dequantize()), @@ -524,8 +508,7 @@ def test_initializes_nvfp4_data_scale_and_amax_buffers(self): buffers, storage = sub_storage.prepare_for_saving() try: assert all( - buffer is None or torch.count_nonzero(buffer).item() == 0 - for buffer in buffers + buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers ) finally: assert storage.restore_from_saved(buffers) == [] @@ -554,9 +537,7 @@ def test_initializes_float8_block_data_and_scale_buffers(self): result = source.new_zeros(source.shape) assert type(result.rowwise_sub_storage) is type(source.rowwise_sub_storage) - assert type(result.columnwise_sub_storage) is type( - source.columnwise_sub_storage - ) + assert type(result.columnwise_sub_storage) is type(source.columnwise_sub_storage) torch.testing.assert_close( result.dequantize(), torch.zeros_like(result.dequantize()), @@ -570,8 +551,7 @@ def test_initializes_float8_block_data_and_scale_buffers(self): buffers, storage = sub_storage.prepare_for_saving() try: assert all( - buffer is None or torch.count_nonzero(buffer).item() == 0 - for buffer in buffers + buffer is None or torch.count_nonzero(buffer).item() == 0 for buffer in buffers ) finally: assert storage.restore_from_saved(buffers) == [] @@ -595,12 +575,8 @@ def test_split_preserves_hybrid_type(self, hybrid_param): dim0 = hybrid_param.shape[0] chunk_size = dim0 // 2 pieces = torch.split(hybrid_param, chunk_size, dim=0) - expected_rowwise = torch.split( - hybrid_param.rowwise_sub_storage, chunk_size, dim=0 - ) - expected_columnwise = torch.split( - hybrid_param.columnwise_sub_storage, chunk_size, dim=0 - ) + expected_rowwise = torch.split(hybrid_param.rowwise_sub_storage, chunk_size, dim=0) + expected_columnwise = torch.split(hybrid_param.columnwise_sub_storage, chunk_size, dim=0) assert len(pieces) == len(expected_rowwise) == len(expected_columnwise) assert len(pieces) >= 2 @@ -712,8 +688,7 @@ def test_copy_between_mismatched_usages_raises_atomically(self, hybrid_param): dst = quantizer.quantize(torch.zeros_like(hybrid_param.dequantize())) src.update_usage(columnwise_usage=False) dst_before = tuple( - None if tensor is None else tensor.clone() - for tensor in _as_data_tensor_tuple(dst) + None if tensor is None else tensor.clone() for tensor in _as_data_tensor_tuple(dst) ) with pytest.raises( @@ -738,9 +713,7 @@ def test_copy_from_bf16_to_hybrid(self, hybrid_param): aten.copy_.default(param, bf16_data) assert isinstance(param, HybridQuantizedTensor) _assert_hybrid_tensor_exact(param, expected, context="BF16 copy_") - torch.testing.assert_close( - param.dequantize(), expected.dequantize(), rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(param.dequantize(), expected.dequantize(), rtol=0.0, atol=0.0) def test_new_zeros_returns_hybrid(self, hybrid_param): """new_zeros returns initialized storage that remains FSDP-copyable.""" @@ -753,12 +726,8 @@ def test_new_zeros_returns_hybrid(self, hybrid_param): assert result.shape == hybrid_param.shape assert result.rowwise_sub_storage is not None assert result.columnwise_sub_storage is not None - assert type(result.rowwise_sub_storage) is type( - hybrid_param.rowwise_sub_storage - ) - assert type(result.columnwise_sub_storage) is type( - hybrid_param.columnwise_sub_storage - ) + assert type(result.rowwise_sub_storage) is type(hybrid_param.rowwise_sub_storage) + assert type(result.columnwise_sub_storage) is type(hybrid_param.columnwise_sub_storage) torch.testing.assert_close( result.dequantize(), torch.zeros_like(result.dequantize()), @@ -771,9 +740,7 @@ def test_new_zeros_returns_hybrid(self, hybrid_param): torch.testing.assert_close( result.dequantize(), hybrid_param.dequantize(), rtol=0.0, atol=0.0 ) - _assert_hybrid_tensor_exact( - result, hybrid_param, context="new_zeros then copy_" - ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="new_zeros then copy_") def test_empty_like_returns_hybrid(self, hybrid_param): """empty_like must return a HybridQuantizedTensor.""" @@ -976,9 +943,7 @@ def test_post_all_gather_buffer_reuse(self, hybrid_param): assert ( second_result is first_result ), "Buffer reuse: post_all_gather(out=prev) should return the same object" - _assert_hybrid_tensor_exact( - second_result, hybrid_param, context="post-all-gather reuse" - ) + _assert_hybrid_tensor_exact(second_result, hybrid_param, context="post-all-gather reuse") def test_post_all_gather_dequantize_matches_original(self, hybrid_param): """Reconstructed tensor should dequantize close to the original.""" @@ -1020,9 +985,7 @@ def test_post_all_gather_sub_storage_types_correct(self, hybrid_param): out=None, ) assert type(result.rowwise_sub_storage) is orig_row_type - _assert_hybrid_tensor_exact( - result, hybrid_param, context="post-all-gather storage types" - ) + _assert_hybrid_tensor_exact(result, hybrid_param, context="post-all-gather storage types") assert type(result.columnwise_sub_storage) is orig_col_type @@ -1093,9 +1056,7 @@ def test_pre_post_roundtrip_buffer_reuse_preserves_data(self, hybrid_param): torch.testing.assert_close( hybrid_param.dequantize(), second_result.dequantize(), rtol=0.0, atol=0.0 ) - _assert_hybrid_tensor_exact( - second_result, hybrid_param, context="FSDP roundtrip reuse" - ) + _assert_hybrid_tensor_exact(second_result, hybrid_param, context="FSDP roundtrip reuse") @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_scale_refresh_across_iterations(self): @@ -1170,9 +1131,7 @@ def test_scale_refresh_across_iterations(self): rtol=0.0, atol=0.0, ) - _assert_hybrid_tensor_exact( - gathered_refreshed, hybrid_param, context="FSDP scale refresh" - ) + _assert_hybrid_tensor_exact(gathered_refreshed, hybrid_param, context="FSDP scale refresh") # And the magnitude really did change (sanity: this test would pass # vacuously if update_quantized didn't actually change anything). assert gathered_refreshed.dequantize().abs().max() > 10.0, ( @@ -1256,9 +1215,7 @@ def test_make_like_preserves_sub_storages(self): assert type(copy.rowwise_sub_storage) is type(param.rowwise_sub_storage) assert type(copy.columnwise_sub_storage) is type(param.columnwise_sub_storage) _assert_hybrid_tensor_exact(copy, param, context="make_like") - torch.testing.assert_close( - copy.dequantize(), param.dequantize(), rtol=0.0, atol=0.0 - ) + torch.testing.assert_close(copy.dequantize(), param.dequantize(), rtol=0.0, atol=0.0) def test_make_like_is_independent(self): """make_like result should not share the same tensor identity.""" @@ -1374,9 +1331,7 @@ def test_columnwise_only_float8_fsdp_extract_uses_m_major_transport(self): assert len(buffers) == 1 assert buffers[0].shape == out.shape assert buffers[0].is_contiguous() - torch.testing.assert_close( - buffers[0], out._transpose.movedim(0, -1).contiguous() - ) + torch.testing.assert_close(buffers[0], out._transpose.movedim(0, -1).contiguous()) assert metadata == { "field_names": ("_transpose",), "transport_layout": "columnwise_m_major", @@ -1394,12 +1349,8 @@ def test_columnwise_only_float8_fsdp_assign_restores_columnwise_layout(self): assert rebuilt._data is None assert rebuilt._transpose.shape == torch.Size((64, 64)) assert not rebuilt._transpose_invalid - torch.testing.assert_close( - rebuilt._transpose, gathered.movedim(-1, 0).contiguous() - ) - torch.testing.assert_close( - rebuilt.dequantize(), torch.cat((expected, expected), dim=0) - ) + torch.testing.assert_close(rebuilt._transpose, gathered.movedim(-1, 0).contiguous()) + torch.testing.assert_close(rebuilt.dequantize(), torch.cat((expected, expected), dim=0)) def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): shards = [self._make_hybrid_shard(value) for value in (1.0, 2.0)] @@ -1419,9 +1370,7 @@ def test_hybrid_fsdp_two_rank_gather_and_buffer_reuse(self): ) expected = torch.cat((shards[0][1], shards[1][1]), dim=0) - rebuilt, _ = shards[0][0].fsdp_post_all_gather( - gathered, extracted[0][1], torch.bfloat16 - ) + rebuilt, _ = shards[0][0].fsdp_post_all_gather(gathered, extracted[0][1], torch.bfloat16) assert rebuilt.shape == torch.Size((64, 64)) assert rebuilt._columnwise_storage.shape == torch.Size((64, 64)) @@ -1542,9 +1491,7 @@ def test_iter2_invalidates_stale_transpose_on_rowwise_substorage(self): # Set up a fake stale transpose (non-None, marked invalid by # the mismatching shape would catch nothing, so just plant # a tensor and clear the invalid flag to "valid"). - out._rowwise_storage._transpose = torch.zeros_like( - out._rowwise_storage._data - ).t() + out._rowwise_storage._transpose = torch.zeros_like(out._rowwise_storage._data).t() out._rowwise_storage._transpose_invalid = False stale_transpose_id = id(out._rowwise_storage._transpose) @@ -1556,9 +1503,7 @@ def test_iter2_invalidates_stale_transpose_on_rowwise_substorage(self): module=None, mp_policy=None, ) - out2, _ = param.fsdp_post_all_gather( - sharded_tensors, metadata, param.dtype, out=out - ) + out2, _ = param.fsdp_post_all_gather(sharded_tensors, metadata, param.dtype, out=out) # After fsdp_post_all_gather, the rowwise sub-quantizer is pinned # columnwise=False, so update_usage(rowwise=True, columnwise=False) diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index c07614d87e7..eb01ccc71dd 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -1072,17 +1072,11 @@ def _make_identity_module(module_name, seed=1234, dtype=torch.bfloat16, bias=Tru if module_name == "Linear": return te.Linear(hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() if module_name == "LayerNormLinear": - return te.LayerNormLinear( - hidden_size, hidden_size, bias=bias, params_dtype=dtype - ).cuda() + return te.LayerNormLinear(hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() if module_name == "LayerNormMLP": - return te.LayerNormMLP( - hidden_size, ffn_hidden_size, bias=bias, params_dtype=dtype - ).cuda() + return te.LayerNormMLP(hidden_size, ffn_hidden_size, bias=bias, params_dtype=dtype).cuda() if module_name == "GroupedLinear": - return te.GroupedLinear( - 2, hidden_size, hidden_size, bias=bias, params_dtype=dtype - ).cuda() + return te.GroupedLinear(2, hidden_size, hidden_size, bias=bias, params_dtype=dtype).cuda() if module_name == "TransformerLayer": return te.TransformerLayer( hidden_size, From a50300935c4e5e3370577c218f3b3bd134425b95 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 22 Jul 2026 12:43:13 +0000 Subject: [PATCH 58/97] Fix delayed-scaling hybrid FP8 FSDP guard Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 37 +++++++++++++--------- transformer_engine/pytorch/tensor/utils.py | 21 ++++++------ 2 files changed, 34 insertions(+), 24 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index c03101a9857..5d78984277e 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -4439,7 +4439,7 @@ def _hybrid_recipe_fp8_current(): ) -def _make_delayed_quantizer(fp8_dtype=None): +def _make_delayed_quantizer(fp8_dtype=None, *, rowwise=True, columnwise=True): """Construct a ``Float8Quantizer`` (delayed scaling) with locally-allocated scale/amax buffers for single-shot unit tests. @@ -4457,6 +4457,8 @@ def _make_delayed_quantizer(fp8_dtype=None): scale=torch.ones(1, dtype=torch.float32, device="cuda"), amax=torch.zeros(1, dtype=torch.float32, device="cuda"), fp8_dtype=fp8_dtype, + rowwise=rowwise, + columnwise=columnwise, ) @@ -4701,24 +4703,29 @@ def test_fp8_current_original_partial_master_raises(self): is_non_tn_fp8_gemm_supported(), reason="Hopper-only: Blackwell supports columnwise per-tensor FP8 FSDP updates", ) - def test_fp8_current_fsdp_transpose_only_raises_before_mutation(self): - """Reject Hopper FSDP updates that would materialize rowwise column storage.""" + @pytest.mark.parametrize("scaling", ("current", "delayed")) + def test_per_tensor_fp8_fsdp_transpose_only_raises_before_mutation(self, scaling): + """Reject Hopper FSDP updates that cannot flatten columnwise-only storage.""" from transformer_engine.pytorch.tensor.utils import quantize_master_weights shape = (4, 8) shard_shape = (2, 8) - row_quantizer = Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, - device="cuda", - rowwise=True, - columnwise=False, - ) - col_quantizer = Float8CurrentScalingQuantizer( - tex.DType.kFloat8E4M3, - device="cuda", - rowwise=False, - columnwise=True, - ) + if scaling == "current": + row_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=True, + columnwise=False, + ) + col_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + rowwise=False, + columnwise=True, + ) + else: + row_quantizer = _make_delayed_quantizer(rowwise=True, columnwise=False) + col_quantizer = _make_delayed_quantizer(rowwise=False, columnwise=True) quantizer = HybridQuantizer( rowwise_quantizer=row_quantizer, columnwise_quantizer=col_quantizer, diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 6695509327a..a2a0d37498e 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -253,7 +253,7 @@ def quantize_master_weights( for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( model_weights, master_weights, start_offsets, fsdp_shard_model_weights ): - _validate_fp8_current_scaling_fsdp_hopper_policy( + _validate_per_tensor_fp8_fsdp_hopper_policy( model_weight, fsdp_shard_model_weight, ) @@ -1170,6 +1170,9 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # hybrid sub-storage IS a single-direction Float8Tensor, so we route them as two # independent entries (into the same bucket for same-format, or into different # buckets for cross-format Float8 — e.g. delayed row + current col). +# The FSDP-sharded columnwise-only representation on Hopper is the exception: +# neither per-tensor cast helper can safely flatten its transpose-only payload, +# so `_validate_per_tensor_fp8_fsdp_hopper_policy` rejects it before mutation. # # Identity routes to an exact copy bucket. Single-direction hybrid (only one # sub-storage populated) routes the present direction only. Both-None hybrids @@ -1219,17 +1222,17 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # --------------------------------------------------------------------------------------------- -def _validate_fp8_current_scaling_fsdp_hopper_policy( +def _validate_per_tensor_fp8_fsdp_hopper_policy( model_weight, fsdp_shard_model_weight, ): - """Reject unsupported transpose-only current-scaling FSDP updates on Hopper. + """Reject unsupported transpose-only per-tensor FP8 FSDP updates on Hopper. The FSDP shard path can receive a ``Float8Tensor`` and pass it directly to - ``Float8Quantizer.update_quantized``. On Hopper, doing that for a - columnwise-only current-scaling tensor materializes rowwise ``_data`` and - violates the direction pinned by ``HybridQuantizer``. Reject the whole - batch before any scale, cache, or storage mutation instead. + a per-tensor FP8 cast helper. On Hopper, both delayed and current scaling + use a transpose-only representation for columnwise-only tensors that the + sharded update path cannot flatten safely. Reject the whole batch before + any scale, cache, or storage mutation instead. """ if fsdp_shard_model_weight is None or is_non_tn_fp8_gemm_supported(): return @@ -1246,12 +1249,12 @@ def _validate_fp8_current_scaling_fsdp_hopper_policy( for storage, quantizer in candidates: if ( storage is not None - and isinstance(quantizer, Float8CurrentScalingQuantizer) + and isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)) and _is_float8_transpose_only(storage) ): raise NotImplementedError( "Columnwise-only per-tensor FP8 quantization is not implemented for " - "FSDP current-scaling updates on this architecture." + "FSDP updates on this architecture." ) From 8942e70d0ec0dc69c43f963d71232c61a86e58b8 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 24 Jul 2026 13:59:48 +0000 Subject: [PATCH 59/97] Resolve comments Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 655 ++++++++++++++++-- tests/pytorch/test_identity_quantizer.py | 44 +- .../pytorch/cpp_extensions/gemm.py | 118 ++-- transformer_engine/pytorch/distributed.py | 19 +- transformer_engine/pytorch/module/_common.py | 31 + transformer_engine/pytorch/module/base.py | 8 - .../pytorch/module/grouped_linear.py | 363 +++++----- .../pytorch/module/layernorm_linear.py | 19 +- .../pytorch/module/layernorm_mlp.py | 27 +- transformer_engine/pytorch/module/linear.py | 77 +- transformer_engine/pytorch/ops/_common.py | 27 +- .../pytorch/ops/basic/layer_norm.py | 20 +- .../pytorch/ops/basic/rmsnorm.py | 11 +- .../pytorch/quantized_tensor.py | 12 +- .../pytorch/tensor/float8_blockwise_tensor.py | 4 + .../pytorch/tensor/float8_tensor.py | 4 + .../pytorch/tensor/hybrid_tensor.py | 117 ++-- .../pytorch/tensor/identity_tensor.py | 14 +- .../pytorch/tensor/mxfp8_tensor.py | 4 + .../pytorch/tensor/nvfp4_tensor.py | 4 + .../tensor/storage/hybrid_tensor_storage.py | 2 +- transformer_engine/pytorch/tensor/utils.py | 20 +- 22 files changed, 1093 insertions(+), 507 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 5d78984277e..5bc10e24306 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -5,6 +5,7 @@ """Tests for hybrid quantization (mixed rowwise/columnwise formats).""" import io +import warnings import pytest import torch @@ -59,9 +60,10 @@ QuantizedTensor, ) from transformer_engine.pytorch.cpp_extensions.gemm import ( - _unwrap_hybrid_A, - _unwrap_hybrid_B, + _unwrap_tensor, + _validate_native_gemm_output_quantizer, ) +from transformer_engine.pytorch.quantized_tensor import Quantizer from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported _fp8_col_factory = _fp8_row_factory @@ -108,6 +110,48 @@ ) +def test_native_gemm_output_quantizer_support_is_opt_in(): + """Unknown quantizers must be rejected before entering the native C++ path.""" + _validate_native_gemm_output_quantizer(None) + + unknown_quantizer = Quantizer(rowwise=True, columnwise=False) + with pytest.raises( + NotImplementedError, + match="Quantizer is not supported as a native GEMM output quantizer", + ): + _validate_native_gemm_output_quantizer(unknown_quantizer) + + +def test_hybrid_storage_snapshots_parent_quantizer(): + """Caller mutations must not change an existing Hybrid tensor's behavior.""" + quantizer = HybridQuantizer( + rowwise_quantizer=IdentityQuantizer(), + columnwise_quantizer=IdentityQuantizer(), + ) + tensor = quantizer(torch.ones((2, 2))) + + assert tensor._quantizer is not quantizer + assert tensor._quantizer.rowwise_quantizer is not quantizer.rowwise_quantizer + assert tensor._quantizer.columnwise_quantizer is not quantizer.columnwise_quantizer + + # Make the two Identity directions independent so a skipped update is observable. + tensor._rowwise_storage._hp_data = tensor._rowwise_storage._hp_data.clone() + tensor._columnwise_storage._hp_data = tensor._columnwise_storage._hp_data.clone() + + quantizer.set_usage(rowwise=False, columnwise=True) + tensor.copy_(torch.full(tensor.shape, 2.0)) + + assert tensor._quantizer.get_usages() == {"rowwise": True, "columnwise": True} + torch.testing.assert_close( + tensor._rowwise_storage.dequantize(), + torch.full(tensor.shape, 2.0), + ) + torch.testing.assert_close( + tensor._columnwise_storage.dequantize(), + torch.full(tensor.shape, 2.0), + ) + + def _make_hybrid_quantizer_fp4_row_fp8_col(): """NVFP4 rowwise + FP8 columnwise (reversed direction).""" return HybridQuantizer( @@ -146,14 +190,14 @@ def _snapshot_model_parameters(model): class _CountingIdentityQuantizer(IdentityQuantizer): - """Identity quantizer that counts quantize calls for regression tests.""" + """Replay-safe identity quantizer that counts quantize calls.""" def __init__(self, counter, *, dtype=None, rowwise=True, columnwise=True): super().__init__(dtype=dtype, rowwise=rowwise, columnwise=columnwise) self.counter = counter def copy(self): - quantizer = _CountingIdentityQuantizer( + quantizer = type(self)( self.counter, dtype=self.dtype, rowwise=self.rowwise_usage, @@ -168,6 +212,41 @@ def quantize_impl(self, tensor): return super().quantize_impl(tensor) +class _CountingUnsafeIdentityQuantizer(_CountingIdentityQuantizer): + """Non-replay-safe identity quantizer that counts quantize calls.""" + + def is_requantization_safe(self): + self.counter["safety_calls"] = self.counter.get("safety_calls", 0) + 1 + return False + + +class _CountingPythonQuantizer(Quantizer): + """Custom Python quantizer used to verify fallback dispatch.""" + + def __init__(self, calls, *, rowwise=True, columnwise=True): + super().__init__(rowwise=rowwise, columnwise=columnwise) + self.calls = calls + + def copy(self): + quantizer = type(self)( + self.calls, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm + return quantizer + + def quantize_impl(self, tensor): + self.calls.append(tensor) + fallback = IdentityQuantizer( + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + fallback.internal = self.internal + return fallback.quantize_impl(tensor) + + @requires_mxfp8_and_nvfp4 class TestComposerStyleFactory: """Composer 2-style row-scaled NVFP4 forward + MXFP8 backward dispatch.""" @@ -246,7 +325,7 @@ def _factory(role): @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) @pytest.mark.parametrize("tensor_type", ["input", "weight"]) - def test_forward_roles_require_saved_forward_quantized_value(self, module_type, tensor_type): + def test_forward_roles_are_requantization_safe(self, module_type, tensor_type): from transformer_engine.pytorch.quantization import QuantizerRole quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type=tensor_type)) @@ -255,16 +334,16 @@ def test_forward_roles_require_saved_forward_quantized_value(self, module_type, assert quantizer.columnwise_source == "rowwise_dequantized" assert isinstance(quantizer.rowwise_quantizer, MXFP8Quantizer) assert isinstance(quantizer.columnwise_quantizer, IdentityQuantizer) - assert quantizer.allows_save_original_input_for_backward() is False + assert quantizer.is_requantization_safe() is True @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) - def test_grad_output_role_allows_save_original_input(self, module_type): + def test_grad_output_role_is_requantization_safe(self, module_type): from transformer_engine.pytorch.quantization import QuantizerRole quantizer = self._factory(QuantizerRole(module_type=module_type, tensor_type="grad_output")) assert isinstance(quantizer, IdentityQuantizer) - assert quantizer.allows_save_original_input_for_backward() is True + assert quantizer.is_requantization_safe() is True @requires_nvfp4 @@ -477,17 +556,71 @@ def test_invalid_columnwise_source_raises(self): columnwise_source="dequantized", ) - def test_default_quantizer_allows_save_original_input_for_backward(self): - assert IdentityQuantizer().allows_save_original_input_for_backward() is True + def test_default_quantizer_is_requantization_safe(self): + assert IdentityQuantizer().is_requantization_safe() is True - @pytest.mark.parametrize( - "columnwise_source,allowed", - [("original", True), ("rowwise_dequantized", False)], - ) - def test_hybrid_save_original_input_policy(self, columnwise_source, allowed): + @pytest.mark.parametrize("columnwise_source", ("original", "rowwise_dequantized")) + def test_hybrid_requantization_safe_with_safe_sub_quantizers(self, columnwise_source): hq = self._make_quantizer(columnwise_source=columnwise_source) - assert hq.allows_save_original_input_for_backward() is allowed + assert hq.is_requantization_safe() is True + + def test_hybrid_requantization_checks_requested_sub_quantizers(self): + from transformer_engine.pytorch.module._common import ( + can_reconstruct_wgrad_input_from_original, + ) + + hq = HybridQuantizer( + rowwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_source="original", + ) + + assert hq.is_requantization_safe() is False + assert can_reconstruct_wgrad_input_from_original(hq) is True + + @pytest.mark.parametrize("columnwise_source", ("original", "rowwise_dequantized")) + @pytest.mark.parametrize("rowwise_safe", (True, False)) + @pytest.mark.parametrize("columnwise_safe", (True, False)) + def test_requantization_and_wgrad_reconstruction_policy_matrix( + self, + columnwise_source, + rowwise_safe, + columnwise_safe, + ): + from transformer_engine.pytorch.module._common import ( + can_reconstruct_wgrad_input_from_original, + ) + + rowwise_cls = ( + _CountingIdentityQuantizer if rowwise_safe else _CountingUnsafeIdentityQuantizer + ) + columnwise_cls = ( + _CountingIdentityQuantizer if columnwise_safe else _CountingUnsafeIdentityQuantizer + ) + hq = HybridQuantizer( + rowwise_quantizer=rowwise_cls({"calls": 0}), + columnwise_quantizer=columnwise_cls({"calls": 0}), + columnwise_source=columnwise_source, + ) + + # General requantization reproduces both requested representations. + assert hq.is_requantization_safe() is (rowwise_safe and columnwise_safe) + # Wgrad reconstruction only repeats the rowwise stage for a + # rowwise-dequantized columnwise source. + expected_reconstruction = columnwise_source == "original" or rowwise_safe + assert ( + can_reconstruct_wgrad_input_from_original(hq) is expected_reconstruction + ) + + def test_hybrid_rowwise_source_requires_safe_rowwise_quantizer(self): + hq = HybridQuantizer( + rowwise_quantizer=_CountingUnsafeIdentityQuantizer({"calls": 0}), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + + assert hq.is_requantization_safe() is False def test_copy_preserves_columnwise_source(self): hq = self._make_quantizer(columnwise_source="rowwise_dequantized") @@ -565,14 +698,118 @@ def test_update_quantized_uses_updated_rowwise_storage(self, input_tensor): class TestHybridSaveOriginalInputPolicy: """Module-level save_original_input policy for hybrid qfactory inputs.""" + def test_linear_high_precision_override_skips_requantization_safety_veto(self): + counters = {"rowwise": {"calls": 0}, "columnwise": {"calls": 0}} + custom_recipe = self._counting_identity_hybrid_recipe(counters) + custom_recipe.backward_override = "high_precision" + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp) + out.float().sum().backward() + + assert counters["rowwise"].get("safety_calls", 0) == 0 + + def test_linear_custom_delayed_scaling_disables_save_original_input(self): + from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( + delayed_scaling_quantizer_factory, + ) + + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + custom_recipe = recipe.CustomRecipe(qfactory=delayed_scaling_quantizer_factory) + + with pytest.warns( + UserWarning, + match="save_original_input is incompatible with delayed-scaling quantizers", + ): + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp) + out.float().sum().backward() + + def test_linear_builtin_delayed_scaling_rejects_save_original_input(self): + model = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with pytest.raises( + ValueError, + match="DelayedScaling recipe is not supported with save_original_input", + ): + with autocast(enabled=True, recipe=recipe.DelayedScaling()): + model(inp) + + def test_grouped_linear_classifies_requantization_safety_once_per_generation(self): + counters = [{"calls": 0}, {"calls": 0}] + input_quantizers = [ + _CountingUnsafeIdentityQuantizer(counter) for counter in counters + ] + generation = [] + for input_quantizer in input_quantizers: + generation.extend((input_quantizer, IdentityQuantizer(), IdentityQuantizer())) + + module = GroupedLinear( + 2, + 16, + 16, + bias=False, + device="meta", + ) + module.quantizers["scaling_fwd"] = generation + + module._validate_quantizer_generation(fwd=True) + assert module._unsafe_requantization_input_quantizer is input_quantizers[0] + assert [counter.get("safety_calls", 0) for counter in counters] == [1, 0] + + # The generation list is stable between forwards, so the O(1) identity + # guard must avoid re-running capability checks. + module._validate_quantizer_generation(fwd=True) + assert [counter.get("safety_calls", 0) for counter in counters] == [1, 0] + @staticmethod - def _counting_identity_hybrid_recipe(counter): + def _counting_identity_hybrid_recipe( + counters, + *, + columnwise_source="rowwise_dequantized", + rowwise_safe=False, + columnwise_safe=True, + ): def factory(role): if role is not None and role.module_type == "linear" and role.tensor_type == "input": + rowwise_cls = ( + _CountingIdentityQuantizer + if rowwise_safe + else _CountingUnsafeIdentityQuantizer + ) + columnwise_cls = ( + _CountingIdentityQuantizer + if columnwise_safe + else _CountingUnsafeIdentityQuantizer + ) return HybridQuantizer( - rowwise_quantizer=_CountingIdentityQuantizer(counter), - columnwise_quantizer=IdentityQuantizer(), - columnwise_source="rowwise_dequantized", + rowwise_quantizer=rowwise_cls(counters["rowwise"]), + columnwise_quantizer=columnwise_cls(counters["columnwise"]), + columnwise_source=columnwise_source, ) return IdentityQuantizer() @@ -580,7 +817,7 @@ def factory(role): def test_linear_save_original_input_veto_uses_saved_forward_quantized_input(self): torch.manual_seed(205) - counter = {"calls": 0} + counters = {"rowwise": {"calls": 0}, "columnwise": {"calls": 0}} model = Linear( 128, 128, @@ -593,16 +830,105 @@ def test_linear_save_original_input_veto_uses_saved_forward_quantized_input(self with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): with autocast( enabled=True, - recipe=self._counting_identity_hybrid_recipe(counter), + recipe=self._counting_identity_hybrid_recipe(counters), ): out = model(inp) - calls_after_forward = counter["calls"] + calls_after_forward = counters["rowwise"]["calls"] assert calls_after_forward > 0 out.float().sum().backward() - assert counter["calls"] == calls_after_forward + assert counters["rowwise"]["calls"] == calls_after_forward + + @pytest.mark.parametrize( + "columnwise_source,rowwise_safe,columnwise_safe,optimization_enabled," + "expected_rowwise_calls", + ( + ("original", True, True, True, 1), + ("original", False, False, True, 1), + ("rowwise_dequantized", True, True, True, 2), + ("rowwise_dequantized", True, False, True, 2), + ("rowwise_dequantized", False, True, False, 1), + ), + ) + def test_linear_save_original_input_policy_is_bitwise_exact( + self, + columnwise_source, + rowwise_safe, + columnwise_safe, + optimization_enabled, + expected_rowwise_calls, + ): + torch.manual_seed(206) + in_features, out_features, batch = 128, 128, 32 + reference = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=False, + ).cuda() + candidate = Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + save_original_input=True, + ).cuda() + candidate.load_state_dict(reference.state_dict()) + + base_input = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + reference_input = base_input.clone().detach().requires_grad_(True) + candidate_input = base_input.clone().detach().requires_grad_(True) + reference_counters = { + "rowwise": {"calls": 0}, + "columnwise": {"calls": 0}, + } + candidate_counters = { + "rowwise": {"calls": 0}, + "columnwise": {"calls": 0}, + } + recipe_kwargs = { + "columnwise_source": columnwise_source, + "rowwise_safe": rowwise_safe, + "columnwise_safe": columnwise_safe, + } + + with autocast( + enabled=True, + recipe=self._counting_identity_hybrid_recipe(reference_counters, **recipe_kwargs), + ): + reference_output = reference(reference_input) + + candidate_recipe = self._counting_identity_hybrid_recipe( + candidate_counters, + **recipe_kwargs, + ) + if optimization_enabled: + with autocast(enabled=True, recipe=candidate_recipe): + candidate_output = candidate(candidate_input) + else: + with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): + with autocast(enabled=True, recipe=candidate_recipe): + candidate_output = candidate(candidate_input) + + assert torch.equal(reference_output, candidate_output) + + reference_output.float().sum().backward() + candidate_output.float().sum().backward() + + assert reference_input.grad is not None and candidate_input.grad is not None + assert torch.equal(reference_input.grad, candidate_input.grad) + reference_weight_grad = dict(reference.named_parameters())["weight"].grad + candidate_weight_grad = dict(candidate.named_parameters())["weight"].grad + assert reference_weight_grad is not None and candidate_weight_grad is not None + assert torch.equal(reference_weight_grad, candidate_weight_grad) + + # The columnwise quantizer is always invoked exactly once. Only a + # rowwise-dequantized reconstruction safely repeats the rowwise stage. + assert candidate_counters["rowwise"]["calls"] == expected_rowwise_calls + assert candidate_counters["columnwise"]["calls"] == 1 @requires_fp8_and_nvfp4 @@ -1026,6 +1352,19 @@ def test_detach_shares_underlying_buffers(self, input_tensor): continue assert a is b, "detach() must share buffer tensors, not copy them" + @pytest.mark.parametrize("direction", ("rowwise", "columnwise")) + def test_detach_rejects_storage_only_sub_storage(self, input_tensor, direction): + """Storage-only children have no supported wrapper-copy protocol.""" + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + getattr(hq, f"{direction}_quantizer").internal = True + ht = hq.quantize(input_tensor) + + with pytest.raises( + NotImplementedError, + match=rf"storage-only {direction} sub-storage", + ): + ht.detach() + @requires_fp8_and_nvfp4 class TestHybridSaveRestore: @@ -1109,6 +1448,17 @@ def test_make_empty_has_sub_storages(self): assert empty.rowwise_sub_storage is not None assert empty.columnwise_sub_storage is not None + def test_make_empty_internal_returns_storage(self): + hq = _make_hybrid_quantizer_fp8_row_fp4_col() + hq.internal = True + + empty = hq.make_empty((128, 256), dtype=torch.bfloat16, device="cuda") + + assert isinstance(empty, HybridQuantizedTensorStorage) + assert not isinstance(empty, HybridQuantizedTensor) + assert empty.size() == torch.Size((128, 256)) + assert empty.dequantize().dtype == torch.bfloat16 + @requires_fp8_and_nvfp4 class TestHybridUsageFlagsRespected: @@ -1608,9 +1958,8 @@ def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwis out_ref = model_ref(inp_ref) qfactory_recipe = recipe.CustomRecipe(qfactory=mxfp8_fwd_dequantized_bwd_quantizer_factory) - with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): - with autocast(enabled=True, recipe=qfactory_recipe): - out_qfactory = model_qfactory(inp_qfactory) + with autocast(enabled=True, recipe=qfactory_recipe): + out_qfactory = model_qfactory(inp_qfactory) assert torch.equal( out_ref, out_qfactory @@ -2641,12 +2990,8 @@ def mixed_factory(role): @requires_fp8_and_nvfp4 -class TestUnwrapHybridDirection: - """Test per-operand unwrap selects the correct sub-storage. - - Operand A: transposed (layout[0]=='T') → rowwise, else → columnwise - Operand B: not-transposed (layout[1]=='N') → rowwise, else → columnwise - """ +class TestUnwrapTensor: + """Test GEMM input dispatch by rowwise or columnwise usage.""" @pytest.fixture def hybrid_tensor(self): @@ -2655,49 +3000,57 @@ def hybrid_tensor(self): hq = _make_hybrid_quantizer_fp8_row_fp4_col() return hq.quantize(inp) - def test_A_tn_returns_rowwise(self, hybrid_tensor): - assert _unwrap_hybrid_A(hybrid_tensor, "TN") is hybrid_tensor.rowwise_sub_storage - - def test_A_nn_returns_columnwise(self, hybrid_tensor): - assert _unwrap_hybrid_A(hybrid_tensor, "NN") is hybrid_tensor.columnwise_sub_storage - - def test_A_nt_returns_columnwise(self, hybrid_tensor): - assert _unwrap_hybrid_A(hybrid_tensor, "NT") is hybrid_tensor.columnwise_sub_storage - - def test_B_tn_returns_rowwise(self, hybrid_tensor): - assert _unwrap_hybrid_B(hybrid_tensor, "TN") is hybrid_tensor.rowwise_sub_storage + def test_hybrid_rowwise_returns_rowwise_sub_storage(self, hybrid_tensor): + assert _unwrap_tensor(hybrid_tensor, "rowwise") is hybrid_tensor.rowwise_sub_storage - def test_B_nn_returns_rowwise(self, hybrid_tensor): - assert _unwrap_hybrid_B(hybrid_tensor, "NN") is hybrid_tensor.rowwise_sub_storage + def test_hybrid_columnwise_returns_columnwise_sub_storage(self, hybrid_tensor): + assert _unwrap_tensor(hybrid_tensor, "columnwise") is hybrid_tensor.columnwise_sub_storage - def test_B_nt_returns_columnwise(self, hybrid_tensor): - assert _unwrap_hybrid_B(hybrid_tensor, "NT") is hybrid_tensor.columnwise_sub_storage - - def test_tn_sub_storage_type(self, hybrid_tensor): + def test_rowwise_sub_storage_type(self, hybrid_tensor): assert isinstance( - _unwrap_hybrid_A(hybrid_tensor, "TN"), + _unwrap_tensor(hybrid_tensor, "rowwise"), (Float8TensorStorage, Float8Tensor), ) - def test_nt_sub_storage_type(self, hybrid_tensor): + def test_columnwise_sub_storage_type(self, hybrid_tensor): assert isinstance( - _unwrap_hybrid_B(hybrid_tensor, "NT"), + _unwrap_tensor(hybrid_tensor, "columnwise"), (NVFP4TensorStorage, NVFP4Tensor), ) def test_non_hybrid_passthrough(self): plain = torch.randn(4, 4, device="cuda") - for layout in ("TN", "NN", "NT"): - assert _unwrap_hybrid_A(plain, layout) is plain - assert _unwrap_hybrid_B(plain, layout) is plain + for usage in ("rowwise", "columnwise"): + assert _unwrap_tensor(plain, usage) is plain def test_fp8_tensor_passthrough(self): quantizer = _make_fp8_quantizer() inp = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") fp8 = quantizer.quantize(inp) - for layout in ("TN", "NN", "NT"): - assert _unwrap_hybrid_A(fp8, layout) is fp8 - assert _unwrap_hybrid_B(fp8, layout) is fp8 + for usage in ("rowwise", "columnwise"): + assert _unwrap_tensor(fp8, usage) is fp8 + + def test_identity_falls_back_to_high_precision(self): + inp = torch.randn(4, 4, device="cuda") + identity = IdentityQuantizer()(inp) + + result = _unwrap_tensor(identity, "rowwise") + + assert isinstance(result, torch.Tensor) + assert not isinstance(result, QuantizedTensor) + assert torch.equal(result, inp) + + def test_custom_tensor_passthrough(self): + from transformer_engine.pytorch.custom_recipes.quantization_ref_current_scaling import ( + CurrentScalingTensorRef, + ) + + custom = CurrentScalingTensorRef() + assert _unwrap_tensor(custom, "rowwise") is custom + + def test_invalid_usage_raises(self): + with pytest.raises(ValueError, match="Unsupported GEMM tensor usage"): + _unwrap_tensor(torch.empty(0), "diagonal") @requires_fp8 @@ -3591,6 +3944,44 @@ def factory(role): return factory +@requires_fp8 +@pytest.mark.parametrize("norm_cls", (te.ops.LayerNorm, te.ops.RMSNorm)) +def test_fusible_norm_hybrid_output_falls_back_to_explicit_quantize(norm_cls): + """Unsupported fused Hybrid output is quantized by the following operation.""" + + def qfactory(_role): + return HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=IdentityQuantizer(), + ) + + hidden_size = 128 + norm = norm_cls( + hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + forward = te.ops.Sequential(norm, te.ops.Quantize()) + inp = torch.randn( + 16, + hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=qfactory)): + out = forward(inp) + + assert isinstance(out, HybridQuantizedTensor) + out.backward(torch.randn_like(inp)) + assert inp.grad is not None + assert norm.weight.grad is not None + + @requires_fp8 @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 class TestHybridAllModules: @@ -3782,6 +4173,67 @@ def test_uniform_lists_validate(self, quantizers): _validate_grouped_quantizer_list(quantizers, operand_name="input") + def test_plain_custom_quantizer_uses_python_split_fallback(self, monkeypatch): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + monkeypatch.setattr( + grouped_linear.tex, + "split_quantize", + lambda *args, **kwargs: pytest.fail("entered native split_quantize"), + ) + calls = [] + tensor = torch.randn((4, 8)) + quantizers = [_CountingPythonQuantizer(calls) for _ in range(2)] + + out = grouped_linear._split_quantize_non_hybrid( + tensor, + [2, 2], + quantizers, + tensor.dtype, + ) + + assert len(calls) == 2 + for actual, expected in zip(out, torch.split(tensor, [2, 2])): + torch.testing.assert_close(actual.dequantize(), expected, rtol=0.0, atol=0.0) + + @pytest.mark.parametrize("direction", ("rowwise", "columnwise")) + def test_hybrid_custom_child_uses_python_split_fallback( + self, + monkeypatch, + direction, + ): + import transformer_engine.pytorch.module.grouped_linear as grouped_linear + + monkeypatch.setattr( + grouped_linear.tex, + "split_quantize", + lambda *args, **kwargs: pytest.fail("entered native split_quantize"), + ) + calls = [] + quantizers = [] + for _ in range(2): + rowwise = IdentityQuantizer() + columnwise = IdentityQuantizer() + if direction == "rowwise": + rowwise = _CountingPythonQuantizer(calls) + else: + columnwise = _CountingPythonQuantizer(calls) + quantizer = HybridQuantizer( + rowwise_quantizer=rowwise, + columnwise_quantizer=columnwise, + ) + quantizers.append(quantizer) + + out = grouped_linear._split_quantize_hybrid( + torch.randn((4, 8)), + [2, 2], + quantizers, + ) + + assert len(calls) == 2 + assert all(result.rowwise_sub_storage is not None for result in out) + assert all(result.columnwise_sub_storage is not None for result in out) + def test_mixed_hybrid_and_plain_raises(self): from transformer_engine.pytorch.module.grouped_linear import ( _validate_grouped_quantizer_list, @@ -3851,7 +4303,7 @@ def test_distinct_delayed_scaling_state_is_allowed(self): ) def test_hybrid_split_quantize_respects_parent_usage_flags(self, usage, expected): from transformer_engine.pytorch.module.grouped_linear import ( - _hybrid_split_quantize, + _split_quantize_hybrid, ) tensor = torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") @@ -3865,7 +4317,7 @@ def test_hybrid_split_quantize_respects_parent_usage_flags(self, usage, expected for quantizer in quantizers: quantizer.set_usage(rowwise=usage[0], columnwise=usage[1]) - out = _hybrid_split_quantize(tensor, [16, 16], quantizers) + out = _split_quantize_hybrid(tensor, [16, 16], quantizers) assert [storage.get_usages() for storage in out] == [expected, expected] @@ -3894,7 +4346,7 @@ def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): for quantizer in quantizers: quantizer.set_usage(rowwise=False, columnwise=True) - out = grouped_linear._hybrid_split_quantize(tensor, [16, 16], quantizers) + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) assert len(calls) == 2 expected_columnwise_source = torch.cat( @@ -3933,7 +4385,7 @@ def tracked_split_quantize(tensor, splits, quantizers, **kwargs): for _ in m_splits ] - out = grouped_linear._hybrid_split_quantize(tensor, m_splits, quantizers) + out = grouped_linear._split_quantize_hybrid(tensor, m_splits, quantizers) assert input_dtypes == [tensor.dtype, tensor.dtype] assert len(out) == len(m_splits) @@ -3961,7 +4413,7 @@ def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): for quantizer in quantizers: quantizer.set_usage(rowwise=True, columnwise=False) - out = grouped_linear._hybrid_split_quantize(tensor, [16, 16], quantizers) + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) assert calls == [tensor] assert all(storage.rowwise_sub_storage is not None for storage in out) @@ -3989,7 +4441,7 @@ def tracked_split_quantize(tensor, m_splits, quantizers, **kwargs): for _ in range(2) ] - out = grouped_linear._hybrid_split_quantize(tensor, [16, 16], quantizers) + out = grouped_linear._split_quantize_hybrid(tensor, [16, 16], quantizers) assert calls == [tensor, tensor] assert all(storage.rowwise_sub_storage is not None for storage in out) @@ -4107,7 +4559,7 @@ def mixed_source_qfactory(role): def test_hybrid_split_quantize_honors_rowwise_dequantized_source(self): """NVFP4 column data must derive from the actual grouped row result.""" from transformer_engine.pytorch.module.grouped_linear import ( - _hybrid_split_quantize, + _split_quantize_hybrid, ) torch.manual_seed(3598) @@ -4123,7 +4575,7 @@ def make_quantizer(): ) quantizers = [make_quantizer(), make_quantizer()] - actual = _hybrid_split_quantize( + actual = _split_quantize_hybrid( tensor, [64, 64], quantizers, @@ -4571,8 +5023,6 @@ class TestHybridQuantizeMasterWeights: sub-storage(s) dequantize close to the master weight after the cast: * Float8CurrentScaling on both directions (same-format, full master) - * Float8CurrentScaling on both directions (DP-sharded master, non-zero - start_offset) * Float8 delayed scaling on both directions (same-format) * Float8 delayed row + Float8 current col (cross-format; row -> delayed bucket, col -> current bucket) @@ -4585,6 +5035,8 @@ class TestHybridQuantizeMasterWeights: * MXFP8 as a hybrid sub-quantizer (rowwise OR columnwise) * NVFP4 as a hybrid sub-quantizer (rowwise OR columnwise) * Float8Blockwise as a hybrid sub-quantizer + * A live ``rowwise_dequantized`` column (deferred to #3158) + * A partial ``original``-source master without a two-direction FSDP shard * Both sub-storages dropped (caller bug: nothing left to cast) """ @@ -4641,6 +5093,63 @@ def _reference_fp8_bytes(master, scale, dtype=torch.bfloat16): quantizer.update_quantized(master.reshape(1, -1), temp) return temp._data.reshape(-1) + @pytest.mark.parametrize("partial_master", (False, True)) + def test_rowwise_dequantized_master_update_raises_before_mutation( + self, + partial_master, + ): + """Defer row-to-column sequencing to the #3158 follow-up.""" + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + group = _ensure_single_rank_dp_group() + shape = (4, 8) + quantizer = HybridQuantizer( + rowwise_quantizer=Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + device="cuda", + ), + columnwise_quantizer=IdentityQuantizer(), + columnwise_source="rowwise_dequantized", + ) + source = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + weight = quantizer(source) + state_before = { + "rowwise": _snapshot_storage_tensor_metadata( + weight._rowwise_storage, + clone=True, + ), + "columnwise": _snapshot_storage_tensor_metadata( + weight._columnwise_storage, + clone=True, + ), + } + start_offset = shape[-1] if partial_master else 0 + master = torch.randn(weight.numel(), dtype=torch.float32, device="cuda") + master = master[start_offset:].contiguous() + + with pytest.raises( + NotImplementedError, + match="rowwise update/all-gather.*#3158", + ): + quantize_master_weights( + [weight], + [master], + [start_offset], + group=group, + ) + + state_after = { + "rowwise": _snapshot_storage_tensor_metadata( + weight._rowwise_storage, + clone=True, + ), + "columnwise": _snapshot_storage_tensor_metadata( + weight._columnwise_storage, + clone=True, + ), + } + _assert_nested_state_exact(state_after, state_before) + def test_fp8_current_original_partial_master_raises(self): """Reject an independently sourced Hybrid column from a partial master shard. @@ -4689,7 +5198,7 @@ def test_fp8_current_original_partial_master_raises(self): ValueError, match=( "partial master shard.*columnwise_source='original'.*" - "columnwise_source='rowwise_dequantized'" + "full-master data" ), ): quantize_master_weights( @@ -6578,12 +7087,12 @@ def _build_and_run(use_checkpoint): def _run_grouped_linear(self, recipe_obj, *, checkpoint_fn=None): """Build a GroupedLinear, run forward+backward with optional activation checkpointing around the module. Exercises the - ``_hybrid_split_quantize`` code path under recompute. + ``_split_quantize_hybrid`` code path under recompute. GroupedLinear is the MoE token-dispatch kernel: a single batch is split along dim-0 into ``num_gemms`` chunks and each chunk goes through its own weight matrix. Under hybrid quantization, - ``_hybrid_split_quantize`` (``module/grouped_linear.py``) runs + ``_split_quantize_hybrid`` (``module/grouped_linear.py``) runs ``tex.split_quantize`` twice (once per sub-quantizer direction) and zips the results into a list of ``HybridQuantizedTensor`` chunks — save-for-backward then receives a *list* of hybrid @@ -6614,7 +7123,7 @@ def _run_grouped_linear(self, recipe_obj, *, checkpoint_fn=None): @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 def test_te_checkpoint_reentrant_grouped_linear_fp8_bitwise(self): """GroupedLinear + te.checkpoint(reentrant) under same-format FP8 - hybrid. Exercises the MoE ``_hybrid_split_quantize`` + list-of- + hybrid. Exercises the MoE ``_split_quantize_hybrid`` + list-of- hybrid-tensors save-for-backward path under recompute.""" import transformer_engine.pytorch as te_pytorch diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index eb01ccc71dd..3f924b23614 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -245,16 +245,27 @@ def test_internal_returns_storage(self): assert isinstance(out, IdentityTensorStorage) assert not isinstance(out, IdentityTensor) + def test_make_empty_internal_returns_storage(self): + q = IdentityQuantizer() + q.internal = True + + out = q.make_empty((8, 16), dtype=torch.bfloat16, device="cuda") + + assert isinstance(out, IdentityTensorStorage) + assert not isinstance(out, IdentityTensor) + assert out.size() == torch.Size((8, 16)) + assert out.dequantize().dtype == torch.bfloat16 + def test_grouped_split_all_identity_uses_plain_tensor_views(self): from transformer_engine.pytorch.module.grouped_linear import ( - _split_quantize_with_identity_fallback, + _split_quantize_non_hybrid, ) x = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) m_splits = [3, 5] quantizers = [IdentityQuantizer(), IdentityQuantizer()] - out = _split_quantize_with_identity_fallback( + out = _split_quantize_non_hybrid( x, m_splits, quantizers, activation_dtype=torch.bfloat16 ) @@ -267,7 +278,7 @@ def test_grouped_split_all_identity_uses_plain_tensor_views(self): IdentityQuantizer(dtype=torch.float32), IdentityQuantizer(dtype=torch.float32), ] - cast_out = _split_quantize_with_identity_fallback( + cast_out = _split_quantize_non_hybrid( x, m_splits, cast_quantizers, activation_dtype=torch.bfloat16 ) assert all(isinstance(t, IdentityTensorStorage) for t in cast_out) @@ -306,7 +317,7 @@ def test_grouped_split_rejects_mixed_identity_and_quantized_operands(self): def test_hybrid_split_forwards_disable_bulk_allocation_to_both_directions(self, monkeypatch): import transformer_engine.pytorch.module.grouped_linear as grouped_linear - from transformer_engine.pytorch.module.grouped_linear import _hybrid_split_quantize + from transformer_engine.pytorch.module.grouped_linear import _split_quantize_hybrid calls = [] @@ -318,6 +329,11 @@ def fake_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocation ] monkeypatch.setattr(grouped_linear.tex, "split_quantize", fake_split_quantize) + monkeypatch.setattr( + grouped_linear, + "_supports_native_split_quantize", + lambda quantizer: True, + ) x = torch.randn(8, 16, dtype=torch.bfloat16) m_splits = [3, 5] quantizers = [ @@ -328,7 +344,7 @@ def fake_split_quantize(tensor, m_splits, quantizers, *, disable_bulk_allocation for _ in m_splits ] - out = _hybrid_split_quantize( + out = _split_quantize_hybrid( x, m_splits, quantizers, @@ -357,7 +373,7 @@ def qfactory(role): calls = [] - def fake_hybrid_split_quantize( + def fake_split_quantize_hybrid( tensor, m_splits, quantizers, *, disable_bulk_allocation=False, **kwargs ): del tensor, m_splits, quantizers, kwargs @@ -365,7 +381,9 @@ def fake_hybrid_split_quantize( raise StopAfterFlagCapture("captured hybrid split kwargs") monkeypatch.setattr(grouped_linear, "is_cpu_offload_enabled", lambda: True) - monkeypatch.setattr(grouped_linear, "_hybrid_split_quantize", fake_hybrid_split_quantize) + monkeypatch.setattr( + grouped_linear, "_split_quantize_hybrid", fake_split_quantize_hybrid + ) model = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) @@ -1486,15 +1504,9 @@ def test_matches_base_recipe_bitwise(self, case_name, module_name): y_ref, dx_ref, grads_ref = _fwd_bwd_zoo_dequantized_module( module_name, ref, inp, ref_recipe ) - if module_name in ("Linear", "GroupedLinear", "LayerNormLinearLinear"): - with pytest.warns(UserWarning, match="Ignoring save_original_input=True"): - y_test, dx_test, grads_test = _fwd_bwd_zoo_dequantized_module( - module_name, test, inp, qfactory_recipe - ) - else: - y_test, dx_test, grads_test = _fwd_bwd_zoo_dequantized_module( - module_name, test, inp, qfactory_recipe - ) + y_test, dx_test, grads_test = _fwd_bwd_zoo_dequantized_module( + module_name, test, inp, qfactory_recipe + ) torch.testing.assert_close(y_test, y_ref, rtol=0.0, atol=0.0) torch.testing.assert_close(dx_test, dx_ref, rtol=0.0, atol=0.0) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 753a76d1fb1..e14a4cfadac 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -4,7 +4,7 @@ """Python interface for GEMM extensions""" -from typing import Iterable, Optional, Tuple, Union, List +from typing import Iterable, Literal, Optional, Tuple, Union, List import os import functools import torch @@ -12,19 +12,21 @@ from ..constants import TE_DType, DType from ..utils import get_sm_count, _empty_tensor -from ..quantized_tensor import Quantizer +from ..quantized_tensor import QuantizedTensorStorage, Quantizer +from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer +from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer +from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage +from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage +from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage +from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom -from ..tensor.hybrid_tensor import HybridQuantizer -from ..tensor.identity_tensor import IdentityQuantizer -from ..tensor.storage.hybrid_tensor_storage import HybridQuantizedTensorStorage -from ..tensor.storage.identity_tensor_storage import IdentityTensorStorage from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizer - __all__ = [ "general_gemm", "general_grouped_gemm", @@ -106,60 +108,62 @@ def _nvfp4_row_scaled_gemm_inputs( ) -def _unwrap_hybrid_A(tensor, layout): - """Extract the direction-appropriate native sub-storage for GEMM operand A. +_NATIVE_GEMM_INPUT_STORAGES = ( + Float8TensorStorage, + MXFP8TensorStorage, + Float8BlockwiseQTensorStorage, + NVFP4TensorStorage, +) - Operand A's data direction is determined by its transpose flag (layout[0]): - T (transposed) → rowwise sub-storage (.data consumed by C++) - N (not-transposed) → columnwise sub-storage (.columnwise_data consumed by C++) - For non-hybrid tensors this is a no-op passthrough. - """ - if not isinstance(tensor, HybridQuantizedTensorStorage): + +def _unwrap_tensor( + tensor: Union[torch.Tensor, QuantizedTensorStorage], + usage: Literal["rowwise", "columnwise"], +) -> Union[torch.Tensor, QuantizedTensorStorage]: + """Prepare a tensor for native or custom GEMM dispatch.""" + if usage not in ("rowwise", "columnwise"): + raise ValueError(f"Unsupported GEMM tensor usage ({usage})") + + # Plain PyTorch tensor + if not isinstance(tensor, QuantizedTensorStorage): return tensor - if layout[0] == "T": - return tensor.rowwise_sub_storage - return tensor.columnwise_sub_storage + # Select the direction of a hybrid tensor, then process its sub-storage. + if isinstance(tensor, HybridQuantizedTensorStorage): + sub_storage = ( + tensor.rowwise_sub_storage if usage == "rowwise" else tensor.columnwise_sub_storage + ) + return _unwrap_tensor(sub_storage, usage) -def _unwrap_hybrid_B(tensor, layout): - """Extract the direction-appropriate native sub-storage for GEMM operand B. + # Preserve custom tensors for custom_gemm dispatch. + if is_custom(tensor): + return tensor - Operand B's data direction is determined by its transpose flag (layout[1]): - N (not-transposed) → rowwise sub-storage (.data consumed by C++) - T (transposed) → columnwise sub-storage (.columnwise_data consumed by C++) - For non-hybrid tensors this is a no-op passthrough. - """ - if not isinstance(tensor, HybridQuantizedTensorStorage): + # Quantized tensor formats with native GEMM support + if isinstance(tensor, _NATIVE_GEMM_INPUT_STORAGES): return tensor - if layout[1] == "N": - return tensor.rowwise_sub_storage - return tensor.columnwise_sub_storage + # Fall back to high-precision GEMM for other quantized tensor formats. + return tensor.dequantize() -def _unwrap_identity_tensor(tensor): - """Replace an :class:`IdentityTensorStorage` operand with its plain tensor. - Identity (high-precision passthrough) operands carry an unquantized tensor; - unwrapping it here routes the matmul through the standard high-precision - GEMM path. Non-identity operands pass through unchanged. Called after the - hybrid unwrap, so a high-precision *direction* of a hybrid tensor is handled - too. - """ - if isinstance(tensor, IdentityTensorStorage): - return tensor.dequantize() - return tensor - - -def _reject_unsupported_output_quantizer(quantization_params): - """Reject output quantizers that the native C++ GEMM path cannot convert.""" - if isinstance(quantization_params, (HybridQuantizer, IdentityQuantizer)): - quantizer_name = type(quantization_params).__name__ - # TODO(#3158): Lower HybridQuantizer output to its native rowwise - # sub-quantizer, and IdentityQuantizer output to an unquantized/no-op - # path, once the returned tensor contract is defined for these boundary - # roles. +_NATIVE_GEMM_OUTPUT_QUANTIZERS = ( + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, +) + + +def _validate_native_gemm_output_quantizer(quantization_params): + """Validate that the native C++ GEMM path can convert an output quantizer.""" + if quantization_params is not None and not isinstance( + quantization_params, _NATIVE_GEMM_OUTPUT_QUANTIZERS + ): raise NotImplementedError( - f"{quantizer_name} is not supported as a native GEMM output quantizer. " + f"{type(quantization_params).__name__} is not supported as a native " + "GEMM output quantizer. " "Return a TE-native quantizer for output/grad_input roles or disable " "quantized GEMM output for this boundary." ) @@ -191,8 +195,8 @@ def general_gemm( transa = layout[0] == "T" transb = layout[1] == "T" - A = _unwrap_identity_tensor(_unwrap_hybrid_A(A, layout)) - B = _unwrap_identity_tensor(_unwrap_hybrid_B(B, layout)) + A = _unwrap_tensor(A, "rowwise" if transa else "columnwise") + B = _unwrap_tensor(B, "columnwise" if transb else "rowwise") alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) @@ -239,7 +243,7 @@ def general_gemm( A = A.get_tensor(not transa) B = B.get_tensor(transb) - _reject_unsupported_output_quantizer(quantization_params) + _validate_native_gemm_output_quantizer(quantization_params) # Use bfloat16 as default bias_dtype bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] @@ -363,12 +367,12 @@ def general_grouped_gemm( """ num_gemms = len(A) - A = [_unwrap_identity_tensor(_unwrap_hybrid_A(a, layout)) for a in A] - B = [_unwrap_identity_tensor(_unwrap_hybrid_B(b, layout)) for b in B] - transa = layout[0] == "T" transb = layout[1] == "T" + A = [_unwrap_tensor(a, "rowwise" if transa else "columnwise") for a in A] + B = [_unwrap_tensor(b, "columnwise" if transb else "rowwise") for b in B] + empty_tensor = _empty_tensor() empty_tensors = [empty_tensor] * num_gemms diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 010f891d3fa..c050f26869f 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -43,7 +43,6 @@ from .tensor.mxfp8_tensor import MXFP8Quantizer from .tensor.nvfp4_tensor import NVFP4Quantizer from .tensor.float8_blockwise_tensor import Float8BlockQuantizer -from .tensor.hybrid_tensor import HybridQuantizer from .quantized_tensor import QuantizedTensorStorage, QuantizedTensor, Quantizer from .tensor.storage.float8_tensor_storage import Float8TensorStorage from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -1764,23 +1763,7 @@ def gather_along_first_dim( memory_format=torch.contiguous_format, ) torch.distributed.all_gather_into_tensor(out, inp, group=process_group) - # Hybrid override: callers drop columnwise before AG, expecting to - # synthesize it post-AG via ``update_usage(columnwise_usage=True)`` - # (native FP8's ``_create_transpose``). Hybrid has no synthesis path — - # that update_usage is a no-op — so re-quantize with both directions, - # mirroring what the planned native hybrid AG dispatch would produce - # (see the TODO in - # :meth:`HybridQuantizer.supports_only_rowwise_all_gather`); once - # native AG lands, hybrid won't reach this fallback. - if isinstance(quantizer, HybridQuantizer): - prev_row, prev_col = quantizer.rowwise_usage, quantizer.columnwise_usage - quantizer.set_usage(rowwise=True, columnwise=True) - try: - out = quantizer(out) - finally: - quantizer.set_usage(rowwise=prev_row, columnwise=prev_col) - else: - out = quantizer(out) + out = quantizer(out) return out, None # Dequantize quantized tensor if not supported diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 9b914e2d6cc..634bdd83d1c 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -13,6 +13,7 @@ from .. import cpp_extensions as tex from ..constants import TE_DType from ..export import is_in_onnx_export_mode +from ..tensor.hybrid_tensor import HybridQuantizer from ..utils import get_default_init_method @@ -31,6 +32,36 @@ def set_quantizer_amax_reduction_group(quantizer, amax_reduction_group) -> None: target.amax_reduction_group = amax_reduction_group +def set_quantizer_usage_for_wgrad_all_gather(quantizer) -> None: + """Configure an all-gather output for consumption by wgrad.""" + if quantizer is None: + return + + target = getattr(quantizer, "parent_quantizer", quantizer) + + # Hybrid currently gathers in high precision, then quantizes the full + # result, so request the columnwise representation consumed by wgrad. + if isinstance(target, HybridQuantizer): + target.set_usage(rowwise=False, columnwise=True) + elif target.supports_only_rowwise_all_gather(): + # Per-tensor FP8 gathers rowwise data and synthesizes its transpose. + target.set_usage(rowwise=True, columnwise=False) + else: + target.set_usage(rowwise=False, columnwise=True) + + +def can_reconstruct_wgrad_input_from_original(quantizer) -> bool: + """Whether wgrad input can be reconstructed from a saved original tensor.""" + target = getattr(quantizer, "parent_quantizer", quantizer) + if target is None: + target = quantizer + if isinstance(target, HybridQuantizer): + if target.columnwise_source == "original": + return True + return target.rowwise_quantizer.is_requantization_safe() + return target.is_requantization_safe() + + def _get_normalization_func(normalization: str, forward: bool): fwd_normalization_funcs = { "LayerNorm": tex.layernorm_fwd, diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index f2c2d77e68b..ebb802b2ae7 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -610,14 +610,6 @@ def fill_userbuffers_buffer_for_all_gather( local_tensor = quantizer(local_tensor) return local_tensor, local_tensor - # TODO(#3158): Support HybridQuantizer with standard Userbuffers. - if isinstance(quantizer, HybridQuantizer): - raise NotImplementedError( - "Standard Userbuffers does not support HybridQuantizer yet. " - "Disable Userbuffers all-gather overlap or use a supported " - "single-format quantizer. See #3158." - ) - # Tensor dimensions local_shape = local_tensor.size() if not local_shape: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index e6bb5852d28..429fb28eae7 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -4,7 +4,7 @@ """GroupedLinear API""" -from typing import Union, Optional, Callable, Tuple, List +from typing import Union, Optional, Callable, Tuple, List, Sequence from itertools import chain import os import warnings @@ -28,7 +28,7 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from ._common import WeightGradStore +from ._common import can_reconstruct_wgrad_input_from_original, WeightGradStore from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( divide, @@ -56,6 +56,7 @@ from ..triton.grouped_dbias_dscales import compute_grouped_dbias from ..tensor import ( + Float8BlockQuantizer, Float8CurrentScalingQuantizer, Float8Quantizer, HybridQuantizer, @@ -73,6 +74,22 @@ from ...debug.pytorch.debug_state import TEDebugState +_NATIVE_SPLIT_QUANTIZER_TYPES = frozenset( + { + Float8Quantizer, + Float8CurrentScalingQuantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, + } +) + + +def _supports_native_split_quantize(quantizer): + """Whether ``tex.split_quantize`` has an exact converter for this quantizer.""" + return type(quantizer) in _NATIVE_SPLIT_QUANTIZER_TYPES + + def _uses_identity_quantizer(quantizer): """Whether a quantizer, including a hybrid sub-quantizer, is Identity-backed.""" if quantizer is None: @@ -223,17 +240,18 @@ def _validate_grouped_quantizer_list(quantizers, *, operand_name="operand") -> N ) -def _split_quantize_with_identity_fallback( +def _split_quantize_non_hybrid( tensor, m_splits, quantizers, activation_dtype, *, disable_bulk_allocation=False, + allow_identity_views=True, ): - """Split+quantize a list validated as uniform for this recipe generation.""" + """Split and quantize one homogeneous, non-Hybrid quantizer list.""" reference = quantizers[0] - if not _uses_identity_quantizer(reference): + if _supports_native_split_quantize(reference): return tex.split_quantize( tensor, m_splits, @@ -242,7 +260,7 @@ def _split_quantize_with_identity_fallback( ) tensor = cast_if_needed(tensor, activation_dtype) - if isinstance(reference, IdentityQuantizer) and ( + if allow_identity_views and type(reference) is IdentityQuantizer and ( reference.dtype is None or reference.dtype == activation_dtype ): return torch.split(tensor, m_splits) @@ -253,7 +271,7 @@ def _split_quantize_with_identity_fallback( ] -def _hybrid_split_quantize( +def _split_quantize_hybrid( tensor, m_splits, quantizers, @@ -274,11 +292,13 @@ def _hybrid_split_quantize( columnwise_enabled and columnwise_source == "rowwise_dequantized" ) row_results = ( - tex.split_quantize( + _split_quantize_non_hybrid( tensor, m_splits, rowwise_quantizers, + tensor.dtype, disable_bulk_allocation=disable_bulk_allocation, + allow_identity_views=False, ) if needs_rowwise_result else [None] * len(quantizers) @@ -293,11 +313,13 @@ def _hybrid_split_quantize( dim=0, ) col_results = ( - tex.split_quantize( + _split_quantize_non_hybrid( columnwise_src, m_splits, columnwise_quantizers, + tensor.dtype, disable_bulk_allocation=disable_bulk_allocation, + allow_identity_views=False, ) if columnwise_enabled else [None] * len(quantizers) @@ -318,6 +340,99 @@ def _hybrid_split_quantize( ] +def _split_quantize( + tensor: torch.Tensor, + split_sizes: List[int], + with_quantized_output: bool, + quantizers: Optional[List[Quantizer]], + dtype: torch.dtype, + with_debug_quantizers: bool, + disable_bulk_allocation: bool, +) -> Sequence[Union[torch.Tensor, QuantizedTensorStorage]]: + """Split a tensor and quantize each part if needed.""" + if not with_quantized_output: + return torch.split(cast_if_needed(tensor, dtype), split_sizes) + + if quantizers is None or quantizers[0] is None: + raise ValueError("Quantizers are required for quantized split output") + + if with_debug_quantizers: + return DebugQuantizer.multi_tensor_quantize(tensor, quantizers, split_sizes, dtype) + + reference = quantizers[0] + if isinstance(reference, HybridQuantizer): + return _split_quantize_hybrid( + tensor, + split_sizes, + quantizers, + disable_bulk_allocation=disable_bulk_allocation, + ) + + return _split_quantize_non_hybrid( + tensor, + split_sizes, + quantizers, + dtype, + disable_bulk_allocation=disable_bulk_allocation, + ) + + +def _split_quantize_and_bias( + tensor: torch.Tensor, + split_sizes: List[int], + *, + fp8: bool, + debug: bool, + quantizers: Optional[List[Quantizer]], + dtype: torch.dtype, + use_bias: bool, + recipe: Recipe, + disable_bulk_allocation: bool, +) -> Tuple[ + Sequence[Union[torch.Tensor, QuantizedTensorStorage]], + List[Optional[torch.Tensor]], +]: + """Split grad output, quantize if needed, and compute unfused bias gradients.""" + num_splits = len(split_sizes) + grad_biases = [None] * num_splits + reference = quantizers[0] + identity = _uses_identity_quantizer(reference) + hybrid = isinstance(reference, HybridQuantizer) and not identity + + use_native_bgrad_quantize = ( + fp8 + and not debug + and not hybrid + and use_bias + and not identity + and (recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8()) + ) + if use_native_bgrad_quantize: + outputs = [None] * num_splits + for i, tensor_part in enumerate(torch.split(tensor, split_sizes)): + grad_biases[i], outputs[i] = tex.bgrad_quantize(tensor_part, quantizers[i]) + return outputs, grad_biases + + with_quantized_output = fp8 or debug + if with_quantized_output and (use_bias or debug): + for i, tensor_part in enumerate(torch.split(tensor, split_sizes)): + grad_biases[i] = tensor_part.sum(dim=0) + + # Preserve the existing CPU-offload policy: only Hybrid split-quantize + # disables bulk allocation in backward. + disable_bulk_allocation = disable_bulk_allocation if hybrid else False + outputs = _split_quantize( + tensor, + split_sizes, + with_quantized_output=with_quantized_output, + quantizers=quantizers, + dtype=dtype, + with_debug_quantizers=debug, + disable_bulk_allocation=disable_bulk_allocation, + ) + return outputs, grad_biases + + __all__ = ["GroupedLinear"] @@ -735,6 +850,8 @@ def forward( cache_weight, skip_fp8_weight_update, save_original_input, + delayed_scaling_input_quantizer, + unsafe_requantization_input_quantizer, debug, ) = non_tensor_args if fp8: @@ -752,38 +869,36 @@ def forward( device = inp.device weight_requires_grad = weights[0].requires_grad backward_needs_input = is_grad_enabled and weight_requires_grad - if save_original_input and backward_needs_input: - disallowed_input_quantizer = next( - ( - q - for q in input_quantizers - if q is not None and not q.allows_save_original_input_for_backward() - ), - None, - ) - if disallowed_input_quantizer is not None: + if backward_override is None and save_original_input and backward_needs_input: + if delayed_scaling_input_quantizer is not None: + if FP8GlobalStateManager.get_fp8_recipe().custom(): + warnings.warn( + "save_original_input is incompatible with delayed-scaling quantizers " + "(Float8Quantizer). Disabling save_original_input for this module.", + stacklevel=2, + ) + save_original_input = False + else: + raise ValueError( + "DelayedScaling recipe is not supported with save_original_input" + ) + + # Megatron-Core may enable this automatically to reuse an activation + # already retained by an upstream operation. The resolved quantizer + # generation is classified once in ``_validate_quantizer_generation``. + if ( + save_original_input + and unsafe_requantization_input_quantizer is not None + ): warnings.warn( - "Ignoring save_original_input=True because the input quantizer requires " - "the forward quantized activation for backward " - f"({disallowed_input_quantizer}).", + "Ignoring save_original_input=True because the input quantizer cannot " + "safely reconstruct the backward operand from the original input " + f"({unsafe_requantization_input_quantizer}).", stacklevel=2, ) save_original_input = False # Configure quantizers - if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): - if FP8GlobalStateManager.get_fp8_recipe().custom(): - # Custom recipe factory may produce DS quantizers unknown to caller. - # TODO(negvet): fix on Megatron side — guard should also exclude 'custom', or - # better: check at runtime whether quantizers are DS-based. - warnings.warn( - "save_original_input is incompatible with delayed-scaling quantizers " - "(Float8Quantizer). Disabling save_original_input for this module.", - stacklevel=2, - ) - save_original_input = False - else: - raise ValueError("DelayedScaling recipe is not supported with save_original_input") if input_quantizers[0] is not None: for input_quantizer in input_quantizers: input_quantizer.set_usage( @@ -865,35 +980,18 @@ def forward( m_splits = m_splits.tolist() inp_view = inp.reshape(-1, in_features) - inputmats: list - input_reference = input_quantizers[0] - input_hybrid = isinstance( - input_reference, HybridQuantizer - ) and not _uses_identity_quantizer(input_reference) - if fp8 and not debug and not input_hybrid: - # Disable bulk allocation when CPU offloading is active: offloading skips small - # tensors (like scales), but bulk allocation shares storage across all tensors, - # so if scales can't be offloaded, nothing in the group can be offloaded. - inputmats = _split_quantize_with_identity_fallback( - inp_view, - m_splits, - input_quantizers, - activation_dtype, - disable_bulk_allocation=cpu_offloading, - ) - elif fp8 and input_hybrid: - inputmats = _hybrid_split_quantize( - inp_view, - m_splits, - input_quantizers, - disable_bulk_allocation=cpu_offloading, - ) - elif debug: - inputmats = DebugQuantizer.multi_tensor_quantize( - inp_view, input_quantizers, m_splits, activation_dtype - ) - else: - inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) + # Disable bulk allocation when CPU offloading is active: offloading skips small + # tensors (like scales), but bulk allocation shares storage across all tensors, + # so if scales can't be offloaded, nothing in the group can be offloaded. + inputmats = _split_quantize( + inp_view, + m_splits, + with_quantized_output=fp8 or debug, + quantizers=input_quantizers, + dtype=activation_dtype, + with_debug_quantizers=debug, + disable_bulk_allocation=cpu_offloading, + ) if cpu_offloading: start_offload(*inputmats) @@ -1291,8 +1389,6 @@ def backward( # Preprocess grad output grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) - grad_output = [None] * ctx.num_gemms - grad_biases = [None] * ctx.num_gemms grad_output_reference = ctx.grad_output_quantizers[0] if ctx.fp8 and isinstance(grad_output_reference, HybridQuantizer): # Usage is a runtime decision, not part of generation validation. @@ -1303,69 +1399,17 @@ def backward( rowwise=ctx.requires_dgrad, columnwise=ctx.weights_requires_grad, ) - grad_output_identity = _uses_identity_quantizer(grad_output_reference) - grad_output_hybrid = ( - isinstance(grad_output_reference, HybridQuantizer) and not grad_output_identity + grad_output, grad_biases = _split_quantize_and_bias( + grad_output_view, + ctx.m_splits, + fp8=ctx.fp8, + debug=ctx.debug, + quantizers=ctx.grad_output_quantizers, + dtype=ctx.activation_dtype, + use_bias=ctx.use_bias, + recipe=ctx.fp8_recipe, + disable_bulk_allocation=ctx.cpu_offloading, ) - if ctx.fp8 and not ctx.debug and not grad_output_hybrid: - if ctx.use_bias: - grad_output_mats = torch.split(grad_output_view, ctx.m_splits) - recipe = ctx.fp8_recipe - if not grad_output_identity and ( - recipe.delayed() or recipe.float8_current_scaling() or recipe.mxfp8() - ): - # Fused bias grad + quantize kernel - for i in range(ctx.num_gemms): - grad_biases[i], grad_output[i] = tex.bgrad_quantize( - grad_output_mats[i], - ctx.grad_output_quantizers[i], - ) - else: - # Unfused bias grad and multi-tensor quantize - for i in range(ctx.num_gemms): - grad_biases[i] = grad_output_mats[i].sum(dim=0) - grad_output = _split_quantize_with_identity_fallback( - grad_output_view, - ctx.m_splits, - ctx.grad_output_quantizers, - ctx.activation_dtype, - ) - else: - # Multi-tensor quantize - grad_output = _split_quantize_with_identity_fallback( - grad_output_view, - ctx.m_splits, - ctx.grad_output_quantizers, - ctx.activation_dtype, - ) - elif ctx.fp8 and grad_output_hybrid: - if ctx.use_bias: - grad_output_mats = torch.split(grad_output_view, ctx.m_splits) - for i in range(ctx.num_gemms): - grad_biases[i] = grad_output_mats[i].sum(dim=0) - grad_output = _hybrid_split_quantize( - grad_output_view, - ctx.m_splits, - ctx.grad_output_quantizers, - disable_bulk_allocation=ctx.cpu_offloading, - ) - elif ctx.debug: - grad_output_mats = torch.split(grad_output_view, ctx.m_splits) - for i in range(ctx.num_gemms): - grad_biases[i] = grad_output_mats[i].sum(dim=0) - grad_output = DebugQuantizer.multi_tensor_quantize( - grad_output_view, - ctx.grad_output_quantizers, - ctx.m_splits, - ctx.activation_dtype, - ) - else: - # Only split grad output. Grad bias is fused with - # wgrad GEMM. - grad_output = torch.split( - cast_if_needed(grad_output_view, ctx.activation_dtype), - ctx.m_splits, - ) if ctx.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( @@ -1458,36 +1502,15 @@ def backward( input_quantizer.set_usage(rowwise=True, columnwise=True) else: input_quantizer.set_usage(rowwise=False, columnwise=True) - inputmats: list - input_reference = ctx.input_quantizers[0] - input_hybrid = isinstance( - input_reference, HybridQuantizer - ) and not _uses_identity_quantizer(input_reference) - if ctx.fp8 and not ctx.debug and not input_hybrid: - inputmats = _split_quantize_with_identity_fallback( - inp_view, - ctx.m_splits, - ctx.input_quantizers, - ctx.activation_dtype, - ) - elif ctx.fp8 and input_hybrid: - inputmats = _hybrid_split_quantize( - inp_view, - ctx.m_splits, - ctx.input_quantizers, - disable_bulk_allocation=ctx.cpu_offloading, - ) - elif ctx.debug: - inputmats = DebugQuantizer.multi_tensor_quantize( - inp_view, - ctx.input_quantizers, - ctx.m_splits, - ctx.activation_dtype, - ) - else: - inputmats = torch.split( - cast_if_needed(inp_view, ctx.activation_dtype), ctx.m_splits - ) + inputmats = _split_quantize( + inp_view, + ctx.m_splits, + with_quantized_output=ctx.fp8 or ctx.debug, + quantizers=ctx.input_quantizers, + dtype=ctx.activation_dtype, + with_debug_quantizers=ctx.debug, + disable_bulk_allocation=ctx.cpu_offloading, + ) elif ctx.backward_override == "dequantized": inputmats_dequant = [] for inputmat in inputmats: @@ -1630,7 +1653,8 @@ class GroupedLinear(TransformerEngineBaseModule): If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. - Cannot work with FP8 DelayedScaling recipe. + Requires input quantizers that can safely reproduce their results from the + original input. Cannot work with FP8 DelayedScaling recipe. single_grouped_weight : bool, default = False If set to ``True``, grouped weights are stored as a single grouped parameter instead of one parameter per GEMM. @@ -1716,6 +1740,8 @@ def __init__( "bwd": 2, } self._validated_quantizer_generations = {} + self._delayed_scaling_input_quantizer = None + self._unsafe_requantization_input_quantizer = None if tp_group is None: self.tp_size = tp_size @@ -1829,6 +1855,23 @@ def _validate_quantizer_generation(self, fwd: bool) -> None: ) _validate_grouped_quantizer_list(input_quantizers, operand_name="input") _validate_grouped_quantizer_list(weight_quantizers, operand_name="weight") + delayed_scaling_input_quantizer = next( + (q for q in input_quantizers if isinstance(q, Float8Quantizer)), + None, + ) + unsafe_requantization_input_quantizer = next( + ( + q + for q in input_quantizers + if q is not None + and not can_reconstruct_wgrad_input_from_original(q) + ), + None, + ) + self._delayed_scaling_input_quantizer = delayed_scaling_input_quantizer + self._unsafe_requantization_input_quantizer = ( + unsafe_requantization_input_quantizer + ) else: stride = self._num_fp8_tensors_per_gemm["bwd"] grad_output_quantizers = tuple( @@ -2253,6 +2296,8 @@ def forward( cache_weight, skip_fp8_weight_update, self.save_original_input, + self._delayed_scaling_input_quantizer, + self._unsafe_requantization_input_quantizer, debug, ) out, new_workspaces = linear_fn( diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 9c457053543..e69dc68c784 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -63,6 +63,7 @@ apply_normalization, noop_cat, set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, WeightGradStore, ) from ..quantized_tensor import ( @@ -752,12 +753,7 @@ def backward( quantizer = None if ctx.input_quantizer is not None and ctx.fp8: quantizer = ctx.input_quantizer - if quantizer.supports_only_rowwise_all_gather(): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if ctx.ub_bulk_dgrad: ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_dgrad, @@ -914,7 +910,7 @@ def backward( and ctx.ub_obj_gradout.with_cublasmp() ): if ctx.grad_output_quantizer is not None: - ctx.grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather(ctx.grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, ctx.tp_group, @@ -1580,15 +1576,6 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: recipe = FP8GlobalStateManager.get_fp8_recipe() if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) - # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP - # activations are gathered in high precision and re-quantized whole, so - # every rank already sees the same global amax. - # TODO(#3158): once native quantized all-gather lands (see - # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path - # quantizes per-shard, needing a hybrid branch here that mirrors the - # current-scaling / NVFP4 SP reduction above: - # elif recipe.custom(): - # ... # enable SP amax reduction on the hybrid input/grad quantizer def get_quantizer_roles( self, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 6c5ae2e776a..0874343446f 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -76,7 +76,12 @@ from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.hybrid_tensor import HybridQuantizer from ..tensor.identity_tensor import IdentityQuantizer -from ._common import apply_normalization, set_quantizer_amax_reduction_group, WeightGradStore +from ._common import ( + apply_normalization, + set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, + WeightGradStore, +) from ..cpu_offload import ( is_cpu_offload_enabled, start_offload, @@ -1177,12 +1182,7 @@ def backward( quantizer = None if ctx.fp8 or ctx.debug: quantizer = ctx.fc1_input_quantizer - if isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if ctx.ub_bulk_dgrad: ub_obj_fc1_dgrad = get_ub("fc1_dgrad", ctx.fp8) ln_out_total, _ = fill_userbuffers_buffer_for_all_gather( @@ -1300,7 +1300,9 @@ def backward( and ctx.ub_obj_gradout.with_cublasmp() ): if ctx.fc2_grad_output_quantizer is not None: - ctx.fc2_grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather( + ctx.fc2_grad_output_quantizer + ) grad_output, _ = gather_along_first_dim( grad_output, ctx.tp_group, @@ -2185,15 +2187,6 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: recipe = FP8GlobalStateManager.get_fp8_recipe() if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) - # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP - # activations are gathered in high precision and re-quantized whole, so - # every rank already sees the same global amax. - # TODO(#3158): once native quantized all-gather lands (see - # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path - # quantizes per-shard, needing a hybrid branch here that mirrors the - # current-scaling / NVFP4 SP reduction above: - # elif recipe.custom(): - # ... # enable SP amax reduction on the hybrid input/grad quantizer def get_quantizer_roles( self, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 9efc0317fa5..6e54a8bf961 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -30,7 +30,13 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from ._common import noop_cat, set_quantizer_amax_reduction_group, WeightGradStore +from ._common import ( + can_reconstruct_wgrad_input_from_original, + noop_cat, + set_quantizer_amax_reduction_group, + set_quantizer_usage_for_wgrad_all_gather, + WeightGradStore, +) from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( cast_if_needed, @@ -287,10 +293,40 @@ def _linear_forward_impl( debug = args.debug backward_override = args.backward_override is_fsdp2 = args.is_fsdp2 + backward_needs_input = is_grad_enabled and weight.requires_grad if backward_override == "high_precision": save_original_input = True elif backward_override == "dequantized": save_original_input = False + if ( + backward_override is None + and save_original_input + and backward_needs_input + and input_quantizer is not None + ): + # Megatron-Core enables this automatically for attention output projections + # to reuse the high-precision DPA output saved by attention backward. Validate + # the resolved quantizer since this is not necessarily an informed user opt-in. + if isinstance(input_quantizer, Float8Quantizer): + if FP8GlobalStateManager.get_fp8_recipe().custom(): + warnings.warn( + "save_original_input is incompatible with delayed-scaling quantizers " + "(Float8Quantizer). Disabling save_original_input for this module.", + stacklevel=2, + ) + save_original_input = False + else: + raise ValueError( + "DelayedScaling recipe is not supported with save_original_input" + ) + elif not can_reconstruct_wgrad_input_from_original(input_quantizer): + warnings.warn( + "Ignoring save_original_input=True because the input quantizer cannot " + "safely reconstruct the backward operand from the original input " + f"({input_quantizer}).", + stacklevel=2, + ) + save_original_input = False # NVTX label for profiling nvtx_label = "transformer_engine._Linear.forward" @@ -303,20 +339,6 @@ def _linear_forward_impl( # Configure tensor-parallel communication tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and weight.requires_grad - if ( - save_original_input - and backward_needs_input - and input_quantizer is not None - and not input_quantizer.allows_save_original_input_for_backward() - ): - warnings.warn( - "Ignoring save_original_input=True because the input quantizer requires " - "the forward quantized activation for backward " - f"({input_quantizer}).", - stacklevel=2, - ) - save_original_input = False with_input_all_gather_nccl = ( parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop ) @@ -347,10 +369,6 @@ def _linear_forward_impl( own_quantized_input = False if fp8: assert_dim_for_fp8_exec(inputmat, weight) - if save_original_input: - assert not isinstance( - input_quantizer, Float8Quantizer - ), "DelayedScaling recipe is not supported with save_original_input" if with_input_all_gather_nccl or ub_overlap_ag_fprop: # All-gather input tensor @@ -942,12 +960,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. quantizer = None if bwd_args.fp8 or bwd_args.debug: quantizer = input_quantizer - if quantizer.supports_only_rowwise_all_gather(): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) + set_quantizer_usage_for_wgrad_all_gather(quantizer) if bwd_args.ub_bulk_dgrad: inputmat_total, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_dgrad, @@ -1103,7 +1116,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. and bwd_args.ub_obj_gradout.with_cublasmp() ): if grad_output_quantizer is not None: - grad_output_quantizer.set_usage(rowwise=True, columnwise=False) + set_quantizer_usage_for_wgrad_all_gather(grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, bwd_args.tp_group, @@ -1494,7 +1507,8 @@ class Linear(TransformerEngineBaseModule): If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. - Cannot work with FP8 DelayedScaling recipe. + Requires an input quantizer that can safely reproduce its result from the + original input. Cannot work with FP8 DelayedScaling recipe. """ def __init__( @@ -1785,15 +1799,6 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: recipe = FP8GlobalStateManager.get_fp8_recipe() if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) - # Hybrid (CustomRecipe) needs no SP amax-reduction setup today: its SP - # activations are gathered in high precision and re-quantized whole, so - # every rank already sees the same global amax. - # TODO(#3158): once native quantized all-gather lands (see - # supports_only_rowwise_all_gather / gather_along_first_dim) the SP path - # quantizes per-shard, needing a hybrid branch here that mirrors the - # current-scaling / NVFP4 SP reduction above: - # elif recipe.custom(): - # ... # enable SP amax reduction on the hybrid input/grad quantizer def reset_parameters(self, defer_init=False): super().reset_parameters(defer_init=defer_init) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 607346ce300..f39115c4c62 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -13,11 +13,36 @@ from transformer_engine_torch import FP8TensorMeta from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager +from ..quantized_tensor import QuantizedTensorStorage, Quantizer +from ..tensor import ( + Float8BlockQuantizer, + Float8CurrentScalingQuantizer, + Float8Quantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) from ..tensor.float8_tensor import Float8Tensor -from ..quantized_tensor import QuantizedTensorStorage from ..utils import canonicalize_dtype +def get_fused_normalization_quantizer( + quantizer: Optional[Quantizer], +) -> Optional[Quantizer]: + """Return a quantizer supported by fused normalization kernels.""" + if isinstance( + quantizer, + ( + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + Float8BlockQuantizer, + NVFP4Quantizer, + ), + ): + return quantizer + return None + + def validate_or_alloc_output( buffer: Optional[torch.Tensor], shape: tuple[int, ...] | list[int], diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index a85ad483ad6..58601731669 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -16,7 +16,7 @@ from ...constants import TE_DType from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...export import is_in_onnx_export_mode -from ...tensor import HybridQuantizer, Quantizer +from ...tensor import Quantizer from ...utils import ( canonicalize_device, canonicalize_dtype, @@ -24,7 +24,11 @@ devices_match, ) from ..op import BasicOperation, OperationContext -from .._common import maybe_autocast_dtype, maybe_dequantize +from .._common import ( + get_fused_normalization_quantizer, + maybe_autocast_dtype, + maybe_dequantize, +) class LayerNorm(BasicOperation): @@ -183,14 +187,10 @@ def op_forward( if is_in_onnx_export_mode(): return self.op_onnx_forward(input_) - # TODO(#3158): Support producing both directional representations at - # the fused normalization boundary. - if isinstance(next_op_input_quantizer, HybridQuantizer): - raise NotImplementedError( - "te.ops.LayerNorm does not support HybridQuantizer output yet. " - "Use transformer_engine.pytorch.LayerNormMLP or a single-format " - "quantizer instead. See #3158." - ) + # Fall back to a high-precision output when fused quantization is unsupported. + next_op_input_quantizer = get_fused_normalization_quantizer( + next_op_input_quantizer + ) # Check tensor dims weight = self.weight diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 1d8d8be9718..13801d92c97 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -24,7 +24,11 @@ devices_match, ) from ..op import BasicOperation, OperationContext -from .._common import maybe_autocast_dtype, maybe_dequantize +from .._common import ( + get_fused_normalization_quantizer, + maybe_autocast_dtype, + maybe_dequantize, +) class RMSNorm(BasicOperation): @@ -166,6 +170,11 @@ def op_forward( if is_in_onnx_export_mode(): return self.op_onnx_forward(input_) + # Fall back to a high-precision output when fused quantization is unsupported. + next_op_input_quantizer = get_fused_normalization_quantizer( + next_op_input_quantizer + ) + # Check tensor dims weight = self.weight weight_dims = tuple(weight.size()) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index f8a447ae4b9..11224516bbe 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -453,15 +453,13 @@ def supports_only_rowwise_all_gather(self) -> bool: """Returns True if the quantizer supports only rowwise all-gather""" return False - def allows_save_original_input_for_backward(self) -> bool: - """Whether ``save_original_input`` preserves backward operand semantics. + def is_requantization_safe(self) -> bool: + """Whether repeated quantization of the same input reproduces the same value. - Modules may use ``save_original_input`` as a memory optimization by saving - the high-precision input and re-preparing the backward operand later. Some - quantizers require the exact forward-produced quantized value for backward, - so replacing it with the original tensor would change recipe semantics. + Stateful or stochastic quantizers should return ``False``. This lets callers + decide whether a quantized value may be discarded and reconstructed later. """ - return True + return False def is_quantizable(self, inp: torch.Tensor) -> bool: # pylint: disable=unused-argument """Whether tensor supports quantized all-gather diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 09fde86f176..6d26ba384e4 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -122,6 +122,10 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) + def is_requantization_safe(self) -> bool: + """Block-FP8 scales are derived deterministically from each input.""" + return True + def get_scale_shape(self, shape: Iterable[int], columnwise: bool) -> Tuple[int, int]: """Scaling tensor shape. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index e21d283eb48..075561a30ed 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -399,6 +399,10 @@ def supports_only_rowwise_all_gather(self) -> bool: """ return True + def is_requantization_safe(self) -> bool: + """Current scaling is derived deterministically from each input.""" + return True + register_value_opaque_quantizer(Float8CurrentScalingQuantizer) diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 2f94ec24c3b..8f6051c5f83 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -5,7 +5,7 @@ """Tensor and quantizer classes for composed rowwise and columnwise representations.""" from __future__ import annotations -from typing import Any, Dict, Iterable, Literal, Optional, Tuple +from typing import Any, Dict, Iterable, Literal, Optional, Tuple, Union import torch @@ -235,7 +235,7 @@ def make_empty( device: Optional[torch.device] = None, requires_grad: bool = False, pin_memory: bool = False, - ) -> HybridQuantizedTensor: + ) -> Union["HybridQuantizedTensor", HybridQuantizedTensorStorage]: # Mirror ``quantize_impl``: invoke each sub-quantizer with its own # ``internal`` setting (no toggle), so the produced sub-storages have # the same type that ``quantize_impl`` would produce via @@ -255,6 +255,14 @@ def make_empty( else None ) + if self.internal: + return HybridQuantizedTensorStorage( + rowwise_storage=rowwise_empty, + columnwise_storage=columnwise_empty, + quantizer=self, + fake_dtype=dtype, + ) + return HybridQuantizedTensor( shape=shape, dtype=dtype, @@ -298,67 +306,13 @@ def update_quantized( return dst def supports_only_rowwise_all_gather(self) -> bool: - """Whether TP activation all-gather must preserve rowwise data. - - Used by ``_linear_forward_impl`` / ``_linear_backward`` to decide - which direction of the saved activation to keep for the backward - input-AG: ``True`` keeps rowwise (drops columnwise), - ``False`` keeps columnwise (drops rowwise, default for block- - scaled formats whose columnwise is directly consumable by wgrad). - - Why hybrid needs a custom rule - ------------------------------ - ``gather_along_first_dim`` has no hybrid-specific dispatch, so - hybrid falls through to the generic BF16 fallback:: - - inp.dequantize() → all_gather BF16 → quantizer(out) - - The direction we preserve must therefore be one the hybrid can - dequantize. Two cases force rowwise preservation: - - 1. The rowwise sub-quantizer itself declares rowwise-only AG - (e.g. Float8 delayed / current scaling). Propagating keeps - hybrid consistent with its component semantics. - 2. The columnwise sub-quantizer is :class:`NVFP4Quantizer`: - ``NVFP4TensorStorage`` has no columnwise dequantize - (``_FromNVFP4Func.forward`` raises for ``is_colwise=True``), - so a columnwise-only NVFP4 sub-storage cannot traverse the - BF16 fallback. Rowwise preservation routes the fallback - through NVFP4's working rowwise dequantize instead. - - For MXFP8 / Float8Block / Float8CurrentScaling columnwise sub- - quantizers, columnwise dequantize works and the default - (``False``) keeps the smaller, wgrad-ready columnwise shard - saved — which is the more efficient memory choice. - - TODO(#3158): Add native hybrid dispatch to - ``gather_along_first_dim`` to remove the BF16 detour. - - * **Scope.** Branch at the top of ``gather_along_first_dim`` that - detects ``HybridQuantizedTensorStorage`` / ``HybridQuantizer``, - extracts ``rowwise_sub_storage`` and ``columnwise_sub_storage`` - with their sub-quantizers, dispatches each to its native - ``_all_gather_{fp8,mxfp8,nvfp4,fp8_blockwise}`` path, and wraps - the gathered sub-storages back into a ``HybridQuantizedTensor``. - Each per-format AG routine already supports rowwise-only or - columnwise-only input natively (including NVFP4 columnwise — - it gathers packed FP4 bytes without dequantize). - - * **Impact.** Replaces 2×–4× BF16 bandwidth cost with native - quantized AG. Mirrors the FSDP2 native-AG pattern we already - ship on ``fsdp_pre_all_gather`` / ``fsdp_post_all_gather``. - Once it lands, the ``NVFP4Quantizer`` branch in this method - can be removed (columnwise NVFP4 AG works natively), leaving - only the rowwise-sub-quantizer propagation. - - * **Implementation notes.** Compose async handles across the two - per-direction AG calls into a single handle object with a - ``.wait()`` that waits on both. Pass ``out_shape=None`` to the - recursive calls so each format computes its own packed shape. - Preserve FP8 current / delayed rowwise-only semantics on - Hopper / L40 (``_all_gather_fp8`` reads ``inp._data`` which - may be ``None`` for a columnwise-only FP8 sub-storage on - those architectures). + """Whether all-gather requires a rowwise-dequantizable source. + + Hybrid tensors currently use the high-precision fallback, which + dequantizes the local shard before communication. Preserve rowwise + data when required by the rowwise sub-quantizer or when the + columnwise sub-quantizer is NVFP4, whose columnwise-only storage + cannot be dequantized. """ if self.rowwise_quantizer.supports_only_rowwise_all_gather(): return True @@ -366,14 +320,21 @@ def supports_only_rowwise_all_gather(self) -> bool: # (nvfp4_tensor → quantized_tensor → hybrid_tensor at module import). from .nvfp4_tensor import NVFP4Quantizer # noqa: PLC0415 - if isinstance(self.columnwise_quantizer, NVFP4Quantizer): - return True - return False - - def allows_save_original_input_for_backward(self) -> bool: - # TODO(#3158): Add an explicit recompute-from-original policy for deterministic - # quantizers that can trade backward compute for saved activation memory. - return self.columnwise_source == "original" + return isinstance(self.columnwise_quantizer, NVFP4Quantizer) + + def is_requantization_safe(self) -> bool: + """Whether repeated quantization reproduces all requested representations.""" + if self.rowwise_usage and not self.rowwise_quantizer.is_requantization_safe(): + return False + if self.columnwise_usage and not self.columnwise_quantizer.is_requantization_safe(): + return False + if ( + self.columnwise_usage + and self.columnwise_source == "rowwise_dequantized" + and not self.rowwise_quantizer.is_requantization_safe() + ): + return False + return True def _get_compatible_recipe(self): # HybridQuantizer is only reachable via CustomRecipe (the qfactory @@ -494,18 +455,20 @@ def detach(self) -> HybridQuantizedTensor: if hasattr(row_cls, "make_like"): row = row_cls.make_like(self._rowwise_storage) else: - # Storage-only sub-storages (HybridQuantizer.internal=True - # path) don't have make_like; the cpu_offload_v2 path does - # not hit this branch, but keep the behaviour safe by - # sharing the reference as before. - row = self._rowwise_storage + raise NotImplementedError( + "HybridQuantizedTensor.detach() does not support storage-only " + f"rowwise sub-storage {row_cls.__name__}" + ) col = None if self._columnwise_storage is not None: col_cls = type(self._columnwise_storage) if hasattr(col_cls, "make_like"): col = col_cls.make_like(self._columnwise_storage) else: - col = self._columnwise_storage + raise NotImplementedError( + "HybridQuantizedTensor.detach() does not support storage-only " + f"columnwise sub-storage {col_cls.__name__}" + ) return HybridQuantizedTensor( shape=self.shape, dtype=self.dtype, diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py index e418f9e7216..9fb980a7559 100644 --- a/transformer_engine/pytorch/tensor/identity_tensor.py +++ b/transformer_engine/pytorch/tensor/identity_tensor.py @@ -14,7 +14,7 @@ """ from __future__ import annotations -from typing import Any, Iterable, Optional, Tuple +from typing import Any, Iterable, Optional, Tuple, Union import torch from torch.ops import aten @@ -96,6 +96,10 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensorStorage: device=data.device, ) + def is_requantization_safe(self) -> bool: + """Identity quantization is deterministic.""" + return True + def make_empty( self, shape: Iterable[int], @@ -104,12 +108,18 @@ def make_empty( device: Optional[torch.device] = None, requires_grad: bool = False, pin_memory: bool = False, - ) -> "IdentityTensor": + ) -> Union["IdentityTensor", IdentityTensorStorage]: if device is None: device = torch.device("cuda") device = torch.device(device) data_dtype = self.dtype if self.dtype is not None else dtype data = torch.empty(tuple(shape), dtype=data_dtype, device=device, pin_memory=pin_memory) + if self.internal: + return IdentityTensorStorage( + hp_data=data, + fake_dtype=data_dtype, + quantizer=self, + ) return IdentityTensor( data.shape, data_dtype, diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 9f27a104d6a..f30108aecc9 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -180,6 +180,10 @@ def onnx_dequantize(self, tensor: Union[MXFP8TensorStorage, MXFP8Tensor]) -> tor def _get_compatible_recipe(self) -> Union[type[Recipe], None]: return MXFP8BlockScaling + def is_requantization_safe(self) -> bool: + """MXFP8 block scales are derived deterministically from each input.""" + return True + register_value_opaque_quantizer(MXFP8Quantizer) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index ccf06ac1667..3f8a11d2e17 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -254,6 +254,10 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) + def is_requantization_safe(self) -> bool: + """NVFP4 quantization is replay-safe unless stochastic rounding is enabled.""" + return not self.stochastic_rounding + def is_quantizable(self, inp: torch.Tensor) -> bool: """Returns whether or not given inp can be quantized""" if self.row_scaled_nvfp4: diff --git a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py index fb08dfabcea..822bb342fe8 100644 --- a/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/hybrid_tensor_storage.py @@ -58,7 +58,7 @@ def __new__( instance._rowwise_storage = rowwise_storage instance._columnwise_storage = columnwise_storage - instance._quantizer = quantizer + instance._quantizer = quantizer.copy() return instance @property diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index a2a0d37498e..e35d57b3630 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1264,16 +1264,20 @@ def _validate_hybrid_partial_master_policy( start_offset, fsdp_shard_model_weight, ): - """Reject an unsafe independently sourced Hybrid column before distopt mutation.""" + """Reject unsupported Hybrid column-source updates before distopt mutation.""" row_sub = model_weight._rowwise_storage col_sub = model_weight._columnwise_storage quantizer = model_weight._get_quantizer() - if ( - master_weight is None - or row_sub is None - or col_sub is None - or quantizer.columnwise_source != "original" - ): + + if col_sub is not None and quantizer.columnwise_source == "rowwise_dequantized": + raise NotImplementedError( + "quantize_master_weights does not support HybridQuantizer with a live " + "columnwise representation and " + "columnwise_source='rowwise_dequantized'. The column must be derived " + "after the rowwise update/all-gather. See #3158." + ) + + if master_weight is None or row_sub is None or col_sub is None: return shard_has_both_directions = ( @@ -1293,7 +1297,7 @@ def _validate_hybrid_partial_master_policy( "HybridQuantizedTensor from a partial master shard when " "columnwise_source='original': the one-payload distributed-optimizer " "path cannot preserve an independently quantized columnwise value. " - "Use columnwise_source='rowwise_dequantized' or provide full-master data." + "Provide full-master data or retain only the rowwise representation." ) From d8d3de1cd441c7f005f325bbbe60c97806835da6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:13:43 +0000 Subject: [PATCH 60/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 17 ++++------------- tests/pytorch/test_identity_quantizer.py | 8 ++------ .../pytorch/module/grouped_linear.py | 18 +++++++----------- .../pytorch/module/layernorm_mlp.py | 4 +--- transformer_engine/pytorch/module/linear.py | 4 +--- .../pytorch/ops/basic/layer_norm.py | 4 +--- .../pytorch/ops/basic/rmsnorm.py | 4 +--- 7 files changed, 17 insertions(+), 42 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 5bc10e24306..84036ae84f3 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -609,9 +609,7 @@ def test_requantization_and_wgrad_reconstruction_policy_matrix( # Wgrad reconstruction only repeats the rowwise stage for a # rowwise-dequantized columnwise source. expected_reconstruction = columnwise_source == "original" or rowwise_safe - assert ( - can_reconstruct_wgrad_input_from_original(hq) is expected_reconstruction - ) + assert can_reconstruct_wgrad_input_from_original(hq) is expected_reconstruction def test_hybrid_rowwise_source_requires_safe_rowwise_quantizer(self): hq = HybridQuantizer( @@ -761,9 +759,7 @@ def test_linear_builtin_delayed_scaling_rejects_save_original_input(self): def test_grouped_linear_classifies_requantization_safety_once_per_generation(self): counters = [{"calls": 0}, {"calls": 0}] - input_quantizers = [ - _CountingUnsafeIdentityQuantizer(counter) for counter in counters - ] + input_quantizers = [_CountingUnsafeIdentityQuantizer(counter) for counter in counters] generation = [] for input_quantizer in input_quantizers: generation.extend((input_quantizer, IdentityQuantizer(), IdentityQuantizer())) @@ -797,9 +793,7 @@ def _counting_identity_hybrid_recipe( def factory(role): if role is not None and role.module_type == "linear" and role.tensor_type == "input": rowwise_cls = ( - _CountingIdentityQuantizer - if rowwise_safe - else _CountingUnsafeIdentityQuantizer + _CountingIdentityQuantizer if rowwise_safe else _CountingUnsafeIdentityQuantizer ) columnwise_cls = ( _CountingIdentityQuantizer @@ -5196,10 +5190,7 @@ def test_fp8_current_original_partial_master_raises(self): with pytest.raises( ValueError, - match=( - "partial master shard.*columnwise_source='original'.*" - "full-master data" - ), + match="partial master shard.*columnwise_source='original'.*full-master data", ): quantize_master_weights( [weight], diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 3f924b23614..422b3fd9757 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -265,9 +265,7 @@ def test_grouped_split_all_identity_uses_plain_tensor_views(self): m_splits = [3, 5] quantizers = [IdentityQuantizer(), IdentityQuantizer()] - out = _split_quantize_non_hybrid( - x, m_splits, quantizers, activation_dtype=torch.bfloat16 - ) + out = _split_quantize_non_hybrid(x, m_splits, quantizers, activation_dtype=torch.bfloat16) assert all(isinstance(t, torch.Tensor) for t in out) assert not any(isinstance(t, IdentityTensorStorage) for t in out) @@ -381,9 +379,7 @@ def fake_split_quantize_hybrid( raise StopAfterFlagCapture("captured hybrid split kwargs") monkeypatch.setattr(grouped_linear, "is_cpu_offload_enabled", lambda: True) - monkeypatch.setattr( - grouped_linear, "_split_quantize_hybrid", fake_split_quantize_hybrid - ) + monkeypatch.setattr(grouped_linear, "_split_quantize_hybrid", fake_split_quantize_hybrid) model = te.GroupedLinear(2, 64, 64, params_dtype=torch.bfloat16).cuda() x = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 96d701cb5ab..3f3963c5493 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -266,8 +266,10 @@ def _split_quantize_non_hybrid( ) tensor = cast_if_needed(tensor, activation_dtype) - if allow_identity_views and type(reference) is IdentityQuantizer and ( - reference.dtype is None or reference.dtype == activation_dtype + if ( + allow_identity_views + and type(reference) is IdentityQuantizer + and (reference.dtype is None or reference.dtype == activation_dtype) ): return torch.split(tensor, m_splits) @@ -917,10 +919,7 @@ def forward( # Megatron-Core may enable this automatically to reuse an activation # already retained by an upstream operation. The resolved quantizer # generation is classified once in ``_validate_quantizer_generation``. - if ( - save_original_input - and unsafe_requantization_input_quantizer is not None - ): + if save_original_input and unsafe_requantization_input_quantizer is not None: warnings.warn( "Ignoring save_original_input=True because the input quantizer cannot " "safely reconstruct the backward operand from the original input " @@ -1922,15 +1921,12 @@ def _validate_quantizer_generation(self, fwd: bool) -> None: ( q for q in input_quantizers - if q is not None - and not can_reconstruct_wgrad_input_from_original(q) + if q is not None and not can_reconstruct_wgrad_input_from_original(q) ), None, ) self._delayed_scaling_input_quantizer = delayed_scaling_input_quantizer - self._unsafe_requantization_input_quantizer = ( - unsafe_requantization_input_quantizer - ) + self._unsafe_requantization_input_quantizer = unsafe_requantization_input_quantizer else: stride = self._num_fp8_tensors_per_gemm["bwd"] grad_output_quantizers = tuple( diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 8a8658fa1a4..6a2ad31461c 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -1300,9 +1300,7 @@ def backward( and ctx.ub_obj_gradout.with_cublasmp() ): if ctx.fc2_grad_output_quantizer is not None: - set_quantizer_usage_for_wgrad_all_gather( - ctx.fc2_grad_output_quantizer - ) + set_quantizer_usage_for_wgrad_all_gather(ctx.fc2_grad_output_quantizer) grad_output, _ = gather_along_first_dim( grad_output, ctx.tp_group, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 151efe3d097..56622db5e6b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -323,9 +323,7 @@ def _linear_forward_impl( ) save_original_input = False else: - raise ValueError( - "DelayedScaling recipe is not supported with save_original_input" - ) + raise ValueError("DelayedScaling recipe is not supported with save_original_input") elif not can_reconstruct_wgrad_input_from_original(input_quantizer): warnings.warn( "Ignoring save_original_input=True because the input quantizer cannot " diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index 58601731669..3372952add7 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -188,9 +188,7 @@ def op_forward( return self.op_onnx_forward(input_) # Fall back to a high-precision output when fused quantization is unsupported. - next_op_input_quantizer = get_fused_normalization_quantizer( - next_op_input_quantizer - ) + next_op_input_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) # Check tensor dims weight = self.weight diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 13801d92c97..96851c8db43 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -171,9 +171,7 @@ def op_forward( return self.op_onnx_forward(input_) # Fall back to a high-precision output when fused quantization is unsupported. - next_op_input_quantizer = get_fused_normalization_quantizer( - next_op_input_quantizer - ) + next_op_input_quantizer = get_fused_normalization_quantizer(next_op_input_quantizer) # Check tensor dims weight = self.weight From 7e15ddcd9140742c11d2c716f4b4ea20800a85b1 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 29 Jul 2026 12:51:41 +0000 Subject: [PATCH 61/97] Quantizer config per recipe + keys to manage callables Signed-off-by: Evgeny --- transformer_engine/common/recipe/__init__.py | 247 +++++++++++++++++- .../quantization_factory_base.py | 7 + .../quantization_factory_zoo.py | 9 + 3 files changed, 260 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index a89ddba917b..103b01e9dce 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -7,7 +7,7 @@ import abc import os from enum import Enum -from typing import Any, Literal, Optional, Union, Callable, NamedTuple +from typing import Any, Callable, Hashable, Literal, NamedTuple, Optional, Union from dataclasses import field from pydantic.dataclasses import dataclass @@ -104,6 +104,73 @@ def __repr__(self) -> str: return self._cached_repr +def _qparams_config(params: QParams) -> tuple: + """Return the immutable semantic configuration for quantization parameters.""" + return ( + ("power_2_scale", params.power_2_scale), + ("amax_epsilon", params.amax_epsilon), + ("random_hadamard_transform", params.random_hadamard_transform), + ("stochastic_rounding", params.stochastic_rounding), + ("fp4_2d_quantization", params.fp4_2d_quantization), + ) + + +def _mmparams_config(params: MMParams) -> tuple: + """Return the immutable semantic configuration for GEMM parameters.""" + return (("use_split_accumulator", params.use_split_accumulator),) + + +def _algorithm_config( + algorithm: Any, + semantic_key: Optional[Hashable], + *, + field_name: str, +) -> tuple: + """Return identity-independent configuration for a literal or callable algorithm.""" + if not callable(algorithm): + return (("literal", algorithm), None) + if semantic_key is None: + raise ValueError( + f"{field_name} is callable, so {field_name}_key must provide an explicit semantic key" + ) + return (("callable",), semantic_key) + + +def _validate_policy_key(key: Any, *, source: str) -> Hashable: + """Validate and return an immutable semantic policy key.""" + if key is None: + raise ValueError(f"{source} must not be None") + try: + hash(key) + except TypeError as exc: + raise TypeError(f"{source} must be hashable") from exc + return key + + +def quantizer_policy(*, key: Hashable) -> Callable: + """Attach a semantic policy key to a quantizer factory without wrapping or calling it. + + The key, rather than Python callable identity, represents the factory's behavior in + :class:`CustomRecipe` semantic configuration. Change the key whenever the factory's + quantizer-selection behavior changes. + """ + policy_key = _validate_policy_key(key, source="quantizer_policy key") + + def decorator(qfactory: Callable) -> Callable: + if not callable(qfactory): + raise TypeError("quantizer_policy can only decorate a callable") + try: + setattr(qfactory, "policy_key", policy_key) + except (AttributeError, TypeError) as exc: + raise TypeError( + "quantizer_policy could not attach metadata to this callable; " + "pass policy_key directly to CustomRecipe instead" + ) from exc + return qfactory + + return decorator + + class Recipe: """ Base recipe class. @@ -114,12 +181,51 @@ class Recipe: # changes. This makes repeated ``str(recipe)`` calls much cheaper _cached_repr: Optional[str] = None + # Cached semantic configuration. Lazily populated by ``quantizer_config`` + # and invalidated together with the cached repr on recipe mutation. + _cached_quantizer_config: Any = None + _has_cached_quantizer_config: bool = False + def __setattr__(self, name: str, value: Any) -> None: - # Invalidate the cached repr on any attribute mutation. - if name != "_cached_repr": + # Invalidate derived caches on any recipe mutation. Updating either + # cache itself must not invalidate the other cache. + if name not in ( + "_cached_repr", + "_cached_quantizer_config", + "_has_cached_quantizer_config", + ): object.__setattr__(self, "_cached_repr", None) + object.__setattr__(self, "_cached_quantizer_config", None) + object.__setattr__(self, "_has_cached_quantizer_config", False) object.__setattr__(self, name, value) + def quantizer_config(self) -> Hashable: + """Return the immutable semantic configuration used to build quantizers. + + Recipe implementations must include every value that affects quantizer construction or + quantized execution behavior. The returned value must be hashable and must not depend on + recipe object identity, callable identity, or quantizer construction. Subclasses implement + :meth:`_make_quantizer_config` rather than overriding this cached accessor. + """ + if not self._has_cached_quantizer_config: + config = self._make_quantizer_config() + try: + hash(config) + except TypeError as exc: + raise TypeError( + f"{self.__class__.__name__}._make_quantizer_config() " + "must return a hashable value" + ) from exc + object.__setattr__(self, "_cached_quantizer_config", config) + object.__setattr__(self, "_has_cached_quantizer_config", True) + return self._cached_quantizer_config + + def _make_quantizer_config(self) -> Hashable: + """Build the immutable semantic configuration returned by :meth:`quantizer_config`.""" + raise NotImplementedError( + f"{self.__class__.__name__} must implement Recipe.quantizer_config()" + ) + @abc.abstractmethod def _make_repr(self) -> str: """Build the string representation for this recipe. @@ -235,6 +341,13 @@ def scaling_factor_compute(amax: Tensor, `LayerNormLinear (FP8 output) -> FP8 DPA -> Linear`. backward_override : {None, 'high_precision', 'dequantized'}, default = None Backward precision mode. Delayed scaling only supports None. + amax_compute_algo_key : Hashable, default = None + Required semantic key when ``amax_compute_algo`` is callable. + Callables with equivalent behavior must use the same key, and a + behavior change must use a different key. + scaling_factor_compute_algo_key : Hashable, default = None + Required semantic key when + ``scaling_factor_compute_algo`` is callable. Notes ----- @@ -259,6 +372,8 @@ def scaling_factor_compute(amax: Tensor, fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + amax_compute_algo_key: Optional[Hashable] = None + scaling_factor_compute_algo_key: Optional[Hashable] = None def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." @@ -269,6 +384,32 @@ def __post_init__(self) -> None: self.backward_override is None ), "Delayed scaling only supports backward_override=None." + def _make_quantizer_config(self) -> Hashable: + amax_compute_algo, amax_compute_algo_key = _algorithm_config( + self.amax_compute_algo, + self.amax_compute_algo_key, + field_name="amax_compute_algo", + ) + scaling_factor_compute_algo, scaling_factor_compute_algo_key = _algorithm_config( + self.scaling_factor_compute_algo, + self.scaling_factor_compute_algo_key, + field_name="scaling_factor_compute_algo", + ) + return ( + ("recipe_type", "DelayedScaling"), + ("margin", self.margin), + ("fp8_format", self.fp8_format.name), + ("amax_history_len", self.amax_history_len), + ("amax_compute_algo", amax_compute_algo), + ("scaling_factor_compute_algo", scaling_factor_compute_algo), + ("reduce_amax", self.reduce_amax), + ("fp8_dpa", self.fp8_dpa), + ("fp8_mha", self.fp8_mha), + ("backward_override", self.backward_override), + ("amax_compute_algo_key", amax_compute_algo_key), + ("scaling_factor_compute_algo_key", scaling_factor_compute_algo_key), + ) + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -317,6 +458,22 @@ def __post_init__(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + def _make_quantizer_config(self) -> Hashable: + return ( + ("recipe_type", "Float8CurrentScaling"), + ("use_power_2_scales", self.use_power_2_scales), + ("fp8_format", self.fp8_format.name), + ("fp8_quant_fwd_inp", _qparams_config(self.fp8_quant_fwd_inp)), + ("fp8_quant_fwd_weight", _qparams_config(self.fp8_quant_fwd_weight)), + ("fp8_quant_bwd_grad", _qparams_config(self.fp8_quant_bwd_grad)), + ("fp8_gemm_fprop", _mmparams_config(self.fp8_gemm_fprop)), + ("fp8_gemm_dgrad", _mmparams_config(self.fp8_gemm_dgrad)), + ("fp8_gemm_wgrad", _mmparams_config(self.fp8_gemm_wgrad)), + ("fp8_dpa", self.fp8_dpa), + ("fp8_mha", self.fp8_mha), + ("backward_override", self.backward_override), + ) + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -375,6 +532,16 @@ def __post_init__(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + def _make_quantizer_config(self) -> Hashable: + return ( + ("recipe_type", "MXFP8BlockScaling"), + ("margin", self.margin), + ("fp8_format", self.fp8_format.name), + ("fp8_dpa", self.fp8_dpa), + ("fp8_mha", self.fp8_mha), + ("backward_override", self.backward_override), + ) + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -462,6 +629,25 @@ def __post_init__(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + def _make_quantizer_config(self) -> Hashable: + return ( + ("recipe_type", "Float8BlockScaling"), + ("use_f32_scales", self.use_f32_scales), + ("fp8_format", self.fp8_format.name), + ("fp8_quant_fwd_inp", _qparams_config(self.fp8_quant_fwd_inp)), + ("fp8_quant_fwd_weight", _qparams_config(self.fp8_quant_fwd_weight)), + ("fp8_quant_bwd_grad", _qparams_config(self.fp8_quant_bwd_grad)), + ("x_block_scaling_dim", self.x_block_scaling_dim), + ("w_block_scaling_dim", self.w_block_scaling_dim), + ("grad_block_scaling_dim", self.grad_block_scaling_dim), + ("fp8_gemm_fprop", _mmparams_config(self.fp8_gemm_fprop)), + ("fp8_gemm_dgrad", _mmparams_config(self.fp8_gemm_dgrad)), + ("fp8_gemm_wgrad", _mmparams_config(self.fp8_gemm_wgrad)), + ("fp8_dpa", self.fp8_dpa), + ("fp8_mha", self.fp8_mha), + ("backward_override", self.backward_override), + ) + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -604,6 +790,26 @@ def __post_init__(self) -> None: fp4_2d_quantization=False, ) + def _make_quantizer_config(self) -> Hashable: + return ( + ("recipe_type", "NVFP4BlockScaling"), + ("disable_rht", self.disable_rht), + ("disable_stochastic_rounding", self.disable_stochastic_rounding), + ("disable_2d_quantization", self.disable_2d_quantization), + ("row_scaled_activation", self.row_scaled_activation), + ("nvfp4_4over6", self.nvfp4_4over6), + ("nvfp4_4over6_e4m3_use_256", self.nvfp4_4over6_e4m3_use_256), + ("nvfp4_4over6_err_mode", self.nvfp4_4over6_err_mode), + ("fp4_format", self.fp4_format.name), + ("fp8_format", self.fp8_format.name), + ("fp4_quant_fwd_inp", _qparams_config(self.fp4_quant_fwd_inp)), + ("fp4_quant_fwd_weight", _qparams_config(self.fp4_quant_fwd_weight)), + ("fp4_quant_bwd_grad", _qparams_config(self.fp4_quant_bwd_grad)), + ("fp8_dpa", self.fp8_dpa), + ("fp8_mha", self.fp8_mha), + ("backward_override", self.backward_override), + ) + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -674,6 +880,11 @@ class CustomRecipe(Recipe): formats. It can be lowered when the factory's full output space is known; for example, a factory restricted to MXFP8 may use 32. Automatic padding reads this value without invoking ``qfactory``. + policy_key : Hashable, default = None + Semantic identifier for the complete behavior of ``qfactory``. Pass it explicitly or + attach it to the factory with ``@quantizer_policy(key=...)``. An explicit value takes + precedence over attached factory metadata. Equivalent factory behavior must use equal + keys, and any behavior change requires a different key. """ qfactory: Callable[..., Any] @@ -687,6 +898,7 @@ class CustomRecipe(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) quantization_alignment: int = 128 + policy_key: Optional[Hashable] = None def __post_init__(self) -> None: assert ( @@ -694,11 +906,40 @@ def __post_init__(self) -> None: ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." if self.quantization_alignment <= 0: raise ValueError("CustomRecipe quantization_alignment must be positive.") + if self.policy_key is None: + attached_key = getattr(self.qfactory, "policy_key", None) + if attached_key is not None: + self.policy_key = _validate_policy_key( + attached_key, + source="qfactory.policy_key", + ) + else: + self.policy_key = _validate_policy_key( + self.policy_key, + source="CustomRecipe policy_key", + ) + + def _make_quantizer_config(self) -> Hashable: + if self.policy_key is None: + raise ValueError( + "CustomRecipe requires a semantic policy key. Pass policy_key=... or decorate " + "qfactory with @quantizer_policy(key=...)." + ) + return ( + ("recipe_type", "CustomRecipe"), + ("policy_key", self.policy_key), + ("fp8_format", self.fp8_format.name), + ("fp8_dpa", self.fp8_dpa), + ("fp8_mha", self.fp8_mha), + ("backward_override", self.backward_override), + ("quantization_alignment", self.quantization_alignment), + ) def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"qfactory={self.qfactory}, " + f"policy_key={self.policy_key}, " f"backward_override={self.backward_override}, " f"quantization_alignment={self.quantization_alignment}" ) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py index e5125fd92fd..79725efa5de 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py @@ -28,10 +28,12 @@ from typing import Optional import torch +from transformer_engine.common.recipe import quantizer_policy from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType +@quantizer_policy(key=("transformer_engine", "high_precision", 1)) def high_precision_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "IdentityQuantizer": @@ -50,6 +52,7 @@ def high_precision_factory( return IdentityQuantizer() +@quantizer_policy(key=("transformer_engine", "delayed_scaling", 1)) def delayed_scaling_quantizer_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "DelayedScalingRequest": @@ -69,6 +72,7 @@ def delayed_scaling_quantizer_factory( return DelayedScalingRequest(fp8_format=Format.HYBRID) +@quantizer_policy(key=("transformer_engine", "current_scaling", 1)) def current_scaling_quantizer_factory( role: Optional[QuantizerRole], ) -> "Float8CurrentScalingQuantizer": @@ -92,6 +96,7 @@ def current_scaling_quantizer_factory( ) +@quantizer_policy(key=("transformer_engine", "mxfp8", 1)) def mxfp8_quantizer_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "MXFP8Quantizer": @@ -107,6 +112,7 @@ def mxfp8_quantizer_factory( ) +@quantizer_policy(key=("transformer_engine", "float8_block_scaling", 1)) def float8_block_scaling_quantizer_factory( role: Optional[QuantizerRole], ) -> "Float8BlockQuantizer": @@ -137,6 +143,7 @@ def float8_block_scaling_quantizer_factory( ) +@quantizer_policy(key=("transformer_engine", "nvfp4", 1)) def nvfp4_quantizer_factory( role: Optional[QuantizerRole], ) -> "NVFP4Quantizer": diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index 8b238970507..67e146a82ef 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -70,6 +70,7 @@ from typing import Optional +from transformer_engine.common.recipe import quantizer_policy from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType from .quantization_factory_base import mxfp8_quantizer_factory, nvfp4_quantizer_factory @@ -79,6 +80,7 @@ # ----------------------------------------------------------------------------- +@quantizer_policy(key=("transformer_engine", "high_precision_fwd_mxfp8_bwd", 1)) def high_precision_fwd_mxfp8_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -102,6 +104,7 @@ def high_precision_fwd_mxfp8_bwd_quantizer_factory( ) +@quantizer_policy(key=("transformer_engine", "mxfp8_fwd_nvfp4_bwd", 1)) def mxfp8_fwd_nvfp4_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -144,6 +147,7 @@ def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): ) +@quantizer_policy(key=("transformer_engine", "nvfp4_1d_double_quantized_weight", 1)) def nvfp4_1d_double_quantized_weight_quantizer_factory( role: Optional[QuantizerRole], ): @@ -182,6 +186,7 @@ def nvfp4_1d_double_quantized_weight_quantizer_factory( # ----------------------------------------------------------------------------- +@quantizer_policy(key=("transformer_engine", "mxfp8_fwd_dequantized_bwd", 1)) def mxfp8_fwd_dequantized_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -218,6 +223,7 @@ def mxfp8_fwd_dequantized_bwd_quantizer_factory( return mxfp8_quantizer_factory(role) +@quantizer_policy(key=("transformer_engine", "nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd", 1)) def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -279,6 +285,7 @@ def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( return mxfp8_quantizer_factory(role) +@quantizer_policy(key=("transformer_engine", "nvfp4_row_scaled_fwd_dequantized_bwd", 1)) def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -330,6 +337,7 @@ def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( # ----------------------------------------------------------------------------- +@quantizer_policy(key=("transformer_engine", "nvfp4_linear_fp8_dpa", 1)) def nvfp4_linear_fp8_dpa_factory( role: Optional[QuantizerRole], ): @@ -408,6 +416,7 @@ def nvfp4_linear_fp8_dpa_factory( return nvfp4_quantizer_factory(role) +@quantizer_policy(key=("transformer_engine", "nvfp4_linear_mxfp8_dpa", 1)) def nvfp4_linear_mxfp8_dpa_factory( role: Optional[QuantizerRole], ): From 5ee018dd746052e9bae5b2cbfe5725dd0f2c3ac7 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 31 Jul 2026 15:44:05 +0000 Subject: [PATCH 62/97] refactor custom recipe Signed-off-by: Evgeny --- .../distributed/fsdp2_tests/fsdp2_utils.py | 11 +- tests/pytorch/distributed/run_hybrid_tp_sp.py | 8 + .../pytorch/distributed/run_numerics_exact.py | 14 +- tests/pytorch/hybrid_quantization_utils.py | 58 +++--- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 8 +- .../nvfp4/test_nvfp4_group_quantize.py | 4 +- .../test_nvfp4_group_quantize_graph_safe.py | 4 +- .../pytorch/nvfp4/test_nvfp4_module_exact.py | 16 +- .../nvfp4/test_nvfp4_quantize_exact.py | 12 +- .../nvfp4/test_nvfp4_rht_quantize_exact.py | 6 +- tests/pytorch/test_custom_recipe.py | 72 +++---- .../test_float8_current_scaling_exact.py | 4 +- tests/pytorch/test_hybrid_quantization.py | 157 +++++++++------ tests/pytorch/test_identity_quantizer.py | 14 +- .../pytorch/cpp_extensions/gemm.py | 19 +- .../pytorch/custom_recipes/gemm.py | 24 ++- .../pytorch/custom_recipes/quantization.py | 29 --- ...factory_base.py => quantizer_factories.py} | 46 ++--- ...actory_zoo.py => quantizer_factory_zoo.py} | 179 ++++++++++-------- ...caling.py => reference_current_scaling.py} | 14 +- ...zation_ref_nvfp4.py => reference_nvfp4.py} | 34 ++-- .../{utils.py => reference_utils.py} | 2 +- .../pytorch/tensor/hybrid_tensor.py | 2 +- 23 files changed, 402 insertions(+), 335 deletions(-) delete mode 100644 transformer_engine/pytorch/custom_recipes/quantization.py rename transformer_engine/pytorch/custom_recipes/{quantization_factory_base.py => quantizer_factories.py} (93%) rename transformer_engine/pytorch/custom_recipes/{quantization_factory_zoo.py => quantizer_factory_zoo.py} (76%) rename transformer_engine/pytorch/custom_recipes/{quantization_ref_current_scaling.py => reference_current_scaling.py} (97%) rename transformer_engine/pytorch/custom_recipes/{quantization_ref_nvfp4.py => reference_nvfp4.py} (98%) rename transformer_engine/pytorch/custom_recipes/{utils.py => reference_utils.py} (83%) diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 042333843b5..81cc8731f5c 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -4,6 +4,15 @@ """Shared utility functions for FSDP2 distributed tests.""" +import sys +from pathlib import Path + +# FSDP2 files are collected directly by subprocess pytest invocations, which +# put this directory (rather than tests/pytorch) on sys.path. +_TEST_ROOT = str(Path(__file__).resolve().parents[2]) +if _TEST_ROOT not in sys.path: + sys.path.append(_TEST_ROOT) + import transformer_engine.common.recipe from transformer_engine.pytorch import QuantizedTensor @@ -40,7 +49,7 @@ def get_hybrid_recipe_from_string(recipe): """Build a CustomRecipe wrapping a module-level (picklable) hybrid qfactory. Each hybrid qfactory composes one or two role-aware base factories from - ``quantization_factory_base`` per direction; per-role behavior is delegated + ``quantizer_factories`` per direction; per-role behavior is delegated to the base factory and the hybrid layer only decides the direction pairing. Supported values: diff --git a/tests/pytorch/distributed/run_hybrid_tp_sp.py b/tests/pytorch/distributed/run_hybrid_tp_sp.py index 6094b6cc519..b4dfb646ef6 100644 --- a/tests/pytorch/distributed/run_hybrid_tp_sp.py +++ b/tests/pytorch/distributed/run_hybrid_tp_sp.py @@ -10,6 +10,14 @@ import datetime import os import sys +from pathlib import Path + +# This file is launched directly with torchrun, so Python only adds the +# distributed test directory to sys.path. Add tests/pytorch for the shared +# hybrid recipe helpers used by both local and distributed tests. +_TEST_ROOT = str(Path(__file__).resolve().parent.parent) +if _TEST_ROOT not in sys.path: + sys.path.append(_TEST_ROOT) import torch import torch.distributed as dist diff --git a/tests/pytorch/distributed/run_numerics_exact.py b/tests/pytorch/distributed/run_numerics_exact.py index 6af992dac28..46782122b29 100644 --- a/tests/pytorch/distributed/run_numerics_exact.py +++ b/tests/pytorch/distributed/run_numerics_exact.py @@ -22,8 +22,8 @@ ) from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.constants import NVFP4_BLOCK_SCALING_SIZE -from transformer_engine.pytorch.custom_recipes import quantization_ref_nvfp4 -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes import reference_nvfp4 +from transformer_engine.pytorch.custom_recipes import reference_utils from run_layer_with_overlap import _compare_tensors @@ -58,7 +58,7 @@ def get_nvfp4_quantizer_factory(): Mirrors the canonical "branch on what we care about, default fall-through" pattern from - ``transformer_engine.pytorch.custom_recipes.quantization_factory_base``; + ``transformer_engine.pytorch.custom_recipes.quantizer_factories``; every slot gets a real :class:`NVFP4QuantizerRef` (``CustomRecipeState`` rejects ``None`` returns). @@ -73,14 +73,14 @@ def factory(role): and role.tensor_type == "weight" ) if is_weight: - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(16, 16), pow_2_scales=False, with_rht=False, ) - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(1, 16), pow_2_scales=False, with_rht=True, diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py index dfbe8d4c344..9444cc4bcce 100644 --- a/tests/pytorch/hybrid_quantization_utils.py +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -13,11 +13,11 @@ import transformer_engine.pytorch as te from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - current_scaling_quantizer_factory, - float8_block_scaling_quantizer_factory, - mxfp8_quantizer_factory, - nvfp4_quantizer_factory, +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + float8_block_scaling_factory, + mxfp8_factory, + nvfp4_factory, ) _LINEAR_MODULE_TYPES = ("linear", "grouped_linear") @@ -215,10 +215,10 @@ def hybrid_fp8_current_qfactory(role): """FP8 current scaling in both hybrid directions for forward tensor roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: return te.HybridQuantizer( - rowwise_quantizer=current_scaling_quantizer_factory(role), - columnwise_quantizer=current_scaling_quantizer_factory(role), + rowwise_quantizer=current_scaling_factory(role), + columnwise_quantizer=current_scaling_factory(role), ) - return current_scaling_quantizer_factory(role) + return current_scaling_factory(role) def hybrid_fp8_current_e5m2_grads_qfactory(role): @@ -237,20 +237,20 @@ def hybrid_mxfp8_qfactory(role): """MXFP8 in both hybrid directions for forward tensor roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: return te.HybridQuantizer( - rowwise_quantizer=mxfp8_quantizer_factory(role), - columnwise_quantizer=mxfp8_quantizer_factory(role), + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=mxfp8_factory(role), ) - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) def hybrid_float8_block_qfactory(role): """Float8 block scaling in both hybrid directions for forward roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: return te.HybridQuantizer( - rowwise_quantizer=float8_block_scaling_quantizer_factory(role), - columnwise_quantizer=float8_block_scaling_quantizer_factory(role), + rowwise_quantizer=float8_block_scaling_factory(role), + columnwise_quantizer=float8_block_scaling_factory(role), ) - return float8_block_scaling_quantizer_factory(role) + return float8_block_scaling_factory(role) def hybrid_block_fp8_e4m3_qfactory(role): @@ -279,34 +279,34 @@ def hybrid_mixed_mxfp8_fp8_qfactory(role): """MXFP8 rowwise plus FP8 current-scaling columnwise.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: return te.HybridQuantizer( - rowwise_quantizer=mxfp8_quantizer_factory(role), - columnwise_quantizer=current_scaling_quantizer_factory(role), + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=current_scaling_factory(role), ) - return current_scaling_quantizer_factory(role) + return current_scaling_factory(role) def hybrid_fp8_current_identity_qfactory(role): """FP8 current-scaling forward plus Identity backward.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: return te.HybridQuantizer( - rowwise_quantizer=current_scaling_quantizer_factory(role), + rowwise_quantizer=current_scaling_factory(role), columnwise_quantizer=te.IdentityQuantizer(), ) if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: return te.IdentityQuantizer() - return current_scaling_quantizer_factory(role) + return current_scaling_factory(role) def hybrid_mxfp8_identity_qfactory(role): """MXFP8 forward plus Identity backward.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: return te.HybridQuantizer( - rowwise_quantizer=mxfp8_quantizer_factory(role), + rowwise_quantizer=mxfp8_factory(role), columnwise_quantizer=te.IdentityQuantizer(), ) if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: return te.IdentityQuantizer() - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) def identity_qfactory(role): # pylint: disable=unused-argument @@ -318,19 +318,19 @@ def hybrid_nvfp4_qfactory(role): """NVFP4 in both hybrid directions for forward tensor roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: return te.HybridQuantizer( - rowwise_quantizer=nvfp4_quantizer_factory(role), - columnwise_quantizer=nvfp4_quantizer_factory(role), + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=nvfp4_factory(role), ) - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) def hybrid_tp_mxfp8_nvfp4_qfactory(role): """TP/SP MXFP8 rowwise plus NVFP4 columnwise, including boundary roles.""" if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) return te.HybridQuantizer( - rowwise_quantizer=mxfp8_quantizer_factory(role), - columnwise_quantizer=nvfp4_quantizer_factory(role), + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=nvfp4_factory(role), ) @@ -351,8 +351,8 @@ def hybrid_mxfp8_nvfp4_qfactory(role): if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: return te.HybridQuantizer( rowwise_quantizer=_make_mxfp8(), - columnwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_factory(role), ) if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) return _make_mxfp8() diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 67dbe1d1d52..eb480060e2b 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -9,8 +9,8 @@ from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage @@ -128,7 +128,7 @@ def check_nvfp4_gemm_versus_reference( # Create reference quantizer for reference GEMM x_ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=not row_scaled_nvfp4, pow_2_scales=False, @@ -140,7 +140,7 @@ def check_nvfp4_gemm_versus_reference( nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) w_ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=True, pow_2_scales=False, diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 148ad83772b..23d49cd77fc 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -13,8 +13,8 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.common.recipe import NVFP4BlockScaling import pytest diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py index b78c66befdf..38bd1b31a0a 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -6,8 +6,8 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.common.recipe import NVFP4BlockScaling from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor diff --git a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py index 66ac0181e38..9fdabf68dd0 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py @@ -6,8 +6,8 @@ import torch import transformer_engine.pytorch as te from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes import quantization_ref_nvfp4 -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes import reference_nvfp4 +from transformer_engine.pytorch.custom_recipes import reference_utils recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) @@ -103,8 +103,8 @@ def get_nvfp4_quantizer_factory( # qfactory, so we return a valid quantizer for those slots; it is harmless because # the GEMM outputs in the high-precision activation dtype, not in NVFP4. def _default_quantizer(): - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(1, 16), pow_2_scales=False, with_rht=with_rht, @@ -116,16 +116,16 @@ def factory(role): if role.tensor_type == "input": # Only the forward activation is row-scaled, mirroring the production # wiring in quantization.py (mode == "forward" and tensor_type != "weight"). - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(1, 16), pow_2_scales=False, with_rht=with_rht, row_scaled_nvfp4=row_scaled, ) if role.tensor_type == "weight": - return quantization_ref_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + return reference_nvfp4.NVFP4QuantizerRef( + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(16, 16) if with_2d_quantization else (1, 16), pow_2_scales=False, with_rht=False, diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 4a3d4d34136..fe1a04334ee 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -11,8 +11,8 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.common.recipe import NVFP4BlockScaling @@ -185,7 +185,7 @@ def check_quantization_nvfp4_versus_reference( # Reference quantization quant_tile_shape = (1, 16) if not with_2d_quantization else (16, 16) ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=return_transpose, pow_2_scales=False, @@ -394,7 +394,7 @@ def test_nvfp4_quantization_extrema_versus_reference( qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=return_transpose, pow_2_scales=False, @@ -539,7 +539,7 @@ def test_nvfp4_quantization_boundary_values( qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=return_transpose, pow_2_scales=False, @@ -670,7 +670,7 @@ def test_nvfp4_quantization_noncontiguous_inputs( qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=True, columnwise=return_transpose, pow_2_scales=False, diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index a65010ef025..aa1a911f4c5 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -12,8 +12,8 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.common.recipe import NVFP4BlockScaling import pytest @@ -98,7 +98,7 @@ def check_quantization_nvfp4_versus_reference( # Reference quantization using NVFP4QuantizerRef with built-in RHT ref_quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, rowwise=return_rowwise, columnwise=return_transpose, pow_2_scales=False, diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index e88ef93c958..9caa33b18be 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -24,19 +24,19 @@ get_align_size_for_quantization, ) import transformer_engine.pytorch.ops as te_ops -from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - current_scaling_quantizer_factory, - mxfp8_quantizer_factory, - float8_block_scaling_quantizer_factory, - nvfp4_quantizer_factory, - delayed_scaling_quantizer_factory, +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + mxfp8_factory, + float8_block_scaling_factory, + nvfp4_factory, + delayed_scaling_factory, high_precision_factory, ) -from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - mxfp8_fwd_nvfp4_bwd_quantizer_factory, +from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_factory, ) -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import ( - nvfp4_ref_rht_2d_quantizer_factory, +from transformer_engine.pytorch.custom_recipes.reference_nvfp4 import ( + nvfp4_ref_rht_2d_factory, ) @@ -67,7 +67,7 @@ def test_custom_recipe_sanity_modules_nvfp4(module_type): inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) # Use NVFP4 quantizer factory - custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_factory) # Execute with custom recipe with autocast(enabled=True, recipe=custom_recipe): @@ -398,7 +398,7 @@ def _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus): def test_factory_matches_delayed_scaling(): - """delayed_scaling_quantizer_factory should produce bit-identical results + """delayed_scaling_factory should produce bit-identical results to the built-in DelayedScaling recipe.""" available, reason = te.is_fp8_available(return_reason=True) if not torch.cuda.is_available() or not available: @@ -408,13 +408,13 @@ def test_factory_matches_delayed_scaling(): out_ref, grad_ref, pgrads_ref = _run_linear_fwd_bwd(model_ref, inp_ref, recipe.DelayedScaling()) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=delayed_scaling_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=delayed_scaling_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) def test_factory_matches_current_scaling(): - """current_scaling_quantizer_factory should produce bit-identical results + """current_scaling_factory should produce bit-identical results to the built-in Float8CurrentScaling recipe.""" available, reason = te.is_fp8_available(return_reason=True) if not torch.cuda.is_available() or not available: @@ -426,13 +426,13 @@ def test_factory_matches_current_scaling(): model_ref, inp_ref, recipe.Float8CurrentScaling() ) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=current_scaling_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=current_scaling_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) def test_factory_matches_mxfp8(): - """mxfp8_quantizer_factory should produce bit-identical results + """mxfp8_factory should produce bit-identical results to the built-in MXFP8BlockScaling recipe.""" available, reason = te.is_mxfp8_available(return_reason=True) if not torch.cuda.is_available() or not available: @@ -444,13 +444,13 @@ def test_factory_matches_mxfp8(): model_ref, inp_ref, recipe.MXFP8BlockScaling() ) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=mxfp8_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=mxfp8_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) def test_factory_matches_block_scaling(): - """float8_block_scaling_quantizer_factory should produce bit-identical results + """float8_block_scaling_factory should produce bit-identical results to the built-in Float8BlockScaling recipe.""" available = te.is_fp8_block_scaling_available() if not torch.cuda.is_available() or not available: @@ -462,13 +462,13 @@ def test_factory_matches_block_scaling(): model_ref, inp_ref, recipe.Float8BlockScaling() ) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=float8_block_scaling_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=float8_block_scaling_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) def test_factory_matches_nvfp4(): - """nvfp4_quantizer_factory should produce bit-identical results + """nvfp4_factory should produce bit-identical results to the built-in NVFP4BlockScaling recipe.""" available = te.is_nvfp4_available() if not torch.cuda.is_available() or not available: @@ -480,7 +480,7 @@ def test_factory_matches_nvfp4(): model_ref, inp_ref, recipe.NVFP4BlockScaling() ) out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( - model_cus, inp_cus, recipe.CustomRecipe(qfactory=nvfp4_quantizer_factory) + model_cus, inp_cus, recipe.CustomRecipe(qfactory=nvfp4_factory) ) _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) @@ -611,30 +611,30 @@ def targeting_factory(role): recorded_roles.append(role) if role is None: - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) assert isinstance(role, QuantizerRole), f"Expected QuantizerRole, got {type(role)}" # Layer 0 (tl0.*): all MXFP8 if role.name.startswith("tl0"): - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) # Layer 1 (tl1.*): NVFP4 default, but fc2 overridden to MXFP8 if role.name == "tl1.layernorm_mlp.fc2": - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) # Layer 2: block scaling for qkv and fc1, rest falls through to default if role.name == "tl2.self_attention.layernorm_linear_qkv": - return float8_block_scaling_quantizer_factory(role) + return float8_block_scaling_factory(role) if role.name == "tl2.layernorm_mlp.fc1": - return float8_block_scaling_quantizer_factory(role) + return float8_block_scaling_factory(role) # Layer 3: current-scaling for proj, rest falls through to default if role.name == "tl3.proj": - return current_scaling_quantizer_factory(role) + return current_scaling_factory(role) # Default: NVFP4 - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) custom_recipe = recipe.CustomRecipe(qfactory=targeting_factory) @@ -1133,7 +1133,7 @@ def test_custom_recipe_dpa_fp8(): Float8Quantizer, Float8CurrentScalingQuantizer, ) - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) @@ -1260,7 +1260,7 @@ def test_custom_recipe_dpa_mxfp8(): from transformer_engine.pytorch.quantization import CustomRecipeState from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_mxfp8_dpa_factory, ) @@ -1407,7 +1407,7 @@ def test_custom_recipe_debug_tool_compat(): in_features, out_features, params_dtype=torch.bfloat16, name="layer" ).cuda() - custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_factory) assert TEDebugState.debug_enabled, "Debug mode should be active" @@ -1905,7 +1905,7 @@ def test_custom_recipe_quantization_alignment_contract(): def qfactory(role): calls.append(role) - return current_scaling_quantizer_factory(role) + return current_scaling_factory(role) custom_recipe = recipe.CustomRecipe(qfactory=qfactory) assert get_align_size_for_quantization(custom_recipe) == 128 @@ -1923,7 +1923,7 @@ def qfactory(role): "qfactory", [ pytest.param( - current_scaling_quantizer_factory, + current_scaling_factory, marks=pytest.mark.skipif( not _alignment_fp8_available, reason=_alignment_fp8_reason, @@ -1931,7 +1931,7 @@ def qfactory(role): id="fp8_current_scaling", ), pytest.param( - mxfp8_quantizer_factory, + mxfp8_factory, marks=pytest.mark.skipif( not _alignment_mxfp8_available, reason=_alignment_mxfp8_reason, @@ -1939,7 +1939,7 @@ def qfactory(role): id="mxfp8", ), pytest.param( - nvfp4_quantizer_factory, + nvfp4_factory, marks=pytest.mark.skipif( not _alignment_nvfp4_available, reason=_alignment_nvfp4_reason, @@ -1947,7 +1947,7 @@ def qfactory(role): id="nvfp4", ), pytest.param( - mxfp8_fwd_nvfp4_bwd_quantizer_factory, + mxfp8_fwd_nvfp4_bwd_factory, marks=pytest.mark.skipif( not (_alignment_mxfp8_available and _alignment_nvfp4_available), reason=f"MXFP8: {_alignment_mxfp8_reason}; NVFP4: {_alignment_nvfp4_reason}", diff --git a/tests/pytorch/test_float8_current_scaling_exact.py b/tests/pytorch/test_float8_current_scaling_exact.py index cac7ccb66bf..1d45d1919af 100644 --- a/tests/pytorch/test_float8_current_scaling_exact.py +++ b/tests/pytorch/test_float8_current_scaling_exact.py @@ -13,8 +13,8 @@ from transformer_engine.common.recipe import Float8CurrentScaling from transformer_engine.pytorch.quantization import autocast, get_fp8_torch_dtype from transformer_engine.pytorch.constants import TE_DType -from transformer_engine.pytorch.custom_recipes.quantization import MMParams -from transformer_engine.pytorch.custom_recipes.quantization_ref_current_scaling import ( +from transformer_engine.pytorch.custom_recipes.gemm import MMParams +from transformer_engine.pytorch.custom_recipes.reference_current_scaling import ( CurrentScalingQuantizerRef, ) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 84036ae84f3..7f9bc16b96a 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -30,11 +30,11 @@ snapshot_storage_tensor_metadata as _snapshot_storage_tensor_metadata, ) from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - nvfp4_quantizer_factory, +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + nvfp4_factory, ) -from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - mxfp8_fwd_nvfp4_bwd_quantizer_factory, +from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_factory, ) from transformer_engine.pytorch import ( autocast, @@ -89,6 +89,16 @@ ), ) +_XFAIL_SELECTIVE_ATTENTION_RECOMPUTE = pytest.mark.xfail( + condition=is_non_tn_fp8_gemm_supported(), + raises=AttributeError, + strict=True, + reason=( + "Selective attention recompute currently consumes LayerNormLinear's saved " + "tensor objects twice; this also fails with BF16 and built-in FP8 recipes." + ), +) + requires_fp8 = pytest.mark.skipif( not fp8_available, reason=f"FP8: {reason_for_no_fp8}", @@ -253,11 +263,11 @@ class TestComposerStyleFactory: @staticmethod def _factory(role): - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_row_scaled_fwd_mxfp8_bwd_factory, ) - return nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory(role) + return nvfp4_row_scaled_fwd_mxfp8_bwd_factory(role) @pytest.mark.parametrize( "tensor_type,row_scaled", @@ -312,16 +322,16 @@ def test_non_linear_roles_fall_back_to_mxfp8(self): @pytest.mark.skipif(not mxfp8_available, reason=f"MXFP8: {reason_for_no_mxfp8}") -class TestMXFP8FwdDequantizedBwdFactory: +class TestMXFP8FwdHighPrecisionBwdFactory: """MXFP8 forward + dequantized backward qfactory dispatch.""" @staticmethod def _factory(role): - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - mxfp8_fwd_dequantized_bwd_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_high_precision_bwd_factory, ) - return mxfp8_fwd_dequantized_bwd_quantizer_factory(role) + return mxfp8_fwd_high_precision_bwd_factory(role) @pytest.mark.parametrize("module_type", ["linear", "grouped_linear"]) @pytest.mark.parametrize("tensor_type", ["input", "weight"]) @@ -347,16 +357,16 @@ def test_grad_output_role_is_requantization_safe(self, module_type): @requires_nvfp4 -class TestNVFP4WeightDoubleQuantFactory: +class TestNVFP41DWeightFactory: """Base NVFP4 recipe with W.T sourced from dequantized 1D W.""" @staticmethod def _factory(role): - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - nvfp4_1d_double_quantized_weight_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_1d_weight_factory, ) - return nvfp4_1d_double_quantized_weight_quantizer_factory(role) + return nvfp4_1d_weight_factory(role) @staticmethod def _assert_plain_1d_nvfp4(quantizer): @@ -399,25 +409,25 @@ def _assert_matches_base_nvfp4_recipe(quantizer, expected): @pytest.mark.parametrize("tensor_type", ["input", "output", "grad_output", "grad_input"]) def test_non_weight_linear_roles_match_base_nvfp4_recipe(self, module_type, tensor_type): from transformer_engine.pytorch.quantization import QuantizerRole - from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - nvfp4_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + nvfp4_factory, ) role = QuantizerRole(module_type=module_type, tensor_type=tensor_type) quantizer = self._factory(role) - expected = nvfp4_quantizer_factory(role) + expected = nvfp4_factory(role) self._assert_matches_base_nvfp4_recipe(quantizer, expected) def test_non_linear_roles_match_base_nvfp4_recipe(self): from transformer_engine.pytorch.quantization import QuantizerRole - from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - nvfp4_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + nvfp4_factory, ) role = QuantizerRole(module_type="dpa", tensor_type="qkv") quantizer = self._factory(role) - expected = nvfp4_quantizer_factory(role) + expected = nvfp4_factory(role) self._assert_matches_base_nvfp4_recipe(quantizer, expected) @@ -718,8 +728,8 @@ def test_linear_high_precision_override_skips_requantization_safety_veto(self): assert counters["rowwise"].get("safety_calls", 0) == 0 def test_linear_custom_delayed_scaling_disables_save_original_input(self): - from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - delayed_scaling_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + delayed_scaling_factory, ) model = Linear( @@ -730,7 +740,7 @@ def test_linear_custom_delayed_scaling_disables_save_original_input(self): save_original_input=True, ).cuda() inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) - custom_recipe = recipe.CustomRecipe(qfactory=delayed_scaling_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=delayed_scaling_factory) with pytest.warns( UserWarning, @@ -1919,8 +1929,8 @@ def hybrid_mxfp8_factory(role): ) def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwise(self): - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - mxfp8_fwd_dequantized_bwd_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_high_precision_bwd_factory, ) torch.manual_seed(204) @@ -1951,7 +1961,7 @@ def test_dequantized_bwd_qfactory_save_original_input_matches_base_recipe_bitwis with autocast(enabled=True, recipe=ref_recipe): out_ref = model_ref(inp_ref) - qfactory_recipe = recipe.CustomRecipe(qfactory=mxfp8_fwd_dequantized_bwd_quantizer_factory) + qfactory_recipe = recipe.CustomRecipe(qfactory=mxfp8_fwd_high_precision_bwd_factory) with autocast(enabled=True, recipe=qfactory_recipe): out_qfactory = model_qfactory(inp_qfactory) @@ -2056,15 +2066,15 @@ def fake_infer(*_args, **_kwargs): @pytest.mark.parametrize( "factory_name,expected", [ - ("current_scaling_quantizer_factory", (True, False)), - ("delayed_scaling_quantizer_factory", (False, False)), + ("current_scaling_factory", (True, False)), + ("delayed_scaling_factory", (False, False)), ], ) def test_qkv_capabilities_reuse_canonical_quantizer(self, factory_name, expected): """Capability queries must not call qfactory outside recipe-state setup.""" - from transformer_engine.pytorch.custom_recipes import quantization_factory_base + from transformer_engine.pytorch.custom_recipes import quantizer_factories - base_qfactory = getattr(quantization_factory_base, factory_name) + base_qfactory = getattr(quantizer_factories, factory_name) calls = [] def counting_qfactory(role): @@ -2119,11 +2129,11 @@ def rebuilt_qfactory(role): def test_mha_forward_does_not_probe_qfactory(self, monkeypatch, fp8_mha): """Repeated MHA forwards must not create discarded DPA quantizers.""" from transformer_engine.pytorch.attention import multi_head_attention as mha_module - from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - current_scaling_quantizer_factory, - delayed_scaling_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + delayed_scaling_factory, ) - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) from transformer_engine.pytorch.utils import get_device_compute_capability @@ -2141,7 +2151,7 @@ def valid_fp8_dpa_qfactory(role): # complete MHA recipe so that all requests share one coherent # delayed-scaling state, including the DPA boundary tensors. if cc < (10, 0): - return delayed_scaling_quantizer_factory(role) + return delayed_scaling_factory(role) is_dpa = role is not None and role.module_type == "dpa" is_dpa_boundary = ( role is not None @@ -2150,7 +2160,7 @@ def valid_fp8_dpa_qfactory(role): ) if is_dpa or is_dpa_boundary: return nvfp4_linear_fp8_dpa_factory(role) - return current_scaling_quantizer_factory(role) + return current_scaling_factory(role) def counting_qfactory(role): calls.append(role) @@ -2327,7 +2337,7 @@ def test_linear_dpa_linear_matches_native_env_recipe_bitwise( if case_name == "mxfp8_dpa" and not mxfp8_available: pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") - from transformer_engine.pytorch.custom_recipes import quantization_factory_zoo + from transformer_engine.pytorch.custom_recipes import quantizer_factory_zoo from transformer_engine.pytorch.utils import get_device_compute_capability cc = get_device_compute_capability() @@ -2360,7 +2370,7 @@ def test_linear_dpa_linear_matches_native_env_recipe_bitwise( grad = torch.randn_like(base_inp) native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) - qfactory = getattr(quantization_factory_zoo, qfactory_name) + qfactory = getattr(quantizer_factory_zoo, qfactory_name) qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) native_out, native_dx, native_grads, native_local_recipes = self._run_model( @@ -2445,7 +2455,7 @@ def test_multihead_attention_matches_native_env_recipe_bitwise( pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") from transformer_engine.pytorch.attention import multi_head_attention as mha_module - from transformer_engine.pytorch.custom_recipes import quantization_factory_zoo + from transformer_engine.pytorch.custom_recipes import quantizer_factory_zoo from transformer_engine.pytorch.utils import get_device_compute_capability cc = get_device_compute_capability() @@ -2503,7 +2513,7 @@ def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) grad = torch.randn_like(base_inp) native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) - qfactory = getattr(quantization_factory_zoo, qfactory_name) + qfactory = getattr(quantizer_factory_zoo, qfactory_name) qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) native_out, native_dx, native_grads = self._run_mha_model( @@ -2601,7 +2611,7 @@ def test_mxfp8_qfactory_uses_plain_bf16_mha_boundaries(self, monkeypatch): from transformer_engine.pytorch.attention.dot_product_attention import ( backends as dpa_backends, ) - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_mxfp8_dpa_factory, ) from transformer_engine.pytorch.utils import get_device_compute_capability @@ -2642,10 +2652,38 @@ def _hook(grad): return _hook - def _record_dpa_inputs(_module, inputs): - for name, tensor in zip(("q", "k", "v"), inputs[:3]): - boundary_tensors[name] = tensor - tensor.register_hook(_record_grad(f"d{name}")) + def _record_dpa_inputs(_module, inputs, kwargs): + q, k, v = inputs[:3] + packed_qkv = kwargs.get("qkv_layer") + packed_kv = kwargs.get("kv_layer") + interleave_dim = kwargs.get("qkv_interleave_dim", -3) + + if packed_qkv is not None: + q, k, v = (packed_qkv.select(interleave_dim, i) for i in range(3)) + + def _record_packed_qkv_grad(grad): + for name, tensor in zip( + ("dq", "dk", "dv"), + (grad.select(interleave_dim, i) for i in range(3)), + ): + boundary_tensors[name] = tensor + + packed_qkv.register_hook(_record_packed_qkv_grad) + else: + q.register_hook(_record_grad("dq")) + if packed_kv is not None: + k, v = (packed_kv.select(interleave_dim, i) for i in range(2)) + + def _record_packed_kv_grad(grad): + boundary_tensors["dk"] = grad.select(interleave_dim, 0) + boundary_tensors["dv"] = grad.select(interleave_dim, 1) + + packed_kv.register_hook(_record_packed_kv_grad) + else: + k.register_hook(_record_grad("dk")) + v.register_hook(_record_grad("dv")) + + boundary_tensors.update(q=q, k=k, v=v) def _record_dpa_output(_module, _inputs, output): boundary_tensors["o"] = output @@ -2656,7 +2694,7 @@ def _record_projection_input(_module, inputs): boundary_tensors["o_is_proj_input"] = boundary_tensors["o"] is inputs[0] handles = ( - model.core_attention.register_forward_pre_hook(_record_dpa_inputs), + model.core_attention.register_forward_pre_hook(_record_dpa_inputs, with_kwargs=True), model.core_attention.register_forward_hook(_record_dpa_output), model.proj.register_forward_pre_hook(_record_projection_input), ) @@ -2800,10 +2838,10 @@ def hybrid_nvfp4_factory(role): and role.module_type in ("linear", "grouped_linear") and role.tensor_type == "grad_output" ): - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) return HybridQuantizer( - rowwise_quantizer=nvfp4_quantizer_factory(role), - columnwise_quantizer=nvfp4_quantizer_factory(role), + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=nvfp4_factory(role), ) hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_factory) @@ -2852,8 +2890,8 @@ def test_linear_fwd_bwd_all_roles_hybrid(self): def hybrid_nvfp4_all_roles_factory(role): return HybridQuantizer( - rowwise_quantizer=nvfp4_quantizer_factory(role), - columnwise_quantizer=nvfp4_quantizer_factory(role), + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=nvfp4_factory(role), ) hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_all_roles_factory) @@ -3035,7 +3073,7 @@ def test_identity_falls_back_to_high_precision(self): assert torch.equal(result, inp) def test_custom_tensor_passthrough(self): - from transformer_engine.pytorch.custom_recipes.quantization_ref_current_scaling import ( + from transformer_engine.pytorch.custom_recipes.reference_current_scaling import ( CurrentScalingTensorRef, ) @@ -6102,10 +6140,10 @@ def qfactory(role): if is_linear and role.tensor_type in ("input", "weight", "output"): return HybridQuantizer( rowwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), - columnwise_quantizer=nvfp4_quantizer_factory(role), + columnwise_quantizer=nvfp4_factory(role), ) if is_linear and role.tensor_type == "grad_output": - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) hybrid_recipe = recipe.CustomRecipe(qfactory=qfactory) @@ -6221,10 +6259,10 @@ def _hybrid_nvfp4_qfactory(role): """Hybrid NVFP4 (E2M1 both dirs, base role behavior).""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type == "grad_output": - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) return HybridQuantizer( - rowwise_quantizer=nvfp4_quantizer_factory(role), - columnwise_quantizer=nvfp4_quantizer_factory(role), + rowwise_quantizer=nvfp4_factory(role), + columnwise_quantizer=nvfp4_factory(role), ) @@ -6462,7 +6500,7 @@ class TestHybridCheckpoint: """Test state_dict save/load round-trips for models with hybrid quantized params.""" def _hybrid_checkpoint_recipe(self): - return recipe.CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_quantizer_factory) + return recipe.CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_factory) def test_state_dict_save_load_roundtrip(self): """state_dict → save → load → same model should produce identical outputs.""" @@ -7143,6 +7181,7 @@ def fn(model, inp, m_splits): # ----- Selective attention recompute ---------------------------- @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 + @_XFAIL_SELECTIVE_ATTENTION_RECOMPUTE def test_selective_attention_recompute_transformer_layer_fp8_bitwise(self): """``TransformerLayer(..., checkpoint_core_attention=True)`` — the Megatron default memory-savings pattern. diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 422b3fd9757..cb0785e3e0e 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -1420,15 +1420,15 @@ def _mxfp8_all_qfactory(role): # pylint: disable=unused-argument def _zoo_dequantized_qfactory(case_name): - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - mxfp8_fwd_dequantized_bwd_quantizer_factory, - nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_high_precision_bwd_factory, + nvfp4_row_scaled_fwd_high_precision_bwd_factory, ) if case_name == "mxfp8": - return mxfp8_fwd_dequantized_bwd_quantizer_factory + return mxfp8_fwd_high_precision_bwd_factory if case_name == "nvfp4_row_scaled": - return nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory + return nvfp4_row_scaled_fwd_high_precision_bwd_factory raise ValueError(case_name) @@ -1487,8 +1487,8 @@ def _assert_zoo_layernorm_mlp_role_semantics(case_name): assert isinstance(qfactory(role), IdentityQuantizer) -class TestZooDequantizedBackwardFactoryModuleCoverage: - """Zoo dequantized-backward recipes should match base recipes across TE modules.""" +class TestZooHighPrecisionBackwardFactoryModuleCoverage: + """Zoo high-precision-backward factories should match base recipes across TE modules.""" @pytest.mark.parametrize("case_name", _ZOO_DEQUANTIZED_CASES) @pytest.mark.parametrize("module_name", _ZOO_DEQUANTIZED_MODULES) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 13d259935ec..2341c56bd9e 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -25,7 +25,7 @@ from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm -from ...debug.pytorch.debug_quantization import DebugQuantizer +from ...debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer __all__ = [ "general_gemm", @@ -132,6 +132,11 @@ def _unwrap_tensor( if not isinstance(tensor, QuantizedTensorStorage): return tensor + # Select the requested representation from a debug wrapper, then process + # the wrapped tensor normally (it may itself be hybrid or TE-native). + if isinstance(tensor, DebugQuantizedTensor): + return _unwrap_tensor(tensor.get_tensor(usage == "columnwise"), usage) + # Select the direction of a hybrid tensor, then process its sub-storage. if isinstance(tensor, HybridQuantizedTensorStorage): sub_storage = ( @@ -199,6 +204,11 @@ def general_gemm( transa = layout[0] == "T" transb = layout[1] == "T" + debug_quantizer = None + if isinstance(quantization_params, DebugQuantizer): + debug_quantizer = quantization_params + quantization_params = quantization_params.parent_quantizer + A = _unwrap_tensor(A, "rowwise" if transa else "columnwise") B = _unwrap_tensor(B, "columnwise" if transb else "rowwise") @@ -240,13 +250,6 @@ def general_gemm( grad, ) - debug_quantizer = None - if isinstance(quantization_params, DebugQuantizer): - debug_quantizer = quantization_params - quantization_params = quantization_params.parent_quantizer - A = A.get_tensor(not transa) - B = B.get_tensor(transb) - _validate_native_gemm_output_quantizer(quantization_params) # Use bfloat16 as default bias_dtype diff --git a/transformer_engine/pytorch/custom_recipes/gemm.py b/transformer_engine/pytorch/custom_recipes/gemm.py index 3d1e1cc43e0..c2901c80236 100644 --- a/transformer_engine/pytorch/custom_recipes/gemm.py +++ b/transformer_engine/pytorch/custom_recipes/gemm.py @@ -4,18 +4,34 @@ """GEMM API that enables custom GEMM logic for custom quantization recipes.""" +import dataclasses +import enum from typing import Iterable, Optional import torch -from transformer_engine.pytorch.custom_recipes.quantization import ( - MMParams, - GEMMType, -) from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer from transformer_engine.pytorch.tensor.utils import is_custom +@enum.unique +class GEMMType(enum.Enum): + """Type of GEMM operation being performed.""" + + FPROP = "fprop" + DGRAD = "dgrad" + WGRAD = "wgrad" + + +@dataclasses.dataclass(frozen=True) +class MMParams: + """Matrix multiplication parameters.""" + + out_dtype: torch.dtype | None = None + # Use split accumulator for more accurate FP8 GEMM + use_split_accumulator: bool = True + + def custom_gemm( A: QuantizedTensorStorage, B: QuantizedTensorStorage, diff --git a/transformer_engine/pytorch/custom_recipes/quantization.py b/transformer_engine/pytorch/custom_recipes/quantization.py deleted file mode 100644 index 85920f50327..00000000000 --- a/transformer_engine/pytorch/custom_recipes/quantization.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Quantization API for experimental middleware between Transformer Engine and Kitchen.""" - -from __future__ import annotations -import dataclasses -import enum - -import torch - - -@enum.unique -class GEMMType(enum.Enum): - """Type of GEMM operation being performed.""" - - FPROP = "fprop" - DGRAD = "dgrad" - WGRAD = "wgrad" - - -@dataclasses.dataclass(frozen=True) -class MMParams: - """Matrix multiplication parameters.""" - - out_dtype: torch.dtype | None = None - # Use split accumulator for more accurate FP8 GEMM - use_split_accumulator: bool = True diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py similarity index 93% rename from transformer_engine/pytorch/custom_recipes/quantization_factory_base.py rename to transformer_engine/pytorch/custom_recipes/quantizer_factories.py index e5125fd92fd..250ce8c9393 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py @@ -3,7 +3,7 @@ # See LICENSE for license information. """ -Quantizer factory examples using native Transformer Engine quantizers. +Quantizer factories that mirror Transformer Engine's built-in recipes. For TE ``Linear`` and ``GroupedLinear`` roles, each factory below mirrors the nominal defaults of a built-in recipe through the ``CustomRecipe`` + @@ -14,11 +14,11 @@ from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_base import ( - nvfp4_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + nvfp4_factory, ) - recipe = CustomRecipe(qfactory=nvfp4_quantizer_factory) + recipe = CustomRecipe(qfactory=nvfp4_factory) with autocast(recipe=recipe): output = model(input) """ @@ -50,7 +50,7 @@ def high_precision_factory( return IdentityQuantizer() -def delayed_scaling_quantizer_factory( +def delayed_scaling_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "DelayedScalingRequest": """Factory that mirrors :class:`DelayedScaling` recipe defaults. @@ -69,7 +69,7 @@ def delayed_scaling_quantizer_factory( return DelayedScalingRequest(fp8_format=Format.HYBRID) -def current_scaling_quantizer_factory( +def current_scaling_factory( role: Optional[QuantizerRole], ) -> "Float8CurrentScalingQuantizer": """Factory that mirrors :class:`Float8CurrentScaling` recipe defaults. @@ -92,22 +92,7 @@ def current_scaling_quantizer_factory( ) -def mxfp8_quantizer_factory( - role: Optional[QuantizerRole], # pylint: disable=unused-argument -) -> "MXFP8Quantizer": - """Factory that mirrors :class:`MXFP8BlockScaling` recipe defaults. - - * E4M3 by default for all tensors - * Block size 32, power-of-2 (E8M0) scales - """ - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - return MXFP8Quantizer( - fp8_dtype=DType.kFloat8E4M3, - ) - - -def float8_block_scaling_quantizer_factory( +def float8_block_scaling_factory( role: Optional[QuantizerRole], ) -> "Float8BlockQuantizer": """Factory that mirrors :class:`Float8BlockScaling` recipe defaults. @@ -137,7 +122,22 @@ def float8_block_scaling_quantizer_factory( ) -def nvfp4_quantizer_factory( +def mxfp8_factory( + role: Optional[QuantizerRole], # pylint: disable=unused-argument +) -> "MXFP8Quantizer": + """Factory that mirrors :class:`MXFP8BlockScaling` recipe defaults. + + * E4M3 by default for all tensors + * Block size 32, power-of-2 (E8M0) scales + """ + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer( + fp8_dtype=DType.kFloat8E4M3, + ) + + +def nvfp4_factory( role: Optional[QuantizerRole], ) -> "NVFP4Quantizer": """Factory that mirrors :class:`NVFP4BlockScaling` recipe defaults. diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py similarity index 76% rename from transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py rename to transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 8b238970507..1b123a474a6 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -3,14 +3,14 @@ # See LICENSE for license information. """ -Quantizer factory zoo. +Example quantizer factories for custom and mixed quantization recipes. -A collection of composed/mixed-recipe factories built on top of the single-format -building blocks. They demonstrate how to use the ``CustomRecipe`` + ``qfactory`` -interface to apply *different* quantization recipes to different -module/tensor types/instances within the same model, and how to use -``HybridQuantizer`` when rowwise and columnwise tensor directions should use -different representations or sources. +A collection of composed/mixed-recipe factories. They demonstrate how to use +the ``CustomRecipe`` + ``qfactory`` interface to apply *different* quantization +recipes to different module/tensor types/instances within the same model. +Factories may return native Transformer Engine quantizers, custom quantizers, +or ``HybridQuantizer`` instances when tensor directions should use different +representations or sources. Within the Linear/GroupedLinear and RL-oriented families, examples are roughly ordered by increasingly aggressive forward quantization. This is an @@ -18,16 +18,17 @@ These factories demonstrate what can be composed; they are not necessarily tuned for end-to-end performance. -When the columnwise/backward representation is higher precision than -the rowwise/forward representation, consider -``columnwise_source="rowwise_dequantized"`` as a starting point. It constructs -backward operands from the value obtained during forward quantization, -improving forward/backward representation consistency. +When forward operands (inputs and weights) are quantized and use a different +representation in backward, consider setting +``columnwise_source="rowwise_dequantized"`` on their hybrid quantizers. This +applies whether backward uses another low-precision format or high precision. +It constructs backward operands from the value obtained during forward +quantization, improving forward/backward representation consistency. This +source choice should not be applied to gradient tensors. Note: dequantization does not recover information discarded during forward quantization. Organization: - * Linear / grouped-linear recipes (pre-training). Favor more precision - on the forward pass. + * Pre-training-oriented recipes: Favor more precision on the forward pass. * RL-oriented recipes: Favor more precision in backward GEMMs. * Linear + attention recipes: factories that also cover ``DotProductAttention`` roles and require ``CustomRecipe(..., fp8_dpa=True)``. @@ -46,13 +47,13 @@ from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( - mxfp8_fwd_nvfp4_bwd_quantizer_factory, + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + mxfp8_fwd_nvfp4_bwd_factory, nvfp4_linear_mxfp8_dpa_factory, ) # Linear-only recipe (no attention quantization): the qfactory is the only knob. - recipe = CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_quantizer_factory) + recipe = CustomRecipe(qfactory=mxfp8_fwd_nvfp4_bwd_factory) with autocast(recipe=recipe): output = model(input) @@ -72,14 +73,14 @@ from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType -from .quantization_factory_base import mxfp8_quantizer_factory, nvfp4_quantizer_factory +from .quantizer_factories import mxfp8_factory, nvfp4_factory # ----------------------------------------------------------------------------- -# Linear / GroupedLinear Recipes (pre-training) +# Pre-training-Oriented Recipes # ----------------------------------------------------------------------------- -def high_precision_fwd_mxfp8_bwd_quantizer_factory( +def high_precision_fwd_mxfp8_bwd_factory( role: Optional[QuantizerRole], ): """Quantizer factory: high-precision forward, MXFP8 backward. @@ -93,16 +94,30 @@ def high_precision_fwd_mxfp8_bwd_quantizer_factory( is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type == "grad_output": - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) # fprop consumes rowwise high precision; dgrad / wgrad consume columnwise MXFP8. return HybridQuantizer( rowwise_quantizer=IdentityQuantizer(), - columnwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=mxfp8_factory(role), ) -def mxfp8_fwd_nvfp4_bwd_quantizer_factory( +def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): + """NVFP4 quantizer without RHT, stochastic rounding, or 2D scaling.""" + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + return NVFP4Quantizer( + fp4_dtype=DType.kFloat4E2M1, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=row_scaled_nvfp4, + ) + + +def mxfp8_fwd_nvfp4_bwd_factory( role: Optional[QuantizerRole], ): """Quantizer factory: MXFP8 forward, NVFP4 backward. @@ -112,53 +127,50 @@ def mxfp8_fwd_nvfp4_bwd_quantizer_factory( * dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)`` * wgrad: ``input.col(NVFP4) x grad_output.col(NVFP4)`` - The NVFP4 side mirrors :func:`nvfp4_quantizer_factory` role semantics: - weights use 2D scaling, activations/grads use the base 1D RHT/SR settings. - ``HybridQuantizer`` then pins each sub-quantizer to the direction that is - actually consumed. + Every backward operand uses 1D NVFP4 scaling. Inputs and gradients mirror + :func:`nvfp4_factory` semantics: RHT is applied only to the columnwise + representations consumed by wgrad, and gradients use stochastic rounding. + The dgrad weight uses plain 1D NVFP4 without RHT or stochastic rounding. + + The backward input and weight representations are quantized to NVFP4 + directly from the original high-precision tensors rather than from the + dequantized MXFP8 forward representations. In ``HybridQuantizer`` terms, + this source choice is expressed with ``columnwise_source="original"``, + which is the default. To derive them from the forward representations + instead, use ``columnwise_source="rowwise_dequantized"``. """ from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight"): + if role.tensor_type == "weight": + backward_quantizer = _plain_nvfp4_quantizer() + else: + backward_quantizer = nvfp4_factory(role) return HybridQuantizer( - rowwise_quantizer=mxfp8_quantizer_factory(role), - columnwise_quantizer=nvfp4_quantizer_factory(role), + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=backward_quantizer, ) if is_linear and role.tensor_type == "grad_output": - return nvfp4_quantizer_factory(role) - return mxfp8_quantizer_factory(role) + return nvfp4_factory(role) + return mxfp8_factory(role) -def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): - """NVFP4 quantizer without RHT, stochastic rounding, or 2D scaling.""" - from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - - return NVFP4Quantizer( - fp4_dtype=DType.kFloat4E2M1, - with_rht=False, - with_post_rht_amax=False, - with_2d_quantization=False, - stochastic_rounding=False, - row_scaled_nvfp4=row_scaled_nvfp4, - ) - - -def nvfp4_1d_double_quantized_weight_quantizer_factory( +def nvfp4_1d_weight_factory( role: Optional[QuantizerRole], ): - """Quantizer factory: NVFP4 recipe with 1D weight double quantization. + """Quantizer factory: NVFP4 recipe with 1D weight scaling. Dispatch logic: * ``linear`` / ``grouped_linear`` ``weight`` -> ``Hybrid(rowwise=plain 1D NVFP4, columnwise=plain 1D NVFP4, columnwise_source="rowwise_dequantized")`` - * everything else -> :func:`nvfp4_quantizer_factory` + * everything else -> :func:`nvfp4_factory` - ``W.T`` is quantized from the dequantized forward/rowwise ``W`` value - instead of directly from the original high-precision weight. In - ``HybridQuantizer`` terms, that source choice is expressed with - ``columnwise_source="rowwise_dequantized"``. + The backward weight representation (``W.T``) is quantized to NVFP4 from + the dequantized NVFP4 forward weight rather than directly from the original + high-precision weight. In ``HybridQuantizer`` terms, this source choice is + expressed with ``columnwise_source="rowwise_dequantized"``. All non-weight roles keep the standard NVFP4 factory behavior, including RHT for inputs and stochastic rounding for gradients. The weight override uses @@ -174,7 +186,7 @@ def nvfp4_1d_double_quantized_weight_quantizer_factory( columnwise_quantizer=_plain_nvfp4_quantizer(), columnwise_source="rowwise_dequantized", ) - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) # ----------------------------------------------------------------------------- @@ -182,10 +194,10 @@ def nvfp4_1d_double_quantized_weight_quantizer_factory( # ----------------------------------------------------------------------------- -def mxfp8_fwd_dequantized_bwd_quantizer_factory( +def mxfp8_fwd_high_precision_bwd_factory( role: Optional[QuantizerRole], ): - """Quantizer factory: MXFP8 forward, high-precision dequantized backward. + """Quantizer factory: MXFP8 forward, high-precision backward. This expresses the linear/grouped-linear equivalent of ``backward_override="dequantized"`` through per-direction quantizers: @@ -195,8 +207,11 @@ def mxfp8_fwd_dequantized_bwd_quantizer_factory( * ``grad_output`` -> ``IdentityQuantizer`` * everything else -> MXFP8 - Backward GEMMs use high-precision operands dequantized from the MXFP8 - forward payload, avoiding gradient quantization. + The backward input and weight representations are high-precision values + dequantized from the MXFP8 forward representations rather than the original + high-precision tensors. In ``HybridQuantizer`` terms, this source choice is + expressed with ``columnwise_source="rowwise_dequantized"``. The gradient + output independently remains in high precision. This recipe targets RL-style training use cases and is motivated by NVIDIA/TransformerEngine#2644, where ``backward_override="dequantized"`` @@ -209,19 +224,19 @@ def mxfp8_fwd_dequantized_bwd_quantizer_factory( is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("input", "weight"): return HybridQuantizer( - rowwise_quantizer=mxfp8_quantizer_factory(role), + rowwise_quantizer=mxfp8_factory(role), columnwise_quantizer=IdentityQuantizer(), columnwise_source="rowwise_dequantized", ) if is_linear and role.tensor_type == "grad_output": return IdentityQuantizer() - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) -def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( +def nvfp4_row_scaled_fwd_mxfp8_bwd_factory( role: Optional[QuantizerRole], ): - """Quantizer factory: row-scaled NVFP4 forward, dequantized-source MXFP8 backward. + """Quantizer factory: row-scaled NVFP4 forward, MXFP8 backward. This RL-related recipe is inspired by the Composer 2 MoE grouped-GEMM recipe described in arXiv:2603.24477. @@ -245,13 +260,13 @@ def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( * ``grad_output`` -> MXFP8 * everything else -> MXFP8 - Row-scaled NVFP4 is fprop-only, so the forward NVFP4 quantizers avoid - RHT, stochastic rounding, and 2D scaling. This sample assumes MXFP8 - backward operands are quantized from the dequantized rowwise NVFP4 forward - value, expressed with ``columnwise_source="rowwise_dequantized"``. To - quantize MXFP8 backward operands from the original high-precision tensor - instead, construct the same ``HybridQuantizer`` while omitting - ``columnwise_source="rowwise_dequantized"``. + Row-scaled NVFP4 is fprop-only, so the forward NVFP4 quantizers avoid RHT, + stochastic rounding, and 2D scaling. The backward input and weight + representations are quantized to MXFP8 from the dequantized NVFP4 forward + representations rather than directly from the original high-precision + tensors. In ``HybridQuantizer`` terms, this source choice is expressed with + ``columnwise_source="rowwise_dequantized"``. To use the original tensors + instead, use ``columnwise_source="original"``. Composer 2 Technical Report: https://arxiv.org/abs/2603.24477 @@ -263,26 +278,26 @@ def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( if is_grouped_linear and role.tensor_type == "input": return HybridQuantizer( rowwise_quantizer=_plain_nvfp4_quantizer(row_scaled_nvfp4=True), - columnwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=mxfp8_factory(role), columnwise_source="rowwise_dequantized", ) if is_grouped_linear and role.tensor_type == "weight": return HybridQuantizer( rowwise_quantizer=_plain_nvfp4_quantizer(), - columnwise_quantizer=mxfp8_quantizer_factory(role), + columnwise_quantizer=mxfp8_factory(role), columnwise_source="rowwise_dequantized", ) if is_grouped_linear and role.tensor_type == "grad_output": - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) if is_linear: - return mxfp8_quantizer_factory(role) - return mxfp8_quantizer_factory(role) + return mxfp8_factory(role) + return mxfp8_factory(role) -def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( +def nvfp4_row_scaled_fwd_high_precision_bwd_factory( role: Optional[QuantizerRole], ): - """Quantizer factory: row-scaled NVFP4 forward, dequantized backward. + """Quantizer factory: row-scaled NVFP4 forward, high-precision backward. This expresses a linear/grouped-linear variant of ``NVFP4BlockScaling(row_scaled_activation=True, @@ -300,6 +315,12 @@ def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( Row-scaled NVFP4 is fprop-only, so the forward quantizers avoid RHT, stochastic rounding, and 2D scaling. + The backward input and weight representations are high-precision values + dequantized from the NVFP4 forward representations rather than the original + high-precision tensors. In ``HybridQuantizer`` terms, this source choice is + expressed with ``columnwise_source="rowwise_dequantized"``. The gradient + output independently remains in high precision. + This recipe targets RL-style training use cases and builds on NVIDIA/TransformerEngine#2931, which introduced row-scaled NVFP4: https://github.com/NVIDIA/TransformerEngine/pull/2931 @@ -367,7 +388,7 @@ def nvfp4_linear_fp8_dpa_factory( from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) @@ -405,7 +426,7 @@ def nvfp4_linear_fp8_dpa_factory( fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) def nvfp4_linear_mxfp8_dpa_factory( @@ -448,7 +469,7 @@ def nvfp4_linear_mxfp8_dpa_factory( from transformer_engine.common.recipe import CustomRecipe from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantization_factory_zoo import ( + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_mxfp8_dpa_factory, ) @@ -475,4 +496,4 @@ def nvfp4_linear_mxfp8_dpa_factory( fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 return MXFP8Quantizer(fp8_dtype=fp8_dtype) - return nvfp4_quantizer_factory(role) + return nvfp4_factory(role) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py similarity index 97% rename from transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py rename to transformer_engine/pytorch/custom_recipes/reference_current_scaling.py index ecbb667ecfa..4016d4b0f22 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py @@ -10,12 +10,12 @@ import torch -from transformer_engine.pytorch.custom_recipes import quantization -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes import gemm +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer -def current_scaling_ref_quantizer_factory(role): +def current_scaling_ref_factory(role): """Factory function for current scaling reference quantizer. Receives a :class:`~transformer_engine.pytorch.quantization.QuantizerRole`. @@ -24,7 +24,7 @@ def current_scaling_ref_quantizer_factory(role): Usage with CustomRecipe and autocast:: - custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_ref_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_ref_factory) with autocast(recipe=custom_recipe): output = model(input) """ @@ -340,7 +340,7 @@ def quantize( **kwargs, # pylint: disable=unused-argument ) -> CurrentScalingTensorRef: # sanity checks - assert tensor.dtype in utils.HIGH_PRECISION_FLOAT_DTYPES, "Unsupported input dtype." + assert tensor.dtype in reference_utils.HIGH_PRECISION_FLOAT_DTYPES, "Unsupported input dtype." # Make it work with 3D tensors original_shape = tensor.shape @@ -374,14 +374,14 @@ def qgemm( self, qx: torch.Tensor, qw: torch.Tensor, - m_params: quantization.MMParams, + m_params: gemm.MMParams, out_dtype: torch.dtype, sx: torch.Tensor, sw: torch.Tensor, bias: torch.Tensor | None = None, out: torch.Tensor | None = None, accumulate: bool = False, - gemm_type: quantization.GEMMType = quantization.GEMMType.FPROP, # pylint: disable=unused-argument + gemm_type: gemm.GEMMType = gemm.GEMMType.FPROP, # pylint: disable=unused-argument qresult_x: QuantizedTensorStorage | None = None, # pylint: disable=unused-argument qresult_w: QuantizedTensorStorage | None = None, # pylint: disable=unused-argument ) -> torch.Tensor: diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py similarity index 98% rename from transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py rename to transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index a2b21465af0..3aae1e9d82c 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -9,12 +9,12 @@ import torch -from transformer_engine.pytorch.custom_recipes import quantization -from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes import gemm +from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer -def nvfp4_ref_rht_2d_quantizer_factory(role): +def nvfp4_ref_rht_2d_factory(role): """ Quantizer factory for NVFP4 recipe reference implementation (RHT and 2D quantization for weights). @@ -22,7 +22,7 @@ def nvfp4_ref_rht_2d_quantizer_factory(role): Usage with CustomRecipe and autocast:: - custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_quantizer_factory) + custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_factory) with autocast(recipe=custom_recipe): output = model(input) """ @@ -33,13 +33,13 @@ def nvfp4_ref_rht_2d_quantizer_factory(role): ) if is_weight_tensor_in_gemm: # 2D quantization for weights in GEMM-based modules return NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(16, 16), pow_2_scales=False, with_rht=False, ) return NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, + dtype=reference_utils.Fp4Formats.E2M1, quant_tile_shape=(1, 16), pow_2_scales=False, with_rht=True, @@ -207,7 +207,7 @@ class NVFP4TensorRef(QuantizedTensorStorage): nominal tensor datatype. device: torch.device device of the tensor. - quant_dtype: Union[utils.Fp4Formats, torch.dtype] + quant_dtype: Union[reference_utils.Fp4Formats, torch.dtype] low precision tensor datatype. original_shape: Tuple[int, ...] original shape of the tensor. @@ -226,7 +226,7 @@ class NVFP4TensorRef(QuantizedTensorStorage): dtype: Optional[torch.dtype] = None device: Optional[torch.device] = None - quant_dtype: Optional[Union[utils.Fp4Formats, torch.dtype]] = None + quant_dtype: Optional[Union[reference_utils.Fp4Formats, torch.dtype]] = None original_shape: Optional[Tuple[int, ...]] = None _quantizer: Optional[Quantizer] = None @@ -345,7 +345,7 @@ class NVFP4QuantizerRef(Quantizer): def __init__( self, - dtype: utils.Fp4Formats, + dtype: reference_utils.Fp4Formats, rowwise: bool = True, columnwise: bool = True, pow_2_scales: bool = False, @@ -452,9 +452,9 @@ def _recover_swizzled_scales( ) -> torch.Tensor: if not swizzled_scale: return scale - rounded_m = utils.roundup_div(m, 128) * 128 - scale_n = utils.roundup_div(n, block_length) - rounded_n = utils.roundup_div(scale_n, 4) * 4 + rounded_m = reference_utils.roundup_div(m, 128) * 128 + scale_n = reference_utils.roundup_div(n, block_length) + rounded_n = reference_utils.roundup_div(scale_n, 4) * 4 # Recover swizzled scaling factor layout -> linear layout tmp = torch.reshape(scale, (rounded_m // 128, rounded_n // 4, 32, 4, 4)) # after permutation, the layout is [rounded_m // 128, 4, 32, rounded_n // 4, 4] @@ -963,10 +963,10 @@ def quantize( **kwargs, # pylint: disable=unused-argument ) -> NVFP4TensorRef: # sanity checks - if tensor.dtype not in utils.HIGH_PRECISION_FLOAT_DTYPES: + if tensor.dtype not in reference_utils.HIGH_PRECISION_FLOAT_DTYPES: raise TypeError( f"Unsupported input dtype {tensor.dtype}, expected one of" - f" {utils.HIGH_PRECISION_FLOAT_DTYPES}" + f" {reference_utils.HIGH_PRECISION_FLOAT_DTYPES}" ) # Make it work with 3D tensors @@ -1076,14 +1076,14 @@ def qgemm( self, qx: torch.Tensor, qw: torch.Tensor, - m_params: quantization.MMParams, # pylint: disable=unused-argument + m_params: gemm.MMParams, # pylint: disable=unused-argument out_dtype: torch.dtype, sx: torch.Tensor, sw: torch.Tensor, bias: torch.Tensor | None = None, out: torch.Tensor | None = None, accumulate: bool = False, - gemm_type: quantization.GEMMType = quantization.GEMMType.FPROP, + gemm_type: gemm.GEMMType = gemm.GEMMType.FPROP, qresult_x: QuantizedTensorStorage | None = None, qresult_w: QuantizedTensorStorage | None = None, ) -> torch.Tensor: @@ -1169,7 +1169,7 @@ def qgemm( fp8_max_w = 448.0 factor = 6.0 * 6.0 * fp8_max_x * fp8_max_w - if gemm_type == quantization.GEMMType.WGRAD: + if gemm_type == gemm.GEMMType.WGRAD: partial_alpha = qresult_x.global_amax_col * qresult_w.global_amax_col # A row-scaled operand contributes a per-output-column (N) vector # here, so broadcast along the last axis. Selecting the axis from diff --git a/transformer_engine/pytorch/custom_recipes/utils.py b/transformer_engine/pytorch/custom_recipes/reference_utils.py similarity index 83% rename from transformer_engine/pytorch/custom_recipes/utils.py rename to transformer_engine/pytorch/custom_recipes/reference_utils.py index 3e23661f147..5bd25e35b21 100644 --- a/transformer_engine/pytorch/custom_recipes/utils.py +++ b/transformer_engine/pytorch/custom_recipes/reference_utils.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Utility functions for experimental middleware between Transformer Engine and Kitchen.""" +"""Shared utilities for custom recipe reference implementations.""" import enum diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py index 8f6051c5f83..dc65c9894b3 100644 --- a/transformer_engine/pytorch/tensor/hybrid_tensor.py +++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py @@ -67,7 +67,7 @@ class HybridQuantizer(Quantizer): specialization. For example, a factory can return a ``HybridQuantizer`` only for ``role.tensor_type == "weight"`` and use regular quantizers for inputs and gradients. See - ``custom_recipes.quantization_factory_zoo.nvfp4_1d_double_quantized_weight_quantizer_factory`` + ``custom_recipes.quantizer_factory_zoo.nvfp4_1d_weight_factory`` for a weight-only double-quantization example. """ From aebd39265b266c4e74b71ef73257ecdd8c9646ac Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:46:07 +0000 Subject: [PATCH 63/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pytorch/custom_recipes/reference_current_scaling.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py index 4016d4b0f22..4e5292019ac 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py @@ -340,7 +340,9 @@ def quantize( **kwargs, # pylint: disable=unused-argument ) -> CurrentScalingTensorRef: # sanity checks - assert tensor.dtype in reference_utils.HIGH_PRECISION_FLOAT_DTYPES, "Unsupported input dtype." + assert ( + tensor.dtype in reference_utils.HIGH_PRECISION_FLOAT_DTYPES + ), "Unsupported input dtype." # Make it work with 3D tensors original_shape = tensor.shape From c59bc39091c13cff70f14a50838630664de00f56 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 31 Jul 2026 16:06:04 +0000 Subject: [PATCH 64/97] update policy keys and add state.quantizer_config + tests Signed-off-by: Evgeny --- tests/pytorch/hybrid_quantization_utils.py | 4 +- tests/pytorch/test_custom_recipe.py | 36 +++ tests/pytorch/test_hybrid_quantization.py | 173 +++++++++-- tests/pytorch/test_identity_quantizer.py | 50 ++- tests/pytorch/test_recipe.py | 288 +++++++++++++++++- transformer_engine/common/recipe/__init__.py | 4 +- transformer_engine/pytorch/__init__.py | 1 + .../dot_product_attention.py | 2 + .../quantization_factory_base.py | 12 +- .../quantization_factory_zoo.py | 16 +- .../quantization_ref_current_scaling.py | 2 + .../custom_recipes/quantization_ref_nvfp4.py | 2 + transformer_engine/pytorch/quantization.py | 28 +- 13 files changed, 563 insertions(+), 55 deletions(-) diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py index dfbe8d4c344..4b0281a4236 100644 --- a/tests/pytorch/hybrid_quantization_utils.py +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -192,7 +192,7 @@ def make_role_aware_quantizer(factory, role): return quantizer -def hybrid_custom_recipe(row_factory, col_factory, grad_factory=None): +def hybrid_custom_recipe(row_factory, col_factory, grad_factory=None, *, policy_key): """Build a CustomRecipe with hybrid forward and configurable grad quantizers.""" if grad_factory is None: grad_factory = col_factory @@ -208,7 +208,7 @@ def qfactory(role): return make_role_aware_quantizer(grad_factory, role) return make_role_aware_quantizer(row_factory, role) - return recipe.CustomRecipe(qfactory=qfactory) + return recipe.CustomRecipe(qfactory=qfactory, policy_key=policy_key) def hybrid_fp8_current_qfactory(role): diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index e88ef93c958..19871ea81f9 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -40,6 +40,30 @@ ) +@pytest.mark.parametrize( + "qfactory", + [ + current_scaling_quantizer_factory, + mxfp8_quantizer_factory, + float8_block_scaling_quantizer_factory, + nvfp4_quantizer_factory, + delayed_scaling_quantizer_factory, + high_precision_factory, + mxfp8_fwd_nvfp4_bwd_quantizer_factory, + nvfp4_ref_rht_2d_quantizer_factory, + ], +) +def test_first_party_qfactories_have_canonical_policy_keys(qfactory): + """First-party factories can be used without repeating policy_key in CustomRecipe.""" + policy_key = getattr(qfactory, "policy_key") + hash(policy_key) + assert isinstance(policy_key, tuple) + assert len(policy_key) == 2 + assert isinstance(policy_key[0], str) and policy_key[0] + assert isinstance(policy_key[1], int) and policy_key[1] > 0 + assert recipe.CustomRecipe(qfactory=qfactory).policy_key == policy_key + + @pytest.mark.parametrize("module_type", ["Linear", "LayerNormLinear", "OpsLinear"]) def test_custom_recipe_sanity_modules_nvfp4(module_type): """Test modules with NVFP4 custom recipe support""" @@ -107,6 +131,7 @@ def test_custom_recipe_sanity(module_type): inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) # Single factory: map roles to quantizers + @recipe.quantizer_policy(key=("test_current_scaling_sanity", 1)) def quantizer_factory(role): if role is None: return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -144,6 +169,7 @@ def test_custom_recipe_grouped_linear_sanity(): model = GroupedLinear(num_gemms, in_features, out_features, params_dtype=torch.bfloat16).cuda() inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) + @recipe.quantizer_policy(key=("test_current_scaling_grouped_linear", 1)) def quantizer_factory(role): if role is None: return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -206,6 +232,7 @@ def test_custom_recipe_matches_current_scaling(): loss_ref.backward() # Custom: single factory returning quantizers per role to match Float8CurrentScaling + @recipe.quantizer_policy(key=("test_current_scaling_parity", 1)) def quantizer_factory(role): if role is None: return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -263,6 +290,7 @@ def test_custom_recipe_ops_linear_2_1_layout(): op = te_ops.Linear(in_features, out_features, device="cuda", dtype=torch.bfloat16) inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) + @recipe.quantizer_policy(key=("test_current_scaling_ops_linear", 1)) def quantizer_factory(role): if role is None: return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -306,6 +334,7 @@ def test_custom_recipe_factory_invocation_counts_and_cycling(): None: 0, } + @recipe.quantizer_policy(key=("test_factory_invocation_counts", 1)) def quantizer_factory(role): if role is None: counts[None] += 1 @@ -607,6 +636,7 @@ def test_custom_recipe_quantization_targets(): # ------------------------------------------------------------------ recorded_roles = [] + @recipe.quantizer_policy(key=("test_per_module_targeting", 1)) def targeting_factory(role): recorded_roles.append(role) @@ -764,6 +794,7 @@ def test_grouped_linear_module_type_dispatch(): recorded_roles = [] + @recipe.quantizer_policy(key=("test_grouped_linear_role_recording", 1)) def recording_factory(role): recorded_roles.append(role) return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -804,6 +835,7 @@ def test_delayed_scaling_request_wiring(): from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.common.recipe import Format + @recipe.quantizer_policy(key=("test_delayed_scaling_request", 1)) def ds_factory(role): return DelayedScalingRequest(fp8_format=Format.HYBRID, amax_history_len=16) @@ -864,6 +896,7 @@ def test_custom_recipe_mixed_ds_and_stateless(): from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.common.recipe import Format + @recipe.quantizer_policy(key=("test_mixed_delayed_current_scaling", 1)) def mixed_factory(role): # Only weight gets delayed scaling, rest get current scaling if role is not None and role.tensor_type == "weight": @@ -908,6 +941,7 @@ def test_custom_recipe_ds_multi_step(): from transformer_engine.pytorch.quantization import DelayedScalingRequest from transformer_engine.common.recipe import Format + @recipe.quantizer_policy(key=("test_delayed_scaling_multi_step", 1)) def ds_factory(role): return DelayedScalingRequest(fp8_format=Format.HYBRID) @@ -1031,6 +1065,7 @@ def test_role_change_preserves_custom_delayed_scaling_state(): ) from transformer_engine.common.recipe import Format + @recipe.quantizer_policy(key=("test_delayed_scaling_role_rebuild", 1)) def ds_factory(role): return DelayedScalingRequest(fp8_format=Format.HYBRID, amax_history_len=8) @@ -1903,6 +1938,7 @@ def test_custom_recipe_quantization_alignment_contract(): """The alignment contract is safe, configurable, and does not invoke qfactory.""" calls = [] + @recipe.quantizer_policy(key=("test_quantization_alignment", 1)) def qfactory(role): calls.append(role) return current_scaling_quantizer_factory(role) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 84036ae84f3..bf2885d065d 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -807,7 +807,16 @@ def factory(role): ) return IdentityQuantizer() - return recipe.CustomRecipe(qfactory=factory) + return recipe.CustomRecipe( + qfactory=factory, + policy_key=( + "test_counting_identity_hybrid", + 1, + columnwise_source, + rowwise_safe, + columnwise_safe, + ), + ) def test_linear_save_original_input_veto_uses_saved_forward_quantized_input(self): torch.manual_seed(205) @@ -1686,6 +1695,7 @@ def test_te_linear_inference_workspace_rowwise_only(self): hybrid_recipe = _hybrid_custom_recipe( row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), + policy_key=("test_inference_workspace_hybrid_fp8", 1), ) torch.manual_seed(2026) model = Linear(128, 256, bias=False, params_dtype=torch.bfloat16).cuda() @@ -1831,7 +1841,10 @@ def hybrid_fp8_factory(role): device="cuda", ) - hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_fp8_factory) + hybrid_recipe = recipe.CustomRecipe( + qfactory=hybrid_fp8_factory, + policy_key=("test_hybrid_fp8_parity", 1), + ) with autocast(enabled=True, recipe=hybrid_recipe): out_hybrid = model_hybrid(inp_hybrid) loss_hybrid = out_hybrid.float().sum() @@ -1897,7 +1910,10 @@ def hybrid_mxfp8_factory(role): columnwise_quantizer=MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3), ) - hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_mxfp8_factory) + hybrid_recipe = recipe.CustomRecipe( + qfactory=hybrid_mxfp8_factory, + policy_key=("test_hybrid_mxfp8_parity", 1), + ) with autocast(enabled=True, recipe=hybrid_recipe): out_hybrid = model_hybrid(inp_hybrid) out_hybrid.float().sum().backward() @@ -1989,7 +2005,10 @@ def test_inference_runs_once_per_quantizer_state_and_clears_stale_labels(self, m from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import FP8GlobalStateManager - custom_recipe = recipe.CustomRecipe(qfactory=lambda _role: IdentityQuantizer()) + custom_recipe = recipe.CustomRecipe( + qfactory=lambda _role: IdentityQuantizer(), + policy_key=("test_dpa_identity", 1), + ) monkeypatch.setattr( FP8GlobalStateManager, "get_fp8_recipe", @@ -2078,6 +2097,7 @@ def counting_qfactory(role): qfactory=counting_qfactory, fp8_dpa=True, fp8_mha=True, + policy_key=("test_dpa_capability_counting", factory_name, 1), ) dpa = te.DotProductAttention( num_attention_heads=2, @@ -2106,6 +2126,7 @@ def rebuilt_qfactory(role): qfactory=rebuilt_qfactory, fp8_dpa=True, fp8_mha=True, + policy_key=("test_dpa_capability_rebuilt", factory_name, 1), ) with autocast(enabled=True, recipe=rebuilt_recipe): rebuilt = dpa.get_qkv_quantization_capabilities() @@ -2161,6 +2182,7 @@ def counting_qfactory(role): qfactory=counting_qfactory, fp8_dpa=True, fp8_mha=fp8_mha, + policy_key=("test_mha_factory_invocation_count", fp8_mha, 1), ) model = te.MultiheadAttention( hidden_size=128, @@ -2745,7 +2767,10 @@ def hybrid_block_fp8_factory(role): ), ) - hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_block_fp8_factory) + hybrid_recipe = recipe.CustomRecipe( + qfactory=hybrid_block_fp8_factory, + policy_key=("test_hybrid_block_fp8_parity", 1), + ) with autocast(enabled=True, recipe=hybrid_recipe): out_hybrid = model_hybrid(inp_hybrid) out_hybrid.float().sum().backward() @@ -2806,7 +2831,10 @@ def hybrid_nvfp4_factory(role): columnwise_quantizer=nvfp4_quantizer_factory(role), ) - hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_factory) + hybrid_recipe = recipe.CustomRecipe( + qfactory=hybrid_nvfp4_factory, + policy_key=("test_hybrid_nvfp4_parity", 1), + ) torch.manual_seed(1202) torch.cuda.manual_seed_all(1202) with autocast(enabled=True, recipe=hybrid_recipe): @@ -2856,7 +2884,10 @@ def hybrid_nvfp4_all_roles_factory(role): columnwise_quantizer=nvfp4_quantizer_factory(role), ) - hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_nvfp4_all_roles_factory) + hybrid_recipe = recipe.CustomRecipe( + qfactory=hybrid_nvfp4_all_roles_factory, + policy_key=("test_hybrid_nvfp4_all_roles", 1), + ) torch.manual_seed(1203) torch.cuda.manual_seed_all(1203) with autocast(enabled=True, recipe=hybrid_recipe): @@ -2918,7 +2949,10 @@ def mixed_factory(role): return _make_nvfp4_quantizer() return _make_fp8_quantizer() - mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + mixed_recipe = recipe.CustomRecipe( + qfactory=mixed_factory, + policy_key=("test_mixed_fp8_nvfp4", 1), + ) with autocast(enabled=True, recipe=mixed_recipe): out = model(inp) @@ -2968,7 +3002,10 @@ def mixed_factory(role): ) return _make_fp8_quantizer() - mixed_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + mixed_recipe = recipe.CustomRecipe( + qfactory=mixed_factory, + policy_key=("test_mixed_fp8_nvfp4_numerics", 1), + ) with torch.no_grad(): with autocast(enabled=True, recipe=mixed_recipe): out_mixed = model(inp) @@ -3078,7 +3115,7 @@ def factory(role): ), ) - return factory + return recipe.quantizer_policy(key=("test_uniform_hybrid_fp8", 1))(factory) def test_bias_grad_matches_vanilla_fp8(self): torch.manual_seed(456) @@ -3169,7 +3206,13 @@ def factory(role): return _make_nvfp4_quantizer() return _make_fp8_quantizer() - with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + with autocast( + enabled=True, + recipe=recipe.CustomRecipe( + qfactory=factory, + policy_key=("test_matching_columnwise_formats", 1), + ), + ): out = model(inp) out.float().sum().backward() assert inp.grad is not None @@ -3201,7 +3244,13 @@ def factory(role): ) return _make_fp8_quantizer() - with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=factory)): + with autocast( + enabled=True, + recipe=recipe.CustomRecipe( + qfactory=factory, + policy_key=("test_mismatched_columnwise_formats", 1), + ), + ): out = model(inp) with pytest.raises(RuntimeError, match="scaling_mode"): out.float().sum().backward() @@ -3236,7 +3285,10 @@ def factory(role): ) return _make_nvfp4_quantizer() - mixed_recipe = recipe.CustomRecipe(qfactory=factory) + mixed_recipe = recipe.CustomRecipe( + qfactory=factory, + policy_key=("test_nvfp4_row_fp8_column_forward", 1), + ) with torch.no_grad(): with autocast(enabled=True, recipe=mixed_recipe): out = model(inp) @@ -3273,7 +3325,10 @@ def factory(role): return _make_fp8_quantizer() return _make_nvfp4_quantizer() - mixed_recipe = recipe.CustomRecipe(qfactory=factory) + mixed_recipe = recipe.CustomRecipe( + qfactory=factory, + policy_key=("test_nvfp4_row_fp8_column", 1), + ) with autocast(enabled=True, recipe=mixed_recipe): out = model(inp) @@ -3330,7 +3385,10 @@ def mixed_factory(role): model_hybrid, base_input, grad_output, - recipe.CustomRecipe(qfactory=mixed_factory), + recipe.CustomRecipe( + qfactory=mixed_factory, + policy_key=("test_single_hybrid_role", hybrid_role, 1), + ), seed=seed + 100, ) native_result = _run_linear_forward_backward( @@ -3504,7 +3562,9 @@ def factory(role): return _make_role_aware_quantizer(grad_factory, role) return _make_role_aware_quantizer(operand_factory, role) - return factory + return recipe.quantizer_policy( + key=("test_plain_linear", operand_factory.__name__, grad_factory.__name__, 1) + )(factory) def _assert_linear_results_exact(actual, expected, *, output, input_grad, param_grads): @@ -3563,7 +3623,10 @@ def hybrid_factory(role): return _make_role_aware_quantizer(make_col_grad, role) return _make_role_aware_quantizer(make_row_operand, role) - hybrid_recipe = recipe.CustomRecipe(qfactory=hybrid_factory) + hybrid_recipe = recipe.CustomRecipe( + qfactory=hybrid_factory, + policy_key=("test_cross_format", row_name, col_name, 1), + ) fprop_ref_recipe = recipe.CustomRecipe( qfactory=_plain_linear_qfactory(make_row_operand, make_row_grad) ) @@ -3835,7 +3898,17 @@ def hybrid_factory(role): model_hybrid, base_input, grad_output, - recipe.CustomRecipe(qfactory=hybrid_factory), + recipe.CustomRecipe( + qfactory=hybrid_factory, + policy_key=( + "test_three_format", + make_fprop.__name__, + make_dgrad.__name__, + make_wgrad.__name__, + plain_grad_output, + 1, + ), + ), seed=seed + 100, ) fprop_ref = _run_linear_forward( @@ -3935,7 +4008,7 @@ def factory(role): device="cuda", ) - return factory + return recipe.quantizer_policy(key=("test_hybrid_fp8_all_modules", 1))(factory) @requires_fp8 @@ -3967,7 +4040,13 @@ def qfactory(_role): requires_grad=True, ) - with autocast(enabled=True, recipe=recipe.CustomRecipe(qfactory=qfactory)): + with autocast( + enabled=True, + recipe=recipe.CustomRecipe( + qfactory=qfactory, + policy_key=("test_fusible_norm_hybrid_output", 1), + ), + ): out = forward(inp) assert isinstance(out, HybridQuantizedTensor) @@ -4483,7 +4562,9 @@ def qfactory(_role): columnwise_source=columnwise_source, ) - return qfactory + return recipe.quantizer_policy( + key=("test_grouped_linear_columnwise_source", columnwise_source, 1) + )(qfactory) model = GroupedLinear(2, 128, 128, bias=False, params_dtype=torch.bfloat16).cuda() tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") @@ -4536,7 +4617,10 @@ def mixed_source_qfactory(role): columnwise_source=source, ) - mixed_recipe = recipe.CustomRecipe(qfactory=mixed_source_qfactory) + mixed_recipe = recipe.CustomRecipe( + qfactory=mixed_source_qfactory, + policy_key=("test_mixed_columnwise_sources", 1), + ) # A failed generation is never marked validated. Base metadata can then # early-return on retry, so the O(1) guard must validate it again. for _ in range(2): @@ -4608,6 +4692,7 @@ def _hybrid_fp8_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_quantized_model_init_hybrid_fp8", 1), ) def test_linear_weight_is_hybrid_quantized_tensor(self): @@ -4711,6 +4796,7 @@ def _hybrid_fp8_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_weight_workspace_hybrid_fp8", 1), ) @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 @@ -4816,6 +4902,7 @@ def _hybrid_fp8_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_weight_quantizer_update_hybrid_fp8", 1), ) def test_quantized_param_survives_multiple_forward_passes(self): @@ -4882,6 +4969,7 @@ def _hybrid_recipe_fp8_current(): row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), + policy_key=("test_hybrid_fp8_current", 1), ) @@ -4914,6 +5002,7 @@ def _hybrid_recipe_fp8_delayed(): row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), + policy_key=("test_hybrid_fp8_delayed", 1), ) @@ -4930,6 +5019,7 @@ def _hybrid_recipe_fp8_delayed_row_current_col(): # grad_factory matches the columnwise direction so the wgrad GEMM's # grad_output sub-quantizer pairs with the input/weight col format. grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), + policy_key=("test_hybrid_fp8_delayed_row_current_column", 1), ) @@ -4943,6 +5033,7 @@ def _hybrid_recipe_fp8_current_row_delayed_col(): row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), + policy_key=("test_hybrid_fp8_current_row_delayed_column", 1), ) @@ -4952,6 +5043,7 @@ def _hybrid_recipe_mxfp8(): row_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), grad_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E5M2), + policy_key=("test_hybrid_mxfp8", 1), ) @@ -4967,6 +5059,7 @@ def _hybrid_recipe_blockwise(): grad_factory=lambda: Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True ), + policy_key=("test_hybrid_blockwise", 1), ) @@ -5526,6 +5619,7 @@ def test_mxfp8_columnwise_raises(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_master_weight_mxfp8_columnwise", 1), ) # Shape must be a multiple of MXFP8 block size (32) on both axes. weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) @@ -5559,6 +5653,7 @@ def test_nvfp4_rowwise_raises(self): grad_factory=lambda: NVFP4Quantizer( fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False ), + policy_key=("test_master_weight_nvfp4_rowwise", 1), ) weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) master_flat = hp_master.view(-1).contiguous() @@ -5580,6 +5675,7 @@ def test_nvfp4_columnwise_raises(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_master_weight_nvfp4_columnwise", 1), ) weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) master_flat = hp_master.view(-1).contiguous() @@ -5625,6 +5721,7 @@ def test_blockwise_columnwise_raises(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_master_weight_blockwise_columnwise", 1), ) weight, hp_master = _build_hybrid_linear_weight(128, 128, hybrid_recipe) master_flat = hp_master.view(-1).contiguous() @@ -5783,6 +5880,7 @@ def _hybrid_fp8_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_recipe_correspondence_hybrid_fp8", 1), ) def test_hybrid_param_with_matching_recipe_does_not_raise(self): @@ -5882,6 +5980,7 @@ def _build_hybrid_model(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_fused_adam_hybrid_fp8", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -5985,6 +6084,7 @@ def _build_model_and_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), + policy_key=("test_end_to_end_hybrid_fp8", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -6108,7 +6208,10 @@ def qfactory(role): return nvfp4_quantizer_factory(role) return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) - hybrid_recipe = recipe.CustomRecipe(qfactory=qfactory) + hybrid_recipe = recipe.CustomRecipe( + qfactory=qfactory, + policy_key=("test_mxfp8_nvfp4_quantized_params", 1), + ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() return model, hybrid_recipe @@ -6396,7 +6499,10 @@ def _vanilla_recipe(self): return recipe.Float8CurrentScaling() def _hybrid_recipe(self): - return recipe.CustomRecipe(qfactory=_hybrid_fp8_current_qfactory) + return recipe.CustomRecipe( + qfactory=_hybrid_fp8_current_qfactory, + policy_key=("test_hybrid_fp8_current", 1), + ) def test_equivalence(self): self._test_equivalence() @@ -6410,7 +6516,10 @@ def _vanilla_recipe(self): return recipe.MXFP8BlockScaling() def _hybrid_recipe(self): - return recipe.CustomRecipe(qfactory=_hybrid_mxfp8_qfactory) + return recipe.CustomRecipe( + qfactory=_hybrid_mxfp8_qfactory, + policy_key=("test_hybrid_mxfp8", 1), + ) def test_equivalence(self): self._test_equivalence() @@ -6424,7 +6533,10 @@ def _vanilla_recipe(self): return recipe.Float8BlockScaling() def _hybrid_recipe(self): - return recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) + return recipe.CustomRecipe( + qfactory=_hybrid_block_fp8_qfactory, + policy_key=("test_hybrid_block_fp8", 1), + ) def test_equivalence(self): self._test_equivalence() @@ -6441,7 +6553,10 @@ def _vanilla_recipe(self): return recipe.NVFP4BlockScaling() def _hybrid_recipe(self): - return recipe.CustomRecipe(qfactory=_hybrid_nvfp4_qfactory) + return recipe.CustomRecipe( + qfactory=_hybrid_nvfp4_qfactory, + policy_key=("test_hybrid_nvfp4", 1), + ) def test_equivalence(self): self._test_equivalence() @@ -6747,6 +6862,7 @@ def _same_format_fp8_recipe(self): row_factory=_fp8_row_factory, col_factory=_fp8_col_factory, grad_factory=_fp8_grad_factory, + policy_key=("test_recompute_hybrid_fp8", 1), ) def _same_format_mxfp8_recipe(self): @@ -6756,6 +6872,7 @@ def _same_format_mxfp8_recipe(self): row_factory=_mxfp8_factory, col_factory=_mxfp8_factory, grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), + policy_key=("test_recompute_hybrid_mxfp8", 1), ) def _cross_format_fp8_mxfp8_recipe(self): @@ -6767,6 +6884,7 @@ def _cross_format_fp8_mxfp8_recipe(self): row_factory=_fp8_row_factory, col_factory=_mxfp8_factory, grad_factory=_mxfp8_factory, + policy_key=("test_recompute_fp8_mxfp8", 1), ) def _run_linear(self, recipe_obj, *, checkpoint_fn=None): @@ -7292,6 +7410,7 @@ def test_te_checkpoint_linear_all_stateless_formats_bitwise(self, format_name, r row_factory=row_factory, col_factory=row_factory, grad_factory=grad_factory, + policy_key=("test_checkpoint_hybrid", format_name, 1), ) def fn(model, inp): diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 422b3fd9757..117366defdd 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -15,7 +15,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex -from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.common.recipe import CustomRecipe, quantizer_policy from transformer_engine.pytorch import ( Float8BlockQuantizer, Float8CurrentScalingQuantizer, @@ -56,6 +56,7 @@ # ── Module-level qfactories (picklable / autocast-friendly) ────────── +@quantizer_policy(key=("test_identity_all", 1)) def identity_all_factory(role): # pylint: disable=unused-argument """Whole layer in high precision: Identity for every slot.""" return IdentityQuantizer() @@ -134,6 +135,7 @@ def _format_quantizer(format_name): def _hybrid_quantized_fwd_identity_bwd_factory(format_name): + @quantizer_policy(key=(f"test_{format_name}_fwd_identity_bwd", 1)) def qfactory(role): is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("grad_output", "grad_input"): @@ -146,6 +148,7 @@ def qfactory(role): return qfactory +@quantizer_policy(key=("test_high_precision_fwd_fp8_bwd", 1)) def fwd_hp_bwd_fp8_factory(role): """High-precision forward, FP8 backward (per-direction via hybrid).""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -157,6 +160,7 @@ def fwd_hp_bwd_fp8_factory(role): ) +@quantizer_policy(key=("test_fp8_fwd_high_precision_bwd", 1)) def fwd_fp8_bwd_hp_factory(role): """FP8 forward, high-precision backward (per-direction via hybrid).""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -168,6 +172,7 @@ def fwd_fp8_bwd_hp_factory(role): ) +@quantizer_policy(key=("test_fp8_fwd_dequantized_high_precision_bwd", 1)) def fwd_fp8_bwd_rowwise_dequantized_hp_factory(role): """FP8 forward, high-precision backward from dequantized fprop values.""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -180,6 +185,7 @@ def fwd_fp8_bwd_rowwise_dequantized_hp_factory(role): ) +@quantizer_policy(key=("test_hybrid_identity_all", 1)) def hybrid_all_identity_factory(role): """All directions high precision, expressed through the hybrid container. @@ -196,6 +202,7 @@ def hybrid_all_identity_factory(role): ) +@quantizer_policy(key=("test_fp8_all", 1)) def fp8_fwd_factory(role): """Plain FP8 current scaling for every slot (E4M3 fwd, E5M2 grad). @@ -386,7 +393,13 @@ def fake_split_quantize_hybrid( m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) with pytest.raises(StopAfterFlagCapture): - with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + with te.autocast( + enabled=True, + recipe=CustomRecipe( + qfactory=qfactory, + policy_key=("test_grouped_linear_cpu_offload_hybrid", 1), + ), + ): model(x, m_splits=m_splits) assert calls == [True] @@ -410,7 +423,13 @@ def qfactory(role): m_splits = torch.tensor([32, 32], device="cuda", dtype=torch.int32) with pytest.raises(ValueError, match="mix Identity-backed and quantized"): - with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + with te.autocast( + enabled=True, + recipe=CustomRecipe( + qfactory=qfactory, + policy_key=("test_grouped_linear_mixed_identity_weights", 1), + ), + ): model(x, m_splits=m_splits) def test_identity_contiguous_preserves_wrapper_and_values(self): @@ -498,7 +517,10 @@ def qfactory(role): # pylint: disable=unused-argument torch.manual_seed(1701) ref = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) - custom_recipe = CustomRecipe(qfactory=qfactory) + custom_recipe = CustomRecipe( + qfactory=qfactory, + policy_key=("test_ops_basic_linear_hybrid_identity_weight", 1), + ) torch.manual_seed(1702) with te.quantized_model_init(enabled=True, recipe=custom_recipe): test = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) @@ -549,7 +571,13 @@ def test_te_ops_quantize_then_gelu_accepts_identity_backed_tensors(self, qfactor x_test = x.detach().clone().requires_grad_(True) y_ref = ref(x_ref) - with te.autocast(enabled=True, recipe=CustomRecipe(qfactory=qfactory)): + with te.autocast( + enabled=True, + recipe=CustomRecipe( + qfactory=qfactory, + policy_key=("test_ops_quantize_gelu_identity_backed", 1), + ), + ): y_test = test(x_test) torch.manual_seed(1801) @@ -1161,7 +1189,10 @@ def test_identity_recipe_matches_bf16_bitwise(self, module_name, qfactory): # wgrad in the same order on every architecture. ref, test = _make_identity_module_pair(module_name, seed=7300, bias=False) x = _identity_module_input(module_name) - recipe = CustomRecipe(qfactory=qfactory) + recipe = CustomRecipe( + qfactory=qfactory, + policy_key=("test_quantized_primary_hybrid_identity", 1), + ) y_ref, dx_ref, wg_ref = _fwd_bwd_module(module_name, ref, x, recipe=None) y_id, dx_id, wg_id = _fwd_bwd_module(module_name, test, x, recipe=recipe) @@ -1204,7 +1235,11 @@ def run(model, x, recipe): y_bo, dx_bo, wg_bo = run( ref, x, - CustomRecipe(qfactory=mxfp8_all_factory, backward_override="high_precision"), + CustomRecipe( + qfactory=mxfp8_all_factory, + policy_key=("test_mxfp8_all", 1), + backward_override="high_precision", + ), ) y_id, dx_id, wg_id = run( test, @@ -1415,6 +1450,7 @@ def _fwd_bwd_zoo_dequantized_module(module_name, module, inp, recipe): return out.detach().clone(), inp.grad.detach().clone(), param_grads +@quantizer_policy(key=("test_mxfp8_all", 1)) def _mxfp8_all_qfactory(role): # pylint: disable=unused-argument return _mxfp8(tex.DType.kFloat8E4M3) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index ccef104a332..12c47de8953 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. +from dataclasses import fields from typing import Optional import pickle @@ -40,6 +41,7 @@ MXFP8BlockScaling, NVFP4BlockScaling, Recipe, + quantizer_policy, ) from transformer_engine.pytorch._extra_state import ( CheckpointExtraStatePolicy, @@ -57,6 +59,288 @@ fp4_available, reason_for_no_fp4 = te.is_nvfp4_available(return_reason=True) +def test_recipe_quantizer_config_cache_and_invalidation(): + """Semantic configuration is built once and rebuilt after recipe mutation.""" + + config_builds = [] + + class CountingRecipe(Recipe): + def __init__(self, value: int) -> None: + self.value = value + + def _make_quantizer_config(self): + config_builds.append(self.value) + return ("counting_recipe", self.value) + + def _make_repr(self) -> str: + return f"counting_recipe={self.value}" + + recipe = CountingRecipe(1) + assert recipe.quantizer_config() == ("counting_recipe", 1) + assert recipe.quantizer_config() == ("counting_recipe", 1) + assert config_builds == [1] + + # Populating the independent repr cache must not invalidate the config cache. + assert repr(recipe) == "counting_recipe=1" + assert recipe.quantizer_config() == ("counting_recipe", 1) + assert config_builds == [1] + + recipe.value = 2 + assert recipe.quantizer_config() == ("counting_recipe", 2) + assert config_builds == [1, 2] + + +def test_recipe_quantizer_config_must_be_hashable(): + """Reject mutable containers that cannot serve as runtime-key components.""" + + class UnhashableConfigRecipe(Recipe): + def _make_quantizer_config(self): + return ["mutable"] + + def _make_repr(self) -> str: + return "unhashable_config_recipe" + + with pytest.raises( + TypeError, + match=r"UnhashableConfigRecipe\._make_quantizer_config\(\) must return a hashable value", + ): + UnhashableConfigRecipe().quantizer_config() + + +@pytest.mark.parametrize( + "recipe_type,derived_fields", + [ + (DelayedScaling, ()), + ( + Float8CurrentScaling, + ("fp8_quant_fwd_inp", "fp8_quant_fwd_weight", "fp8_quant_bwd_grad"), + ), + (MXFP8BlockScaling, ()), + ( + Float8BlockScaling, + ("fp8_quant_fwd_inp", "fp8_quant_fwd_weight", "fp8_quant_bwd_grad"), + ), + ( + NVFP4BlockScaling, + ("fp4_quant_fwd_inp", "fp4_quant_fwd_weight", "fp4_quant_bwd_grad"), + ), + ], +) +def test_builtin_recipe_quantizer_configs_are_exhaustive( + recipe_type, + derived_fields, +): + """Built-in configs include every declared field and derived parameter bundle.""" + recipe = recipe_type() + config = recipe.quantizer_config() + config_labels = tuple(name for name, _ in config) + expected_labels = { + "recipe_type", + *(field.name for field in fields(recipe_type)), + *derived_fields, + } + + assert len(config_labels) == len(set(config_labels)) + assert set(config_labels) == expected_labels + + config_by_name = dict(config) + for name in derived_fields: + params = getattr(recipe, name) + expected = tuple( + (field.name, getattr(params, field.name)) for field in fields(type(params)) + ) + assert config_by_name[name] == expected + + for name in ("fp8_gemm_fprop", "fp8_gemm_dgrad", "fp8_gemm_wgrad"): + if hasattr(recipe, name): + params = getattr(recipe, name) + expected = tuple( + (field.name, getattr(params, field.name)) for field in fields(type(params)) + ) + assert config_by_name[name] == expected + + +@pytest.mark.parametrize( + "recipe_type", + [ + DelayedScaling, + Float8CurrentScaling, + MXFP8BlockScaling, + Float8BlockScaling, + NVFP4BlockScaling, + ], +) +def test_equal_builtin_recipes_have_equal_quantizer_configs(recipe_type): + """Independent built-in recipes with the same ingredients are semantically equal.""" + first = recipe_type().quantizer_config() + second = recipe_type().quantizer_config() + assert first == second + + +@pytest.mark.parametrize( + "algorithm_field", + ["amax_compute_algo", "scaling_factor_compute_algo"], +) +def test_delayed_scaling_callable_config_uses_explicit_key(algorithm_field): + """Delayed-scaling callable identity is replaced by an explicit semantic key.""" + calls = [] + + def first_algorithm(*args): + calls.append(("first", args)) + + def second_algorithm(*args): + calls.append(("second", args)) + + key_field = f"{algorithm_field}_key" + shared_key = ("test_algorithm", 1) + first = DelayedScaling(**{algorithm_field: first_algorithm, key_field: shared_key}) + second = DelayedScaling(**{algorithm_field: second_algorithm, key_field: shared_key}) + + assert first.quantizer_config() == second.quantizer_config() + assert calls == [] + + second = DelayedScaling(**{algorithm_field: second_algorithm, key_field: ("test_algorithm", 2)}) + assert first.quantizer_config() != second.quantizer_config() + + missing_key = DelayedScaling(**{algorithm_field: first_algorithm}) + with pytest.raises( + ValueError, + match=rf"{algorithm_field} is callable, so {key_field} must provide", + ): + missing_key.quantizer_config() + + +def test_quantizer_policy_contract(): + """The decorator validates and attaches metadata without wrapping or calling.""" + calls = [] + + def factory(role): + calls.append(role) + + decorated = quantizer_policy(key=("test_policy", 1))(factory) + + assert decorated is factory + assert getattr(decorated, "policy_key") == ("test_policy", 1) + assert calls == [] + assert te.quantizer_policy is quantizer_policy + + with pytest.raises(TypeError, match="quantizer_policy key must be hashable"): + quantizer_policy(key=["mutable"]) + + +def test_custom_recipe_policy_key_contract(): + """Custom recipes resolve semantic keys without calling their factories.""" + calls = [] + + def first_factory(role): + calls.append(("first", role)) + + def second_factory(role): + calls.append(("second", role)) + + policy_key = ("hybrid_policy", 2, ("double_quantization", True)) + first = CustomRecipe(qfactory=first_factory, policy_key=policy_key) + second = CustomRecipe(qfactory=second_factory, policy_key=policy_key) + + assert first.quantizer_config() == second.quantizer_config() + assert set(dict(first.quantizer_config())) == { + "recipe_type", + "policy_key", + "fp8_format", + "fp8_dpa", + "fp8_mha", + "backward_override", + "quantization_alignment", + } + assert calls == [] + + @quantizer_policy(key=("attached_policy", 1)) + def factory(role): # pylint: disable=unused-argument + calls.append(("attached", role)) + + attached = CustomRecipe(qfactory=factory) + assert attached.policy_key == ("attached_policy", 1) + assert dict(attached.quantizer_config())["policy_key"] == ("attached_policy", 1) + + explicit = CustomRecipe(qfactory=factory, policy_key=("explicit_policy", 2)) + assert explicit.policy_key == ("explicit_policy", 2) + assert dict(explicit.quantizer_config())["policy_key"] == ("explicit_policy", 2) + + def unkeyed_factory(role): + calls.append(("unkeyed", role)) + + with pytest.raises( + ValueError, + match=r"Pass policy_key=.*@quantizer_policy", + ): + CustomRecipe(qfactory=unkeyed_factory).quantizer_config() + assert calls == [] + + +def test_global_state_caches_active_quantizer_config(): + """Recipe activation publishes the recipe and its semantic configuration together.""" + FP8GlobalStateManager.reset() + active_recipe = Float8CurrentScaling() + + FP8GlobalStateManager.activate_recipe(active_recipe) + original_config = active_recipe.quantizer_config() + assert FP8GlobalStateManager.get_fp8_recipe() is active_recipe + assert FP8GlobalStateManager.get_quantizer_config() is original_config + + # Recipe mutation is requested state until the recipe is activated again. + active_recipe.fp8_dpa = True + assert FP8GlobalStateManager.get_quantizer_config() is original_config + + FP8GlobalStateManager.activate_recipe(active_recipe) + assert FP8GlobalStateManager.get_quantizer_config() == active_recipe.quantizer_config() + assert FP8GlobalStateManager.get_quantizer_config() != original_config + FP8GlobalStateManager.reset() + + +def test_recipe_activation_does_not_call_qfactory_and_is_atomic_on_config_error(): + """Configuration is resolved without probing the factory and before changing active state.""" + FP8GlobalStateManager.reset() + calls = [] + + @quantizer_policy(key=("activation_test", 1)) + def keyed_factory(role): + calls.append(role) + + keyed_recipe = CustomRecipe(qfactory=keyed_factory) + FP8GlobalStateManager.activate_recipe(keyed_recipe) + keyed_config = keyed_recipe.quantizer_config() + assert FP8GlobalStateManager.get_fp8_recipe() is keyed_recipe + assert FP8GlobalStateManager.get_quantizer_config() is keyed_config + assert calls == [] + + def unkeyed_factory(role): + calls.append(role) + + with pytest.raises(ValueError, match="requires a semantic policy key"): + FP8GlobalStateManager.activate_recipe(CustomRecipe(qfactory=unkeyed_factory)) + assert FP8GlobalStateManager.get_fp8_recipe() is keyed_recipe + assert FP8GlobalStateManager.get_quantizer_config() is keyed_config + assert calls == [] + FP8GlobalStateManager.reset() + + +def test_autocast_restores_recipe_and_quantizer_config_together(): + """Leaving autocast restores the previously active recipe/configuration pair.""" + FP8GlobalStateManager.reset() + outer_recipe = Float8CurrentScaling() + inner_recipe = MXFP8BlockScaling() + FP8GlobalStateManager.activate_recipe(outer_recipe) + outer_config = FP8GlobalStateManager.get_quantizer_config() + + with te.autocast(enabled=False, recipe=inner_recipe): + assert FP8GlobalStateManager.get_fp8_recipe() is inner_recipe + assert FP8GlobalStateManager.get_quantizer_config() == inner_recipe.quantizer_config() + + assert FP8GlobalStateManager.get_fp8_recipe() is outer_recipe + assert FP8GlobalStateManager.get_quantizer_config() is outer_config + FP8GlobalStateManager.reset() + + # FP8 per tensor delayed scaling @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) class TestFP8Recipe: @@ -331,7 +615,9 @@ def test_scale_update_numeric_scenarios(self, amax_case, fused_update, fp8_dtype ) ) recipe = transformer_engine.common.recipe.DelayedScaling( - fp8_format=fp8_format, scaling_factor_compute_algo=scaling_factor_compute_algo + fp8_format=fp8_format, + scaling_factor_compute_algo=scaling_factor_compute_algo, + scaling_factor_compute_algo_key=("test_fused_update", 1) if fused_update else None, ) # Setup fp8_meta dictionary diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 103b01e9dce..2dd00a915c9 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -152,7 +152,9 @@ def quantizer_policy(*, key: Hashable) -> Callable: The key, rather than Python callable identity, represents the factory's behavior in :class:`CustomRecipe` semantic configuration. Change the key whenever the factory's - quantizer-selection behavior changes. + quantizer-selection behavior changes. A recommended convention is + ``(policy_name, behavior_revision)``, where ``policy_name`` is a short, stable name and + ``behavior_revision`` is incremented only when the factory's behavior changes. """ policy_key = _validate_policy_key(key, source="quantizer_policy key") diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 4c6b7fc67a6..c01117733cb 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -53,6 +53,7 @@ from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.quantization import QuantizerRequest from transformer_engine.pytorch.quantization import DelayedScalingRequest +from transformer_engine.common.recipe import quantizer_policy from transformer_engine.pytorch.utils import get_cudnn_version from transformer_engine.pytorch.utils import get_device_compute_capability from transformer_engine.pytorch.utils import is_bf16_available diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 9ae0df880a8..0ac06a7f1be 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -132,6 +132,8 @@ def _delayed_scaling_recipe() -> Optional[DelayedScaling]: reduce_amax=ds_recipe.reduce_amax, fp8_dpa=fp8_recipe.fp8_dpa, fp8_mha=fp8_recipe.fp8_mha, + amax_compute_algo_key=ds_recipe.amax_compute_algo_key, + scaling_factor_compute_algo_key=ds_recipe.scaling_factor_compute_algo_key, ) if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py index 79725efa5de..e63302a4329 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_base.py @@ -33,7 +33,7 @@ from ..constants import DType -@quantizer_policy(key=("transformer_engine", "high_precision", 1)) +@quantizer_policy(key=("high_precision", 1)) def high_precision_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "IdentityQuantizer": @@ -52,7 +52,7 @@ def high_precision_factory( return IdentityQuantizer() -@quantizer_policy(key=("transformer_engine", "delayed_scaling", 1)) +@quantizer_policy(key=("delayed_scaling", 1)) def delayed_scaling_quantizer_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "DelayedScalingRequest": @@ -72,7 +72,7 @@ def delayed_scaling_quantizer_factory( return DelayedScalingRequest(fp8_format=Format.HYBRID) -@quantizer_policy(key=("transformer_engine", "current_scaling", 1)) +@quantizer_policy(key=("current_scaling", 1)) def current_scaling_quantizer_factory( role: Optional[QuantizerRole], ) -> "Float8CurrentScalingQuantizer": @@ -96,7 +96,7 @@ def current_scaling_quantizer_factory( ) -@quantizer_policy(key=("transformer_engine", "mxfp8", 1)) +@quantizer_policy(key=("mxfp8", 1)) def mxfp8_quantizer_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "MXFP8Quantizer": @@ -112,7 +112,7 @@ def mxfp8_quantizer_factory( ) -@quantizer_policy(key=("transformer_engine", "float8_block_scaling", 1)) +@quantizer_policy(key=("float8_block_scaling", 1)) def float8_block_scaling_quantizer_factory( role: Optional[QuantizerRole], ) -> "Float8BlockQuantizer": @@ -143,7 +143,7 @@ def float8_block_scaling_quantizer_factory( ) -@quantizer_policy(key=("transformer_engine", "nvfp4", 1)) +@quantizer_policy(key=("nvfp4", 1)) def nvfp4_quantizer_factory( role: Optional[QuantizerRole], ) -> "NVFP4Quantizer": diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py index 67e146a82ef..44d6d72134a 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_zoo.py @@ -80,7 +80,7 @@ # ----------------------------------------------------------------------------- -@quantizer_policy(key=("transformer_engine", "high_precision_fwd_mxfp8_bwd", 1)) +@quantizer_policy(key=("high_precision_fwd_mxfp8_bwd", 1)) def high_precision_fwd_mxfp8_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -104,7 +104,7 @@ def high_precision_fwd_mxfp8_bwd_quantizer_factory( ) -@quantizer_policy(key=("transformer_engine", "mxfp8_fwd_nvfp4_bwd", 1)) +@quantizer_policy(key=("mxfp8_fwd_nvfp4_bwd", 1)) def mxfp8_fwd_nvfp4_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -147,7 +147,7 @@ def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): ) -@quantizer_policy(key=("transformer_engine", "nvfp4_1d_double_quantized_weight", 1)) +@quantizer_policy(key=("nvfp4_1d_weight", 1)) def nvfp4_1d_double_quantized_weight_quantizer_factory( role: Optional[QuantizerRole], ): @@ -186,7 +186,7 @@ def nvfp4_1d_double_quantized_weight_quantizer_factory( # ----------------------------------------------------------------------------- -@quantizer_policy(key=("transformer_engine", "mxfp8_fwd_dequantized_bwd", 1)) +@quantizer_policy(key=("mxfp8_fwd_high_precision_bwd", 1)) def mxfp8_fwd_dequantized_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -223,7 +223,7 @@ def mxfp8_fwd_dequantized_bwd_quantizer_factory( return mxfp8_quantizer_factory(role) -@quantizer_policy(key=("transformer_engine", "nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd", 1)) +@quantizer_policy(key=("composer2", 1)) def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -285,7 +285,7 @@ def nvfp4_row_scaled_fwd_dequantized_mxfp8_bwd_quantizer_factory( return mxfp8_quantizer_factory(role) -@quantizer_policy(key=("transformer_engine", "nvfp4_row_scaled_fwd_dequantized_bwd", 1)) +@quantizer_policy(key=("nvfp4_row_scaled_fwd_high_precision_bwd", 1)) def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( role: Optional[QuantizerRole], ): @@ -337,7 +337,7 @@ def nvfp4_row_scaled_fwd_dequantized_bwd_quantizer_factory( # ----------------------------------------------------------------------------- -@quantizer_policy(key=("transformer_engine", "nvfp4_linear_fp8_dpa", 1)) +@quantizer_policy(key=("nvfp4_linear_fp8_dpa", 1)) def nvfp4_linear_fp8_dpa_factory( role: Optional[QuantizerRole], ): @@ -416,7 +416,7 @@ def nvfp4_linear_fp8_dpa_factory( return nvfp4_quantizer_factory(role) -@quantizer_policy(key=("transformer_engine", "nvfp4_linear_mxfp8_dpa", 1)) +@quantizer_policy(key=("nvfp4_linear_mxfp8_dpa", 1)) def nvfp4_linear_mxfp8_dpa_factory( role: Optional[QuantizerRole], ): diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py b/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py index ecbb667ecfa..6715b9cfff8 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py @@ -10,11 +10,13 @@ import torch +from transformer_engine.common.recipe import quantizer_policy from transformer_engine.pytorch.custom_recipes import quantization from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer +@quantizer_policy(key=("current_scaling_reference", 1)) def current_scaling_ref_quantizer_factory(role): """Factory function for current scaling reference quantizer. diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py index a2b21465af0..568b8c190a3 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py @@ -9,11 +9,13 @@ import torch +from transformer_engine.common.recipe import quantizer_policy from transformer_engine.pytorch.custom_recipes import quantization from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer +@quantizer_policy(key=("nvfp4_reference_rht_2d", 1)) def nvfp4_ref_rht_2d_quantizer_factory(role): """ Quantizer factory for NVFP4 recipe reference implementation (RHT and 2D quantization for weights). diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 07a3b804830..f1fc52a4246 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field from contextlib import contextmanager from collections import deque -from typing import Callable, List, Optional, Dict, Any, Tuple, Union +from typing import Callable, Hashable, List, Optional, Dict, Any, Tuple, Union import torch import transformer_engine_torch as tex @@ -396,6 +396,7 @@ class FP8GlobalState: fp8_enabled: bool = False fp8_calibration: bool = False fp8_recipe: Optional[Recipe] = None + quantizer_config: Optional[Hashable] = None fp8_distributed_group: Optional[dist_group_type] = None fp8_parameters: bool = False high_precision_init_val: bool = False @@ -610,6 +611,23 @@ def get_fp8_recipe(cls) -> Recipe: return cls.quantization_state.fp8_recipe return get_default_fp8_recipe() + @classmethod + def activate_recipe(cls, recipe: Recipe) -> Hashable: + """Make a recipe and its semantic quantizer configuration active together.""" + quantizer_config = recipe.quantizer_config() + qstate = cls.quantization_state + qstate.fp8_recipe = recipe + qstate.quantizer_config = quantizer_config + return quantizer_config + + @classmethod + def get_quantizer_config(cls) -> Hashable: + """Return the active recipe's cached semantic quantizer configuration.""" + quantizer_config = cls.quantization_state.quantizer_config + if quantizer_config is None: + return cls.activate_recipe(cls.get_fp8_recipe()) + return quantizer_config + @classmethod def get_fp8_group(cls) -> Union[dist_group_type, None]: """Return the fp8 group for scale/amax comm""" @@ -623,6 +641,7 @@ def get_autocast_state(cls) -> tuple: qstate.fp8_enabled, qstate.fp8_calibration, qstate.fp8_recipe, + qstate.quantizer_config, qstate.fp8_distributed_group, qstate.is_first_fp8_module, qstate.fp8_graph_capturing, @@ -636,6 +655,7 @@ def set_autocast_state(cls, state: tuple) -> None: qstate.fp8_enabled, qstate.fp8_calibration, qstate.fp8_recipe, + qstate.quantizer_config, qstate.fp8_distributed_group, qstate.is_first_fp8_module, qstate.fp8_graph_capturing, @@ -747,7 +767,7 @@ def autocast_enter( qstate.fp8_enabled = enabled qstate.fp8_calibration = calibrating - qstate.fp8_recipe = fp8_recipe + cls.activate_recipe(fp8_recipe) qstate.fp8_distributed_group = fp8_group qstate.fp8_graph_capturing = _graph @@ -926,15 +946,17 @@ def quantized_model_init( qstate = FP8GlobalStateManager.quantization_state _fp8_parameters = qstate.fp8_parameters _fp8_recipe = qstate.fp8_recipe + _quantizer_config = qstate.quantizer_config _high_precision_init_val = qstate.high_precision_init_val qstate.fp8_parameters = enabled - qstate.fp8_recipe = get_default_fp8_recipe() if recipe is None else recipe + FP8GlobalStateManager.activate_recipe(get_default_fp8_recipe() if recipe is None else recipe) qstate.high_precision_init_val = preserve_high_precision_init_val try: yield finally: qstate.fp8_parameters = _fp8_parameters qstate.fp8_recipe = _fp8_recipe + qstate.quantizer_config = _quantizer_config qstate.high_precision_init_val = _high_precision_init_val From 24907078c0a507bc4bdd72152de5d8095fc69b25 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 3 Aug 2026 15:43:07 +0000 Subject: [PATCH 65/97] Minor changes Signed-off-by: Evgeny --- tests/pytorch/debug/test_sanity.py | 50 +++++++++ tests/pytorch/hybrid_quantization_utils.py | 19 ++++ tests/pytorch/test_custom_recipe.py | 14 +-- tests/pytorch/test_hybrid_quantization.py | 48 ++++++--- .../debug/pytorch/debug_quantization.py | 33 ++++-- .../dot_product_attention.py | 84 +++++++++++++++ .../pytorch/cpp_extensions/gemm.py | 10 ++ .../custom_recipes/quantizer_factory_zoo.py | 100 +++--------------- .../pytorch/module/grouped_linear.py | 4 +- .../pytorch/module/layernorm_mlp.py | 6 +- 10 files changed, 244 insertions(+), 124 deletions(-) diff --git a/tests/pytorch/debug/test_sanity.py b/tests/pytorch/debug/test_sanity.py index 2bc4b355909..e1f255b6246 100644 --- a/tests/pytorch/debug/test_sanity.py +++ b/tests/pytorch/debug/test_sanity.py @@ -7,6 +7,8 @@ import nvdlfw_inspect.api as debug_api import transformer_engine.pytorch as te +from transformer_engine.debug.pytorch.debug_quantization import DebugQuantizedTensor +from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer from test_numerics import create_config_file @@ -58,6 +60,54 @@ """, } + +def test_debug_quantized_tensor_routes_usage_to_distinct_representations(): + """A usage request should only reach the child for that GEMM direction.""" + rowwise = IdentityQuantizer()(torch.ones(2, 2)) + columnwise = IdentityQuantizer()(torch.ones(2, 2)) + rowwise_calls = [] + columnwise_calls = [] + + def record_rowwise(rowwise_usage=None, columnwise_usage=None): + rowwise_calls.append((rowwise_usage, columnwise_usage)) + + def record_columnwise(rowwise_usage=None, columnwise_usage=None): + columnwise_calls.append((rowwise_usage, columnwise_usage)) + + rowwise.update_usage = record_rowwise + columnwise.update_usage = record_columnwise + tensor = DebugQuantizedTensor(rowwise, columnwise, quantizer=None) + + tensor.update_usage(columnwise_usage=True) + assert rowwise_calls == [] + assert columnwise_calls == [(None, True)] + + tensor.update_usage(rowwise_usage=True) + assert rowwise_calls == [(True, None)] + assert columnwise_calls == [(None, True)] + + tensor.update_usage(rowwise_usage=False) + assert tensor.rowwise_gemm_tensor is None + with pytest.raises(RuntimeError, match="Cannot recreate rowwise tensor"): + tensor.update_usage(rowwise_usage=True, columnwise_usage=False) + assert tensor.columnwise_gemm_tensor is columnwise + + +def test_debug_quantized_tensor_updates_shared_representation_once(): + """A shared child receives one combined usage update.""" + shared = IdentityQuantizer()(torch.ones(2, 2)) + calls = [] + + def record_usage(rowwise_usage=None, columnwise_usage=None): + calls.append((rowwise_usage, columnwise_usage)) + + shared.update_usage = record_usage + tensor = DebugQuantizedTensor(shared, shared, quantizer=None) + + tensor.update_usage(rowwise_usage=True, columnwise_usage=False) + assert calls == [(True, False)] + + # Configs that require FP8 to be enabled fp8_required_configs = {"log_fp8"} diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py index 9444cc4bcce..8ef4a7f421b 100644 --- a/tests/pytorch/hybrid_quantization_utils.py +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -37,6 +37,25 @@ def _make_mxfp8(*, fp8_dtype=te.DType.kFloat8E4M3): return te.MXFP8Quantizer(fp8_dtype=fp8_dtype) +def nvfp4_linear_mxfp8_dpa_test_factory(role): + """Test-local reference for the experimental NVFP4 Linear + MXFP8 DPA example.""" + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + + if is_dpa or is_dpa_boundary: + is_bwd_role = (is_dpa and role.tensor_type in ("do", "dp", "dqkv")) or ( + is_dpa_boundary and "dpa_grad_input" in role.name + ) + fp8_dtype = te.DType.kFloat8E5M2 if is_bwd_role else te.DType.kFloat8E4M3 + return te.MXFP8Quantizer(fp8_dtype=fp8_dtype) + + return nvfp4_factory(role) + + def fp8_e4m3_factory(): """Construct the default E4M3 current-scaling test quantizer.""" return te.Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 9caa33b18be..a1a6ef7b373 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -7,6 +7,9 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex +from hybrid_quantization_utils import ( + nvfp4_linear_mxfp8_dpa_test_factory as _nvfp4_linear_mxfp8_dpa_factory, +) from transformer_engine.common import recipe from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx from transformer_engine.pytorch import ( @@ -1233,15 +1236,15 @@ def test_custom_recipe_dpa_fp8(): def test_custom_recipe_dpa_mxfp8(): """DotProductAttention forward+backward with CustomRecipe and MXFP8 attention. - Uses the nvfp4_linear_mxfp8_dpa_factory which dispatches: + Uses a test-local reference factory which dispatches: * DPA roles (QKV/O/S/dO/dP/dQKV) -> MXFP8Quantizer (S/dP later nulled out by ``get_attention_quantizers`` since the MXFP8 fused-attention kernel handles those slots internally) * DPA boundary hints -> MXFP8Quantizer * Linear slots -> NVFP4Quantizer - Mirrors the documented "NVFP4 linear + MXFP8 attention" combo from - ``dot_product_attention.py``'s recipe-combination table. + Mirrors the documented experimental "NVFP4 Linear + MXFP8 attention" + configuration. """ available, reason = te.is_fp8_available(return_reason=True) if not torch.cuda.is_available() or not available: @@ -1260,9 +1263,6 @@ def test_custom_recipe_dpa_mxfp8(): from transformer_engine.pytorch.quantization import CustomRecipeState from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( - nvfp4_linear_mxfp8_dpa_factory, - ) torch.manual_seed(42) @@ -1282,7 +1282,7 @@ def test_custom_recipe_dpa_mxfp8(): out_proj = Linear(H, H, params_dtype=torch.bfloat16, bias=False, name="proj").cuda() custom_recipe = recipe.CustomRecipe( - qfactory=nvfp4_linear_mxfp8_dpa_factory, + qfactory=_nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True, ) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 7f9bc16b96a..74ec0a05ecf 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -27,6 +27,7 @@ make_nvfp4_quantizer as _make_nvfp4_quantizer, make_role_aware_quantizer as _make_role_aware_quantizer, mxfp8_e4m3_factory as _mxfp8_factory, + nvfp4_linear_mxfp8_dpa_test_factory as _nvfp4_linear_mxfp8_dpa_factory, snapshot_storage_tensor_metadata as _snapshot_storage_tensor_metadata, ) from transformer_engine.common import recipe @@ -35,6 +36,7 @@ ) from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( mxfp8_fwd_nvfp4_bwd_factory, + nvfp4_linear_fp8_dpa_factory, ) from transformer_engine.pytorch import ( autocast, @@ -2313,17 +2315,17 @@ def _run_model(self, model, inp, grad, fp8_recipe, seed): ) @pytest.mark.parametrize( - "case_name,native_dpa_recipe,qfactory_name", + "case_name,native_dpa_recipe,qfactory", [ ( "fp8_dpa", "Float8CurrentScaling", - "nvfp4_linear_fp8_dpa_factory", + nvfp4_linear_fp8_dpa_factory, ), ( "mxfp8_dpa", "MXFP8BlockScaling", - "nvfp4_linear_mxfp8_dpa_factory", + _nvfp4_linear_mxfp8_dpa_factory, ), ], ) @@ -2332,12 +2334,11 @@ def test_linear_dpa_linear_matches_native_env_recipe_bitwise( monkeypatch, case_name, native_dpa_recipe, - qfactory_name, + qfactory, ): if case_name == "mxfp8_dpa" and not mxfp8_available: pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") - from transformer_engine.pytorch.custom_recipes import quantizer_factory_zoo from transformer_engine.pytorch.utils import get_device_compute_capability cc = get_device_compute_capability() @@ -2370,7 +2371,6 @@ def test_linear_dpa_linear_matches_native_env_recipe_bitwise( grad = torch.randn_like(base_inp) native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) - qfactory = getattr(quantizer_factory_zoo, qfactory_name) qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) native_out, native_dx, native_grads, native_local_recipes = self._run_model( @@ -2427,18 +2427,18 @@ def _run_mha_model(self, model, inp, grad, fp8_recipe, seed): ) @pytest.mark.parametrize( - "case_name,native_dpa_recipe,qfactory_name,expected_flags", + "case_name,native_dpa_recipe,qfactory,expected_flags", [ ( "fp8_dpa", "Float8CurrentScaling", - "nvfp4_linear_fp8_dpa_factory", + nvfp4_linear_fp8_dpa_factory, (False, False, False), ), ( "mxfp8_dpa", "MXFP8BlockScaling", - "nvfp4_linear_mxfp8_dpa_factory", + _nvfp4_linear_mxfp8_dpa_factory, (False, False, False), ), ], @@ -2448,14 +2448,13 @@ def test_multihead_attention_matches_native_env_recipe_bitwise( monkeypatch, case_name, native_dpa_recipe, - qfactory_name, + qfactory, expected_flags, ): if case_name == "mxfp8_dpa" and not mxfp8_available: pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") from transformer_engine.pytorch.attention import multi_head_attention as mha_module - from transformer_engine.pytorch.custom_recipes import quantizer_factory_zoo from transformer_engine.pytorch.utils import get_device_compute_capability cc = get_device_compute_capability() @@ -2513,7 +2512,6 @@ def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) grad = torch.randn_like(base_inp) native_recipe = recipe.NVFP4BlockScaling(fp8_dpa=True) - qfactory = getattr(quantizer_factory_zoo, qfactory_name) qfactory_recipe = recipe.CustomRecipe(qfactory=qfactory, fp8_dpa=True) native_out, native_dx, native_grads = self._run_mha_model( @@ -2611,9 +2609,6 @@ def test_mxfp8_qfactory_uses_plain_bf16_mha_boundaries(self, monkeypatch): from transformer_engine.pytorch.attention.dot_product_attention import ( backends as dpa_backends, ) - from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( - nvfp4_linear_mxfp8_dpa_factory, - ) from transformer_engine.pytorch.utils import get_device_compute_capability cc = get_device_compute_capability() @@ -2709,7 +2704,7 @@ def _record_projection_input(_module, inputs): ) grad = torch.randn_like(inp) qfactory_recipe = recipe.CustomRecipe( - qfactory=nvfp4_linear_mxfp8_dpa_factory, + qfactory=_nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True, ) try: @@ -3038,6 +3033,27 @@ def test_hybrid_rowwise_returns_rowwise_sub_storage(self, hybrid_tensor): def test_hybrid_columnwise_returns_columnwise_sub_storage(self, hybrid_tensor): assert _unwrap_tensor(hybrid_tensor, "columnwise") is hybrid_tensor.columnwise_sub_storage + @pytest.mark.parametrize( + "available_usage,missing_usage", + [("rowwise", "columnwise"), ("columnwise", "rowwise")], + ) + def test_missing_hybrid_representation_raises(self, available_usage, missing_usage): + inp = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + quantizer = _make_hybrid_quantizer_fp8_row_fp4_col() + quantizer.set_usage( + rowwise=available_usage == "rowwise", + columnwise=available_usage == "columnwise", + ) + hybrid_tensor = quantizer.quantize(inp) + + assert getattr(hybrid_tensor, f"{available_usage}_sub_storage") is not None + assert getattr(hybrid_tensor, f"{missing_usage}_sub_storage") is None + with pytest.raises( + RuntimeError, + match=rf"GEMM requested the {missing_usage} representation, but it is unavailable", + ): + _unwrap_tensor(hybrid_tensor, missing_usage) + def test_rowwise_sub_storage_type(self, hybrid_tensor): assert isinstance( _unwrap_tensor(hybrid_tensor, "rowwise"), diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index ed5fdd4660d..107dd7a3732 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -676,20 +676,31 @@ def size(self, *args): def update_usage(self, rowwise_usage: bool = None, columnwise_usage: bool = None): """Update usage of the tensor.""" - if self.rowwise_gemm_tensor is not self.columnwise_gemm_tensor: - # If the same object is used both for rowwise and columnwise gemms, - # there is no benefit in erasing the usage of one of them. - # And there are scenarios when not deleting the usage of one of them is needed. - # For example when we want to recreate columnwise from rowwise. + if self.rowwise_gemm_tensor is self.columnwise_gemm_tensor: + if isinstance(self.rowwise_gemm_tensor, QuantizedTensor): + self.rowwise_gemm_tensor.update_usage(rowwise_usage, columnwise_usage) + else: + # Each tensor owns only the representation for its GEMM direction. + # Validate requests before dropping either representation so a failed + # mixed drop/enable request does not partially mutate this wrapper. + if rowwise_usage and self.rowwise_gemm_tensor is None: + raise RuntimeError( + "Cannot recreate rowwise tensor from columnwise tensor in debug mode." + ) + if columnwise_usage and self.columnwise_gemm_tensor is None: + raise RuntimeError( + "Cannot recreate columnwise tensor from rowwise tensor in debug mode." + ) + if rowwise_usage is False: self.rowwise_gemm_tensor = None + elif rowwise_usage and isinstance(self.rowwise_gemm_tensor, QuantizedTensor): + self.rowwise_gemm_tensor.update_usage(rowwise_usage=True) + if columnwise_usage is False: self.columnwise_gemm_tensor = None - - if isinstance(self.rowwise_gemm_tensor, QuantizedTensor): - self.rowwise_gemm_tensor.update_usage(rowwise_usage, columnwise_usage) - if isinstance(self.columnwise_gemm_tensor, QuantizedTensor): - self.columnwise_gemm_tensor.update_usage(rowwise_usage, columnwise_usage) + elif columnwise_usage and isinstance(self.columnwise_gemm_tensor, QuantizedTensor): + self.columnwise_gemm_tensor.update_usage(columnwise_usage=True) if rowwise_usage and self.rowwise_gemm_tensor is None: raise RuntimeError( @@ -698,7 +709,7 @@ def update_usage(self, rowwise_usage: bool = None, columnwise_usage: bool = None if columnwise_usage and self.columnwise_gemm_tensor is None: raise RuntimeError( - "Cannot recreate columnwise tensor from rowwise tensor is debug mode." + "Cannot recreate columnwise tensor from rowwise tensor in debug mode." ) @property diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 9ae0df880a8..11d6500ab7e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -393,6 +393,90 @@ class DotProductAttention(TransformerEngineBaseModule): As the FP8 attention support expands from one backend to multiple backends, the location of that key has also shifted (see `FP8 checkpoint compatibility `_). + .. rubric:: Fine-grained Linear and attention recipes + + .. warning:: + + Fine-grained attention configuration through ``CustomRecipe`` and a quantizer factory is + experimental and subject to change. + + A quantizer factory can select different recipes for Linear and DotProductAttention tensors + using ``QuantizerRole``. DotProductAttention itself supports only its fixed recipe families: + FP8 delayed scaling, FP8 current scaling, and MXFP8 block scaling. In particular, a factory may + return NVFP4 quantizers for Linear roles, but it must not return NVFP4 quantizers for DPA roles. + + .. list-table:: Example Linear and attention combinations + :header-rows: 1 + + * - Linear + - Attention + - Configuration + - Status + * - NVFP4 + - FP8 current scaling for QKV/O and delayed scaling for S/dP + - ``nvfp4_linear_fp8_dpa_factory`` + - Validated factory provided by Transformer Engine + * - NVFP4 + - MXFP8 + - User-defined factory shown below + - Experimental example; not broadly validated + + The validated NVFP4 Linear + FP8 attention combination is available from the quantizer factory + zoo:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_fp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + + The following factory demonstrates the experimental NVFP4 Linear + MXFP8 attention + combination. With ``CustomRecipe``, the per-role selection is expressed directly in the + factory, so ``NVTE_DPA_FP8_RECIPE`` is not needed. DPA also issues hint-only roles for its + output boundaries; these must resolve to a DPA-supported quantizer even when the boundary + tensor remains in BF16. For MXFP8 attention, the fused kernel handles the S/dP slots + internally, so their factory-provided quantizers are not consumed:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.constants import DType + from transformer_engine.pytorch.custom_recipes.quantizer_factories import nvfp4_factory + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + def nvfp4_linear_mxfp8_dpa_factory(role): + # NVFP4 for Linear roles and MXFP8 for supported DPA roles. + is_dpa = role is not None and role.module_type == "dpa" + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + + if is_dpa or is_dpa_boundary: + is_bwd_role = ( + is_dpa and role.tensor_type in ("do", "dp", "dqkv") + ) or ( + is_dpa_boundary and "dpa_grad_input" in role.name + ) + fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 + return MXFP8Quantizer(fp8_dtype=fp8_dtype) + + return nvfp4_factory(role) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_mxfp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + Parameters ---------- diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 2341c56bd9e..f3d97b72695 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -128,6 +128,16 @@ def _unwrap_tensor( if usage not in ("rowwise", "columnwise"): raise ValueError(f"Unsupported GEMM tensor usage ({usage})") + # Hybrid and debug wrappers may omit a representation that is not needed + # by their configured GEMMs. Fail here if a caller requests that missing + # direction instead of passing ``None`` deeper into GEMM preparation. + if tensor is None: + raise RuntimeError( + f"GEMM requested the {usage} representation, but it is unavailable. " + f"Ensure {usage}_usage is enabled and the representation has not " + "been dropped by update_usage()." + ) + # Plain PyTorch tensor if not isinstance(tensor, QuantizedTensorStorage): return tensor diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 1b123a474a6..0d5fc0c4419 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -49,7 +49,7 @@ from transformer_engine.pytorch.quantization import autocast from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( mxfp8_fwd_nvfp4_bwd_factory, - nvfp4_linear_mxfp8_dpa_factory, + nvfp4_linear_fp8_dpa_factory, ) # Linear-only recipe (no attention quantization): the qfactory is the only knob. @@ -59,7 +59,7 @@ # Recipe that also quantizes DotProductAttention: set ``fp8_dpa=True`` so the # attention GEMMs request quantizers from the factory (DPA roles) too. - recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True) + recipe = CustomRecipe(qfactory=nvfp4_linear_fp8_dpa_factory, fp8_dpa=True) with autocast(recipe=recipe): output = model(input) @@ -132,24 +132,26 @@ def mxfp8_fwd_nvfp4_bwd_factory( representations consumed by wgrad, and gradients use stochastic rounding. The dgrad weight uses plain 1D NVFP4 without RHT or stochastic rounding. - The backward input and weight representations are quantized to NVFP4 - directly from the original high-precision tensors rather than from the - dequantized MXFP8 forward representations. In ``HybridQuantizer`` terms, - this source choice is expressed with ``columnwise_source="original"``, - which is the default. To derive them from the forward representations - instead, use ``columnwise_source="rowwise_dequantized"``. + The backward weight representation consumed by dgrad is quantized to + NVFP4 from the dequantized MXFP8 forward weight. In ``HybridQuantizer`` + terms, the weight uses ``columnwise_source="rowwise_dequantized"``. The + backward input representation consumed by wgrad remains quantized directly + from the original high-precision input with ``columnwise_source="original"``. """ from transformer_engine.pytorch.tensor.hybrid_tensor import HybridQuantizer is_linear = role is not None and role.module_type in ("linear", "grouped_linear") - if is_linear and role.tensor_type in ("input", "weight"): - if role.tensor_type == "weight": - backward_quantizer = _plain_nvfp4_quantizer() - else: - backward_quantizer = nvfp4_factory(role) + if is_linear and role.tensor_type == "input": return HybridQuantizer( rowwise_quantizer=mxfp8_factory(role), - columnwise_quantizer=backward_quantizer, + columnwise_quantizer=nvfp4_factory(role), + columnwise_source="original", + ) + if is_linear and role.tensor_type == "weight": + return HybridQuantizer( + rowwise_quantizer=mxfp8_factory(role), + columnwise_quantizer=_plain_nvfp4_quantizer(), + columnwise_source="rowwise_dequantized", ) if is_linear and role.tensor_type == "grad_output": return nvfp4_factory(role) @@ -427,73 +429,3 @@ def nvfp4_linear_fp8_dpa_factory( return Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device="cuda") return nvfp4_factory(role) - - -def nvfp4_linear_mxfp8_dpa_factory( - role: Optional[QuantizerRole], -): - """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``DotProductAttention``. - - Mirrors the documented "NVFP4 linear + MXFP8 attention" combo from - :mod:`transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention` - (see the recipe-combination table at the top of that module). With - ``CustomRecipe`` the per-tensor decision is made directly here, so the - ``NVTE_DPA_FP8_RECIPE="MXFP8BlockScaling"`` env override that the - built-in recipes would otherwise need is unnecessary. - - DPA-owned tensor types (``role.module_type == "dpa"``): - - =========== ============================================================ - tensor_type Description - =========== ============================================================ - ``"qkv"`` Query, Key, Value inputs to the first attention GEMM - ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM - ``"do"`` Gradient of the attention output (dO), backward input - ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward - =========== ============================================================ - - Dispatch logic: - * ``role.module_type == "dpa"`` -> MXFP8 - (``Format.HYBRID``: E4M3 fwd, E5M2 bwd) - The MXFP8 fused-attention kernel handles the S/dP slots - internally, so any quantizer returned for those roles is later - nulled out by ``get_attention_quantizers``. Returning MXFP8 is - the simplest valid choice. - * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in - ``role.name``) -> MXFP8 placeholder. The fused attention kernel - requires FP8-compatible quantizers in all DPA slots. - * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) - -> NVFP4 (E2M1), configured per tensor role. - - Usage:: - - from transformer_engine.common.recipe import CustomRecipe - from transformer_engine.pytorch.quantization import autocast - from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( - nvfp4_linear_mxfp8_dpa_factory, - ) - - recipe = CustomRecipe( - qfactory=nvfp4_linear_mxfp8_dpa_factory, - fp8_dpa=True, - ) - with autocast(recipe=recipe): - output = model(input) - """ - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer - - is_dpa = role is not None and role.module_type == "dpa" - is_dpa_boundary = ( - role is not None - and not role.module_type - and ("dpa_output" in role.name or "dpa_grad_input" in role.name) - ) - - if is_dpa or is_dpa_boundary: - is_bwd_role = (is_dpa and role.tensor_type in ("do", "dp", "dqkv")) or ( - is_dpa_boundary and "dpa_grad_input" in role.name - ) - fp8_dtype = DType.kFloat8E5M2 if is_bwd_role else DType.kFloat8E4M3 - return MXFP8Quantizer(fp8_dtype=fp8_dtype) - - return nvfp4_factory(role) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 1049565ce90..12497d0cb0a 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -271,7 +271,9 @@ def _split_quantize_non_hybrid( tensor = cast_if_needed(tensor, activation_dtype) if ( allow_identity_views - and type(reference) is IdentityQuantizer + # Only the base IdentityQuantizer can bypass quantization; subclasses + # may override its behavior and must go through their normal call path. + and type(reference) is IdentityQuantizer # pylint: disable=unidiomatic-typecheck and (reference.dtype is None or reference.dtype == activation_dtype) ): return torch.split(tensor, m_splits) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 6a2ad31461c..3ee0cda50c5 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -66,11 +66,7 @@ from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing -from ..tensor.float8_tensor import ( - Float8CurrentScalingQuantizer, - Float8Quantizer, - Float8Tensor, -) +from ..tensor.float8_tensor import Float8Tensor from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer From 0ac0c0519e96a816427d27442a59ce71e6b91eca Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 4 Aug 2026 14:05:03 +0000 Subject: [PATCH 66/97] Renaming Signed-off-by: Evgeny --- tests/pytorch/hybrid_quantization_utils.py | 4 +- tests/pytorch/test_custom_recipe.py | 42 ++++---- tests/pytorch/test_hybrid_quantization.py | 102 +++++++++--------- tests/pytorch/test_identity_quantizer.py | 30 +++--- tests/pytorch/test_recipe.py | 42 ++++---- transformer_engine/common/recipe/__init__.py | 50 ++++----- transformer_engine/pytorch/__init__.py | 2 +- .../custom_recipes/quantizer_factories.py | 14 +-- .../custom_recipes/quantizer_factory_zoo.py | 16 +-- .../reference_current_scaling.py | 4 +- .../pytorch/custom_recipes/reference_nvfp4.py | 4 +- 11 files changed, 155 insertions(+), 155 deletions(-) diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py index 2ff93f8ab7d..7285e092dc3 100644 --- a/tests/pytorch/hybrid_quantization_utils.py +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -211,7 +211,7 @@ def make_role_aware_quantizer(factory, role): return quantizer -def hybrid_custom_recipe(row_factory, col_factory, grad_factory=None, *, policy_key): +def hybrid_custom_recipe(row_factory, col_factory, grad_factory=None, *, qfactory_key): """Build a CustomRecipe with hybrid forward and configurable grad quantizers.""" if grad_factory is None: grad_factory = col_factory @@ -227,7 +227,7 @@ def qfactory(role): return make_role_aware_quantizer(grad_factory, role) return make_role_aware_quantizer(row_factory, role) - return recipe.CustomRecipe(qfactory=qfactory, policy_key=policy_key) + return recipe.CustomRecipe(qfactory=qfactory, qfactory_key=qfactory_key) def hybrid_fp8_current_qfactory(role): diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index afb6975493c..de62e8eec38 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -56,15 +56,15 @@ nvfp4_ref_rht_2d_factory, ], ) -def test_first_party_qfactories_have_canonical_policy_keys(qfactory): - """First-party factories can be used without repeating policy_key in CustomRecipe.""" - policy_key = getattr(qfactory, "policy_key") - hash(policy_key) - assert isinstance(policy_key, tuple) - assert len(policy_key) == 2 - assert isinstance(policy_key[0], str) and policy_key[0] - assert isinstance(policy_key[1], int) and policy_key[1] > 0 - assert recipe.CustomRecipe(qfactory=qfactory).policy_key == policy_key +def test_first_party_qfactories_have_canonical_qfactory_keys(qfactory): + """First-party factories can be used without repeating qfactory_key in CustomRecipe.""" + qfactory_key = getattr(qfactory, "qfactory_key") + hash(qfactory_key) + assert isinstance(qfactory_key, tuple) + assert len(qfactory_key) == 2 + assert isinstance(qfactory_key[0], str) and qfactory_key[0] + assert isinstance(qfactory_key[1], int) and qfactory_key[1] > 0 + assert recipe.CustomRecipe(qfactory=qfactory).qfactory_key == qfactory_key @pytest.mark.parametrize("module_type", ["Linear", "LayerNormLinear", "OpsLinear"]) @@ -134,7 +134,7 @@ def test_custom_recipe_sanity(module_type): inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) # Single factory: map roles to quantizers - @recipe.quantizer_policy(key=("test_current_scaling_sanity", 1)) + @recipe.quantizer_factory(key=("test_current_scaling_sanity", 1)) def quantizer_factory(role): if role is None: return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -172,7 +172,7 @@ def test_custom_recipe_grouped_linear_sanity(): model = GroupedLinear(num_gemms, in_features, out_features, params_dtype=torch.bfloat16).cuda() inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) - @recipe.quantizer_policy(key=("test_current_scaling_grouped_linear", 1)) + @recipe.quantizer_factory(key=("test_current_scaling_grouped_linear", 1)) def quantizer_factory(role): if role is None: return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -235,7 +235,7 @@ def test_custom_recipe_matches_current_scaling(): loss_ref.backward() # Custom: single factory returning quantizers per role to match Float8CurrentScaling - @recipe.quantizer_policy(key=("test_current_scaling_parity", 1)) + @recipe.quantizer_factory(key=("test_current_scaling_parity", 1)) def quantizer_factory(role): if role is None: return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -293,7 +293,7 @@ def test_custom_recipe_ops_linear_2_1_layout(): op = te_ops.Linear(in_features, out_features, device="cuda", dtype=torch.bfloat16) inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) - @recipe.quantizer_policy(key=("test_current_scaling_ops_linear", 1)) + @recipe.quantizer_factory(key=("test_current_scaling_ops_linear", 1)) def quantizer_factory(role): if role is None: return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -337,7 +337,7 @@ def test_custom_recipe_factory_invocation_counts_and_cycling(): None: 0, } - @recipe.quantizer_policy(key=("test_factory_invocation_counts", 1)) + @recipe.quantizer_factory(key=("test_factory_invocation_counts", 1)) def quantizer_factory(role): if role is None: counts[None] += 1 @@ -639,7 +639,7 @@ def test_custom_recipe_quantization_targets(): # ------------------------------------------------------------------ recorded_roles = [] - @recipe.quantizer_policy(key=("test_per_module_targeting", 1)) + @recipe.quantizer_factory(key=("test_per_module_targeting", 1)) def targeting_factory(role): recorded_roles.append(role) @@ -797,7 +797,7 @@ def test_grouped_linear_module_type_dispatch(): recorded_roles = [] - @recipe.quantizer_policy(key=("test_grouped_linear_role_recording", 1)) + @recipe.quantizer_factory(key=("test_grouped_linear_role_recording", 1)) def recording_factory(role): recorded_roles.append(role) return Float8CurrentScalingQuantizer(te.DType.kFloat8E4M3, device="cuda") @@ -838,7 +838,7 @@ def test_delayed_scaling_request_wiring(): from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.common.recipe import Format - @recipe.quantizer_policy(key=("test_delayed_scaling_request", 1)) + @recipe.quantizer_factory(key=("test_delayed_scaling_request", 1)) def ds_factory(role): return DelayedScalingRequest(fp8_format=Format.HYBRID, amax_history_len=16) @@ -899,7 +899,7 @@ def test_custom_recipe_mixed_ds_and_stateless(): from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer from transformer_engine.common.recipe import Format - @recipe.quantizer_policy(key=("test_mixed_delayed_current_scaling", 1)) + @recipe.quantizer_factory(key=("test_mixed_delayed_current_scaling", 1)) def mixed_factory(role): # Only weight gets delayed scaling, rest get current scaling if role is not None and role.tensor_type == "weight": @@ -944,7 +944,7 @@ def test_custom_recipe_ds_multi_step(): from transformer_engine.pytorch.quantization import DelayedScalingRequest from transformer_engine.common.recipe import Format - @recipe.quantizer_policy(key=("test_delayed_scaling_multi_step", 1)) + @recipe.quantizer_factory(key=("test_delayed_scaling_multi_step", 1)) def ds_factory(role): return DelayedScalingRequest(fp8_format=Format.HYBRID) @@ -1068,7 +1068,7 @@ def test_role_change_preserves_custom_delayed_scaling_state(): ) from transformer_engine.common.recipe import Format - @recipe.quantizer_policy(key=("test_delayed_scaling_role_rebuild", 1)) + @recipe.quantizer_factory(key=("test_delayed_scaling_role_rebuild", 1)) def ds_factory(role): return DelayedScalingRequest(fp8_format=Format.HYBRID, amax_history_len=8) @@ -1938,7 +1938,7 @@ def test_custom_recipe_quantization_alignment_contract(): """The alignment contract is safe, configurable, and does not invoke qfactory.""" calls = [] - @recipe.quantizer_policy(key=("test_quantization_alignment", 1)) + @recipe.quantizer_factory(key=("test_quantization_alignment", 1)) def qfactory(role): calls.append(role) return current_scaling_factory(role) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index a27764ae9fd..3027d2e8c9d 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -821,7 +821,7 @@ def factory(role): return recipe.CustomRecipe( qfactory=factory, - policy_key=( + qfactory_key=( "test_counting_identity_hybrid", 1, columnwise_source, @@ -1707,7 +1707,7 @@ def test_te_linear_inference_workspace_rowwise_only(self): hybrid_recipe = _hybrid_custom_recipe( row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), - policy_key=("test_inference_workspace_hybrid_fp8", 1), + qfactory_key=("test_inference_workspace_hybrid_fp8", 1), ) torch.manual_seed(2026) model = Linear(128, 256, bias=False, params_dtype=torch.bfloat16).cuda() @@ -1855,7 +1855,7 @@ def hybrid_fp8_factory(role): hybrid_recipe = recipe.CustomRecipe( qfactory=hybrid_fp8_factory, - policy_key=("test_hybrid_fp8_parity", 1), + qfactory_key=("test_hybrid_fp8_parity", 1), ) with autocast(enabled=True, recipe=hybrid_recipe): out_hybrid = model_hybrid(inp_hybrid) @@ -1924,7 +1924,7 @@ def hybrid_mxfp8_factory(role): hybrid_recipe = recipe.CustomRecipe( qfactory=hybrid_mxfp8_factory, - policy_key=("test_hybrid_mxfp8_parity", 1), + qfactory_key=("test_hybrid_mxfp8_parity", 1), ) with autocast(enabled=True, recipe=hybrid_recipe): out_hybrid = model_hybrid(inp_hybrid) @@ -2019,7 +2019,7 @@ def test_inference_runs_once_per_quantizer_state_and_clears_stale_labels(self, m custom_recipe = recipe.CustomRecipe( qfactory=lambda _role: IdentityQuantizer(), - policy_key=("test_dpa_identity", 1), + qfactory_key=("test_dpa_identity", 1), ) monkeypatch.setattr( FP8GlobalStateManager, @@ -2109,7 +2109,7 @@ def counting_qfactory(role): qfactory=counting_qfactory, fp8_dpa=True, fp8_mha=True, - policy_key=("test_dpa_capability_counting", factory_name, 1), + qfactory_key=("test_dpa_capability_counting", factory_name, 1), ) dpa = te.DotProductAttention( num_attention_heads=2, @@ -2138,7 +2138,7 @@ def rebuilt_qfactory(role): qfactory=rebuilt_qfactory, fp8_dpa=True, fp8_mha=True, - policy_key=("test_dpa_capability_rebuilt", factory_name, 1), + qfactory_key=("test_dpa_capability_rebuilt", factory_name, 1), ) with autocast(enabled=True, recipe=rebuilt_recipe): rebuilt = dpa.get_qkv_quantization_capabilities() @@ -2194,7 +2194,7 @@ def counting_qfactory(role): qfactory=counting_qfactory, fp8_dpa=True, fp8_mha=fp8_mha, - policy_key=("test_mha_factory_invocation_count", fp8_mha, 1), + qfactory_key=("test_mha_factory_invocation_count", fp8_mha, 1), ) model = te.MultiheadAttention( hidden_size=128, @@ -2802,7 +2802,7 @@ def hybrid_block_fp8_factory(role): hybrid_recipe = recipe.CustomRecipe( qfactory=hybrid_block_fp8_factory, - policy_key=("test_hybrid_block_fp8_parity", 1), + qfactory_key=("test_hybrid_block_fp8_parity", 1), ) with autocast(enabled=True, recipe=hybrid_recipe): out_hybrid = model_hybrid(inp_hybrid) @@ -2866,7 +2866,7 @@ def hybrid_nvfp4_factory(role): hybrid_recipe = recipe.CustomRecipe( qfactory=hybrid_nvfp4_factory, - policy_key=("test_hybrid_nvfp4_parity", 1), + qfactory_key=("test_hybrid_nvfp4_parity", 1), ) torch.manual_seed(1202) torch.cuda.manual_seed_all(1202) @@ -2919,7 +2919,7 @@ def hybrid_nvfp4_all_roles_factory(role): hybrid_recipe = recipe.CustomRecipe( qfactory=hybrid_nvfp4_all_roles_factory, - policy_key=("test_hybrid_nvfp4_all_roles", 1), + qfactory_key=("test_hybrid_nvfp4_all_roles", 1), ) torch.manual_seed(1203) torch.cuda.manual_seed_all(1203) @@ -2984,7 +2984,7 @@ def mixed_factory(role): mixed_recipe = recipe.CustomRecipe( qfactory=mixed_factory, - policy_key=("test_mixed_fp8_nvfp4", 1), + qfactory_key=("test_mixed_fp8_nvfp4", 1), ) with autocast(enabled=True, recipe=mixed_recipe): @@ -3037,7 +3037,7 @@ def mixed_factory(role): mixed_recipe = recipe.CustomRecipe( qfactory=mixed_factory, - policy_key=("test_mixed_fp8_nvfp4_numerics", 1), + qfactory_key=("test_mixed_fp8_nvfp4_numerics", 1), ) with torch.no_grad(): with autocast(enabled=True, recipe=mixed_recipe): @@ -3169,7 +3169,7 @@ def factory(role): ), ) - return recipe.quantizer_policy(key=("test_uniform_hybrid_fp8", 1))(factory) + return recipe.quantizer_factory(key=("test_uniform_hybrid_fp8", 1))(factory) def test_bias_grad_matches_vanilla_fp8(self): torch.manual_seed(456) @@ -3264,7 +3264,7 @@ def factory(role): enabled=True, recipe=recipe.CustomRecipe( qfactory=factory, - policy_key=("test_matching_columnwise_formats", 1), + qfactory_key=("test_matching_columnwise_formats", 1), ), ): out = model(inp) @@ -3302,7 +3302,7 @@ def factory(role): enabled=True, recipe=recipe.CustomRecipe( qfactory=factory, - policy_key=("test_mismatched_columnwise_formats", 1), + qfactory_key=("test_mismatched_columnwise_formats", 1), ), ): out = model(inp) @@ -3341,7 +3341,7 @@ def factory(role): mixed_recipe = recipe.CustomRecipe( qfactory=factory, - policy_key=("test_nvfp4_row_fp8_column_forward", 1), + qfactory_key=("test_nvfp4_row_fp8_column_forward", 1), ) with torch.no_grad(): with autocast(enabled=True, recipe=mixed_recipe): @@ -3381,7 +3381,7 @@ def factory(role): mixed_recipe = recipe.CustomRecipe( qfactory=factory, - policy_key=("test_nvfp4_row_fp8_column", 1), + qfactory_key=("test_nvfp4_row_fp8_column", 1), ) with autocast(enabled=True, recipe=mixed_recipe): out = model(inp) @@ -3441,7 +3441,7 @@ def mixed_factory(role): grad_output, recipe.CustomRecipe( qfactory=mixed_factory, - policy_key=("test_single_hybrid_role", hybrid_role, 1), + qfactory_key=("test_single_hybrid_role", hybrid_role, 1), ), seed=seed + 100, ) @@ -3616,7 +3616,7 @@ def factory(role): return _make_role_aware_quantizer(grad_factory, role) return _make_role_aware_quantizer(operand_factory, role) - return recipe.quantizer_policy( + return recipe.quantizer_factory( key=("test_plain_linear", operand_factory.__name__, grad_factory.__name__, 1) )(factory) @@ -3679,7 +3679,7 @@ def hybrid_factory(role): hybrid_recipe = recipe.CustomRecipe( qfactory=hybrid_factory, - policy_key=("test_cross_format", row_name, col_name, 1), + qfactory_key=("test_cross_format", row_name, col_name, 1), ) fprop_ref_recipe = recipe.CustomRecipe( qfactory=_plain_linear_qfactory(make_row_operand, make_row_grad) @@ -3954,7 +3954,7 @@ def hybrid_factory(role): grad_output, recipe.CustomRecipe( qfactory=hybrid_factory, - policy_key=( + qfactory_key=( "test_three_format", make_fprop.__name__, make_dgrad.__name__, @@ -4062,7 +4062,7 @@ def factory(role): device="cuda", ) - return recipe.quantizer_policy(key=("test_hybrid_fp8_all_modules", 1))(factory) + return recipe.quantizer_factory(key=("test_hybrid_fp8_all_modules", 1))(factory) @requires_fp8 @@ -4098,7 +4098,7 @@ def qfactory(_role): enabled=True, recipe=recipe.CustomRecipe( qfactory=qfactory, - policy_key=("test_fusible_norm_hybrid_output", 1), + qfactory_key=("test_fusible_norm_hybrid_output", 1), ), ): out = forward(inp) @@ -4616,7 +4616,7 @@ def qfactory(_role): columnwise_source=columnwise_source, ) - return recipe.quantizer_policy( + return recipe.quantizer_factory( key=("test_grouped_linear_columnwise_source", columnwise_source, 1) )(qfactory) @@ -4673,7 +4673,7 @@ def mixed_source_qfactory(role): mixed_recipe = recipe.CustomRecipe( qfactory=mixed_source_qfactory, - policy_key=("test_mixed_columnwise_sources", 1), + qfactory_key=("test_mixed_columnwise_sources", 1), ) # A failed generation is never marked validated. Base metadata can then # early-return on retry, so the O(1) guard must validate it again. @@ -4746,7 +4746,7 @@ def _hybrid_fp8_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_quantized_model_init_hybrid_fp8", 1), + qfactory_key=("test_quantized_model_init_hybrid_fp8", 1), ) def test_linear_weight_is_hybrid_quantized_tensor(self): @@ -4850,7 +4850,7 @@ def _hybrid_fp8_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_weight_workspace_hybrid_fp8", 1), + qfactory_key=("test_weight_workspace_hybrid_fp8", 1), ) @_XFAIL_HOPPER_COLUMNWISE_PER_TENSOR_FP8 @@ -4956,7 +4956,7 @@ def _hybrid_fp8_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_weight_quantizer_update_hybrid_fp8", 1), + qfactory_key=("test_weight_quantizer_update_hybrid_fp8", 1), ) def test_quantized_param_survives_multiple_forward_passes(self): @@ -5023,7 +5023,7 @@ def _hybrid_recipe_fp8_current(): row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), col_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), - policy_key=("test_hybrid_fp8_current", 1), + qfactory_key=("test_hybrid_fp8_current", 1), ) @@ -5056,7 +5056,7 @@ def _hybrid_recipe_fp8_delayed(): row_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), - policy_key=("test_hybrid_fp8_delayed", 1), + qfactory_key=("test_hybrid_fp8_delayed", 1), ) @@ -5073,7 +5073,7 @@ def _hybrid_recipe_fp8_delayed_row_current_col(): # grad_factory matches the columnwise direction so the wgrad GEMM's # grad_output sub-quantizer pairs with the input/weight col format. grad_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda"), - policy_key=("test_hybrid_fp8_delayed_row_current_column", 1), + qfactory_key=("test_hybrid_fp8_delayed_row_current_column", 1), ) @@ -5087,7 +5087,7 @@ def _hybrid_recipe_fp8_current_row_delayed_col(): row_factory=lambda: Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda"), col_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E4M3), grad_factory=lambda: _make_delayed_quantizer(tex.DType.kFloat8E5M2), - policy_key=("test_hybrid_fp8_current_row_delayed_column", 1), + qfactory_key=("test_hybrid_fp8_current_row_delayed_column", 1), ) @@ -5097,7 +5097,7 @@ def _hybrid_recipe_mxfp8(): row_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), col_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E4M3), grad_factory=lambda: MXFP8Quantizer(tex.DType.kFloat8E5M2), - policy_key=("test_hybrid_mxfp8", 1), + qfactory_key=("test_hybrid_mxfp8", 1), ) @@ -5113,7 +5113,7 @@ def _hybrid_recipe_blockwise(): grad_factory=lambda: Float8BlockQuantizer( fp8_dtype=tex.DType.kFloat8E5M2, rowwise=True, columnwise=True ), - policy_key=("test_hybrid_blockwise", 1), + qfactory_key=("test_hybrid_blockwise", 1), ) @@ -5673,7 +5673,7 @@ def test_mxfp8_columnwise_raises(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_master_weight_mxfp8_columnwise", 1), + qfactory_key=("test_master_weight_mxfp8_columnwise", 1), ) # Shape must be a multiple of MXFP8 block size (32) on both axes. weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) @@ -5707,7 +5707,7 @@ def test_nvfp4_rowwise_raises(self): grad_factory=lambda: NVFP4Quantizer( fp4_dtype=tex.DType.kFloat4E2M1, with_2d_quantization=False ), - policy_key=("test_master_weight_nvfp4_rowwise", 1), + qfactory_key=("test_master_weight_nvfp4_rowwise", 1), ) weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) master_flat = hp_master.view(-1).contiguous() @@ -5729,7 +5729,7 @@ def test_nvfp4_columnwise_raises(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_master_weight_nvfp4_columnwise", 1), + qfactory_key=("test_master_weight_nvfp4_columnwise", 1), ) weight, hp_master = _build_hybrid_linear_weight(64, 128, hybrid_recipe) master_flat = hp_master.view(-1).contiguous() @@ -5775,7 +5775,7 @@ def test_blockwise_columnwise_raises(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_master_weight_blockwise_columnwise", 1), + qfactory_key=("test_master_weight_blockwise_columnwise", 1), ) weight, hp_master = _build_hybrid_linear_weight(128, 128, hybrid_recipe) master_flat = hp_master.view(-1).contiguous() @@ -5934,7 +5934,7 @@ def _hybrid_fp8_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_recipe_correspondence_hybrid_fp8", 1), + qfactory_key=("test_recipe_correspondence_hybrid_fp8", 1), ) def test_hybrid_param_with_matching_recipe_does_not_raise(self): @@ -6034,7 +6034,7 @@ def _build_hybrid_model(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_fused_adam_hybrid_fp8", 1), + qfactory_key=("test_fused_adam_hybrid_fp8", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -6138,7 +6138,7 @@ def _build_model_and_recipe(self): grad_factory=lambda: Float8CurrentScalingQuantizer( tex.DType.kFloat8E5M2, device="cuda" ), - policy_key=("test_end_to_end_hybrid_fp8", 1), + qfactory_key=("test_end_to_end_hybrid_fp8", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -6264,7 +6264,7 @@ def qfactory(role): hybrid_recipe = recipe.CustomRecipe( qfactory=qfactory, - policy_key=("test_mxfp8_nvfp4_quantized_params", 1), + qfactory_key=("test_mxfp8_nvfp4_quantized_params", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() @@ -6555,7 +6555,7 @@ def _vanilla_recipe(self): def _hybrid_recipe(self): return recipe.CustomRecipe( qfactory=_hybrid_fp8_current_qfactory, - policy_key=("test_hybrid_fp8_current", 1), + qfactory_key=("test_hybrid_fp8_current", 1), ) def test_equivalence(self): @@ -6572,7 +6572,7 @@ def _vanilla_recipe(self): def _hybrid_recipe(self): return recipe.CustomRecipe( qfactory=_hybrid_mxfp8_qfactory, - policy_key=("test_hybrid_mxfp8", 1), + qfactory_key=("test_hybrid_mxfp8", 1), ) def test_equivalence(self): @@ -6589,7 +6589,7 @@ def _vanilla_recipe(self): def _hybrid_recipe(self): return recipe.CustomRecipe( qfactory=_hybrid_block_fp8_qfactory, - policy_key=("test_hybrid_block_fp8", 1), + qfactory_key=("test_hybrid_block_fp8", 1), ) def test_equivalence(self): @@ -6609,7 +6609,7 @@ def _vanilla_recipe(self): def _hybrid_recipe(self): return recipe.CustomRecipe( qfactory=_hybrid_nvfp4_qfactory, - policy_key=("test_hybrid_nvfp4", 1), + qfactory_key=("test_hybrid_nvfp4", 1), ) def test_equivalence(self): @@ -6916,7 +6916,7 @@ def _same_format_fp8_recipe(self): row_factory=_fp8_row_factory, col_factory=_fp8_col_factory, grad_factory=_fp8_grad_factory, - policy_key=("test_recompute_hybrid_fp8", 1), + qfactory_key=("test_recompute_hybrid_fp8", 1), ) def _same_format_mxfp8_recipe(self): @@ -6926,7 +6926,7 @@ def _same_format_mxfp8_recipe(self): row_factory=_mxfp8_factory, col_factory=_mxfp8_factory, grad_factory=lambda: MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E5M2), - policy_key=("test_recompute_hybrid_mxfp8", 1), + qfactory_key=("test_recompute_hybrid_mxfp8", 1), ) def _cross_format_fp8_mxfp8_recipe(self): @@ -6938,7 +6938,7 @@ def _cross_format_fp8_mxfp8_recipe(self): row_factory=_fp8_row_factory, col_factory=_mxfp8_factory, grad_factory=_mxfp8_factory, - policy_key=("test_recompute_fp8_mxfp8", 1), + qfactory_key=("test_recompute_fp8_mxfp8", 1), ) def _run_linear(self, recipe_obj, *, checkpoint_fn=None): @@ -7465,7 +7465,7 @@ def test_te_checkpoint_linear_all_stateless_formats_bitwise(self, format_name, r row_factory=row_factory, col_factory=row_factory, grad_factory=grad_factory, - policy_key=("test_checkpoint_hybrid", format_name, 1), + qfactory_key=("test_checkpoint_hybrid", format_name, 1), ) def fn(model, inp): diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 82c33c4ad77..3cbfab0aa2e 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -15,7 +15,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex -from transformer_engine.common.recipe import CustomRecipe, quantizer_policy +from transformer_engine.common.recipe import CustomRecipe, quantizer_factory from transformer_engine.pytorch import ( Float8BlockQuantizer, Float8CurrentScalingQuantizer, @@ -56,7 +56,7 @@ # ── Module-level qfactories (picklable / autocast-friendly) ────────── -@quantizer_policy(key=("test_identity_all", 1)) +@quantizer_factory(key=("test_identity_all", 1)) def identity_all_factory(role): # pylint: disable=unused-argument """Whole layer in high precision: Identity for every slot.""" return IdentityQuantizer() @@ -135,7 +135,7 @@ def _format_quantizer(format_name): def _hybrid_quantized_fwd_identity_bwd_factory(format_name): - @quantizer_policy(key=(f"test_{format_name}_fwd_identity_bwd", 1)) + @quantizer_factory(key=(f"test_{format_name}_fwd_identity_bwd", 1)) def qfactory(role): is_linear = role is not None and role.module_type in ("linear", "grouped_linear") if is_linear and role.tensor_type in ("grad_output", "grad_input"): @@ -148,7 +148,7 @@ def qfactory(role): return qfactory -@quantizer_policy(key=("test_high_precision_fwd_fp8_bwd", 1)) +@quantizer_factory(key=("test_high_precision_fwd_fp8_bwd", 1)) def fwd_hp_bwd_fp8_factory(role): """High-precision forward, FP8 backward (per-direction via hybrid).""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -160,7 +160,7 @@ def fwd_hp_bwd_fp8_factory(role): ) -@quantizer_policy(key=("test_fp8_fwd_high_precision_bwd", 1)) +@quantizer_factory(key=("test_fp8_fwd_high_precision_bwd", 1)) def fwd_fp8_bwd_hp_factory(role): """FP8 forward, high-precision backward (per-direction via hybrid).""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -172,7 +172,7 @@ def fwd_fp8_bwd_hp_factory(role): ) -@quantizer_policy(key=("test_fp8_fwd_dequantized_high_precision_bwd", 1)) +@quantizer_factory(key=("test_fp8_fwd_dequantized_high_precision_bwd", 1)) def fwd_fp8_bwd_rowwise_dequantized_hp_factory(role): """FP8 forward, high-precision backward from dequantized fprop values.""" is_linear = role is not None and role.module_type in ("linear", "grouped_linear") @@ -185,7 +185,7 @@ def fwd_fp8_bwd_rowwise_dequantized_hp_factory(role): ) -@quantizer_policy(key=("test_hybrid_identity_all", 1)) +@quantizer_factory(key=("test_hybrid_identity_all", 1)) def hybrid_all_identity_factory(role): """All directions high precision, expressed through the hybrid container. @@ -202,7 +202,7 @@ def hybrid_all_identity_factory(role): ) -@quantizer_policy(key=("test_fp8_all", 1)) +@quantizer_factory(key=("test_fp8_all", 1)) def fp8_fwd_factory(role): """Plain FP8 current scaling for every slot (E4M3 fwd, E5M2 grad). @@ -397,7 +397,7 @@ def fake_split_quantize_hybrid( enabled=True, recipe=CustomRecipe( qfactory=qfactory, - policy_key=("test_grouped_linear_cpu_offload_hybrid", 1), + qfactory_key=("test_grouped_linear_cpu_offload_hybrid", 1), ), ): model(x, m_splits=m_splits) @@ -427,7 +427,7 @@ def qfactory(role): enabled=True, recipe=CustomRecipe( qfactory=qfactory, - policy_key=("test_grouped_linear_mixed_identity_weights", 1), + qfactory_key=("test_grouped_linear_mixed_identity_weights", 1), ), ): model(x, m_splits=m_splits) @@ -519,7 +519,7 @@ def qfactory(role): # pylint: disable=unused-argument ref = te_ops.BasicLinear(16, 16, device="cuda", dtype=torch.bfloat16) custom_recipe = CustomRecipe( qfactory=qfactory, - policy_key=("test_ops_basic_linear_hybrid_identity_weight", 1), + qfactory_key=("test_ops_basic_linear_hybrid_identity_weight", 1), ) torch.manual_seed(1702) with te.quantized_model_init(enabled=True, recipe=custom_recipe): @@ -575,7 +575,7 @@ def test_te_ops_quantize_then_gelu_accepts_identity_backed_tensors(self, qfactor enabled=True, recipe=CustomRecipe( qfactory=qfactory, - policy_key=("test_ops_quantize_gelu_identity_backed", 1), + qfactory_key=("test_ops_quantize_gelu_identity_backed", 1), ), ): y_test = test(x_test) @@ -1191,7 +1191,7 @@ def test_identity_recipe_matches_bf16_bitwise(self, module_name, qfactory): x = _identity_module_input(module_name) recipe = CustomRecipe( qfactory=qfactory, - policy_key=("test_quantized_primary_hybrid_identity", 1), + qfactory_key=("test_quantized_primary_hybrid_identity", 1), ) y_ref, dx_ref, wg_ref = _fwd_bwd_module(module_name, ref, x, recipe=None) @@ -1237,7 +1237,7 @@ def run(model, x, recipe): x, CustomRecipe( qfactory=mxfp8_all_factory, - policy_key=("test_mxfp8_all", 1), + qfactory_key=("test_mxfp8_all", 1), backward_override="high_precision", ), ) @@ -1450,7 +1450,7 @@ def _fwd_bwd_zoo_dequantized_module(module_name, module, inp, recipe): return out.detach().clone(), inp.grad.detach().clone(), param_grads -@quantizer_policy(key=("test_mxfp8_all", 1)) +@quantizer_factory(key=("test_mxfp8_all", 1)) def _mxfp8_all_qfactory(role): # pylint: disable=unused-argument return _mxfp8(tex.DType.kFloat8E4M3) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 12c47de8953..3329762f5da 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -41,7 +41,7 @@ MXFP8BlockScaling, NVFP4BlockScaling, Recipe, - quantizer_policy, + quantizer_factory, ) from transformer_engine.pytorch._extra_state import ( CheckpointExtraStatePolicy, @@ -210,25 +210,25 @@ def second_algorithm(*args): missing_key.quantizer_config() -def test_quantizer_policy_contract(): +def test_quantizer_factory_contract(): """The decorator validates and attaches metadata without wrapping or calling.""" calls = [] def factory(role): calls.append(role) - decorated = quantizer_policy(key=("test_policy", 1))(factory) + decorated = quantizer_factory(key=("test_policy", 1))(factory) assert decorated is factory - assert getattr(decorated, "policy_key") == ("test_policy", 1) + assert getattr(decorated, "qfactory_key") == ("test_policy", 1) assert calls == [] - assert te.quantizer_policy is quantizer_policy + assert te.quantizer_factory is quantizer_factory - with pytest.raises(TypeError, match="quantizer_policy key must be hashable"): - quantizer_policy(key=["mutable"]) + with pytest.raises(TypeError, match="quantizer_factory key must be hashable"): + quantizer_factory(key=["mutable"]) -def test_custom_recipe_policy_key_contract(): +def test_custom_recipe_qfactory_key_contract(): """Custom recipes resolve semantic keys without calling their factories.""" calls = [] @@ -238,14 +238,14 @@ def first_factory(role): def second_factory(role): calls.append(("second", role)) - policy_key = ("hybrid_policy", 2, ("double_quantization", True)) - first = CustomRecipe(qfactory=first_factory, policy_key=policy_key) - second = CustomRecipe(qfactory=second_factory, policy_key=policy_key) + qfactory_key = ("hybrid_policy", 2, ("double_quantization", True)) + first = CustomRecipe(qfactory=first_factory, qfactory_key=qfactory_key) + second = CustomRecipe(qfactory=second_factory, qfactory_key=qfactory_key) assert first.quantizer_config() == second.quantizer_config() assert set(dict(first.quantizer_config())) == { "recipe_type", - "policy_key", + "qfactory_key", "fp8_format", "fp8_dpa", "fp8_mha", @@ -254,24 +254,24 @@ def second_factory(role): } assert calls == [] - @quantizer_policy(key=("attached_policy", 1)) + @quantizer_factory(key=("attached_policy", 1)) def factory(role): # pylint: disable=unused-argument calls.append(("attached", role)) attached = CustomRecipe(qfactory=factory) - assert attached.policy_key == ("attached_policy", 1) - assert dict(attached.quantizer_config())["policy_key"] == ("attached_policy", 1) + assert attached.qfactory_key == ("attached_policy", 1) + assert dict(attached.quantizer_config())["qfactory_key"] == ("attached_policy", 1) - explicit = CustomRecipe(qfactory=factory, policy_key=("explicit_policy", 2)) - assert explicit.policy_key == ("explicit_policy", 2) - assert dict(explicit.quantizer_config())["policy_key"] == ("explicit_policy", 2) + explicit = CustomRecipe(qfactory=factory, qfactory_key=("explicit_policy", 2)) + assert explicit.qfactory_key == ("explicit_policy", 2) + assert dict(explicit.quantizer_config())["qfactory_key"] == ("explicit_policy", 2) def unkeyed_factory(role): calls.append(("unkeyed", role)) with pytest.raises( ValueError, - match=r"Pass policy_key=.*@quantizer_policy", + match=r"Pass qfactory_key=.*@quantizer_factory", ): CustomRecipe(qfactory=unkeyed_factory).quantizer_config() assert calls == [] @@ -302,7 +302,7 @@ def test_recipe_activation_does_not_call_qfactory_and_is_atomic_on_config_error( FP8GlobalStateManager.reset() calls = [] - @quantizer_policy(key=("activation_test", 1)) + @quantizer_factory(key=("activation_test", 1)) def keyed_factory(role): calls.append(role) @@ -316,7 +316,7 @@ def keyed_factory(role): def unkeyed_factory(role): calls.append(role) - with pytest.raises(ValueError, match="requires a semantic policy key"): + with pytest.raises(ValueError, match="requires a semantic qfactory key"): FP8GlobalStateManager.activate_recipe(CustomRecipe(qfactory=unkeyed_factory)) assert FP8GlobalStateManager.get_fp8_recipe() is keyed_recipe assert FP8GlobalStateManager.get_quantizer_config() is keyed_config diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 2dd00a915c9..871868b91c0 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -136,8 +136,8 @@ def _algorithm_config( return (("callable",), semantic_key) -def _validate_policy_key(key: Any, *, source: str) -> Hashable: - """Validate and return an immutable semantic policy key.""" +def _validate_qfactory_key(key: Any, *, source: str) -> Hashable: + """Validate and return an immutable semantic quantizer-factory key.""" if key is None: raise ValueError(f"{source} must not be None") try: @@ -147,26 +147,26 @@ def _validate_policy_key(key: Any, *, source: str) -> Hashable: return key -def quantizer_policy(*, key: Hashable) -> Callable: - """Attach a semantic policy key to a quantizer factory without wrapping or calling it. +def quantizer_factory(*, key: Hashable) -> Callable: + """Attach a semantic key to a quantizer factory without wrapping or calling it. The key, rather than Python callable identity, represents the factory's behavior in :class:`CustomRecipe` semantic configuration. Change the key whenever the factory's quantizer-selection behavior changes. A recommended convention is - ``(policy_name, behavior_revision)``, where ``policy_name`` is a short, stable name and + ``(factory_name, behavior_revision)``, where ``factory_name`` is a short, stable name and ``behavior_revision`` is incremented only when the factory's behavior changes. """ - policy_key = _validate_policy_key(key, source="quantizer_policy key") + qfactory_key = _validate_qfactory_key(key, source="quantizer_factory key") def decorator(qfactory: Callable) -> Callable: if not callable(qfactory): - raise TypeError("quantizer_policy can only decorate a callable") + raise TypeError("quantizer_factory can only decorate a callable") try: - setattr(qfactory, "policy_key", policy_key) + setattr(qfactory, "qfactory_key", qfactory_key) except (AttributeError, TypeError) as exc: raise TypeError( - "quantizer_policy could not attach metadata to this callable; " - "pass policy_key directly to CustomRecipe instead" + "quantizer_factory could not attach metadata to this callable; " + "pass qfactory_key directly to CustomRecipe instead" ) from exc return qfactory @@ -882,9 +882,9 @@ class CustomRecipe(Recipe): formats. It can be lowered when the factory's full output space is known; for example, a factory restricted to MXFP8 may use 32. Automatic padding reads this value without invoking ``qfactory``. - policy_key : Hashable, default = None + qfactory_key : Hashable, default = None Semantic identifier for the complete behavior of ``qfactory``. Pass it explicitly or - attach it to the factory with ``@quantizer_policy(key=...)``. An explicit value takes + attach it to the factory with ``@quantizer_factory(key=...)``. An explicit value takes precedence over attached factory metadata. Equivalent factory behavior must use equal keys, and any behavior change requires a different key. """ @@ -900,7 +900,7 @@ class CustomRecipe(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) quantization_alignment: int = 128 - policy_key: Optional[Hashable] = None + qfactory_key: Optional[Hashable] = None def __post_init__(self) -> None: assert ( @@ -908,28 +908,28 @@ def __post_init__(self) -> None: ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." if self.quantization_alignment <= 0: raise ValueError("CustomRecipe quantization_alignment must be positive.") - if self.policy_key is None: - attached_key = getattr(self.qfactory, "policy_key", None) + if self.qfactory_key is None: + attached_key = getattr(self.qfactory, "qfactory_key", None) if attached_key is not None: - self.policy_key = _validate_policy_key( + self.qfactory_key = _validate_qfactory_key( attached_key, - source="qfactory.policy_key", + source="qfactory.qfactory_key", ) else: - self.policy_key = _validate_policy_key( - self.policy_key, - source="CustomRecipe policy_key", + self.qfactory_key = _validate_qfactory_key( + self.qfactory_key, + source="CustomRecipe qfactory_key", ) def _make_quantizer_config(self) -> Hashable: - if self.policy_key is None: + if self.qfactory_key is None: raise ValueError( - "CustomRecipe requires a semantic policy key. Pass policy_key=... or decorate " - "qfactory with @quantizer_policy(key=...)." + "CustomRecipe requires a semantic qfactory key. Pass qfactory_key=... or decorate " + "qfactory with @quantizer_factory(key=...)." ) return ( ("recipe_type", "CustomRecipe"), - ("policy_key", self.policy_key), + ("qfactory_key", self.qfactory_key), ("fp8_format", self.fp8_format.name), ("fp8_dpa", self.fp8_dpa), ("fp8_mha", self.fp8_mha), @@ -941,7 +941,7 @@ def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"qfactory={self.qfactory}, " - f"policy_key={self.policy_key}, " + f"qfactory_key={self.qfactory_key}, " f"backward_override={self.backward_override}, " f"quantization_alignment={self.quantization_alignment}" ) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index c01117733cb..1e8aa6e5e00 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -53,7 +53,7 @@ from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.quantization import QuantizerRequest from transformer_engine.pytorch.quantization import DelayedScalingRequest -from transformer_engine.common.recipe import quantizer_policy +from transformer_engine.common.recipe import quantizer_factory from transformer_engine.pytorch.utils import get_cudnn_version from transformer_engine.pytorch.utils import get_device_compute_capability from transformer_engine.pytorch.utils import is_bf16_available diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py index 3be5121114a..05f23eb1902 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factories.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factories.py @@ -28,12 +28,12 @@ from typing import Optional import torch -from transformer_engine.common.recipe import quantizer_policy +from transformer_engine.common.recipe import quantizer_factory from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType -@quantizer_policy(key=("high_precision", 1)) +@quantizer_factory(key=("high_precision", 1)) def high_precision_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "IdentityQuantizer": @@ -52,7 +52,7 @@ def high_precision_factory( return IdentityQuantizer() -@quantizer_policy(key=("delayed_scaling", 1)) +@quantizer_factory(key=("delayed_scaling", 1)) def delayed_scaling_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "DelayedScalingRequest": @@ -72,7 +72,7 @@ def delayed_scaling_factory( return DelayedScalingRequest(fp8_format=Format.HYBRID) -@quantizer_policy(key=("current_scaling", 1)) +@quantizer_factory(key=("current_scaling", 1)) def current_scaling_factory( role: Optional[QuantizerRole], ) -> "Float8CurrentScalingQuantizer": @@ -96,7 +96,7 @@ def current_scaling_factory( ) -@quantizer_policy(key=("float8_block_scaling", 1)) +@quantizer_factory(key=("float8_block_scaling", 1)) def float8_block_scaling_factory( role: Optional[QuantizerRole], ) -> "Float8BlockQuantizer": @@ -127,7 +127,7 @@ def float8_block_scaling_factory( ) -@quantizer_policy(key=("mxfp8", 1)) +@quantizer_factory(key=("mxfp8", 1)) def mxfp8_factory( role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> "MXFP8Quantizer": @@ -143,7 +143,7 @@ def mxfp8_factory( ) -@quantizer_policy(key=("nvfp4", 1)) +@quantizer_factory(key=("nvfp4", 1)) def nvfp4_factory( role: Optional[QuantizerRole], ) -> "NVFP4Quantizer": diff --git a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py index 8e12d43ffe0..eee8bc94594 100644 --- a/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py +++ b/transformer_engine/pytorch/custom_recipes/quantizer_factory_zoo.py @@ -71,7 +71,7 @@ from typing import Optional -from transformer_engine.common.recipe import quantizer_policy +from transformer_engine.common.recipe import quantizer_factory from transformer_engine.pytorch.quantization import QuantizerRole from ..constants import DType from .quantizer_factories import mxfp8_factory, nvfp4_factory @@ -81,7 +81,7 @@ # ----------------------------------------------------------------------------- -@quantizer_policy(key=("high_precision_fwd_mxfp8_bwd", 1)) +@quantizer_factory(key=("high_precision_fwd_mxfp8_bwd", 1)) def high_precision_fwd_mxfp8_bwd_factory( role: Optional[QuantizerRole], ): @@ -119,7 +119,7 @@ def _plain_nvfp4_quantizer(*, row_scaled_nvfp4: bool = False): ) -@quantizer_policy(key=("mxfp8_fwd_nvfp4_bwd", 1)) +@quantizer_factory(key=("mxfp8_fwd_nvfp4_bwd", 1)) def mxfp8_fwd_nvfp4_bwd_factory( role: Optional[QuantizerRole], ): @@ -161,7 +161,7 @@ def mxfp8_fwd_nvfp4_bwd_factory( return mxfp8_factory(role) -@quantizer_policy(key=("nvfp4_1d_weight", 1)) +@quantizer_factory(key=("nvfp4_1d_weight", 1)) def nvfp4_1d_weight_factory( role: Optional[QuantizerRole], ): @@ -200,7 +200,7 @@ def nvfp4_1d_weight_factory( # ----------------------------------------------------------------------------- -@quantizer_policy(key=("mxfp8_fwd_high_precision_bwd", 1)) +@quantizer_factory(key=("mxfp8_fwd_high_precision_bwd", 1)) def mxfp8_fwd_high_precision_bwd_factory( role: Optional[QuantizerRole], ): @@ -240,7 +240,7 @@ def mxfp8_fwd_high_precision_bwd_factory( return mxfp8_factory(role) -@quantizer_policy(key=("composer2", 1)) +@quantizer_factory(key=("composer2", 1)) def nvfp4_row_scaled_fwd_mxfp8_bwd_factory( role: Optional[QuantizerRole], ): @@ -302,7 +302,7 @@ def nvfp4_row_scaled_fwd_mxfp8_bwd_factory( return mxfp8_factory(role) -@quantizer_policy(key=("nvfp4_row_scaled_fwd_high_precision_bwd", 1)) +@quantizer_factory(key=("nvfp4_row_scaled_fwd_high_precision_bwd", 1)) def nvfp4_row_scaled_fwd_high_precision_bwd_factory( role: Optional[QuantizerRole], ): @@ -360,7 +360,7 @@ def nvfp4_row_scaled_fwd_high_precision_bwd_factory( # ----------------------------------------------------------------------------- -@quantizer_policy(key=("nvfp4_linear_fp8_dpa", 1)) +@quantizer_factory(key=("nvfp4_linear_fp8_dpa", 1)) def nvfp4_linear_fp8_dpa_factory( role: Optional[QuantizerRole], ): diff --git a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py index 3efd0909de5..27fbc5b54f4 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/reference_current_scaling.py @@ -10,13 +10,13 @@ import torch -from transformer_engine.common.recipe import quantizer_policy +from transformer_engine.common.recipe import quantizer_factory from transformer_engine.pytorch.custom_recipes import gemm from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer -@quantizer_policy(key=("current_scaling_reference", 1)) +@quantizer_factory(key=("current_scaling_reference", 1)) def current_scaling_ref_factory(role): """Factory function for current scaling reference quantizer. diff --git a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py index a1e2092026f..50e88808d85 100644 --- a/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/reference_nvfp4.py @@ -9,13 +9,13 @@ import torch -from transformer_engine.common.recipe import quantizer_policy +from transformer_engine.common.recipe import quantizer_factory from transformer_engine.pytorch.custom_recipes import gemm from transformer_engine.pytorch.custom_recipes import reference_utils from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer -@quantizer_policy(key=("nvfp4_reference_rht_2d", 1)) +@quantizer_factory(key=("nvfp4_reference_rht_2d", 1)) def nvfp4_ref_rht_2d_factory(role): """ Quantizer factory for NVFP4 recipe reference implementation (RHT and 2D quantization for weights). From 8d62371ffd7c40833aed72927dcb94a0727b6a89 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 10 Aug 2026 14:19:53 +0000 Subject: [PATCH 67/97] Support mid-training recipe updates in Pytorch modules Signed-off-by: Evgeny --- docs/api/common.rst | 6 +- tests/pytorch/test_custom_recipe.py | 28 +- tests/pytorch/test_hybrid_quantization.py | 21 +- tests/pytorch/test_recipe.py | 426 +++++++++++++++- tests/pytorch/test_recipe_update.py | 462 ++++++++++++++++++ transformer_engine/pytorch/module/base.py | 368 +++++++++----- .../pytorch/module/grouped_linear.py | 195 +++++--- .../pytorch/module/layernorm_linear.py | 38 -- .../pytorch/module/layernorm_mlp.py | 59 --- transformer_engine/pytorch/module/linear.py | 38 -- .../pytorch/ops/basic/basic_linear.py | 12 - .../pytorch/ops/basic/grouped_linear.py | 14 - transformer_engine/pytorch/quantization.py | 215 ++++++-- 13 files changed, 1483 insertions(+), 399 deletions(-) create mode 100644 tests/pytorch/test_recipe_update.py diff --git a/docs/api/common.rst b/docs/api/common.rst index 3ddb66df980..8009dc81e94 100644 --- a/docs/api/common.rst +++ b/docs/api/common.rst @@ -8,7 +8,7 @@ Common API .. autoapiclass:: transformer_engine.common.recipe.Format -.. autoapiclass:: transformer_engine.common.recipe.DelayedScaling(margin=0, fp8_format=Format.HYBRID, amax_history_len=1024, amax_compute_algo="max", scaling_factor_compute_algo=None) +.. autoapiclass:: transformer_engine.common.recipe.DelayedScaling(margin=0, fp8_format=Format.HYBRID, amax_history_len=1024, amax_compute_algo="max", scaling_factor_compute_algo=None, reduce_amax=True, fp8_dpa=False, fp8_mha=False, backward_override=None, amax_compute_algo_key=None, scaling_factor_compute_algo_key=None) .. autoapiclass:: transformer_engine.common.recipe.MXFP8BlockScaling(fp8_format=Format.E4M3) @@ -18,4 +18,6 @@ Common API .. autoapiclass:: transformer_engine.common.recipe.Float8BlockScaling(fp8_format=Format.E4M3) -.. autoapiclass:: transformer_engine.common.recipe.CustomRecipe(qfactory, fp8_dpa=False, fp8_mha=False) +.. autofunction:: transformer_engine.common.recipe.quantizer_factory + +.. autoapiclass:: transformer_engine.common.recipe.CustomRecipe(qfactory, fp8_dpa=False, fp8_mha=False, qfactory_key=None) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index de62e8eec38..36734487dbc 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -982,9 +982,9 @@ def ds_factory(role): # ---------------------------------------------------------------------- # # Setting ``output_quantizer_role`` / ``grad_input_quantizer_role`` to a -# different value flips ``fp8_meta_tensors_initialized = False`` so the -# next ``set_meta_tensor`` call rebuilds the recipe state and quantizers -# with up-to-date roles. That rebuild MUST preserve persistent training +# different value advances the requested role revision without mutating the +# active runtime. The next runtime update rebuilds the recipe state and +# quantizers with up-to-date roles. That rebuild MUST preserve persistent training # buffers (delayed scaling's ``scale`` / ``amax_history``); otherwise # checkpointed amax history is silently destroyed on the first forward # pass after ``load_state_dict`` (when MHA wires boundary roles for the @@ -1023,13 +1023,17 @@ def test_role_change_preserves_delayed_scaling_state(): scale_data_ptr = state_before.scale.data_ptr() amax_data_ptr = state_before.amax_history.data_ptr() - # Trigger role-driven invalidation. Setting a non-None role flips - # ``fp8_meta_tensors_initialized = False`` so the next ``set_meta_tensor`` - # falls through and creates a fresh ``RecipeState``. + # Change the requested role. The active compatibility views remain live + # until the next runtime update commits a complete replacement. + runtime_before = model._quantization_runtime + role_revision_before = model._role_revision model.output_quantizer_role = QuantizerRole( module_type="dpa", tensor_type="qkv", name="downstream" ) - assert not model.fp8_meta_tensors_initialized + assert model.fp8_meta_tensors_initialized + assert model._quantization_runtime is runtime_before + assert model._role_revision == role_revision_before + 1 + assert model.fp8_meta["scaling_fwd"] is state_before # Trigger the rebuild directly (no forward, so we can compare buffers exactly). model.init_fp8_meta_tensors(fp8_recipe) @@ -1091,11 +1095,16 @@ def ds_factory(role): scale_obj_id = id(inner_before.scale) amax_obj_id = id(inner_before.amax_history) - # Trigger role-driven invalidation. + # Change the requested role without invalidating the active runtime. + runtime_before = model._quantization_runtime + role_revision_before = model._role_revision model.output_quantizer_role = QuantizerRole( module_type="dpa", tensor_type="qkv", name="downstream" ) - assert not model.fp8_meta_tensors_initialized + assert model.fp8_meta_tensors_initialized + assert model._quantization_runtime is runtime_before + assert model._role_revision == role_revision_before + 1 + assert model.fp8_meta["scaling_fwd"] is state_before # Rebuild. model.init_fp8_meta_tensors(custom_recipe) @@ -1318,6 +1327,7 @@ def test_custom_recipe_dpa_mxfp8(): custom_recipe = recipe.CustomRecipe( qfactory=_nvfp4_linear_mxfp8_dpa_factory, + qfactory_key=("test_nvfp4_linear_mxfp8_dpa", 1), fp8_dpa=True, ) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 3027d2e8c9d..94787c78033 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -4663,7 +4663,7 @@ def mixed_source_qfactory(role): nonlocal input_count source = "original" if role is not None and role.tensor_type == "input": - source = "original" if input_count == 0 else "rowwise_dequantized" + source = "original" if input_count % 2 == 0 else "rowwise_dequantized" input_count += 1 return HybridQuantizer( rowwise_quantizer=_make_fp8_quantizer(), @@ -4675,13 +4675,28 @@ def mixed_source_qfactory(role): qfactory=mixed_source_qfactory, qfactory_key=("test_mixed_columnwise_sources", 1), ) - # A failed generation is never marked validated. Base metadata can then - # early-return on retry, so the O(1) guard must validate it again. + # A failed candidate is neither committed nor marked validated. The + # active runtime stays behind the requested revision, so each retry + # prepares and validates a fresh candidate. + rebuilt_runtime = model._quantization_runtime + rebuilt_backward_generation = model._validated_quantizer_generations["scaling_bwd"] + rebuilt_delayed_quantizer = model._delayed_scaling_input_quantizer + rebuilt_unsafe_quantizer = model._unsafe_requantization_input_quantizer + failed_validation_call_count = len(validation_calls) for _ in range(2): with pytest.raises(ValueError, match="mixed columnwise source policies"): with torch.no_grad(), autocast(enabled=True, recipe=mixed_recipe): model(tensor, m_splits) + assert model._quantization_runtime is rebuilt_runtime assert model._validated_quantizer_generations["scaling_fwd"] is rebuilt_generation + assert ( + model._validated_quantizer_generations["scaling_bwd"] + is rebuilt_backward_generation + ) + assert model._delayed_scaling_input_quantizer is rebuilt_delayed_quantizer + assert model._unsafe_requantization_input_quantizer is rebuilt_unsafe_quantizer + assert len(validation_calls) > failed_validation_call_count + failed_validation_call_count = len(validation_calls) # Stale invalid recipe metadata must not affect the non-quantized path. with torch.no_grad(): diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 3329762f5da..9e802c0b650 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -2,7 +2,8 @@ # # See LICENSE for license information. -from dataclasses import fields +from dataclasses import fields, replace +from types import SimpleNamespace from typing import Optional import pickle @@ -28,8 +29,12 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.quantization import ( FP8GlobalStateManager, + Float8CurrentScalingRecipeState, NVFP4BlockScalingRecipeState, QuantizerRole, + RecipeState, + _QuantizationRuntimeKey, + _QuantizationRuntime, _amax_and_scale_update, ) import transformer_engine.pytorch.ops as te_ops @@ -107,6 +112,65 @@ def _make_repr(self) -> str: UnhashableConfigRecipe().quantizer_config() +def test_quantization_runtime_key_is_semantic_and_normalizes_slot_roles(): + """Runtime-key equality includes recipe config and each ordered role slot.""" + input_role = QuantizerRole(module_type="linear", tensor_type="input", name="qkv") + weight_role = QuantizerRole(module_type="linear", tensor_type="weight", name="qkv") + grad_role = QuantizerRole(module_type="linear", tensor_type="grad_output", name="qkv") + + key = _QuantizationRuntimeKey( + recipe_config=("recipe", "fp8"), + forward_roles=[input_role, None, weight_role], + backward_roles=[grad_role], + ) + same_request = _QuantizationRuntimeKey( + recipe_config=("recipe", "fp8"), + forward_roles=(input_role, None, weight_role), + backward_roles=(grad_role,), + ) + + assert key == same_request + assert hash(key) == hash(same_request) + assert key.forward_roles == (input_role, None, weight_role) + assert key != _QuantizationRuntimeKey( + recipe_config=("recipe", "fp8"), + forward_roles=(input_role, weight_role), + backward_roles=(grad_role,), + ) + assert key != _QuantizationRuntimeKey( + recipe_config=("recipe", "new-fp8"), + forward_roles=(input_role, None, weight_role), + backward_roles=(grad_role,), + ) + + +def test_quantization_runtime_bundles_both_quantization_directions(): + """A runtime retains the complete active-or-candidate quantizer bundle.""" + runtime_key = _QuantizationRuntimeKey( + recipe_config=("recipe", "fp8"), + forward_roles=(), + backward_roles=(), + ) + forward_quantizers = [] + backward_quantizers = [] + runtime = _QuantizationRuntime( + key=runtime_key, + num_gemms=1, + recipe_config_revision=3, + role_revision=2, + forward_states=(), + backward_states=(), + forward_quantizers=forward_quantizers, + backward_quantizers=backward_quantizers, + delayed_registration_handles=("forward-handle", "backward-handle"), + ) + + assert runtime.key is runtime_key + assert runtime.forward_quantizers is forward_quantizers + assert runtime.backward_quantizers is backward_quantizers + assert runtime.delayed_registration_handles == ("forward-handle", "backward-handle") + + @pytest.mark.parametrize( "recipe_type,derived_fields", [ @@ -177,6 +241,301 @@ def test_equal_builtin_recipes_have_equal_quantizer_configs(recipe_type): assert first == second +def test_same_recipe_mutation_invalidates_config_for_direct_and_nested_parameters(): + """Built-in recipe config changes for same-object and nested-qparam updates.""" + recipe = Float8CurrentScaling() + original_config = recipe.quantizer_config() + + recipe.fp8_dpa = True + direct_mutation_config = recipe.quantizer_config() + assert direct_mutation_config != original_config + assert dict(direct_mutation_config)["fp8_dpa"] is True + + recipe.fp8_quant_fwd_inp = replace(recipe.fp8_quant_fwd_inp, amax_epsilon=0.25) + nested_mutation_config = recipe.quantizer_config() + assert nested_mutation_config != direct_mutation_config + assert dict(nested_mutation_config)["fp8_quant_fwd_inp"] == ( + ("power_2_scale", recipe.fp8_quant_fwd_inp.power_2_scale), + ("amax_epsilon", 0.25), + ("random_hadamard_transform", recipe.fp8_quant_fwd_inp.random_hadamard_transform), + ("stochastic_rounding", recipe.fp8_quant_fwd_inp.stochastic_rounding), + ("fp4_2d_quantization", recipe.fp8_quant_fwd_inp.fp4_2d_quantization), + ) + + +def test_current_scaling_recipe_state_configures_roles_and_preserves_boundary_defaults(): + """Current-scaling construction follows role dispatch and legacy boundary defaults.""" + recipe = Float8CurrentScaling() + recipe.fp8_quant_fwd_inp = replace( + recipe.fp8_quant_fwd_inp, power_2_scale=True, amax_epsilon=0.1 + ) + recipe.fp8_quant_fwd_weight = replace( + recipe.fp8_quant_fwd_weight, power_2_scale=False, amax_epsilon=0.2 + ) + recipe.fp8_quant_bwd_grad = replace( + recipe.fp8_quant_bwd_grad, power_2_scale=True, amax_epsilon=0.3 + ) + + forward_quantizers = Float8CurrentScalingRecipeState( + recipe, + mode="forward", + num_quantizers=4, + device=torch.device("cpu"), + roles=[ + QuantizerRole(tensor_type="input"), + QuantizerRole(tensor_type="weight"), + None, + QuantizerRole(tensor_type="unknown"), + ], + ).make_quantizers() + backward_quantizers = Float8CurrentScalingRecipeState( + recipe, + mode="backward", + num_quantizers=3, + device=torch.device("cpu"), + roles=[ + QuantizerRole(tensor_type="grad_output"), + None, + QuantizerRole(tensor_type="grad_input"), + ], + ).make_quantizers() + + assert [(q.force_pow_2_scales, q.amax_epsilon) for q in forward_quantizers] == [ + (True, 0.1), + (False, 0.2), + (recipe.use_power_2_scales, 0.0), + (True, 0.1), + ] + assert [(q.force_pow_2_scales, q.amax_epsilon) for q in backward_quantizers] == [ + (True, 0.3), + (recipe.use_power_2_scales, 0.0), + (recipe.use_power_2_scales, 0.0), + ] + + +@pytest.mark.parametrize( + "recipe", + [ + DelayedScaling(), + Float8CurrentScaling(), + MXFP8BlockScaling(), + Float8BlockScaling(), + NVFP4BlockScaling(), + CustomRecipe(qfactory=lambda role: role, qfactory_key=("invalid-mode-test", 1)), + ], +) +def test_recipe_state_rejects_invalid_mode(recipe): + """Every RecipeState implementation uses the same mode validation.""" + with pytest.raises(ValueError, match=r"Unexpected recipe mode \(invalid\)"): + RecipeState.create(recipe, mode="invalid", device=torch.device("cpu")) + + +def test_current_scaling_role_layouts_cover_module_and_basic_op_families(): + """Every Phase 2 owner supplies roles that select the expected per-slot qparams.""" + recipe = Float8CurrentScaling() + recipe.fp8_quant_fwd_inp = replace( + recipe.fp8_quant_fwd_inp, power_2_scale=True, amax_epsilon=0.1 + ) + recipe.fp8_quant_fwd_weight = replace( + recipe.fp8_quant_fwd_weight, power_2_scale=False, amax_epsilon=0.2 + ) + recipe.fp8_quant_bwd_grad = replace( + recipe.fp8_quant_bwd_grad, power_2_scale=True, amax_epsilon=0.3 + ) + + boundary = (recipe.use_power_2_scales, 0.0) + inp = (True, 0.1) + weight = (False, 0.2) + grad = (True, 0.3) + + def settings(mode, roles): + quantizers = Float8CurrentScalingRecipeState( + recipe, + mode=mode, + num_quantizers=len(roles), + device=torch.device("cpu"), + roles=roles, + ).make_quantizers() + return [(q.force_pow_2_scales, q.amax_epsilon) for q in quantizers] + + module_owner = SimpleNamespace( + name="test", + _output_quantizer_role=None, + _grad_input_quantizer_role=None, + ) + for module_type in (Linear, LayerNormLinear): + forward_roles = module_type.get_quantizer_roles(module_owner, fwd=True, num_quantizers=3) + backward_roles = module_type.get_quantizer_roles(module_owner, fwd=False, num_quantizers=2) + assert settings("forward", forward_roles) == [inp, weight, boundary] + assert settings("backward", backward_roles) == [grad, boundary] + + mlp_forward_roles = LayerNormMLP.get_quantizer_roles( + module_owner, fwd=True, num_quantizers=6 + ) + mlp_backward_roles = LayerNormMLP.get_quantizer_roles( + module_owner, fwd=False, num_quantizers=4 + ) + assert settings("forward", mlp_forward_roles) == [ + inp, + weight, + inp, + inp, + weight, + boundary, + ] + assert settings("backward", mlp_backward_roles) == [grad, boundary, grad, grad] + + grouped_forward_roles = GroupedLinear.get_quantizer_roles( + module_owner, fwd=True, num_quantizers=6 + ) + grouped_backward_roles = GroupedLinear.get_quantizer_roles( + module_owner, fwd=False, num_quantizers=4 + ) + assert settings("forward", grouped_forward_roles) == [ + inp, + weight, + boundary, + inp, + weight, + boundary, + ] + assert settings("backward", grouped_backward_roles) == [grad, boundary, grad, boundary] + + basic_owner = SimpleNamespace(name="test", num_groups=2) + basic_forward_roles = te_ops.BasicLinear.get_quantizer_roles(basic_owner, "forward") + basic_backward_roles = te_ops.BasicLinear.get_quantizer_roles(basic_owner, "backward") + assert settings("forward", basic_forward_roles) == [inp, weight] + assert settings("backward", basic_backward_roles) == [grad] + + basic_grouped_forward_roles = te_ops.GroupedLinear.get_quantizer_roles( + basic_owner, "forward" + ) + basic_grouped_backward_roles = te_ops.GroupedLinear.get_quantizer_roles( + basic_owner, "backward" + ) + assert settings("forward", basic_grouped_forward_roles) == [inp, weight, inp, weight] + assert settings("backward", basic_grouped_backward_roles) == [grad, grad] + + +def test_current_scaling_owner_configuration_paths_preserve_numerics_and_traits(): + """Owner initialization paths publish identical qparams and retain non-numerical traits.""" + recipe = Float8CurrentScaling() + recipe.fp8_quant_fwd_inp = replace( + recipe.fp8_quant_fwd_inp, power_2_scale=True, amax_epsilon=0.1 + ) + recipe.fp8_quant_fwd_weight = replace( + recipe.fp8_quant_fwd_weight, power_2_scale=False, amax_epsilon=0.2 + ) + recipe.fp8_quant_bwd_grad = replace( + recipe.fp8_quant_bwd_grad, power_2_scale=True, amax_epsilon=0.3 + ) + + boundary = (recipe.use_power_2_scales, 0.0) + inp = (True, 0.1) + weight = (False, 0.2) + grad = (True, 0.3) + + def settings(quantizers): + return [(q.force_pow_2_scales, q.amax_epsilon) for q in quantizers] + + def make_module_owner(module_type, num_gemms): + owner = object.__new__(module_type) + owner.__dict__.update( + name="test", + num_gemms=num_gemms, + fp8_meta={"num_gemms": num_gemms}, + fp8_meta_tensors_initialized=False, + quantizers={"scaling_fwd": [], "scaling_bwd": []}, + _output_quantizer_role=None, + _grad_input_quantizer_role=None, + ) + return owner + + FP8GlobalStateManager.reset() + FP8GlobalStateManager.activate_recipe(recipe) + try: + for module_type in (Linear, LayerNormLinear): + owner = make_module_owner(module_type, 1) + module_type.set_meta_tensor(owner, True, recipe) + module_type.set_meta_tensor(owner, False, recipe) + assert settings(owner.quantizers["scaling_fwd"]) == [inp, weight, boundary] + assert settings(owner.quantizers["scaling_bwd"]) == [grad, boundary] + + mlp = make_module_owner(LayerNormMLP, 2) + LayerNormMLP.set_meta_tensor(mlp, True, recipe) + LayerNormMLP.set_meta_tensor(mlp, False, recipe) + assert settings(mlp.quantizers["scaling_fwd"]) == [ + inp, + weight, + inp, + inp, + weight, + boundary, + ] + assert settings(mlp.quantizers["scaling_bwd"]) == [grad, boundary, grad, grad] + + grouped = make_module_owner(GroupedLinear, 2) + grouped.__dict__.update( + tp_size=1, + _offsets={ + "input": 0, + "weight": 1, + "output": 2, + "grad_output": 0, + "grad_input": 1, + }, + _num_fp8_tensors_per_gemm={"fwd": 3, "bwd": 2}, + _validated_quantizer_generations={}, + _delayed_scaling_input_quantizer=None, + _unsafe_requantization_input_quantizer=None, + ) + GroupedLinear.set_meta_tensor(grouped, True, recipe) + GroupedLinear.set_meta_tensor(grouped, False, recipe) + assert settings(grouped.quantizers["scaling_fwd"]) == [ + inp, + weight, + boundary, + inp, + weight, + boundary, + ] + assert settings(grouped.quantizers["scaling_bwd"]) == [ + grad, + boundary, + grad, + boundary, + ] + + grouped_tp = make_module_owner(GroupedLinear, 2) + grouped_tp.__dict__["tp_size"] = 2 + with pytest.raises( + ValueError, + match="GroupedLinear doesn't support TP > 1 with Float8 current scaling", + ): + GroupedLinear.set_meta_tensor(grouped_tp, True, recipe) + + basic_linear = te_ops.BasicLinear(4, 4, device="meta") + basic_linear.reset_recipe_state(recipe=recipe) + assert settings(basic_linear._quantizers["forward"]) == [inp, weight] + assert settings(basic_linear._quantizers["backward"]) == [grad] + assert basic_linear.get_quantizer("forward", 0).internal + assert basic_linear.get_quantizer("forward", 0).optimize_for_gemm + assert basic_linear.get_quantizer("forward", 1).internal + assert basic_linear.get_quantizer("backward", 0).internal + assert basic_linear.get_quantizer("backward", 0).optimize_for_gemm + + basic_grouped = te_ops.GroupedLinear(2, 4, 4, device="meta") + basic_grouped.reset_recipe_state(recipe=recipe) + assert settings(basic_grouped._quantizers["forward"]) == [inp, weight, inp, weight] + assert settings(basic_grouped._quantizers["backward"]) == [grad, grad] + for quantizer in basic_grouped._quantizers["forward"]: + assert quantizer.internal + for quantizer in basic_grouped._quantizers["backward"]: + assert quantizer.internal + finally: + FP8GlobalStateManager.reset() + + @pytest.mark.parametrize( "algorithm_field", ["amax_compute_algo", "scaling_factor_compute_algo"], @@ -277,23 +636,50 @@ def unkeyed_factory(role): assert calls == [] -def test_global_state_caches_active_quantizer_config(): - """Recipe activation publishes the recipe and its semantic configuration together.""" +def test_custom_recipe_qfactory_key_mutation_changes_semantic_configuration(): + """Changing a CustomRecipe policy key invalidates its cached config without probing.""" + calls = [] + + def factory(role): + calls.append(role) + + recipe = CustomRecipe(qfactory=factory, qfactory_key=("test-policy", 1)) + original_config = recipe.quantizer_config() + recipe.qfactory_key = ("test-policy", 2) + + assert recipe.quantizer_config() != original_config + assert dict(recipe.quantizer_config())["qfactory_key"] == ("test-policy", 2) + assert calls == [] + + +def test_global_state_caches_active_quantizer_config_and_revision(): + """Recipe activation publishes semantic configuration with a manager-owned revision.""" FP8GlobalStateManager.reset() active_recipe = Float8CurrentScaling() + assert FP8GlobalStateManager.get_quantizer_config_revision() == 0 FP8GlobalStateManager.activate_recipe(active_recipe) original_config = active_recipe.quantizer_config() assert FP8GlobalStateManager.get_fp8_recipe() is active_recipe assert FP8GlobalStateManager.get_quantizer_config() is original_config + assert FP8GlobalStateManager.get_quantizer_config_revision() == 1 + + # Equivalent recipe objects update the requested recipe without advancing the revision. + equivalent_recipe = Float8CurrentScaling() + FP8GlobalStateManager.activate_recipe(equivalent_recipe) + assert FP8GlobalStateManager.get_fp8_recipe() is equivalent_recipe + assert FP8GlobalStateManager.get_quantizer_config() is original_config + assert FP8GlobalStateManager.get_quantizer_config_revision() == 1 # Recipe mutation is requested state until the recipe is activated again. active_recipe.fp8_dpa = True assert FP8GlobalStateManager.get_quantizer_config() is original_config + assert FP8GlobalStateManager.get_quantizer_config_revision() == 1 FP8GlobalStateManager.activate_recipe(active_recipe) assert FP8GlobalStateManager.get_quantizer_config() == active_recipe.quantizer_config() assert FP8GlobalStateManager.get_quantizer_config() != original_config + assert FP8GlobalStateManager.get_quantizer_config_revision() == 2 FP8GlobalStateManager.reset() @@ -309,6 +695,7 @@ def keyed_factory(role): keyed_recipe = CustomRecipe(qfactory=keyed_factory) FP8GlobalStateManager.activate_recipe(keyed_recipe) keyed_config = keyed_recipe.quantizer_config() + keyed_revision = FP8GlobalStateManager.get_quantizer_config_revision() assert FP8GlobalStateManager.get_fp8_recipe() is keyed_recipe assert FP8GlobalStateManager.get_quantizer_config() is keyed_config assert calls == [] @@ -320,6 +707,7 @@ def unkeyed_factory(role): FP8GlobalStateManager.activate_recipe(CustomRecipe(qfactory=unkeyed_factory)) assert FP8GlobalStateManager.get_fp8_recipe() is keyed_recipe assert FP8GlobalStateManager.get_quantizer_config() is keyed_config + assert FP8GlobalStateManager.get_quantizer_config_revision() == keyed_revision assert calls == [] FP8GlobalStateManager.reset() @@ -331,13 +719,40 @@ def test_autocast_restores_recipe_and_quantizer_config_together(): inner_recipe = MXFP8BlockScaling() FP8GlobalStateManager.activate_recipe(outer_recipe) outer_config = FP8GlobalStateManager.get_quantizer_config() + outer_revision = FP8GlobalStateManager.get_quantizer_config_revision() + + # Equal independent configurations do not advance the revision on entry or restoration. + with te.autocast(enabled=False, recipe=Float8CurrentScaling()): + assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision with te.autocast(enabled=False, recipe=inner_recipe): assert FP8GlobalStateManager.get_fp8_recipe() is inner_recipe assert FP8GlobalStateManager.get_quantizer_config() == inner_recipe.quantizer_config() + assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + 1 + + assert FP8GlobalStateManager.get_fp8_recipe() is outer_recipe + assert FP8GlobalStateManager.get_quantizer_config() is outer_config + assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + 2 + FP8GlobalStateManager.reset() + + +def test_quantized_model_init_restores_recipe_config_with_monotonic_revision(): + """Model-init recipe restoration republishes config rather than reusing an old revision.""" + FP8GlobalStateManager.reset() + outer_recipe = Float8CurrentScaling() + inner_recipe = MXFP8BlockScaling() + FP8GlobalStateManager.activate_recipe(outer_recipe) + outer_config = FP8GlobalStateManager.get_quantizer_config() + outer_revision = FP8GlobalStateManager.get_quantizer_config_revision() + + with te.quantized_model_init(enabled=False, recipe=inner_recipe): + assert FP8GlobalStateManager.get_fp8_recipe() is inner_recipe + assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + 1 assert FP8GlobalStateManager.get_fp8_recipe() is outer_recipe assert FP8GlobalStateManager.get_quantizer_config() is outer_config + assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + 2 FP8GlobalStateManager.reset() @@ -790,7 +1205,7 @@ def test_dynamic_recipe_update( GroupedLinear, ], ) - def test_quantizer_update(self, module_class): + def test_quantized_primary_recipe_update_is_rejected(self, module_class): in_features = 32 out_features = 32 batch_size = 32 @@ -805,8 +1220,7 @@ def test_quantizer_update(self, module_class): x = torch.randn(batch_size, in_features, device="cuda") recipe = DelayedScaling(amax_history_len=1) with te.autocast(enabled=True, recipe=recipe): - warn_msg = "Quantizer is being updated, this may affect model behavior" - with pytest.warns(UserWarning, match=warn_msg): + with pytest.raises(RuntimeError, match="Recipe mismatch for quantized primary weights"): if module_class == GroupedLinear: y = module(x, [batch_size]) else: diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py new file mode 100644 index 00000000000..16cc9115a19 --- /dev/null +++ b/tests/pytorch/test_recipe_update.py @@ -0,0 +1,462 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Module-local runtime lifecycle tests for mid-training recipe updates.""" + +import pytest +import torch + +from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.pytorch import ( + GroupedLinear, + LayerNormLinear, + LayerNormMLP, + Linear, + autocast, +) +from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA is required" +) + + +def _make_counting_recipe(key, calls, *, fail_on_grad_output=False): + """Build an identity recipe whose factory calls are externally observable.""" + + def qfactory(role): + calls.append(role) + if ( + fail_on_grad_output + and role is not None + and role.tensor_type == "grad_output" + ): + raise RuntimeError("backward factory failure") + return IdentityQuantizer() + + return CustomRecipe(qfactory=qfactory, qfactory_key=key) + + +def _ensure_runtime(module, recipe, revision, *, num_gemms=1): + return ( + module._ensure_active_quantization_runtime( # pylint: disable=protected-access + recipe=recipe, + recipe_config=recipe.quantizer_config(), + recipe_config_revision=revision, + num_gemms=num_gemms, + ) + ) + + +def test_equal_runtime_key_synchronizes_revision_without_factory_calls(): + """Equal independent recipes and missed A -> B -> A updates reuse the runtime.""" + calls = [] + first_recipe = _make_counting_recipe(("runtime-reuse", 1), calls) + equal_recipe = _make_counting_recipe(("runtime-reuse", 1), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + + assert _ensure_runtime(module, first_recipe, revision=1) + active = module._quantization_runtime # pylint: disable=protected-access + assert active is not None + assert len(calls) == 5 + + # Revision 3 models a module that did not run while revision 2 was active. + assert not _ensure_runtime(module, equal_recipe, revision=3) + assert module._quantization_runtime is active # pylint: disable=protected-access + assert active.recipe_config_revision == 3 + assert len(calls) == 5 + + +def test_unchanged_runtime_uses_revision_hot_path(monkeypatch): + """The steady-state check does not resolve roles or invoke the factory.""" + calls = [] + recipe = _make_counting_recipe(("runtime-hot-path", 1), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + _ensure_runtime(module, recipe, revision=1) + + def unexpected_role_resolution(**_kwargs): + raise AssertionError("unchanged runtime resolved quantizer roles") + + monkeypatch.setattr(module, "get_quantizer_roles", unexpected_role_resolution) + assert not _ensure_runtime(module, recipe, revision=1) + assert len(calls) == 5 + + +def test_unchanged_forward_uses_revision_hot_path(monkeypatch): + """Repeated forwards do not enter any runtime-construction cold path.""" + calls = [] + recipe = _make_counting_recipe(("forward-hot-path", 1), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + inp = torch.randn(8, 16, device="cuda") + + with autocast(enabled=True, recipe=recipe): + module(inp) + active = module._quantization_runtime # pylint: disable=protected-access + factory_call_count = len(calls) + assert active is not None + assert factory_call_count > 0 + + def unexpected_cold_path(*_args, **_kwargs): + raise AssertionError("unchanged forward entered the runtime cold path") + + monkeypatch.setattr(module, "get_quantizer_roles", unexpected_cold_path) + monkeypatch.setattr( + module, "_prepare_quantization_runtime", unexpected_cold_path + ) + monkeypatch.setattr( + module, "_validate_quantization_runtime", unexpected_cold_path + ) + monkeypatch.setattr( + module, "_commit_quantization_runtime", unexpected_cold_path + ) + + workspace_sentinel = object() + module._fp8_workspaces["sentinel"] = ( + workspace_sentinel # pylint: disable=protected-access + ) + module(inp) + module(inp) + + assert ( + module._quantization_runtime is active + ) # pylint: disable=protected-access + assert len(calls) == factory_call_count + assert ( + module._fp8_workspaces["sentinel"] is workspace_sentinel + ) # pylint: disable=protected-access + + +def test_same_recipe_object_semantic_mutation_rebuilds_runtime(): + """Runtime matching observes a new semantic config on the same recipe object.""" + calls = [] + recipe = _make_counting_recipe(("same-object-update", 1), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + _ensure_runtime(module, recipe, revision=1) + old_runtime = module._quantization_runtime # pylint: disable=protected-access + + recipe.qfactory_key = ("same-object-update", 2) + assert _ensure_runtime(module, recipe, revision=2) + assert ( + module._quantization_runtime is not old_runtime + ) # pylint: disable=protected-access + assert len(calls) == 10 + + +def test_role_revision_is_requested_until_atomic_runtime_commit(): + """Changing a boundary role does not mutate live state before the next update.""" + calls = [] + recipe = _make_counting_recipe(("role-update", 1), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + _ensure_runtime(module, recipe, revision=1) + + old_runtime = module._quantization_runtime # pylint: disable=protected-access + old_forward_state = module.fp8_meta["scaling_fwd"] + old_forward_quantizers = module.quantizers["scaling_fwd"] + old_role_revision = module._role_revision # pylint: disable=protected-access + role = QuantizerRole(module_type="linear", tensor_type="input", name="consumer") + + module.output_quantizer_role = role + assert ( + module._role_revision == old_role_revision + 1 + ) # pylint: disable=protected-access + assert ( + module._quantization_runtime is old_runtime + ) # pylint: disable=protected-access + assert module.fp8_meta["scaling_fwd"] is old_forward_state + assert module.quantizers["scaling_fwd"] is old_forward_quantizers + + assert _ensure_runtime(module, recipe, revision=1) + new_runtime = module._quantization_runtime # pylint: disable=protected-access + assert new_runtime is not old_runtime + assert new_runtime.key.forward_roles[-1] == role + assert module.fp8_meta["scaling_fwd"] is new_runtime.forward_states[0] + assert module.quantizers["scaling_fwd"] is new_runtime.forward_quantizers + assert len(calls) == 10 + + # Reassigning an equal immutable role is a hot-path no-op. + module.output_quantizer_role = QuantizerRole( + module_type="linear", tensor_type="input", name="consumer" + ) + assert ( + module._role_revision == old_role_revision + 1 + ) # pylint: disable=protected-access + assert not _ensure_runtime(module, recipe, revision=1) + assert len(calls) == 10 + + +def test_backward_factory_failure_keeps_complete_active_runtime(): + """A backward preparation failure must not publish candidate forward state.""" + calls = [] + active_recipe = _make_counting_recipe(("atomic-update", 1), calls) + failing_recipe = _make_counting_recipe( + ("atomic-update", 2), + calls, + fail_on_grad_output=True, + ) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + _ensure_runtime(module, active_recipe, revision=1) + + old_runtime = module._quantization_runtime # pylint: disable=protected-access + old_recipe = module.fp8_meta["recipe"] + old_forward_state = module.fp8_meta["scaling_fwd"] + old_backward_state = module.fp8_meta["scaling_bwd"] + old_forward_quantizers = module.quantizers["scaling_fwd"] + old_backward_quantizers = module.quantizers["scaling_bwd"] + + with pytest.raises(RuntimeError, match="backward factory failure"): + _ensure_runtime(module, failing_recipe, revision=2) + + assert ( + module._quantization_runtime is old_runtime + ) # pylint: disable=protected-access + assert module.fp8_meta["recipe"] is old_recipe + assert module.fp8_meta["scaling_fwd"] is old_forward_state + assert module.fp8_meta["scaling_bwd"] is old_backward_state + assert module.quantizers["scaling_fwd"] is old_forward_quantizers + assert module.quantizers["scaling_bwd"] is old_backward_quantizers + assert old_runtime.recipe_config_revision == 1 + + +@pytest.mark.parametrize( + ("module_factory", "num_gemms", "cache_names"), + ( + pytest.param( + lambda: Linear(16, 16, bias=False, device="cuda", name="linear"), + 1, + ("weight",), + id="linear", + ), + pytest.param( + lambda: LayerNormLinear( + 16, + 16, + bias=False, + device="cuda", + name="layernorm_linear", + ), + 1, + ("weight",), + id="layernorm_linear", + ), + pytest.param( + lambda: LayerNormMLP( + 16, + 32, + bias=False, + device="cuda", + name="layernorm_mlp", + ), + 2, + ("fc1_weight", "fc2_weight"), + id="layernorm_mlp", + ), + pytest.param( + lambda: GroupedLinear( + 2, + 16, + 16, + bias=False, + device="cuda", + name="grouped_linear", + ), + 2, + ("weight0", "weight1"), + id="grouped_linear", + ), + ), +) +def test_runtime_update_workspace_lifecycle( + module_factory, + num_gemms, + cache_names, +): + """Workspace ownership stays atomic across updates in every module family.""" + calls = [] + active_recipe = _make_counting_recipe(("workspace-update", 1), calls) + equal_recipe = _make_counting_recipe(("workspace-update", 1), calls) + failing_recipe = _make_counting_recipe( + ("workspace-update", 2), + calls, + fail_on_grad_output=True, + ) + replacement_recipe = _make_counting_recipe(("workspace-update", 3), calls) + module = module_factory() + assert _ensure_runtime(module, active_recipe, revision=1, num_gemms=num_gemms) + + active_runtime = module._quantization_runtime # pylint: disable=protected-access + active_views = ( + module.fp8_meta["recipe"], + module.fp8_meta["scaling_fwd"], + module.fp8_meta["scaling_bwd"], + module.quantizers["scaling_fwd"], + module.quantizers["scaling_bwd"], + ) + workspaces = {cache_name: object() for cache_name in cache_names} + for cache_name, workspace in workspaces.items(): + module._fp8_workspaces[cache_name] = ( + workspace # pylint: disable=protected-access + ) + + # A global revision change with an equal semantic runtime preserves caches. + assert not _ensure_runtime(module, equal_recipe, revision=2, num_gemms=num_gemms) + assert ( + module._quantization_runtime is active_runtime + ) # pylint: disable=protected-access + for cache_name, workspace in workspaces.items(): + assert ( + module._fp8_workspaces[cache_name] is workspace + ) # pylint: disable=protected-access + + # Candidate construction failure preserves both compatibility views and caches. + with pytest.raises(RuntimeError, match="backward factory failure"): + _ensure_runtime(module, failing_recipe, revision=3, num_gemms=num_gemms) + assert ( + module._quantization_runtime is active_runtime + ) # pylint: disable=protected-access + current_views = ( + module.fp8_meta["recipe"], + module.fp8_meta["scaling_fwd"], + module.fp8_meta["scaling_bwd"], + module.quantizers["scaling_fwd"], + module.quantizers["scaling_bwd"], + ) + assert all( + current is active for current, active in zip(current_views, active_views) + ) + for cache_name, workspace in workspaces.items(): + assert ( + module._fp8_workspaces[cache_name] is workspace + ) # pylint: disable=protected-access + + # Only a fully committed replacement clears cached workspaces. + assert _ensure_runtime(module, replacement_recipe, revision=4, num_gemms=num_gemms) + replacement_runtime = ( + module._quantization_runtime + ) # pylint: disable=protected-access + assert replacement_runtime is not active_runtime + assert not module._fp8_workspaces # pylint: disable=protected-access + + +def test_candidate_validation_failure_keeps_complete_active_runtime(monkeypatch): + """Validation runs before any candidate state is published.""" + calls = [] + active_recipe = _make_counting_recipe(("validation-update", 1), calls) + candidate_recipe = _make_counting_recipe(("validation-update", 2), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + _ensure_runtime(module, active_recipe, revision=1) + + old_runtime = module._quantization_runtime # pylint: disable=protected-access + old_forward_state = module.fp8_meta["scaling_fwd"] + old_backward_state = module.fp8_meta["scaling_bwd"] + old_forward_quantizers = module.quantizers["scaling_fwd"] + old_backward_quantizers = module.quantizers["scaling_bwd"] + + def reject_candidate(_candidate): + assert ( + module._quantization_runtime is old_runtime + ) # pylint: disable=protected-access + raise RuntimeError("candidate validation failure") + + monkeypatch.setattr(module, "_validate_quantization_runtime", reject_candidate) + with pytest.raises(RuntimeError, match="candidate validation failure"): + _ensure_runtime(module, candidate_recipe, revision=2) + + assert ( + module._quantization_runtime is old_runtime + ) # pylint: disable=protected-access + assert module.fp8_meta["scaling_fwd"] is old_forward_state + assert module.fp8_meta["scaling_bwd"] is old_backward_state + assert module.quantizers["scaling_fwd"] is old_forward_quantizers + assert module.quantizers["scaling_bwd"] is old_backward_quantizers + + +@pytest.mark.parametrize("mismatched_tensor_type", ("input", "grad_output")) +def test_grouped_candidate_validation_is_atomic(mismatched_tensor_type): + """Forward and backward grouped validation failures preserve all live state.""" + + class UnsafeIdentityQuantizer(IdentityQuantizer): + def is_requantization_safe(self): + return False + + def make_recipe( + key, + *, + mismatched_role=None, + dtype=torch.bfloat16, + unsafe_inputs=False, + ): + matching_role_count = 0 + + def qfactory(role): + nonlocal matching_role_count + quantizer_dtype = dtype + if role is not None and role.tensor_type == mismatched_role: + quantizer_dtype = ( + torch.bfloat16 if matching_role_count % 2 == 0 else torch.float16 + ) + matching_role_count += 1 + quantizer_type = ( + UnsafeIdentityQuantizer + if unsafe_inputs and role is not None and role.tensor_type == "input" + else IdentityQuantizer + ) + return quantizer_type(dtype=quantizer_dtype) + + return CustomRecipe(qfactory=qfactory, qfactory_key=key) + + module = GroupedLinear(2, 16, 16, bias=False, device="cuda", name="grouped") + active_recipe = make_recipe(("grouped-atomic", "active", mismatched_tensor_type)) + _ensure_runtime(module, active_recipe, revision=1, num_gemms=2) + + old_runtime = module._quantization_runtime # pylint: disable=protected-access + old_forward_state = module.fp8_meta["scaling_fwd"] + old_backward_state = module.fp8_meta["scaling_bwd"] + old_forward_quantizers = module.quantizers["scaling_fwd"] + old_backward_quantizers = module.quantizers["scaling_bwd"] + old_validated_generations = module._validated_quantizer_generations + old_delayed_quantizer = module._delayed_scaling_input_quantizer + old_unsafe_quantizer = module._unsafe_requantization_input_quantizer + + invalid_recipe = make_recipe( + ("grouped-atomic", "invalid", mismatched_tensor_type), + mismatched_role=mismatched_tensor_type, + ) + with pytest.raises(ValueError, match="incompatible plain backend configurations"): + _ensure_runtime(module, invalid_recipe, revision=2, num_gemms=2) + + assert ( + module._quantization_runtime is old_runtime + ) # pylint: disable=protected-access + assert module.fp8_meta["scaling_fwd"] is old_forward_state + assert module.fp8_meta["scaling_bwd"] is old_backward_state + assert module.quantizers["scaling_fwd"] is old_forward_quantizers + assert module.quantizers["scaling_bwd"] is old_backward_quantizers + assert module._validated_quantizer_generations is old_validated_generations + assert module._delayed_scaling_input_quantizer is old_delayed_quantizer + assert module._unsafe_requantization_input_quantizer is old_unsafe_quantizer + + replacement_recipe = make_recipe( + ("grouped-atomic", "replacement", mismatched_tensor_type), + dtype=torch.float16, + unsafe_inputs=True, + ) + assert _ensure_runtime(module, replacement_recipe, revision=3, num_gemms=2) + replacement_runtime = ( + module._quantization_runtime + ) # pylint: disable=protected-access + assert replacement_runtime is not old_runtime + assert ( + module._validated_quantizer_generations["scaling_fwd"] + is replacement_runtime.forward_quantizers + ) + assert ( + module._validated_quantizer_generations["scaling_bwd"] + is replacement_runtime.backward_quantizers + ) + assert module._delayed_scaling_input_quantizer is None + assert ( + module._unsafe_requantization_input_quantizer + is replacement_runtime.forward_quantizers[0] + ) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 4c61f8e68a3..ade11413286 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -10,7 +10,7 @@ import warnings from enum import Enum from abc import ABC, abstractmethod -from typing import Any, Dict, Generator, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, Hashable, List, Optional, Tuple, Union from contextlib import contextmanager from types import MethodType @@ -28,15 +28,11 @@ unsafe_pickle_extra_state_enabled, ) from ..quantization import ( - MXFP8BlockScalingRecipeState, - DelayedScalingRecipeState, - Float8CurrentScalingRecipeState, - Float8BlockScalingRecipeState, - NVFP4BlockScalingRecipeState, - CustomRecipeState, FP8GlobalStateManager, QuantizerRole, RecipeState, + _QuantizationRuntime, + _QuantizationRuntimeKey, _has_delayed_scaling_state, ) from ..distributed import ( @@ -918,6 +914,8 @@ def __init__(self, name: Optional[str] = None) -> None: self.wgrad_store = None self._output_quantizer_role: Optional[QuantizerRole] = None self._grad_input_quantizer_role: Optional[QuantizerRole] = None + self._role_revision = 0 + self._quantization_runtime: Optional[_QuantizationRuntime] = None if not TEDebugState.debug_enabled: TEDebugState.initialize() @@ -954,11 +952,7 @@ def output_quantizer_role(self) -> Optional[QuantizerRole]: @output_quantizer_role.setter def output_quantizer_role(self, role: Optional[QuantizerRole]) -> None: - if role == self._output_quantizer_role: - return - self._output_quantizer_role = role - if self.fp8_meta_tensors_initialized: - self.fp8_meta_tensors_initialized = False + self._set_quantizer_role("_output_quantizer_role", role) @property def grad_input_quantizer_role(self) -> Optional[QuantizerRole]: @@ -970,11 +964,20 @@ def grad_input_quantizer_role(self) -> Optional[QuantizerRole]: @grad_input_quantizer_role.setter def grad_input_quantizer_role(self, role: Optional[QuantizerRole]) -> None: - if role == self._grad_input_quantizer_role: + self._set_quantizer_role("_grad_input_quantizer_role", role) + + def _set_quantizer_role( + self, + attribute: str, + role: Optional[QuantizerRole], + ) -> None: + """Update a requested boundary role without mutating the active runtime.""" + if role is not None and not isinstance(role, QuantizerRole): + raise TypeError(f"{attribute} must be a QuantizerRole or None") + if getattr(self, attribute) == role: return - self._grad_input_quantizer_role = role - if self.fp8_meta_tensors_initialized: - self.fp8_meta_tensors_initialized = False + self.fast_setattr(attribute, role) + self.fast_setattr("_role_revision", getattr(self, "_role_revision", 0) + 1) def _warn_missing_output_quantizer_role( self, @@ -1074,76 +1077,230 @@ def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> meta_key ].amax_history[0] - def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: - """Init scales and amaxes for fwd | bwd.""" - fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd" - - # Return early if recipe state matches recipe - if self.fp8_meta_tensors_initialized: - recipe_state = self.fp8_meta[fp8_meta_tensor_key] - # TODO(#3157): Match built-in recipes by full config, not just RecipeState type, so - # same-class mid-training changes rebuild quantizers/workspaces correctly. - if recipe.delayed() and isinstance(recipe_state, DelayedScalingRecipeState): - self.adjust_amax_history_length(recipe.amax_history_len, fwd=fwd) - return - if recipe.mxfp8() and isinstance(recipe_state, MXFP8BlockScalingRecipeState): - return - if recipe.float8_current_scaling() and isinstance( - recipe_state, Float8CurrentScalingRecipeState - ): - return - if recipe.float8_block_scaling() and isinstance( - recipe_state, Float8BlockScalingRecipeState - ): - return - if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): - return - if recipe.custom() and isinstance(recipe_state, CustomRecipeState): - # TODO(#3157): Compare CustomRecipe/qfactory config here. qfactory changes made - # mid-training on the same recipe object currently do not take effect because - # stale quantizers are reused. - if recipe_state.recipe is recipe: - return - - # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and - # 2 (grad_output and grad_input) for bwd - num_fp8_tensors = self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2 - - # Initialize recipe state and quantizers + def _resolve_quantizer_roles( + self, + *, + fwd: bool, + num_quantizers: int, + ) -> Tuple[ + Tuple[Optional[QuantizerRole], ...], + Optional[List[Optional[QuantizerRole]]], + ]: + """Resolve semantic key roles and the role list passed to ``RecipeState``.""" roles = self.get_quantizer_roles( # pylint: disable=assignment-from-none - fwd=fwd, num_quantizers=num_fp8_tensors + fwd=fwd, + num_quantizers=num_quantizers, ) - if roles is not None: - assert ( - len(roles) == num_fp8_tensors - ), f"Recipe roles must match number of quantizers ({len(roles)=} vs {num_fp8_tensors=})" - recipe_state = RecipeState.create( # pylint: disable=assignment-from-none + if roles is None: + # ``RecipeState`` treats a missing role list as one bare role per + # slot. Encode that effective behavior in the runtime key while + # preserving ``None`` for the state so CustomRecipe keeps its + # existing missing-role warning. + return tuple(QuantizerRole() for _ in range(num_quantizers)), None + if len(roles) != num_quantizers: + raise ValueError( + "Recipe roles must match number of quantizers " + f"({len(roles)=} vs {num_quantizers=})" + ) + state_roles = list(roles) + return tuple(state_roles), state_roles + + def _prepare_quantization_runtime( + self, + *, + recipe: Recipe, + key: _QuantizationRuntimeKey, + num_gemms: int, + recipe_config_revision: int, + role_revision: int, + forward_state_roles: Optional[List[Optional[QuantizerRole]]], + backward_state_roles: Optional[List[Optional[QuantizerRole]]], + ) -> _QuantizationRuntime: + """Build both quantization directions without changing live module state.""" + num_forward_quantizers = num_gemms * 3 + num_backward_quantizers = num_gemms * 2 + + forward_state = RecipeState.create( # pylint: disable=assignment-from-none recipe, - mode=("forward" if fwd else "backward"), - num_quantizers=num_fp8_tensors, - roles=roles, + mode="forward", + num_quantizers=num_forward_quantizers, + roles=forward_state_roles, + ) + backward_state = RecipeState.create( # pylint: disable=assignment-from-none + recipe, + mode="backward", + num_quantizers=num_backward_quantizers, + roles=backward_state_roles, ) - # Reached the rebuild path because ``fp8_meta_tensors_initialized`` - # was flipped to False after first init — most commonly because the - # ``output_quantizer_role`` / ``grad_input_quantizer_role`` setter - # invalidated state when a parent module (e.g. ``MultiheadAttention``) - # wired boundary roles. That setter is recipe-agnostic, so this code - # fires even for built-in recipes that don't consume role information - # in ``make_quantizers``. - # - # Rebuilding the recipe state must preserve persistent training - # buffers (delayed-scaling ``scale`` / ``amax_history``) so the new - # quantizer instances and the ``FP8GlobalStateManager`` reduction - # buffers end up viewing the SAME tensor objects, and so any - # checkpoint-loaded state isn't silently destroyed on the first - # forward after ``load_state_dict``. - old_state = self.fp8_meta.get(fp8_meta_tensor_key) - if old_state is not None: - recipe_state.inherit_state_from(old_state) - - self.fp8_meta[fp8_meta_tensor_key] = recipe_state - self.quantizers[fp8_meta_tensor_key] = recipe_state.make_quantizers() + active = getattr(self, "_quantization_runtime", None) + if active is not None: + old_forward_state = active.forward_states[0] + old_backward_state = active.backward_states[0] + else: + # Checkpoint loading and legacy direct initialization can populate + # compatibility metadata before the first runtime is active. + old_forward_state = self.fp8_meta.get("scaling_fwd") + old_backward_state = self.fp8_meta.get("scaling_bwd") + if old_forward_state is not None: + forward_state.inherit_state_from(old_forward_state) + if old_backward_state is not None: + backward_state.inherit_state_from(old_backward_state) + + # A failure in either direction discards this entire local candidate. + forward_quantizers = forward_state.make_quantizers() + backward_quantizers = backward_state.make_quantizers() + + return _QuantizationRuntime( + key=key, + num_gemms=num_gemms, + recipe_config_revision=recipe_config_revision, + role_revision=role_revision, + forward_states=(forward_state,), + backward_states=(backward_state,), + forward_quantizers=forward_quantizers, + backward_quantizers=backward_quantizers, + delayed_registration_handles=(), + ) + + def _validate_quantization_runtime(self, candidate: _QuantizationRuntime) -> Any: + """Validate a candidate and return optional transient commit data.""" + return None + + def _commit_quantization_runtime( + self, + candidate: _QuantizationRuntime, + *, + recipe: Recipe, + validation_result: Any = None, + ) -> None: + """Publish a fully prepared and validated runtime as one unit.""" + del validation_result + forward_state = candidate.forward_states[0] + backward_state = candidate.backward_states[0] + + # ``fp8_meta`` and ``quantizers`` remain compatibility views. Only + # update them after both directions have been built and validated. + self.fp8_meta["recipe"] = recipe + self.fp8_meta["num_gemms"] = candidate.num_gemms + self.fp8_meta["scaling_fwd"] = forward_state + self.fp8_meta["scaling_bwd"] = backward_state + if hasattr(recipe, "fp8_format"): + self.fp8_meta["fp8_max_fwd"] = recipe.fp8_format.value.max_fwd + self.fp8_meta["fp8_max_bwd"] = recipe.fp8_format.value.max_bwd + self.quantizers["scaling_fwd"] = candidate.forward_quantizers + self.quantizers["scaling_bwd"] = candidate.backward_quantizers + self.fast_setattr("_quantization_runtime", candidate) + self.fast_setattr("fp8_meta_tensors_initialized", True) + + # Clear cached weight workspaces only after the replacement runtime has + # committed fully. A later forward will recreate them with its quantizers. + workspaces = getattr(self, "_fp8_workspaces", None) + if workspaces is not None: + workspaces.clear() + + def _ensure_active_quantization_runtime( + self, + *, + recipe: Recipe, + recipe_config: Hashable, + recipe_config_revision: int, + num_gemms: int, + ) -> bool: + """Ensure the active runtime matches the requested recipe, roles, and slot layout. + + Return whether a new runtime was committed. + """ + active = getattr(self, "_quantization_runtime", None) + role_revision = getattr(self, "_role_revision", 0) + + # Constant-time steady-state path: no role construction, factory + # calls, recipe traversal, or candidate validation. + if ( + active is not None + and active.recipe_config_revision == recipe_config_revision + and active.role_revision == role_revision + and active.num_gemms == num_gemms + ): + return False + + num_forward_quantizers = num_gemms * 3 + num_backward_quantizers = num_gemms * 2 + forward_key_roles, forward_state_roles = self._resolve_quantizer_roles( + fwd=True, + num_quantizers=num_forward_quantizers, + ) + backward_key_roles, backward_state_roles = self._resolve_quantizer_roles( + fwd=False, + num_quantizers=num_backward_quantizers, + ) + requested_key = _QuantizationRuntimeKey( + recipe_config=recipe_config, + forward_roles=forward_key_roles, + backward_roles=backward_key_roles, + ) + + # A module can miss an A -> B -> A sequence. Synchronize its cheap + # revision markers without rebuilding an equal semantic runtime. + if ( + active is not None + and active.key == requested_key + and active.num_gemms == num_gemms + ): + active.recipe_config_revision = recipe_config_revision + active.role_revision = role_revision + return False + + # Until per-weight semantic keys are available, a recipe change may + # not retarget already-quantized primary storage. Role-only changes + # are safe here because the module's weight-slot roles are fixed; the + # requested setters only affect output/grad-input boundary slots. + if ( + active is not None + and getattr(self, "primary_weights_in_fp8", False) + and ( + active.key.recipe_config != requested_key.recipe_config + or active.num_gemms != num_gemms + ) + ): + raise RuntimeError( + "Recipe mismatch for quantized primary weights: mid-training recipe " + "configuration changes require per-weight semantic compatibility and " + "are not supported yet." + ) + + candidate = self._prepare_quantization_runtime( + recipe=recipe, + key=requested_key, + num_gemms=num_gemms, + recipe_config_revision=recipe_config_revision, + role_revision=role_revision, + forward_state_roles=forward_state_roles, + backward_state_roles=backward_state_roles, + ) + validation_result = self._validate_quantization_runtime(candidate) + self._commit_quantization_runtime( + candidate, + recipe=recipe, + validation_result=validation_result, + ) + return True + + def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: + """Initialize both quantization directions atomically. + + ``fwd`` is retained for API compatibility. A runtime is never prepared + or published one direction at a time. + """ + # TODO(negvet): Remove set_meta_tensor after checkpoint, GroupedLinear, and DPA migration. + del fwd + num_gemms = self.fp8_meta.get("num_gemms", 1) + self._ensure_active_quantization_runtime( + recipe=recipe, + recipe_config=recipe.quantizer_config(), + recipe_config_revision=-1, + num_gemms=num_gemms, + ) def get_quantizer_roles( self, @@ -1509,51 +1666,36 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: fp8_enabled = fp8 or fp8_calibration meta["fp8_checkpoint"] = fp8_enabled - _original_recipe = None - - if fp8_parameters or fp8_enabled: - _original_recipe = meta.get("recipe", None) - if self.fp8_initialized and FP8GlobalStateManager.get_fp8_recipe() == _original_recipe: - # FP8 init has already been run and recipe is the same, don't do anything. - return - meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() - else: + if not (fp8_parameters or fp8_enabled): # If fp8 isn't enabled, turn off and return. self.fast_setattr("fp8_initialized", False) return - if fp8_parameters and not self.fp8_initialized: - meta["num_gemms"] = num_gemms - self.init_fp8_meta_tensors(meta["recipe"]) - + original_recipe = meta.get("recipe") + recipe = FP8GlobalStateManager.get_fp8_recipe() + recipe_config = FP8GlobalStateManager.get_quantizer_config() + recipe_config_revision = FP8GlobalStateManager.get_quantizer_config_revision() if fp8_enabled: - # Set FP8 and other FP8 metadata - meta["num_gemms"] = num_gemms meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() - # Set FP8_MAX per tensor according to recipe - if hasattr(meta["recipe"], "fp8_format"): - meta["fp8_max_fwd"] = meta["recipe"].fp8_format.value.max_fwd - meta["fp8_max_bwd"] = meta["recipe"].fp8_format.value.max_bwd - - # Allocate scales and amaxes - self.init_fp8_meta_tensors(meta["recipe"]) + runtime_changed = self._ensure_active_quantization_runtime( + recipe=recipe, + recipe_config=recipe_config, + recipe_config_revision=recipe_config_revision, + num_gemms=num_gemms, + ) + if fp8_enabled: self.fast_setattr("fp8_initialized", True) - meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() - - _current_recipe = meta["recipe"] - if _original_recipe is not None and not ( - issubclass(_current_recipe.__class__, _original_recipe.__class__) - or issubclass(_original_recipe.__class__, _current_recipe.__class__) + if runtime_changed and original_recipe is not None and not ( + issubclass(recipe.__class__, original_recipe.__class__) + or issubclass(original_recipe.__class__, recipe.__class__) ): warnings.warn( - f"Recipe type changed from {_original_recipe.__class__.__name__} " - f"to {_current_recipe.__class__.__name__}. " + f"Recipe type changed from {original_recipe.__class__.__name__} " + f"to {recipe.__class__.__name__}. " "This may affect model behavior." ) - # Clear cached workspaces as they were created with the old recipe/quantizer type - self._fp8_workspaces.clear() def prepare_forward( self, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 12497d0cb0a..29ac818254b 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -32,7 +32,12 @@ _get_high_precision_init_val, ) from ._common import can_reconstruct_wgrad_input_from_original, WeightGradStore -from ..quantization import FP8GlobalStateManager, QuantizerRole +from ..quantization import ( + FP8GlobalStateManager, + Float8CurrentScalingRecipeState, + QuantizerRole, + _QuantizationRuntime, +) from ..utils import ( divide, cast_if_needed, @@ -1886,14 +1891,116 @@ def __init__( def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Init scales and amaxes for fwd | bwd.""" + if recipe.float8_current_scaling() and self.tp_size > 1: + raise ValueError( + "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " + "Because the TP communication is handled outside of this module." + ) super().set_meta_tensor(fwd, recipe) + self._validate_quantizer_generation(fwd) + + def _validate_forward_quantizer_generation( + self, + generation: List[Quantizer], + *, + num_gemms: int, + ) -> Tuple[Optional[Quantizer], Optional[Quantizer]]: + """Validate candidate forward quantizers and derive grouped runtime state.""" + stride = self._num_fp8_tensors_per_gemm["fwd"] + input_quantizers = tuple( + generation[self._offsets["input"] + i * stride] for i in range(num_gemms) + ) + weight_quantizers = tuple( + generation[self._offsets["weight"] + i * stride] for i in range(num_gemms) + ) + _validate_grouped_quantizer_list(input_quantizers, operand_name="input") + _validate_grouped_quantizer_list(weight_quantizers, operand_name="weight") + delayed_scaling_input_quantizer = next( + (q for q in input_quantizers if isinstance(q, Float8Quantizer)), + None, + ) + unsafe_requantization_input_quantizer = next( + ( + q + for q in input_quantizers + if q is not None and not can_reconstruct_wgrad_input_from_original(q) + ), + None, + ) + return ( + delayed_scaling_input_quantizer, + unsafe_requantization_input_quantizer, + ) - # Recipe-specific quantizer configuration - recipe = FP8GlobalStateManager.get_fp8_recipe() - if recipe.float8_current_scaling(): - self._customize_quantizers_float8_current_scaling(fwd, recipe) + def _validate_backward_quantizer_generation( + self, + generation: List[Quantizer], + *, + num_gemms: int, + ) -> None: + """Validate candidate backward quantizers for grouped-kernel compatibility.""" + stride = self._num_fp8_tensors_per_gemm["bwd"] + grad_output_quantizers = tuple( + generation[self._offsets["grad_output"] + i * stride] for i in range(num_gemms) + ) + _validate_grouped_quantizer_list( + grad_output_quantizers, + operand_name="grad_output", + ) - self._validate_quantizer_generation(fwd) + def _validate_quantization_runtime( + self, + candidate: _QuantizationRuntime, + ) -> Tuple[Optional[Quantizer], Optional[Quantizer]]: + """Validate both grouped directions without changing live module state.""" + super()._validate_quantization_runtime(candidate) + if ( + isinstance(candidate.forward_states[0], Float8CurrentScalingRecipeState) + and self.tp_size > 1 + ): + raise ValueError( + "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " + "Because the TP communication is handled outside of this module." + ) + + validation_result = self._validate_forward_quantizer_generation( + candidate.forward_quantizers, + num_gemms=candidate.num_gemms, + ) + self._validate_backward_quantizer_generation( + candidate.backward_quantizers, + num_gemms=candidate.num_gemms, + ) + return validation_result + + def _commit_quantization_runtime( + self, + candidate: _QuantizationRuntime, + *, + recipe: Recipe, + validation_result: Tuple[Optional[Quantizer], Optional[Quantizer]], + ) -> None: + """Publish validated grouped-derived state with the candidate runtime.""" + delayed_scaling_input_quantizer, unsafe_requantization_input_quantizer = validation_result + validated_generations = { + "scaling_fwd": candidate.forward_quantizers, + "scaling_bwd": candidate.backward_quantizers, + } + + super()._commit_quantization_runtime( + candidate, + recipe=recipe, + validation_result=validation_result, + ) + self.fast_setattr( + "_delayed_scaling_input_quantizer", + delayed_scaling_input_quantizer, + ) + self.fast_setattr( + "_unsafe_requantization_input_quantizer", + unsafe_requantization_input_quantizer, + ) + self.fast_setattr("_validated_quantizer_generations", validated_generations) def _validate_quantizer_generation(self, fwd: bool) -> None: """Validate grouped-kernel invariants once per quantizer generation.""" @@ -1909,37 +2016,19 @@ def _validate_quantizer_generation(self, fwd: bool) -> None: return if fwd: - stride = self._num_fp8_tensors_per_gemm["fwd"] - input_quantizers = tuple( - generation[self._offsets["input"] + i * stride] for i in range(self.num_gemms) - ) - weight_quantizers = tuple( - generation[self._offsets["weight"] + i * stride] for i in range(self.num_gemms) - ) - _validate_grouped_quantizer_list(input_quantizers, operand_name="input") - _validate_grouped_quantizer_list(weight_quantizers, operand_name="weight") - delayed_scaling_input_quantizer = next( - (q for q in input_quantizers if isinstance(q, Float8Quantizer)), - None, - ) - unsafe_requantization_input_quantizer = next( - ( - q - for q in input_quantizers - if q is not None and not can_reconstruct_wgrad_input_from_original(q) - ), - None, + ( + delayed_scaling_input_quantizer, + unsafe_requantization_input_quantizer, + ) = self._validate_forward_quantizer_generation( + generation, + num_gemms=self.num_gemms, ) self._delayed_scaling_input_quantizer = delayed_scaling_input_quantizer self._unsafe_requantization_input_quantizer = unsafe_requantization_input_quantizer else: - stride = self._num_fp8_tensors_per_gemm["bwd"] - grad_output_quantizers = tuple( - generation[self._offsets["grad_output"] + i * stride] for i in range(self.num_gemms) - ) - _validate_grouped_quantizer_list( - grad_output_quantizers, - operand_name="grad_output", + self._validate_backward_quantizer_generation( + generation, + num_gemms=self.num_gemms, ) self._validated_quantizer_generations[meta_key] = generation @@ -2453,41 +2542,6 @@ def backward_dw(self): del tensor_list self._trigger_wgrad_accumulation_and_reduce_hooks() - def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None: - """Customize quantizers based on current scaling recipe + linear.""" - - if self.tp_size > 1: - raise ValueError( - "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " - "Because the TP communication is handled outside of this module." - ) - - if fwd: - for i in range(self.num_gemms): - # set configs about amax epsilon and power_2_scale - self.quantizers["scaling_fwd"][ - self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["fwd"] - ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale - self.quantizers["scaling_fwd"][ - self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["fwd"] - ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon - # also set weight quantizer with same amax_epsilon & power_2_scale - self.quantizers["scaling_fwd"][ - self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"] - ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale - self.quantizers["scaling_fwd"][ - self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"] - ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon - else: - for i in range(self.num_gemms): - # set grad_output_quantizer with amax epsilon and power_2_scale - self.quantizers["scaling_bwd"][ - self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["bwd"] - ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale - self.quantizers["scaling_bwd"][ - self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["bwd"] - ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon - def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]: """Get the weight tensors of the module.""" grouped_weight = getattr(self, "weight", None) @@ -2535,9 +2589,8 @@ def _get_weight_quantizers(self) -> List[Quantizer]: def _get_quantizers(self): if self.fp8: - # Normally validated while installing recipe metadata. Keep this - # O(1) generation guard so failed transitions cannot reuse stale - # validation state if base metadata takes an early return on retry. + # Normally validated before candidate commit. Keep this O(1) + # generation guard for legacy/direct quantizer-list replacement. self._validate_quantizer_generation(True) if torch.is_grad_enabled(): self._validate_quantizer_generation(False) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 561e813348c..d80c4a3d7cc 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1595,15 +1595,6 @@ def __init__( if name in self.weight_names or name in self.bias_names: param.skip_backward_post_hook = True - def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: - """Init scales and amaxes for fwd | bwd.""" - super().set_meta_tensor(fwd, recipe) - - # Recipe-specific quantizer configuration - recipe = FP8GlobalStateManager.get_fp8_recipe() - if recipe.float8_current_scaling(): - self._customize_quantizers_float8_current_scaling(fwd, recipe) - def get_quantizer_roles( self, *, @@ -1948,35 +1939,6 @@ def onnx_forward( return output, bias_tensor.to(inp_dtype) return output - def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None: - """Customize quantizers based on current scaling recipe + layernorm_linear.""" - assert ( - recipe.float8_current_scaling() - ), "current scaling recipe quantizer customization here" - if fwd: - # set configs about amax epsilon and power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_INPUT - ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_INPUT - ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon - # also set weight quantizer with same amax_epsilon & power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_WEIGHT - ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_WEIGHT - ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon - else: - # set grad_output_quantizer with amax epsilon and power_2_scale (no amax reduction here) - self.quantizers["scaling_bwd"][ - FP8BwdTensorIdx.GRAD_OUTPUT1 - ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale - self.quantizers["scaling_bwd"][ - FP8BwdTensorIdx.GRAD_OUTPUT1 - ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon - def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]: """Get the weight tensors of the module.""" unfused_weights = [getattr(self, name) for name in self.weight_names] diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 3ee0cda50c5..83080eb14f4 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2173,15 +2173,6 @@ def __init__( self.bwd_ln_sm_margin = int(os.getenv("NVTE_BWD_LAYERNORM_SM_MARGIN", "0")) self.inf_ln_sm_margin = int(os.getenv("NVTE_INF_LAYERNORM_SM_MARGIN", "0")) - def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: - """Init scales and amaxes for fwd | bwd.""" - super().set_meta_tensor(fwd, recipe) - - # Recipe-specific quantizer configuration - recipe = FP8GlobalStateManager.get_fp8_recipe() - if recipe.float8_current_scaling(): - self._customize_quantizers_float8_current_scaling(fwd, recipe) - def get_quantizer_roles( self, *, @@ -2665,56 +2656,6 @@ def make_debug(prefix, offset): return tuple(make_debug("fc1", 0) + make_debug("fc2", 6)) - def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None: - """Customize quantizers based on current scaling recipe + layernorm_mlp.""" - assert ( - recipe.float8_current_scaling() - ), "current scaling recipe quantizer customization here" - if fwd: - # fc1_input_quantizer: set configs about amax epsilon and power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_INPUT - ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_INPUT - ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon - # fc2_input_quantizer - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM2_INPUT - ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM2_INPUT - ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon - # fc1_weight_quantizer: also set numerical configs about weight - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_WEIGHT - ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_WEIGHT - ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon - # fc2_weight_quantizer - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM2_WEIGHT - ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM2_WEIGHT - ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon - else: - # fc2_grad_output_quantizer: set configs about amax epsilon and power_2_scale for fc2_grad_output_quantizer - self.quantizers["scaling_bwd"][ - FP8BwdTensorIdx.GRAD_OUTPUT2 - ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale - self.quantizers["scaling_bwd"][ - FP8BwdTensorIdx.GRAD_OUTPUT2 - ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon - # fc1_grad_output_quantizer: also set numerical configs for fc1_grad_output_quantizer - self.quantizers["scaling_bwd"][ - FP8BwdTensorIdx.GRAD_OUTPUT1 - ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale - self.quantizers["scaling_bwd"][ - FP8BwdTensorIdx.GRAD_OUTPUT1 - ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon - def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]: """Get the weight tensors of the module.""" return [self.fc1_weight, self.fc2_weight] diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 56622db5e6b..f773a499f41 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1839,15 +1839,6 @@ def get_quantizer_roles( ] return [base[i % len(base)] for i in range(num_quantizers)] - def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: - """Init scales and amaxes for fwd | bwd.""" - super().set_meta_tensor(fwd, recipe) - - # Recipe-specific quantizer configuration - recipe = FP8GlobalStateManager.get_fp8_recipe() - if recipe.float8_current_scaling(): - self._customize_quantizers_float8_current_scaling(fwd, recipe) - def reset_parameters(self, defer_init=False): super().reset_parameters(defer_init=defer_init) @@ -2199,35 +2190,6 @@ def onnx_forward( return output - def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None: - """Customize quantizers based on current scaling recipe + linear.""" - assert ( - recipe.float8_current_scaling() - ), "current scaling recipe quantizer customization here" - if fwd: - # set configs about amax epsilon and power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_INPUT - ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_INPUT - ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon - # also set weight quantizer with same amax_epsilon & power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_WEIGHT - ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale - self.quantizers["scaling_fwd"][ - FP8FwdTensorIdx.GEMM1_WEIGHT - ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon - else: - # set grad_output_quantizer with amax epsilon and power_2_scale - self.quantizers["scaling_bwd"][ - FP8BwdTensorIdx.GRAD_OUTPUT1 - ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale - self.quantizers["scaling_bwd"][ - FP8BwdTensorIdx.GRAD_OUTPUT1 - ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon - def _get_weight_quantizers(self) -> List[Quantizer]: """Get the weight quantizers of the module.""" if not self.fp8 and not self.fp8_calibration: diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index cb429055a4f..2ecac158290 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -391,18 +391,6 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: or is_quantized_tensor(weight) ) - # Recipe-specific configuration - # Note: This function may be called in base class constructor, - # before any basic linear attrs have been set. - if recipe is not None: - if recipe.float8_current_scaling(): - input_quantizer.force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale - input_quantizer.amax_epsilon_scales = recipe.fp8_quant_fwd_inp.amax_epsilon - weight_quantizer.force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale - weight_quantizer.amax_epsilon_scales = recipe.fp8_quant_fwd_weight.amax_epsilon - grad_output_quantizer.force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale - grad_output_quantizer.amax_epsilon_scales = recipe.fp8_quant_bwd_grad.amax_epsilon - # Update quantizer in quantized weight tensor if weight_quantizer is not None and is_quantized_tensor(weight): if weight._quantizer is not None: diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index be931829ea1..57e59a278e2 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -738,20 +738,6 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: or weight_is_quantized ) - # Recipe-specific configuration - if recipe is not None: - if recipe.float8_current_scaling(): - input_quantizer.force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale - input_quantizer.amax_epsilon_scales = recipe.fp8_quant_fwd_inp.amax_epsilon - weight_quantizer.force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale - weight_quantizer.amax_epsilon_scales = recipe.fp8_quant_fwd_weight.amax_epsilon - grad_output_quantizer.force_pow_2_scales = ( - recipe.fp8_quant_bwd_grad.power_2_scale - ) - grad_output_quantizer.amax_epsilon_scales = ( - recipe.fp8_quant_bwd_grad.amax_epsilon - ) - # Update quantizer in quantized weight tensor if weight_quantizer is not None and weight_is_quantized: # Get quantizer from weight tensor diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index f1fc52a4246..f99ed65f8d1 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Quantization utilities for TransformerEngine""" + from __future__ import annotations import abc @@ -12,12 +13,13 @@ from dataclasses import dataclass, field from contextlib import contextmanager from collections import deque -from typing import Callable, Hashable, List, Optional, Dict, Any, Tuple, Union +from typing import Callable, Hashable, List, Optional, Dict, Any, Tuple, TYPE_CHECKING, Union import torch import transformer_engine_torch as tex from transformer_engine.common.recipe import ( Recipe, + QParams, DelayedScaling, Format, MXFP8BlockScaling, @@ -31,6 +33,8 @@ from .utils import get_device_compute_capability from .jit import jit_fuser +if TYPE_CHECKING: + from .quantized_tensor import Quantizer __all__ = [ "autocast", @@ -96,6 +100,68 @@ def __str__(self) -> str: return "|".join(parts) if parts else "QuantizerRole()" +@dataclass(frozen=True) +class _QuantizationRuntimeKey: + """Immutable semantic request for one module or operation's quantizers. + + ``recipe_config`` is the recipe-owned semantic configuration cached by + :class:`FP8GlobalStateManager`. The role tuples preserve the ordered + forward and backward quantizer-slot layouts for this particular runtime + owner. ``None`` is a meaningful boundary slot and is therefore retained + rather than filtered out. + + This is intentionally private: runtime owners use it to decide whether a + candidate runtime would construct the same quantizers as the active one. + It must contain only immutable semantic values, never recipe, factory, + state, or quantizer object identity. + """ + + recipe_config: Hashable + forward_roles: Tuple[Optional[QuantizerRole], ...] + backward_roles: Tuple[Optional[QuantizerRole], ...] + + def __post_init__(self) -> None: + """Canonicalize role containers and reject non-semantic slot values.""" + try: + hash(self.recipe_config) + except TypeError as exc: + raise TypeError("_QuantizationRuntimeKey.recipe_config must be hashable") from exc + + for attribute in ("forward_roles", "backward_roles"): + roles = tuple(getattr(self, attribute)) + if any(role is not None and not isinstance(role, QuantizerRole) for role in roles): + raise TypeError( + f"_QuantizationRuntimeKey.{attribute} entries must be QuantizerRole or None" + ) + object.__setattr__(self, attribute, roles) + + +@dataclass +class _QuantizationRuntime: + """Complete active-or-candidate quantization state for one runtime owner. + + A module or operation prepares this bundle away from its active state, + validates it, and then publishes it as one unit. Keeping forward and + backward states and quantizers together prevents a failed update from + exposing a new forward configuration with stale backward state. + + ``delayed_registration_handles`` is intentionally opaque because delayed + scaling registration is owned by :class:`FP8GlobalStateManager`. Runtime + owners replace or unregister those handles only during a successful + lifecycle transition. + """ + + key: _QuantizationRuntimeKey + num_gemms: int + recipe_config_revision: int + role_revision: int + forward_states: Tuple["RecipeState", ...] + backward_states: Tuple["RecipeState", ...] + forward_quantizers: List["Quantizer"] + backward_quantizers: List["Quantizer"] + delayed_registration_handles: Tuple[Any, ...] + + @dataclasses.dataclass(frozen=True) class QuantizerRequest: """Base class for stateful quantizer requests. @@ -397,6 +463,7 @@ class FP8GlobalState: fp8_calibration: bool = False fp8_recipe: Optional[Recipe] = None quantizer_config: Optional[Hashable] = None + quantizer_config_revision: int = 0 fp8_distributed_group: Optional[dist_group_type] = None fp8_parameters: bool = False high_precision_init_val: bool = False @@ -611,13 +678,25 @@ def get_fp8_recipe(cls) -> Recipe: return cls.quantization_state.fp8_recipe return get_default_fp8_recipe() + @classmethod + def _set_recipe_config( + cls, + recipe: Optional[Recipe], + quantizer_config: Optional[Hashable], + ) -> None: + """Publish a recipe/configuration pair and advance its monotonic revision.""" + qstate = cls.quantization_state + config_changed = quantizer_config != qstate.quantizer_config + if config_changed: + qstate.quantizer_config_revision += 1 + qstate.quantizer_config = quantizer_config + qstate.fp8_recipe = recipe + @classmethod def activate_recipe(cls, recipe: Recipe) -> Hashable: """Make a recipe and its semantic quantizer configuration active together.""" quantizer_config = recipe.quantizer_config() - qstate = cls.quantization_state - qstate.fp8_recipe = recipe - qstate.quantizer_config = quantizer_config + cls._set_recipe_config(recipe, quantizer_config) return quantizer_config @classmethod @@ -628,6 +707,11 @@ def get_quantizer_config(cls) -> Hashable: return cls.activate_recipe(cls.get_fp8_recipe()) return quantizer_config + @classmethod + def get_quantizer_config_revision(cls) -> int: + """Return the manager-owned active recipe-configuration revision.""" + return cls.quantization_state.quantizer_config_revision + @classmethod def get_fp8_group(cls) -> Union[dist_group_type, None]: """Return the fp8 group for scale/amax comm""" @@ -649,17 +733,23 @@ def get_autocast_state(cls) -> tuple: @classmethod def set_autocast_state(cls, state: tuple) -> None: - """Restore a previously saved autocast state snapshot.""" + """Restore an autocast snapshot without restoring its old configuration revision.""" qstate = cls.quantization_state ( - qstate.fp8_enabled, - qstate.fp8_calibration, - qstate.fp8_recipe, - qstate.quantizer_config, - qstate.fp8_distributed_group, - qstate.is_first_fp8_module, - qstate.fp8_graph_capturing, + fp8_enabled, + fp8_calibration, + fp8_recipe, + quantizer_config, + fp8_distributed_group, + is_first_fp8_module, + fp8_graph_capturing, ) = state + cls._set_recipe_config(fp8_recipe, quantizer_config) + qstate.fp8_enabled = fp8_enabled + qstate.fp8_calibration = fp8_calibration + qstate.fp8_distributed_group = fp8_distributed_group + qstate.is_first_fp8_module = is_first_fp8_module + qstate.fp8_graph_capturing = fp8_graph_capturing @staticmethod def reduce_tensor_across_group_op_max(tensor: torch.Tensor, group: dist_group_type) -> None: @@ -955,8 +1045,7 @@ def quantized_model_init( yield finally: qstate.fp8_parameters = _fp8_parameters - qstate.fp8_recipe = _fp8_recipe - qstate.quantizer_config = _quantizer_config + FP8GlobalStateManager._set_recipe_config(_fp8_recipe, _quantizer_config) qstate.high_precision_init_val = _high_precision_init_val @@ -1235,6 +1324,12 @@ class RecipeState(abc.ABC): _FWD_DEFAULT_TENSOR_TYPES = ("input", "weight", "output") _BWD_DEFAULT_TENSOR_TYPES = ("grad_output", "grad_input") + @staticmethod + def _validate_mode(mode: str) -> None: + """Validate the quantization direction shared by every recipe state.""" + if mode not in ("forward", "backward"): + raise ValueError(f"Unexpected recipe mode ({mode})") + @staticmethod def _validate_roles( roles: Optional[List[QuantizerRole]], @@ -1313,6 +1408,29 @@ def _slot_tensor_type(self, idx: int) -> str: ) return default_tensor_types[idx % len(default_tensor_types)] + @staticmethod + def _qparams_for_tensor_type( + tensor_type: str, + *, + input_qparams: QParams, + weight_qparams: QParams, + grad_qparams: QParams, + output_uses_input_qparams: bool = True, + grad_input_uses_grad_qparams: bool = True, + ) -> Optional[QParams]: + """Select a recipe qparam bundle for a canonical tensor type.""" + if tensor_type == "weight": + return weight_qparams + if tensor_type == "input" or ( + tensor_type == "output" and output_uses_input_qparams + ): + return input_qparams + if tensor_type == "grad_output" or ( + tensor_type == "grad_input" and grad_input_uses_grad_qparams + ): + return grad_qparams + return None + @staticmethod def create( recipe: Recipe, @@ -1445,6 +1563,7 @@ def __init__( device: Optional[torch.device] = None, roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_mode(mode) self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode @@ -1494,6 +1613,7 @@ def __init__( device: Optional[torch.device] = None, roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_mode(mode) self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode @@ -1507,14 +1627,38 @@ def __init__( self.device = device def make_quantizers(self) -> list: + """Build one current-scaling quantizer per slot, dispatched by tensor type. + + Input, weight, and grad-output slots use their corresponding recipe + qparams. Output and grad-input boundary slots retain the legacy + constructor defaults. Missing or non-canonical roles use the common + positional fallback provided by :meth:`RecipeState._slot_tensor_type`. + """ from .tensor.float8_tensor import Float8CurrentScalingQuantizer - return [ - Float8CurrentScalingQuantizer( - self.dtype, device=self.device, force_pow_2_scales=self.recipe.use_power_2_scales + def _make(tensor_type: str) -> Float8CurrentScalingQuantizer: + qparams = self._qparams_for_tensor_type( + tensor_type, + input_qparams=self.recipe.fp8_quant_fwd_inp, + weight_qparams=self.recipe.fp8_quant_fwd_weight, + grad_qparams=self.recipe.fp8_quant_bwd_grad, + output_uses_input_qparams=False, + grad_input_uses_grad_qparams=False, + ) + + force_pow_2_scales = self.recipe.use_power_2_scales + amax_epsilon = 0.0 + if qparams is not None: + force_pow_2_scales = qparams.power_2_scale + amax_epsilon = qparams.amax_epsilon + return Float8CurrentScalingQuantizer( + self.dtype, + device=self.device, + force_pow_2_scales=force_pow_2_scales, + amax_epsilon=amax_epsilon, ) - for i in range(self.num_quantizers) - ] + + return [_make(self._slot_tensor_type(idx)) for idx in range(self.num_quantizers)] class MXFP8BlockScalingRecipeState(RecipeState): @@ -1537,6 +1681,7 @@ def __init__( device: Optional[torch.device] = None, roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_mode(mode) self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode @@ -1577,6 +1722,7 @@ def __init__( device: Optional[torch.device] = None, roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_mode(mode) self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode @@ -1615,18 +1761,22 @@ def make_quantizers(self) -> list: from .tensor.float8_blockwise_tensor import Float8BlockQuantizer def _make(tensor_type: str) -> Float8BlockQuantizer: + qparams = self._qparams_for_tensor_type( + tensor_type, + input_qparams=self.recipe.fp8_quant_fwd_inp, + weight_qparams=self.recipe.fp8_quant_fwd_weight, + grad_qparams=self.recipe.fp8_quant_bwd_grad, + ) + assert qparams is not None if tensor_type == "weight": - qparams = self.recipe.fp8_quant_fwd_weight fp8_dtype = self.qw_dtype block_scaling_dim = self.recipe.w_block_scaling_dim elif tensor_type in ("grad_output", "grad_input"): - qparams = self.recipe.fp8_quant_bwd_grad fp8_dtype = self.qgrad_dtype block_scaling_dim = self.recipe.grad_block_scaling_dim else: # "input", "output", or any unknown forward type fall back to # the input config, matching the legacy positional behavior. - qparams = self.recipe.fp8_quant_fwd_inp fp8_dtype = self.qx_dtype block_scaling_dim = self.recipe.x_block_scaling_dim return Float8BlockQuantizer( @@ -1638,7 +1788,6 @@ def _make(tensor_type: str) -> Float8BlockQuantizer: block_scaling_dim=block_scaling_dim, ) - assert self.mode in ("forward", "backward"), f"Unexpected mode {self.mode}" return [_make(self._slot_tensor_type(idx)) for idx in range(self.num_quantizers)] @@ -1662,6 +1811,7 @@ def __init__( device: Optional[torch.device] = None, roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_mode(mode) self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode @@ -1696,15 +1846,14 @@ def make_quantizers(self) -> list: """ from .tensor.nvfp4_tensor import NVFP4Quantizer - def _qparams(tensor_type: str): - if tensor_type in ("grad_output", "grad_input"): - return self.recipe.fp4_quant_bwd_grad - if tensor_type == "weight": - return self.recipe.fp4_quant_fwd_weight - return self.recipe.fp4_quant_fwd_inp - def _make(tensor_type: str) -> NVFP4Quantizer: - qparams = _qparams(tensor_type) + qparams = self._qparams_for_tensor_type( + tensor_type, + input_qparams=self.recipe.fp4_quant_fwd_inp, + weight_qparams=self.recipe.fp4_quant_fwd_weight, + grad_qparams=self.recipe.fp4_quant_bwd_grad, + ) + assert qparams is not None nvfp4_use_4over6 = False if tensor_type not in ("grad_output", "grad_input"): if self.recipe.nvfp4_4over6 == "all": @@ -1751,9 +1900,6 @@ def _make(tensor_type: str) -> NVFP4Quantizer: nvfp4_4over6_err_mode=self.recipe.nvfp4_4over6_err_mode, ) - if self.mode not in ("forward", "backward"): - raise RuntimeError(f"Unexpected recipe mode ({self.mode})") - return [_make(self._slot_tensor_type(idx)) for idx in range(self.num_quantizers)] @@ -1896,6 +2042,7 @@ def __init__( device: Optional[torch.device] = None, roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_mode(mode) self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode From 4eac680081c5f2dae8cc88507a42ceb1afed4ff2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:48:37 +0000 Subject: [PATCH 68/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 3 +- tests/pytorch/test_recipe.py | 16 +-- tests/pytorch/test_recipe_update.py | 113 ++++++--------------- transformer_engine/pytorch/module/base.py | 19 ++-- transformer_engine/pytorch/quantization.py | 4 +- 5 files changed, 45 insertions(+), 110 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 94787c78033..9bc46a0c6b3 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -4690,8 +4690,7 @@ def mixed_source_qfactory(role): assert model._quantization_runtime is rebuilt_runtime assert model._validated_quantizer_generations["scaling_fwd"] is rebuilt_generation assert ( - model._validated_quantizer_generations["scaling_bwd"] - is rebuilt_backward_generation + model._validated_quantizer_generations["scaling_bwd"] is rebuilt_backward_generation ) assert model._delayed_scaling_input_quantizer is rebuilt_delayed_quantizer assert model._unsafe_requantization_input_quantizer is rebuilt_unsafe_quantizer diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 9e802c0b650..c3ef8c974a9 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -369,12 +369,8 @@ def settings(mode, roles): assert settings("forward", forward_roles) == [inp, weight, boundary] assert settings("backward", backward_roles) == [grad, boundary] - mlp_forward_roles = LayerNormMLP.get_quantizer_roles( - module_owner, fwd=True, num_quantizers=6 - ) - mlp_backward_roles = LayerNormMLP.get_quantizer_roles( - module_owner, fwd=False, num_quantizers=4 - ) + mlp_forward_roles = LayerNormMLP.get_quantizer_roles(module_owner, fwd=True, num_quantizers=6) + mlp_backward_roles = LayerNormMLP.get_quantizer_roles(module_owner, fwd=False, num_quantizers=4) assert settings("forward", mlp_forward_roles) == [ inp, weight, @@ -407,12 +403,8 @@ def settings(mode, roles): assert settings("forward", basic_forward_roles) == [inp, weight] assert settings("backward", basic_backward_roles) == [grad] - basic_grouped_forward_roles = te_ops.GroupedLinear.get_quantizer_roles( - basic_owner, "forward" - ) - basic_grouped_backward_roles = te_ops.GroupedLinear.get_quantizer_roles( - basic_owner, "backward" - ) + basic_grouped_forward_roles = te_ops.GroupedLinear.get_quantizer_roles(basic_owner, "forward") + basic_grouped_backward_roles = te_ops.GroupedLinear.get_quantizer_roles(basic_owner, "backward") assert settings("forward", basic_grouped_forward_roles) == [inp, weight, inp, weight] assert settings("backward", basic_grouped_backward_roles) == [grad, grad] diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 16cc9115a19..111cde94b36 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -18,9 +18,7 @@ from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="CUDA is required" -) +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def _make_counting_recipe(key, calls, *, fail_on_grad_output=False): @@ -28,11 +26,7 @@ def _make_counting_recipe(key, calls, *, fail_on_grad_output=False): def qfactory(role): calls.append(role) - if ( - fail_on_grad_output - and role is not None - and role.tensor_type == "grad_output" - ): + if fail_on_grad_output and role is not None and role.tensor_type == "grad_output": raise RuntimeError("backward factory failure") return IdentityQuantizer() @@ -40,13 +34,11 @@ def qfactory(role): def _ensure_runtime(module, recipe, revision, *, num_gemms=1): - return ( - module._ensure_active_quantization_runtime( # pylint: disable=protected-access - recipe=recipe, - recipe_config=recipe.quantizer_config(), - recipe_config_revision=revision, - num_gemms=num_gemms, - ) + return module._ensure_active_quantization_runtime( # pylint: disable=protected-access + recipe=recipe, + recipe_config=recipe.quantizer_config(), + recipe_config_revision=revision, + num_gemms=num_gemms, ) @@ -102,26 +94,16 @@ def unexpected_cold_path(*_args, **_kwargs): raise AssertionError("unchanged forward entered the runtime cold path") monkeypatch.setattr(module, "get_quantizer_roles", unexpected_cold_path) - monkeypatch.setattr( - module, "_prepare_quantization_runtime", unexpected_cold_path - ) - monkeypatch.setattr( - module, "_validate_quantization_runtime", unexpected_cold_path - ) - monkeypatch.setattr( - module, "_commit_quantization_runtime", unexpected_cold_path - ) + monkeypatch.setattr(module, "_prepare_quantization_runtime", unexpected_cold_path) + monkeypatch.setattr(module, "_validate_quantization_runtime", unexpected_cold_path) + monkeypatch.setattr(module, "_commit_quantization_runtime", unexpected_cold_path) workspace_sentinel = object() - module._fp8_workspaces["sentinel"] = ( - workspace_sentinel # pylint: disable=protected-access - ) + module._fp8_workspaces["sentinel"] = workspace_sentinel # pylint: disable=protected-access module(inp) module(inp) - assert ( - module._quantization_runtime is active - ) # pylint: disable=protected-access + assert module._quantization_runtime is active # pylint: disable=protected-access assert len(calls) == factory_call_count assert ( module._fp8_workspaces["sentinel"] is workspace_sentinel @@ -138,9 +120,7 @@ def test_same_recipe_object_semantic_mutation_rebuilds_runtime(): recipe.qfactory_key = ("same-object-update", 2) assert _ensure_runtime(module, recipe, revision=2) - assert ( - module._quantization_runtime is not old_runtime - ) # pylint: disable=protected-access + assert module._quantization_runtime is not old_runtime # pylint: disable=protected-access assert len(calls) == 10 @@ -158,12 +138,8 @@ def test_role_revision_is_requested_until_atomic_runtime_commit(): role = QuantizerRole(module_type="linear", tensor_type="input", name="consumer") module.output_quantizer_role = role - assert ( - module._role_revision == old_role_revision + 1 - ) # pylint: disable=protected-access - assert ( - module._quantization_runtime is old_runtime - ) # pylint: disable=protected-access + assert module._role_revision == old_role_revision + 1 # pylint: disable=protected-access + assert module._quantization_runtime is old_runtime # pylint: disable=protected-access assert module.fp8_meta["scaling_fwd"] is old_forward_state assert module.quantizers["scaling_fwd"] is old_forward_quantizers @@ -179,9 +155,7 @@ def test_role_revision_is_requested_until_atomic_runtime_commit(): module.output_quantizer_role = QuantizerRole( module_type="linear", tensor_type="input", name="consumer" ) - assert ( - module._role_revision == old_role_revision + 1 - ) # pylint: disable=protected-access + assert module._role_revision == old_role_revision + 1 # pylint: disable=protected-access assert not _ensure_runtime(module, recipe, revision=1) assert len(calls) == 10 @@ -208,9 +182,7 @@ def test_backward_factory_failure_keeps_complete_active_runtime(): with pytest.raises(RuntimeError, match="backward factory failure"): _ensure_runtime(module, failing_recipe, revision=2) - assert ( - module._quantization_runtime is old_runtime - ) # pylint: disable=protected-access + assert module._quantization_runtime is old_runtime # pylint: disable=protected-access assert module.fp8_meta["recipe"] is old_recipe assert module.fp8_meta["scaling_fwd"] is old_forward_state assert module.fp8_meta["scaling_bwd"] is old_backward_state @@ -295,26 +267,18 @@ def test_runtime_update_workspace_lifecycle( ) workspaces = {cache_name: object() for cache_name in cache_names} for cache_name, workspace in workspaces.items(): - module._fp8_workspaces[cache_name] = ( - workspace # pylint: disable=protected-access - ) + module._fp8_workspaces[cache_name] = workspace # pylint: disable=protected-access # A global revision change with an equal semantic runtime preserves caches. assert not _ensure_runtime(module, equal_recipe, revision=2, num_gemms=num_gemms) - assert ( - module._quantization_runtime is active_runtime - ) # pylint: disable=protected-access + assert module._quantization_runtime is active_runtime # pylint: disable=protected-access for cache_name, workspace in workspaces.items(): - assert ( - module._fp8_workspaces[cache_name] is workspace - ) # pylint: disable=protected-access + assert module._fp8_workspaces[cache_name] is workspace # pylint: disable=protected-access # Candidate construction failure preserves both compatibility views and caches. with pytest.raises(RuntimeError, match="backward factory failure"): _ensure_runtime(module, failing_recipe, revision=3, num_gemms=num_gemms) - assert ( - module._quantization_runtime is active_runtime - ) # pylint: disable=protected-access + assert module._quantization_runtime is active_runtime # pylint: disable=protected-access current_views = ( module.fp8_meta["recipe"], module.fp8_meta["scaling_fwd"], @@ -322,19 +286,13 @@ def test_runtime_update_workspace_lifecycle( module.quantizers["scaling_fwd"], module.quantizers["scaling_bwd"], ) - assert all( - current is active for current, active in zip(current_views, active_views) - ) + assert all(current is active for current, active in zip(current_views, active_views)) for cache_name, workspace in workspaces.items(): - assert ( - module._fp8_workspaces[cache_name] is workspace - ) # pylint: disable=protected-access + assert module._fp8_workspaces[cache_name] is workspace # pylint: disable=protected-access # Only a fully committed replacement clears cached workspaces. assert _ensure_runtime(module, replacement_recipe, revision=4, num_gemms=num_gemms) - replacement_runtime = ( - module._quantization_runtime - ) # pylint: disable=protected-access + replacement_runtime = module._quantization_runtime # pylint: disable=protected-access assert replacement_runtime is not active_runtime assert not module._fp8_workspaces # pylint: disable=protected-access @@ -354,18 +312,14 @@ def test_candidate_validation_failure_keeps_complete_active_runtime(monkeypatch) old_backward_quantizers = module.quantizers["scaling_bwd"] def reject_candidate(_candidate): - assert ( - module._quantization_runtime is old_runtime - ) # pylint: disable=protected-access + assert module._quantization_runtime is old_runtime # pylint: disable=protected-access raise RuntimeError("candidate validation failure") monkeypatch.setattr(module, "_validate_quantization_runtime", reject_candidate) with pytest.raises(RuntimeError, match="candidate validation failure"): _ensure_runtime(module, candidate_recipe, revision=2) - assert ( - module._quantization_runtime is old_runtime - ) # pylint: disable=protected-access + assert module._quantization_runtime is old_runtime # pylint: disable=protected-access assert module.fp8_meta["scaling_fwd"] is old_forward_state assert module.fp8_meta["scaling_bwd"] is old_backward_state assert module.quantizers["scaling_fwd"] is old_forward_quantizers @@ -393,9 +347,7 @@ def qfactory(role): nonlocal matching_role_count quantizer_dtype = dtype if role is not None and role.tensor_type == mismatched_role: - quantizer_dtype = ( - torch.bfloat16 if matching_role_count % 2 == 0 else torch.float16 - ) + quantizer_dtype = torch.bfloat16 if matching_role_count % 2 == 0 else torch.float16 matching_role_count += 1 quantizer_type = ( UnsafeIdentityQuantizer @@ -426,9 +378,7 @@ def qfactory(role): with pytest.raises(ValueError, match="incompatible plain backend configurations"): _ensure_runtime(module, invalid_recipe, revision=2, num_gemms=2) - assert ( - module._quantization_runtime is old_runtime - ) # pylint: disable=protected-access + assert module._quantization_runtime is old_runtime # pylint: disable=protected-access assert module.fp8_meta["scaling_fwd"] is old_forward_state assert module.fp8_meta["scaling_bwd"] is old_backward_state assert module.quantizers["scaling_fwd"] is old_forward_quantizers @@ -443,9 +393,7 @@ def qfactory(role): unsafe_inputs=True, ) assert _ensure_runtime(module, replacement_recipe, revision=3, num_gemms=2) - replacement_runtime = ( - module._quantization_runtime - ) # pylint: disable=protected-access + replacement_runtime = module._quantization_runtime # pylint: disable=protected-access assert replacement_runtime is not old_runtime assert ( module._validated_quantizer_generations["scaling_fwd"] @@ -457,6 +405,5 @@ def qfactory(role): ) assert module._delayed_scaling_input_quantizer is None assert ( - module._unsafe_requantization_input_quantizer - is replacement_runtime.forward_quantizers[0] + module._unsafe_requantization_input_quantizer is replacement_runtime.forward_quantizers[0] ) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 50214663d9e..f28c049373b 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1128,8 +1128,7 @@ def _resolve_quantizer_roles( return tuple(QuantizerRole() for _ in range(num_quantizers)), None if len(roles) != num_quantizers: raise ValueError( - "Recipe roles must match number of quantizers " - f"({len(roles)=} vs {num_quantizers=})" + f"Recipe roles must match number of quantizers ({len(roles)=} vs {num_quantizers=})" ) state_roles = list(roles) return tuple(state_roles), state_roles @@ -1271,11 +1270,7 @@ def _ensure_active_quantization_runtime( # A module can miss an A -> B -> A sequence. Synchronize its cheap # revision markers without rebuilding an equal semantic runtime. - if ( - active is not None - and active.key == requested_key - and active.num_gemms == num_gemms - ): + if active is not None and active.key == requested_key and active.num_gemms == num_gemms: active.recipe_config_revision = recipe_config_revision active.role_revision = role_revision return False @@ -1716,9 +1711,13 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: if fp8_enabled: self.fast_setattr("fp8_initialized", True) - if runtime_changed and original_recipe is not None and not ( - issubclass(recipe.__class__, original_recipe.__class__) - or issubclass(original_recipe.__class__, recipe.__class__) + if ( + runtime_changed + and original_recipe is not None + and not ( + issubclass(recipe.__class__, original_recipe.__class__) + or issubclass(original_recipe.__class__, recipe.__class__) + ) ): warnings.warn( f"Recipe type changed from {original_recipe.__class__.__name__} " diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index f99ed65f8d1..ab06405f31b 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1421,9 +1421,7 @@ def _qparams_for_tensor_type( """Select a recipe qparam bundle for a canonical tensor type.""" if tensor_type == "weight": return weight_qparams - if tensor_type == "input" or ( - tensor_type == "output" and output_uses_input_qparams - ): + if tensor_type == "input" or (tensor_type == "output" and output_uses_input_qparams): return input_qparams if tensor_type == "grad_output" or ( tensor_type == "grad_input" and grad_input_uses_grad_qparams From 160ec058aad952e5d1d9d0f80616053658f453ea Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 11 Aug 2026 10:08:40 +0000 Subject: [PATCH 69/97] CI fixes Signed-off-by: Evgeny --- docs/api/common.rst | 2 +- tests/pytorch/test_recipe_update.py | 2 +- transformer_engine/common/recipe/__init__.py | 3 ++- transformer_engine/pytorch/module/base.py | 8 ++++++-- transformer_engine/pytorch/module/layernorm_linear.py | 1 - transformer_engine/pytorch/module/linear.py | 1 - transformer_engine/pytorch/quantization.py | 5 +++-- 7 files changed, 13 insertions(+), 9 deletions(-) diff --git a/docs/api/common.rst b/docs/api/common.rst index 8009dc81e94..0c849285674 100644 --- a/docs/api/common.rst +++ b/docs/api/common.rst @@ -18,6 +18,6 @@ Common API .. autoapiclass:: transformer_engine.common.recipe.Float8BlockScaling(fp8_format=Format.E4M3) -.. autofunction:: transformer_engine.common.recipe.quantizer_factory +.. autoapifunction:: transformer_engine.common.recipe.quantizer_factory .. autoapiclass:: transformer_engine.common.recipe.CustomRecipe(qfactory, fp8_dpa=False, fp8_mha=False, qfactory_key=None) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 111cde94b36..91c2cbfc6f7 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 871868b91c0..9a32d71d5b9 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -6,8 +6,9 @@ from __future__ import annotations import abc import os +from collections.abc import Hashable from enum import Enum -from typing import Any, Callable, Hashable, Literal, NamedTuple, Optional, Union +from typing import Any, Callable, Literal, NamedTuple, Optional, Union from dataclasses import field from pydantic.dataclasses import dataclass diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index f28c049373b..c05dcf79305 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -8,9 +8,10 @@ import os import pickle import warnings +from collections.abc import Hashable from enum import Enum from abc import ABC, abstractmethod -from typing import Any, Dict, Generator, Hashable, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, List, Optional, Tuple, Union from contextlib import contextmanager from types import MethodType @@ -1193,7 +1194,7 @@ def _prepare_quantization_runtime( def _validate_quantization_runtime(self, candidate: _QuantizationRuntime) -> Any: """Validate a candidate and return optional transient commit data.""" - return None + del candidate def _commit_quantization_runtime( self, @@ -1302,6 +1303,9 @@ def _ensure_active_quantization_runtime( forward_state_roles=forward_state_roles, backward_state_roles=backward_state_roles, ) + # Subclasses may return transient commit data even though the base implementation + # returns None, which pylint cannot infer through dynamic dispatch. + # pylint: disable-next=assignment-from-no-return validation_result = self._validate_quantization_runtime(candidate) self._commit_quantization_runtime( candidate, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index d80c4a3d7cc..427a00e2b70 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -15,7 +15,6 @@ import transformer_engine_torch as tex -from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom from .base import ( diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index f773a499f41..5246c9b3c22 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -14,7 +14,6 @@ import transformer_engine_torch as tex -from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version from .base import ( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index ab06405f31b..e34a1c3bb2f 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -10,10 +10,11 @@ import dataclasses import warnings import os +from collections import deque +from collections.abc import Hashable from dataclasses import dataclass, field from contextlib import contextmanager -from collections import deque -from typing import Callable, Hashable, List, Optional, Dict, Any, Tuple, TYPE_CHECKING, Union +from typing import Callable, List, Optional, Dict, Any, Tuple, TYPE_CHECKING, Union import torch import transformer_engine_torch as tex From fa854319cf25cf1029a77ff4355cbbf537dbfeb8 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 11 Aug 2026 13:28:47 +0000 Subject: [PATCH 70/97] Exclude delayed scaling from recipe update Signed-off-by: Evgeny --- tests/pytorch/test_custom_recipe.py | 70 +++----- tests/pytorch/test_hybrid_quantization.py | 13 +- tests/pytorch/test_recipe.py | 16 +- tests/pytorch/test_recipe_update.py | 198 ++++++++++++++++++++- transformer_engine/pytorch/module/base.py | 55 +++++- transformer_engine/pytorch/quantization.py | 84 +++++++-- 6 files changed, 350 insertions(+), 86 deletions(-) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 36734487dbc..761a3e075e9 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -978,29 +978,18 @@ def ds_factory(role): # ---------------------------------------------------------------------- -# State preservation across role-driven rebuilds +# Stateful recipe boundary for role-driven updates # ---------------------------------------------------------------------- # # Setting ``output_quantizer_role`` / ``grad_input_quantizer_role`` to a # different value advances the requested role revision without mutating the -# active runtime. The next runtime update rebuilds the recipe state and -# quantizers with up-to-date roles. That rebuild MUST preserve persistent training -# buffers (delayed scaling's ``scale`` / ``amax_history``); otherwise -# checkpointed amax history is silently destroyed on the first forward -# pass after ``load_state_dict`` (when MHA wires boundary roles for the -# first time on the freshly-loaded module). The buffers must also be -# preserved by tensor-object identity, not just by value: the -# ``FP8GlobalStateManager`` reduction buffer holds a direct reference to -# the tensor created at first init, so any rebuild that allocates fresh -# tensors would break amax all-reduce. - - -def test_role_change_preserves_delayed_scaling_state(): - """Built-in DelayedScaling: role-driven rebuild preserves scale / amax_history. - - Stashes sentinel values into the buffers, forces a rebuild via the role - setter, and verifies values + tensor-object identity survive. - """ +# active runtime. Stateless recipes rebuild on the next update. Delayed-scaling +# runtimes are frozen in round one, so the update must reject and leave all +# persistent buffers and global registration references untouched. + + +def test_role_change_rejects_delayed_scaling_runtime_update(): + """Built-in DelayedScaling rejects a role-driven runtime rebuild.""" available, reason = te.is_fp8_available(return_reason=True) if not torch.cuda.is_available() or not available: pytest.skip(f"FP8 unsupported: {reason}") @@ -1035,33 +1024,23 @@ def test_role_change_preserves_delayed_scaling_state(): assert model._role_revision == role_revision_before + 1 assert model.fp8_meta["scaling_fwd"] is state_before - # Trigger the rebuild directly (no forward, so we can compare buffers exactly). - model.init_fp8_meta_tensors(fp8_recipe) + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + model.init_fp8_meta_tensors(fp8_recipe) assert model.fp8_meta_tensors_initialized state_after = model.fp8_meta["scaling_fwd"] - assert state_after is not state_before, "state should have been rebuilt" - # Tensor objects must be inherited (not freshly allocated) so the - # FP8GlobalStateManager reduction buffer's reference stays valid. - assert ( - id(state_after.scale) == scale_obj_id - ), "scale tensor object replaced by rebuild; global reduction buffer would dangle" + assert model._quantization_runtime is runtime_before + assert state_after is state_before + assert id(state_after.scale) == scale_obj_id assert id(state_after.amax_history) == amax_obj_id assert state_after.scale.data_ptr() == scale_data_ptr assert state_after.amax_history.data_ptr() == amax_data_ptr - # Sentinel values must be preserved. - assert state_after.scale.eq(3.14).all(), "scale was wiped by role-driven rebuild" - assert state_after.amax_history.eq(2.71).all(), "amax_history was wiped" + assert state_after.scale.eq(3.14).all() + assert state_after.amax_history.eq(2.71).all() -def test_role_change_preserves_custom_delayed_scaling_state(): - """CustomRecipe + DelayedScalingRequest: role-driven rebuild preserves inner DSRS. - - Same property as the built-in case, but for the - ``CustomRecipeState`` -> composed ``DelayedScalingRecipeState`` path. - The inner DS state must be re-used across the rebuild so its - accumulated buffers (and any external references to them) survive. - """ +def test_role_change_rejects_custom_delayed_scaling_runtime_update(): + """CustomRecipe with delayed requests rejects a role-driven rebuild.""" available, reason = te.is_fp8_available(return_reason=True) if not torch.cuda.is_available() or not available: pytest.skip(f"FP8 unsupported: {reason}") @@ -1106,20 +1085,17 @@ def ds_factory(role): assert model._role_revision == role_revision_before + 1 assert model.fp8_meta["scaling_fwd"] is state_before - # Rebuild. - model.init_fp8_meta_tensors(custom_recipe) + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + model.init_fp8_meta_tensors(custom_recipe) assert model.fp8_meta_tensors_initialized state_after = model.fp8_meta["scaling_fwd"] assert isinstance(state_after, CustomRecipeState) - assert state_after is not state_before, "outer CustomRecipeState should have been rebuilt" - assert state_after._has_delayed_scaling, "rebuild lost the inner DS state" + assert model._quantization_runtime is runtime_before + assert state_after is state_before + assert state_after._has_delayed_scaling inner_after = state_after._ds_state - # Inner DSRS object identity is preserved (we reuse the existing inner state), - # which means its buffers' tensor objects are also preserved. - assert ( - inner_after is inner_before - ), "inner DSRS replaced; FP8GlobalStateManager reduction buffer would dangle" + assert inner_after is inner_before assert id(inner_after.scale) == scale_obj_id assert id(inner_after.amax_history) == amax_obj_id # Sentinel values preserved. diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 9bc46a0c6b3..86d43eae114 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -2129,8 +2129,9 @@ def counting_qfactory(role): assert dpa._qkv_capabilities_quantizer is canonical_qkv assert len(calls) == calls_after_first - # A recipe-state rebuild creates a new canonical slot and must - # invalidate the capability cache automatically. + # A supported recipe-state rebuild creates a new canonical slot and + # must invalidate the capability cache automatically. Delayed state is + # frozen and rejects the rebuild before invoking the new factory. def rebuilt_qfactory(role): return counting_qfactory(role) @@ -2140,6 +2141,14 @@ def rebuilt_qfactory(role): fp8_mha=True, qfactory_key=("test_dpa_capability_rebuilt", factory_name, 1), ) + if factory_name == "delayed_scaling_factory": + with autocast(enabled=True, recipe=rebuilt_recipe): + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + dpa.get_qkv_quantization_capabilities() + assert dpa._qkv_capabilities_quantizer is canonical_qkv + assert len(calls) == calls_after_first + return + with autocast(enabled=True, recipe=rebuilt_recipe): rebuilt = dpa.get_qkv_quantization_capabilities() rebuilt_qkv = dpa._qkv_capabilities_quantizer diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index c3ef8c974a9..bd21adc26fe 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -16,8 +16,8 @@ import transformer_engine.pytorch as te from transformer_engine.pytorch import ( Float8BlockQuantizer, + Float8CurrentScalingQuantizer, MXFP8Quantizer, - Float8Quantizer, NVFP4Quantizer, quantized_model_init, Linear, @@ -162,13 +162,11 @@ def test_quantization_runtime_bundles_both_quantization_directions(): backward_states=(), forward_quantizers=forward_quantizers, backward_quantizers=backward_quantizers, - delayed_registration_handles=("forward-handle", "backward-handle"), ) assert runtime.key is runtime_key assert runtime.forward_quantizers is forward_quantizers assert runtime.backward_quantizers is backward_quantizers - assert runtime.delayed_registration_handles == ("forward-handle", "backward-handle") @pytest.mark.parametrize( @@ -1128,14 +1126,14 @@ def test_check_for_weight_tensor_and_recipe_correspondence(self, model_init_reci MXFP8Quantizer, mxfp8_available, reason_for_no_mxfp8, - id="DelayedScaling->MXFP8BlockScaling", + id="Float8CurrentScaling->MXFP8BlockScaling", ), pytest.param( Float8BlockScaling, Float8BlockQuantizer, fp8_block_scaling_available, reason_for_no_fp8_block_scaling, - id="DelayedScaling->Float8BlockScaling", + id="Float8CurrentScaling->Float8BlockScaling", ), ], ) @@ -1149,9 +1147,11 @@ def test_dynamic_recipe_update( out_features = 32 batch_size = 32 linear = Linear(in_features, out_features).cuda() - initial_recipe = DelayedScaling() + initial_recipe = Float8CurrentScaling() - # Run initial iterations with DelayedScaling + # Run initial iterations with a stateless recipe. Delayed-scaling + # enter/leave transitions are intentionally rejected by the lazy + # module-local update path and covered separately. for _ in range(3): x = torch.randn(batch_size, in_features, device="cuda") with te.autocast(enabled=True, recipe=initial_recipe): @@ -1160,7 +1160,7 @@ def test_dynamic_recipe_update( loss.backward() for quantizer in linear.quantizers["scaling_fwd"]: - assert isinstance(quantizer, Float8Quantizer) + assert isinstance(quantizer, Float8CurrentScalingQuantizer) # Change recipe target_recipe = target_recipe_class() diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 91c2cbfc6f7..add0ec3f135 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -7,15 +7,19 @@ import pytest import torch -from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.common.recipe import CustomRecipe, DelayedScaling, Float8CurrentScaling from transformer_engine.pytorch import ( GroupedLinear, LayerNormLinear, LayerNormMLP, Linear, autocast, + is_fp8_available, +) +from transformer_engine.pytorch.quantization import ( + DelayedScalingRequest, + QuantizerRole, ) -from transformer_engine.pytorch.quantization import QuantizerRole from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -42,7 +46,195 @@ def _ensure_runtime(module, recipe, revision, *, num_gemms=1): ) -def test_equal_runtime_key_synchronizes_revision_without_factory_calls(): +def _mixed_delayed_factory(role): + """Mix delayed/plain slots while keeping delayed state in both directions.""" + if role is not None and role.tensor_type in ("input", "weight", "grad_output"): + return DelayedScalingRequest(amax_history_len=4) + return IdentityQuantizer() + + +def _forward_only_delayed_factory(role): + """Request delayed state only in the forward direction.""" + if role is not None and role.tensor_type in ("input", "weight"): + return DelayedScalingRequest(amax_history_len=4) + return IdentityQuantizer() + + +def _runtime_views(module): + """Return identity-bearing views that a rejected update must preserve.""" + runtime = module._quantization_runtime # pylint: disable=protected-access + return ( + runtime, + module.fp8_meta["recipe"], + module.fp8_meta["scaling_fwd"], + module.fp8_meta["scaling_bwd"], + module.quantizers["scaling_fwd"], + module.quantizers["scaling_bwd"], + ) + + +@pytest.mark.parametrize( + "replacement", + ( + pytest.param(DelayedScaling(margin=1), id="margin"), + pytest.param(DelayedScaling(amax_history_len=8), id="history-length"), + pytest.param(Float8CurrentScaling(), id="leave-delayed-scaling"), + ), +) +def test_delayed_runtime_rejects_effective_recipe_updates_atomically(replacement): + """Delayed state is frozen once a module runtime has been initialized.""" + module = Linear(16, 16, bias=False, device="cuda", name="linear") + recipe = DelayedScaling(amax_history_len=4) + assert _ensure_runtime(module, recipe, revision=1) + old_views = _runtime_views(module) + + with pytest.raises( + RuntimeError, + match="Mid-training recipe updates do not support delayed scaling", + ): + _ensure_runtime(module, replacement, revision=2) + + assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + assert old_views[0].recipe_config_revision == 1 + + +def test_delayed_runtime_rejects_role_and_slot_layout_updates(): + """Role and GEMM-layout changes cannot rebuild a delayed runtime.""" + module = Linear(16, 16, bias=False, device="cuda", name="linear") + recipe = DelayedScaling(amax_history_len=4) + assert _ensure_runtime(module, recipe, revision=1) + old_views = _runtime_views(module) + + module.output_quantizer_role = QuantizerRole( + module_type="linear", + tensor_type="input", + name="consumer", + ) + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + _ensure_runtime(module, recipe, revision=1) + assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + _ensure_runtime(module, recipe, revision=1, num_gemms=2) + assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + + """An active stateless runtime cannot acquire delayed state mid-training.""" + module = Linear(16, 16, bias=False, device="cuda", name="linear") + assert _ensure_runtime(module, Float8CurrentScaling(), revision=1) + old_views = _runtime_views(module) + + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + _ensure_runtime(module, DelayedScaling(amax_history_len=4), revision=2) + assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + + +def test_custom_recipe_cannot_introduce_delayed_state(): + """A delayed request discovered in a CustomRecipe candidate cannot commit.""" + calls = [] + active_recipe = _make_counting_recipe(("custom-enter-delayed", 1), calls) + replacement = CustomRecipe( + qfactory=_mixed_delayed_factory, + qfactory_key=("custom-enter-delayed", 2), + ) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + assert _ensure_runtime(module, active_recipe, revision=1) + old_views = _runtime_views(module) + + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + _ensure_runtime(module, replacement, revision=2) + assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + + +def test_asymmetric_custom_delayed_scaling_is_rejected_clearly(): + """This PR does not add a delayed topology unsupported on main.""" + recipe = CustomRecipe( + qfactory=_forward_only_delayed_factory, + qfactory_key=("forward-only-delayed", 1), + ) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + + with pytest.raises( + RuntimeError, + match="This hybrid quantization configuration with delayed scaling is not supported", + ): + _ensure_runtime(module, recipe, revision=1) + assert module._quantization_runtime is None # pylint: disable=protected-access + + +def test_mixed_custom_recipe_is_frozen_when_it_contains_delayed_state(): + """Even a nominally non-delayed CustomRecipe edit is outside the contract.""" + calls = [] + + def active_factory(role): + calls.append(role) + return _mixed_delayed_factory(role) + + active_recipe = CustomRecipe(qfactory=active_factory, qfactory_key=("mixed", 1)) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + assert _ensure_runtime(module, active_recipe, revision=1) + old_views = _runtime_views(module) + active_call_count = len(calls) + + def replacement_factory(_role): + raise AssertionError("frozen delayed runtime invoked the replacement factory") + + replacement = CustomRecipe(qfactory=replacement_factory, qfactory_key=("mixed", 2)) + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + _ensure_runtime(module, replacement, revision=2) + + assert len(calls) == active_call_count + assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + + +def test_same_delayed_recipe_object_mutation_keeps_committed_snapshot(): + """Rejecting caller mutation must not mutate the active reduction recipe.""" + module = Linear(16, 16, bias=False, device="cuda", name="linear") + recipe = DelayedScaling(amax_history_len=4, margin=0) + assert _ensure_runtime(module, recipe, revision=1) + committed_recipe = module.fp8_meta["recipe"] + assert committed_recipe is not recipe + + recipe.margin = 1 + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + _ensure_runtime(module, recipe, revision=2) + + assert module.fp8_meta["recipe"] is committed_recipe + assert committed_recipe.margin == 0 + + +def test_rejected_delayed_update_aborts_autocast_reduction(): + """A caught activation failure cannot update the old delayed tensors on exit.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + module = Linear( + 16, + 16, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="linear", + ) + inp = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) + recipe = DelayedScaling(amax_history_len=4, margin=0) + with autocast(enabled=True, recipe=recipe): + module(inp) + + state = module.fp8_meta["scaling_fwd"] + state.scale.fill_(3) + state.amax_history.fill_(7) + expected_scale = state.scale.clone() + expected_history = state.amax_history.clone() + + recipe.margin = 1 + with autocast(enabled=True, recipe=recipe): + with pytest.raises(RuntimeError, match="do not support delayed scaling"): + module(inp) + + assert torch.equal(state.scale, expected_scale) + assert torch.equal(state.amax_history, expected_history) + """Equal independent recipes and missed A -> B -> A updates reuse the runtime.""" calls = [] first_recipe = _make_counting_recipe(("runtime-reuse", 1), calls) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index c05dcf79305..10c2021aafd 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Base modules and utilities for TransformerEngine PyTorch API""" +import copy import io import math import os @@ -35,6 +36,7 @@ _QuantizationRuntime, _QuantizationRuntimeKey, _has_delayed_scaling_state, + _is_delayed_scaling_state, ) from ..distributed import ( gather_along_first_dim, @@ -1149,14 +1151,20 @@ def _prepare_quantization_runtime( num_forward_quantizers = num_gemms * 3 num_backward_quantizers = num_gemms * 2 + # Recipe objects are caller-owned and mutable. States and delayed + # reduction buckets must retain the exact configuration that this + # candidate validated, even if the caller later mutates and reuses the + # same recipe instance. + runtime_recipe = copy.copy(recipe) + forward_state = RecipeState.create( # pylint: disable=assignment-from-none - recipe, + runtime_recipe, mode="forward", num_quantizers=num_forward_quantizers, roles=forward_state_roles, ) backward_state = RecipeState.create( # pylint: disable=assignment-from-none - recipe, + runtime_recipe, mode="backward", num_quantizers=num_backward_quantizers, roles=backward_state_roles, @@ -1180,6 +1188,12 @@ def _prepare_quantization_runtime( forward_quantizers = forward_state.make_quantizers() backward_quantizers = backward_state.make_quantizers() + if _is_delayed_scaling_state(forward_state) != _is_delayed_scaling_state(backward_state): + FP8GlobalStateManager.abort_current_amax_reduction() + raise RuntimeError( + "This hybrid quantization configuration with delayed scaling is not supported." + ) + return _QuantizationRuntime( key=key, num_gemms=num_gemms, @@ -1189,7 +1203,6 @@ def _prepare_quantization_runtime( backward_states=(backward_state,), forward_quantizers=forward_quantizers, backward_quantizers=backward_quantizers, - delayed_registration_handles=(), ) def _validate_quantization_runtime(self, candidate: _QuantizationRuntime) -> Any: @@ -1276,10 +1289,17 @@ def _ensure_active_quantization_runtime( active.role_revision = role_revision return False + def has_delayed_scaling(runtime: _QuantizationRuntime) -> bool: + return any( + _is_delayed_scaling_state(state) + for state in runtime.forward_states + runtime.backward_states + ) + + active_has_delayed_scaling = active is not None and has_delayed_scaling(active) + # Until per-weight semantic keys are available, a recipe change may - # not retarget already-quantized primary storage. Role-only changes - # are safe here because the module's weight-slot roles are fixed; the - # requested setters only affect output/grad-input boundary slots. + # not retarget already-quantized primary storage. Keep this more + # specific diagnostic ahead of the general delayed-scaling boundary. if ( active is not None and getattr(self, "primary_weights_in_fp8", False) @@ -1288,12 +1308,31 @@ def _ensure_active_quantization_runtime( or active.num_gemms != num_gemms ) ): + if active_has_delayed_scaling or recipe.delayed(): + FP8GlobalStateManager.abort_current_amax_reduction() raise RuntimeError( "Recipe mismatch for quantized primary weights: mid-training recipe " "configuration changes require per-weight semantic compatibility and " "are not supported yet." ) + def reject_delayed_scaling_update() -> None: + FP8GlobalStateManager.abort_current_amax_reduction() + raise RuntimeError( + "Mid-training recipe updates do not support delayed scaling. " + "A delayed-scaling runtime is frozen after initialization; only " + "unchanged execution is supported." + ) + + # Delayed state participates in global reduction registration and is + # intentionally outside this round's migration contract. Reject any + # effective update before invoking a CustomRecipe factory or preparing + # replacement tensors. + if active_has_delayed_scaling: + reject_delayed_scaling_update() + if active is not None and not active_has_delayed_scaling and recipe.delayed(): + reject_delayed_scaling_update() + candidate = self._prepare_quantization_runtime( recipe=recipe, key=requested_key, @@ -1303,13 +1342,15 @@ def _ensure_active_quantization_runtime( forward_state_roles=forward_state_roles, backward_state_roles=backward_state_roles, ) + if active is not None and has_delayed_scaling(candidate): + reject_delayed_scaling_update() # Subclasses may return transient commit data even though the base implementation # returns None, which pylint cannot infer through dynamic dispatch. # pylint: disable-next=assignment-from-no-return validation_result = self._validate_quantization_runtime(candidate) self._commit_quantization_runtime( candidate, - recipe=recipe, + recipe=candidate.forward_states[0].recipe, validation_result=validation_result, ) return True diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index e34a1c3bb2f..be1d6f32129 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -145,11 +145,6 @@ class _QuantizationRuntime: validates it, and then publishes it as one unit. Keeping forward and backward states and quantizers together prevents a failed update from exposing a new forward configuration with stale backward state. - - ``delayed_registration_handles`` is intentionally opaque because delayed - scaling registration is owned by :class:`FP8GlobalStateManager`. Runtime - owners replace or unregister those handles only during a successful - lifecycle transition. """ key: _QuantizationRuntimeKey @@ -160,7 +155,6 @@ class _QuantizationRuntime: backward_states: Tuple["RecipeState", ...] forward_quantizers: List["Quantizer"] backward_quantizers: List["Quantizer"] - delayed_registration_handles: Tuple[Any, ...] @dataclasses.dataclass(frozen=True) @@ -471,6 +465,7 @@ class FP8GlobalState: is_first_fp8_module: bool = False fp8_graph_capturing: bool = False autocast_depth: int = 0 + abort_amax_reduction: bool = False global_amax_buffer: Dict[str, list] = field(default_factory=dict) global_amax_history_buffer: Dict[str, list] = field(default_factory=dict) global_scale_buffer: Dict[str, list] = field(default_factory=dict) @@ -591,6 +586,17 @@ def add_fp8_tensors_to_global_buffer( if not _has_delayed_scaling_state(fp8_meta): return + forward_key = cls.get_meta_tensor_key(forward=True) + backward_key = cls.get_meta_tensor_key(forward=False) + if forward_key in fp8_meta and backward_key in fp8_meta: + forward_delayed = _is_delayed_scaling_state(fp8_meta[forward_key]) + backward_delayed = _is_delayed_scaling_state(fp8_meta[backward_key]) + if forward_delayed != backward_delayed: + cls.abort_current_amax_reduction() + raise RuntimeError( + "This hybrid quantization configuration with delayed scaling is not supported." + ) + # Every module must call this function exactly once since # the amax tensors are static. Ensures that compatibility # with non-graphed modules is maintained. @@ -608,15 +614,25 @@ def add_fp8_tensors_to_global_buffer( state = fp8_meta[fp8_meta_tensor_key] - # Determine recipe + buffers: built-in DS or custom with DS requests + # Determine recipe + buffers: built-in DS or custom with DS requests. if isinstance(state, CustomRecipeState) and state._has_delayed_scaling: - inner_recipe = state._inner_delayed_scaling_recipe - key = cls.get_key_in_buffer(forward, inner_recipe, fp8_meta["fp8_group"]) - # Register inner recipe in autocast_arguments for reduction - autocast_key = cls.get_unique_autocast_key(inner_recipe, fp8_meta["fp8_group"]) - qstate.autocast_arguments[autocast_key] = (inner_recipe, fp8_meta["fp8_group"]) + committed_recipe = state._inner_delayed_scaling_recipe else: - key = cls.get_key_in_buffer(forward, fp8_meta["recipe"], fp8_meta["fp8_group"]) + # The state owns the committed recipe snapshot. Register that + # snapshot rather than the caller-owned object installed by + # ``autocast_enter`` so later caller mutation cannot change an + # existing bucket's reduction behavior. + committed_recipe = state.recipe + + key = cls.get_key_in_buffer(forward, committed_recipe, fp8_meta["fp8_group"]) + autocast_key = cls.get_unique_autocast_key( + committed_recipe, + fp8_meta["fp8_group"], + ) + qstate.autocast_arguments[autocast_key] = ( + committed_recipe, + fp8_meta["fp8_group"], + ) if key not in qstate.global_amax_buffer: qstate.global_amax_buffer[key] = [fp8_meta[fp8_meta_tensor_key].amax_history[0]] @@ -771,6 +787,8 @@ def reduce_and_update_fp8_tensors( """Delayed scaling only. Concatenate, reduce, and split amaxes in the global buffer.""" # global_amax_buffer should only be non-empty for fp8 delayed scaling qstate = cls.quantization_state + if qstate.abort_amax_reduction: + return for ( buffer_key, amax_buffer, @@ -851,9 +869,12 @@ def autocast_enter( fp8_recipe = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe autocast_key = cls.get_unique_autocast_key(fp8_recipe, fp8_group) qstate = cls.quantization_state - qstate.autocast_arguments[autocast_key] = ( - fp8_recipe, - fp8_group, + # Once a delayed bucket is registered, its committed recipe snapshot + # owns the reduction semantics for this key. Do not replace it with a + # caller-owned recipe object merely by entering another autocast. + qstate.autocast_arguments.setdefault( + autocast_key, + (fp8_recipe, fp8_group), ) qstate.fp8_enabled = enabled @@ -863,6 +884,7 @@ def autocast_enter( qstate.fp8_graph_capturing = _graph if qstate.autocast_depth == 0: + qstate.abort_amax_reduction = False qstate.is_first_fp8_module = True qstate.autocast_depth += 1 @@ -879,17 +901,34 @@ def autocast_enter( nvfp4_available, reason_for_no_nvfp4 = cls.is_nvfp4_available() assert nvfp4_available, reason_for_no_nvfp4 + @classmethod + def abort_current_amax_reduction(cls) -> None: + """Prevent delayed-state updates when the active autocast has failed.""" + qstate = cls.quantization_state + if qstate.autocast_depth > 0: + qstate.abort_amax_reduction = True + @classmethod def autocast_exit(cls, enabled: bool, _graph: bool) -> None: """Set state and tracking variables for exit from FP8 region.""" qstate = cls.quantization_state qstate.autocast_depth -= 1 + outermost = qstate.autocast_depth == 0 # Reduce only the non-FP8 weight modules here. # FP8 weight modules are reduced at the end of the optimizer # step after the weight amax is populated. - if enabled and qstate.autocast_depth == 0 and not _graph and torch.is_grad_enabled(): - # delayed scaling only function, for other recipes (current scaling with any granularity), - # this is noop for other recipes because cls.global_amax_buffer is empty list + should_reduce = ( + enabled + and outermost + and not qstate.abort_amax_reduction + and not _graph + and torch.is_grad_enabled() + ) + if outermost: + qstate.abort_amax_reduction = False + if should_reduce: + # Delayed scaling only function. For other recipes this is a + # no-op because the global amax buffer is empty. cls.reduce_and_update_fp8_tensors(forward=True) @classmethod @@ -1989,6 +2028,13 @@ def _handle_delayed_scaling_requests( return dsrs +def _is_delayed_scaling_state(state: Any) -> bool: + """Return whether one direction owns delayed scale/history tensors.""" + return isinstance(state, DelayedScalingRecipeState) or ( + isinstance(state, CustomRecipeState) and state._has_delayed_scaling + ) + + def _has_delayed_scaling_state(fp8_meta: Dict[str, Any]) -> bool: """Check if fp8_meta has delayed scaling state (built-in or custom).""" if fp8_meta["recipe"].delayed(): From e33937d4d804c652bdf41450fc860af11848a2ec Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 11 Aug 2026 14:50:57 +0000 Subject: [PATCH 71/97] DPA support Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 130 ++++++++++++++++++ tests/pytorch/test_recipe.py | 3 + .../dot_product_attention.py | 30 ++++ transformer_engine/pytorch/module/base.py | 4 +- .../pytorch/module/grouped_linear.py | 2 - transformer_engine/pytorch/quantization.py | 5 +- 6 files changed, 169 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 86d43eae114..b0525dda622 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -134,6 +134,136 @@ def test_native_gemm_output_quantizer_support_is_opt_in(): _validate_native_gemm_output_quantizer(unknown_quantizer) +@requires_fp8 +class TestDPARuntimeRecipeUpdate: + """CustomRecipe DPA uses the normal module runtime update path.""" + + @staticmethod + def _make_dpa(): + return te.DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + name="recipe_update_dpa", + ).cuda() + + def test_custom_recipe_update_replaces_both_directions(self): + """A qfactory change commits one complete DPA runtime.""" + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + ) + + dpa = self._make_dpa() + + def first_factory(role): + return current_scaling_factory(role) + + first_recipe = recipe.CustomRecipe( + qfactory=first_factory, + qfactory_key=("dpa-runtime-update", 1), + fp8_dpa=True, + ) + + with autocast(enabled=True, recipe=first_recipe): + first_capabilities = dpa.get_qkv_quantization_capabilities() + first_runtime = dpa._quantization_runtime + first_forward_quantizers = dpa.quantizers["scaling_fwd"] + first_backward_quantizers = dpa.quantizers["scaling_bwd"] + first_capabilities_quantizer = dpa._qkv_capabilities_quantizer + assert dpa.get_qkv_quantization_capabilities() == first_capabilities + + assert dpa._quantization_runtime is first_runtime + assert len(first_runtime.forward_quantizers) == 9 + assert len(first_runtime.backward_quantizers) == 6 + assert first_runtime.forward_quantizers is first_forward_quantizers + assert first_runtime.backward_quantizers is first_backward_quantizers + + dpa._fp8_workspaces["old"] = object() + + def second_factory(role): + return current_scaling_factory(role) + + second_recipe = recipe.CustomRecipe( + qfactory=second_factory, + qfactory_key=("dpa-runtime-update", 2), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=second_recipe): + second_capabilities = dpa.get_qkv_quantization_capabilities() + + second_runtime = dpa._quantization_runtime + assert second_runtime is not first_runtime + assert dpa.fp8_meta["scaling_fwd"] is second_runtime.forward_states[0] + assert dpa.fp8_meta["scaling_bwd"] is second_runtime.backward_states[0] + assert dpa.quantizers["scaling_fwd"] is second_runtime.forward_quantizers + assert dpa.quantizers["scaling_bwd"] is second_runtime.backward_quantizers + assert second_capabilities == first_capabilities == (True, False) + assert dpa._qkv_capabilities_quantizer is not first_capabilities_quantizer + assert any( + dpa._qkv_capabilities_quantizer is quantizer + for quantizer in second_runtime.forward_quantizers + ) + assert not dpa._fp8_workspaces + + def test_candidate_failure_preserves_active_runtime_and_caches(self): + """A failed candidate validation does not partially update DPA.""" + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + ) + + dpa = self._make_dpa() + active_recipe = recipe.CustomRecipe( + qfactory=current_scaling_factory, + qfactory_key=("dpa-runtime-valid", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=active_recipe): + dpa.get_qkv_quantization_capabilities() + + old_runtime = dpa._quantization_runtime + old_views = ( + dpa.fp8_meta["recipe"], + dpa.fp8_meta["scaling_fwd"], + dpa.fp8_meta["scaling_bwd"], + dpa.quantizers["scaling_fwd"], + dpa.quantizers["scaling_bwd"], + dpa._qkv_capabilities_quantizer, + dpa._qkv_capabilities_cache, + ) + workspace = object() + dpa._fp8_workspaces["old"] = workspace + + invalid_recipe = recipe.CustomRecipe( + qfactory=lambda _role: IdentityQuantizer(), + qfactory_key=("invalid-dpa-update", 1), + fp8_dpa=True, + ) + with autocast( + enabled=True, + recipe=invalid_recipe, + ): + with pytest.raises(TypeError, match="FP8 attention requires FP8-compatible quantizers"): + dpa.get_qkv_quantization_capabilities() + + assert dpa._quantization_runtime is old_runtime + assert all( + current is old + for current, old in zip( + ( + dpa.fp8_meta["recipe"], + dpa.fp8_meta["scaling_fwd"], + dpa.fp8_meta["scaling_bwd"], + dpa.quantizers["scaling_fwd"], + dpa.quantizers["scaling_bwd"], + dpa._qkv_capabilities_quantizer, + dpa._qkv_capabilities_cache, + ), + old_views, + ) + ) + assert dpa._fp8_workspaces["old"] is workspace + + def test_hybrid_storage_snapshots_parent_quantizer(): """Caller mutations must not change an existing Hybrid tensor's behavior.""" quantizer = HybridQuantizer( diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index bd21adc26fe..bdfecb3983a 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -153,8 +153,10 @@ def test_quantization_runtime_bundles_both_quantization_directions(): ) forward_quantizers = [] backward_quantizers = [] + runtime_recipe = Float8CurrentScaling() runtime = _QuantizationRuntime( key=runtime_key, + recipe=runtime_recipe, num_gemms=1, recipe_config_revision=3, role_revision=2, @@ -165,6 +167,7 @@ def test_quantization_runtime_bundles_both_quantization_directions(): ) assert runtime.key is runtime_key + assert runtime.recipe is runtime_recipe assert runtime.forward_quantizers is forward_quantizers assert runtime.backward_quantizers is backward_quantizers diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 608d29e067d..edf0571d7d4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -30,6 +30,7 @@ MXFP8BlockScalingRecipeState, Float8CurrentScalingRecipeState, Float8BlockScalingRecipeState, + _QuantizationRuntime, ) from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -882,7 +883,10 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # global recipe set in autocast() fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): + old_runtime = getattr(self, "_quantization_runtime", None) super().init_fp8_metadata(num_gemms=num_gemms) + if self._quantization_runtime is not old_runtime: + self._invalidate_dpa_runtime_caches() fwd_quantizers = self.quantizers.get("scaling_fwd", ()) cache_key = ( id(self.fp8_meta.get("scaling_fwd")), @@ -1109,6 +1113,32 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # Clear cached workspaces as they were created with the old recipe/quantizer type self._fp8_workspaces.clear() + def _invalidate_dpa_runtime_caches(self) -> None: + """Invalidate DPA caches after a committed quantizer replacement.""" + self._custom_dpa_local_recipes_cache_key = None + self._custom_dpa_local_recipes_cache = None + self._qkv_capabilities_quantizer = None + self._qkv_capabilities_cache = None + _attention_backends["backend_selection_requires_update"] = True + + def _validate_quantization_runtime(self, candidate: _QuantizationRuntime) -> Any: + """Validate an FP8-attention candidate before publishing it.""" + super()._validate_quantization_runtime(candidate) + if not candidate.recipe.fp8_dpa: + return None + if len(candidate.forward_quantizers) != 9 or len(candidate.backward_quantizers) != 6: + raise ValueError( + "FP8 DotProductAttention requires 9 forward and 6 backward quantizer slots." + ) + dpa_utils.get_attention_quantizers( + True, + { + "scaling_fwd": candidate.forward_quantizers, + "scaling_bwd": candidate.backward_quantizers, + }, + ) + return None + def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: """Return MHA boundary capabilities from the canonical QKV quantizer. diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 10c2021aafd..de16379b60d 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1196,6 +1196,7 @@ def _prepare_quantization_runtime( return _QuantizationRuntime( key=key, + recipe=runtime_recipe, num_gemms=num_gemms, recipe_config_revision=recipe_config_revision, role_revision=role_revision, @@ -1213,11 +1214,11 @@ def _commit_quantization_runtime( self, candidate: _QuantizationRuntime, *, - recipe: Recipe, validation_result: Any = None, ) -> None: """Publish a fully prepared and validated runtime as one unit.""" del validation_result + recipe = candidate.recipe forward_state = candidate.forward_states[0] backward_state = candidate.backward_states[0] @@ -1350,7 +1351,6 @@ def reject_delayed_scaling_update() -> None: validation_result = self._validate_quantization_runtime(candidate) self._commit_quantization_runtime( candidate, - recipe=candidate.forward_states[0].recipe, validation_result=validation_result, ) return True diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 4b4fbbda0b3..0dcf7de0777 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1999,7 +1999,6 @@ def _commit_quantization_runtime( self, candidate: _QuantizationRuntime, *, - recipe: Recipe, validation_result: Tuple[Optional[Quantizer], Optional[Quantizer]], ) -> None: """Publish validated grouped-derived state with the candidate runtime.""" @@ -2011,7 +2010,6 @@ def _commit_quantization_runtime( super()._commit_quantization_runtime( candidate, - recipe=recipe, validation_result=validation_result, ) self.fast_setattr( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index be1d6f32129..2250aa4ccdc 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -144,10 +144,13 @@ class _QuantizationRuntime: A module or operation prepares this bundle away from its active state, validates it, and then publishes it as one unit. Keeping forward and backward states and quantizers together prevents a failed update from - exposing a new forward configuration with stale backward state. + exposing a new forward configuration with stale backward state. ``recipe`` + is the runtime-owned snapshot used to construct those states; it is not the + mutable recipe object owned by the caller. """ key: _QuantizationRuntimeKey + recipe: Recipe num_gemms: int recipe_config_revision: int role_revision: int From ea6ef5a6667ac871471409ef0cc4640c2c844ed9 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 12 Aug 2026 13:21:39 +0000 Subject: [PATCH 72/97] Fusible ops and graphs handling Signed-off-by: Evgeny --- tests/pytorch/test_cuda_graphs.py | 67 ++++++++++++++++++++++ tests/pytorch/test_fusible_ops.py | 47 +++++++++++++++ transformer_engine/pytorch/graph.py | 31 ++++++++++ transformer_engine/pytorch/ops/fuser.py | 17 ++++++ transformer_engine/pytorch/ops/op.py | 38 +++++++++++- transformer_engine/pytorch/quantization.py | 6 ++ 6 files changed, 205 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 5a848dc0e80..980b2d2d12f 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -751,6 +751,73 @@ def test_make_graphed_callables_with_kwargs( assert_all_equal(outputs, graph_outputs) +@pytest.mark.skipif(not fp8_available, reason="FP8 is not available") +def test_graphed_callable_rejects_changed_quantization_recipe(monkeypatch) -> None: + """Validate the capture recipe once per revision and reject mismatches.""" + model_config = model_configs["small"] + dtype = torch.bfloat16 + captured_recipe = recipe.Float8CurrentScaling() + equal_recipe = copy.deepcopy(captured_recipe) + model = Linear( + model_config.hidden_size, + model_config.hidden_size, + device="cuda", + params_dtype=dtype, + ) + for param in model.parameters(): + param.grad = torch.empty_like(param) + model = make_graphed_callables( + model, + (generate_data(model_config, dtype, warmup=True),), + enabled=True, + recipe=captured_recipe, + ) + + get_config_calls = 0 + get_quantizer_config = FP8GlobalStateManager.get_quantizer_config + + def counted_get_quantizer_config(cls): + nonlocal get_config_calls + get_config_calls += 1 + return get_quantizer_config() + + monkeypatch.setattr( + FP8GlobalStateManager, + "get_quantizer_config", + classmethod(counted_get_quantizer_config), + ) + + try: + with autocast(enabled=True, recipe=equal_recipe): + for _ in range(2): + output = model(generate_data(model_config, dtype)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + assert get_config_calls == 1 + + # Re-entering after autocast restores the previous recipe advances the + # revision. Validate the equal semantic configuration once again, then + # stay on the integer-only fast path. + with autocast(enabled=True, recipe=equal_recipe): + for _ in range(2): + output = model(generate_data(model_config, dtype)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) + assert get_config_calls == 2 + + # Mutating and reusing the caller-owned recipe must not mutate the + # immutable semantic configuration saved at capture time. + captured_recipe.fp8_quant_fwd_inp = recipe.QParams(amax_epsilon=0.25) + with autocast(enabled=True, recipe=captured_recipe): + for _ in range(2): + with pytest.raises( + RuntimeError, + match="Recapture the graph with the new recipe", + ): + model(generate_data(model_config, dtype)) + assert get_config_calls == 3 + finally: + reset_graphs(model) + + def test_make_graphed_callables_returns_owned_parameter_grads() -> None: """Parameter grads returned from graph replay must not alias static graph buffers.""" reset_rng_states() diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d81253..6af9a0e9e20 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -5,6 +5,7 @@ from __future__ import annotations from collections.abc import Iterable, Sequence +import copy import io import os import math @@ -24,6 +25,7 @@ GRAD_INPUT_BUFFER_KEY, ) from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV +from transformer_engine.pytorch.quantization import FP8GlobalStateManager from transformer_engine.pytorch.ops.fused import ( BackwardActivationBias, @@ -145,6 +147,51 @@ def maybe_skip_quantization( pytest.skip("NVFP4 quantization is only supported with BF16 data") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("persistent_fuser", (False, True)) +def test_fusible_operation_rejects_unsupported_same_class_recipe_update( + persistent_fuser: bool, +) -> None: + """Fusible ops must not silently retain stale same-class quantizers.""" + op = te_ops.Quantize() + model = te_ops.Sequential(op) if persistent_fuser else op + x = torch.randn(16, 16, dtype=torch.bfloat16, device="cuda") + initial_recipe = transformer_engine.common.recipe.Float8CurrentScaling() + equal_recipe = copy.deepcopy(initial_recipe) + changed_recipe = copy.deepcopy(initial_recipe) + changed_recipe.fp8_quant_fwd_inp = transformer_engine.common.recipe.QParams(amax_epsilon=0.25) + + try: + with torch.no_grad(), te.autocast(recipe=initial_recipe): + model(x) + with torch.no_grad(), te.autocast(recipe=equal_recipe): + model(x) + with pytest.raises(RuntimeError, match="not supported for fusible operations"): + with torch.no_grad(), te.autocast(recipe=changed_recipe): + model(x) + finally: + FP8GlobalStateManager.reset() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_fusible_operation_preserves_delayed_history_resize() -> None: + """The explicit legacy delayed-history update remains supported.""" + op = te_ops.Quantize() + model = te_ops.Sequential(op) + x = torch.randn(16, 16, dtype=torch.bfloat16, device="cuda") + initial_recipe = transformer_engine.common.recipe.DelayedScaling(amax_history_len=2) + resized_recipe = transformer_engine.common.recipe.DelayedScaling(amax_history_len=4) + + try: + with torch.no_grad(), te.autocast(recipe=initial_recipe): + model(x) + with torch.no_grad(), te.autocast(recipe=resized_recipe): + model(x) + assert op._fp8_metas["forward"]["scaling_fwd"].amax_history.shape[0] == 4 + finally: + FP8GlobalStateManager.reset() + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index b298b3d8fff..418c984b95b 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -159,6 +159,8 @@ def _make_graphed_callables( pre_warmup_hook: Optional[Callable] = None, post_warmup_hook: Optional[Callable] = None, capture_time_hooks: List[Dict[str, Dict]], + module_uses_fp8: Dict[int, bool], + captured_recipe_config: Any, ) -> SingleOrTuple[Callable]: """ Helper method for `make_graphed_callables` @@ -1197,6 +1199,23 @@ def reset(): return backward_dw, reset + # Cache CUDA-graph recipe validation by the manager-owned configuration + # revision. Semantic configuration equality is only evaluated when recipe + # activation advances the revision, while unchanged replays only compare + # integers. + checked_recipe_config_revision = None + checked_recipe_config_matches = False + + def active_recipe_config_matches_capture() -> bool: + nonlocal checked_recipe_config_revision, checked_recipe_config_matches + recipe_config_revision = FP8GlobalStateManager.get_quantizer_config_revision() + if recipe_config_revision != checked_recipe_config_revision: + checked_recipe_config_matches = ( + FP8GlobalStateManager.get_quantizer_config() == captured_recipe_config + ) + checked_recipe_config_revision = recipe_config_revision + return checked_recipe_config_matches + # Put together the final graphed callables ret = [] for i in range(len(sample_args)): @@ -1223,6 +1242,15 @@ def new_fwd(*user_args, **user_kwargs): # If the module's training-or-eval state matches what we graphed, # run the graph, otherwise run the original forward method if func.training == graph_training_state: + if ( + module_uses_fp8.get(id(func), False) + and FP8GlobalStateManager.is_fp8_enabled() + and not active_recipe_config_matches_capture() + ): + raise RuntimeError( + "The active quantization recipe differs from the CUDA graph " + "capture recipe. Recapture the graph with the new recipe." + ) # Set the FP8 group from global amax reduction. if FP8GlobalStateManager.is_fp8_enabled(): fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() @@ -1552,6 +1580,7 @@ def make_graphed_callables( recipe = get_default_fp8_recipe() elif not any(enabled): recipe = None + captured_recipe_config = recipe.quantizer_config() if recipe is not None else None module_uses_fp8 = dict(zip((id(m) for m in modules), enabled)) # Canonicalize capture_time_hooks kwarg. @@ -1625,6 +1654,8 @@ def call_func(self, *args, **kwargs): pre_warmup_hook=pre_warmup_hook, post_warmup_hook=post_warmup_hook, capture_time_hooks=capture_time_hooks, + module_uses_fp8=module_uses_fp8, + captured_recipe_config=captured_recipe_config, ) # Ensures warmup does not affect numerics for ops such as dropout. diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004ddd..1dc3f59c0ff 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -18,6 +18,7 @@ FusibleOperation, FusedOperation, OperationContext, + _only_delayed_history_length_changed, ) @@ -350,6 +351,7 @@ def __init__( # Cache and detect change of state relevant for fusing operations self.recipe_type = None + self.recipe_config = None self.first_op_requiring_backward = 0 self.backward_override = None self._last_amax_history_len = 0 @@ -439,6 +441,20 @@ def maybe_fuse_ops( # Early exit if fusion parameters haven't changed need_reset = False recipe_type = type(recipe) + recipe_config = recipe.quantizer_config() if recipe is not None else None + if ( + recipe_type is self.recipe_type + and self.recipe_config != recipe_config + and not ( + isinstance(recipe, DelayedScaling) + and _only_delayed_history_length_changed(self.recipe_config, recipe_config) + ) + ): + FP8GlobalStateManager.abort_current_amax_reduction() + raise RuntimeError( + "Mid-training recipe updates are not supported for fusible operations. " + "Recreate the fusible operation or operation pipeline with the new recipe." + ) backward_override = recipe.backward_override if recipe is not None else None fusion_params = (recipe_type, first_op_requiring_backward, backward_override) if fusion_params != ( @@ -494,6 +510,7 @@ def maybe_fuse_ops( # Save current fusion params self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params + self.recipe_config = recipe_config # Save amax history length if isinstance(recipe, DelayedScaling): diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 5106ec9e0af..93da268b378 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -6,7 +6,7 @@ from __future__ import annotations import abc -from collections.abc import Iterable +from collections.abc import Hashable, Iterable import dataclasses import pickle from typing import Any, Optional @@ -24,6 +24,18 @@ from ..tensor import Quantizer +def _only_delayed_history_length_changed( + old_config: Hashable, + new_config: Hashable, +) -> bool: + """Whether delayed-scaling configs differ only in history length.""" + + def without_history_length(config: Hashable) -> tuple: + return tuple(item for item in config if item[0] != "amax_history_len") + + return without_history_length(old_config) == without_history_length(new_config) + + @dataclasses.dataclass class OperationContext: """State needed to apply an operation @@ -190,6 +202,8 @@ def __init__(self) -> None: # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None + self._recipe_type: Optional[type[Recipe]] = None + self._recipe_config: Optional[Hashable] = None @property def is_fused_op(self) -> bool: @@ -243,8 +257,27 @@ def reset_recipe_state( if recipe is None: self._fp8_metas = None self._quantizers = None + self._recipe_type = None + self._recipe_config = None return + recipe_type = type(recipe) + recipe_config = recipe.quantizer_config() + if ( + self._recipe_type is recipe_type + and self._recipe_config is not None + and self._recipe_config != recipe_config + and not ( + recipe.delayed() + and _only_delayed_history_length_changed(self._recipe_config, recipe_config) + ) + ): + FP8GlobalStateManager.abort_current_amax_reduction() + raise RuntimeError( + "Mid-training recipe updates are not supported for fusible operations. " + "Recreate the fusible operation or operation pipeline with the new recipe." + ) + # Communication group for FP8 amax reductions fp8_group = FP8GlobalStateManager.get_fp8_group() @@ -360,6 +393,9 @@ def reset_recipe_state( self._fp8_metas[mode], ) + self._recipe_type = recipe_type + self._recipe_config = recipe_config + def get_quantizer( self, mode: str, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 2250aa4ccdc..79498cb9094 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1146,6 +1146,12 @@ class autocast: module more than once inside an `autocast` region overrides the amax tensors before reduction can occur. + .. note:: + + Mid-training recipe changes are not supported within an existing ``torch.compile`` + graph. Change the recipe outside the compiled region and reset or recompile the callable + before continuing. + Parameters ---------- enabled : bool, default = True From 35ac69dbe237faee41a7a40090c0d7f89d687df7 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 13 Aug 2026 11:30:26 +0000 Subject: [PATCH 73/97] te.apply_recipe() Signed-off-by: Evgeny --- docs/api/pytorch.rst | 2 + tests/pytorch/test_hybrid_quantization.py | 112 ++-- tests/pytorch/test_recipe_update.py | 615 +++++++++++++++++- transformer_engine/pytorch/__init__.py | 3 +- .../dot_product_attention.py | 85 ++- transformer_engine/pytorch/module/base.py | 250 ++++--- .../pytorch/module/grouped_linear.py | 8 +- .../pytorch/module/layernorm_mlp.py | 4 + transformer_engine/pytorch/quantization.py | 118 ++++ 9 files changed, 1040 insertions(+), 157 deletions(-) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 5fac0a89a6f..0e1bab3a126 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -40,6 +40,8 @@ PyTorch .. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None) +.. autoapifunction:: transformer_engine.pytorch.apply_recipe + .. autoapifunction:: transformer_engine.pytorch.quantized_model_init .. autoapifunction:: transformer_engine.pytorch.checkpoint diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index b0525dda622..fb1ac1ccf1b 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -205,6 +205,52 @@ def second_factory(role): ) assert not dpa._fp8_workspaces + def test_direct_runtime_commit_updates_dpa_derived_state(self): + """DPA cache and recipe labels are part of its runtime commit hook.""" + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + ) + + dpa = self._make_dpa() + first_recipe = recipe.CustomRecipe( + qfactory=current_scaling_factory, + qfactory_key=("dpa-direct-commit", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=first_recipe): + dpa.get_qkv_quantization_capabilities() + + old_runtime = dpa._quantization_runtime + old_local_recipes = object() + dpa.fp8_meta["local_recipes"] = old_local_recipes + old_cache_key = dpa._custom_dpa_local_recipes_cache_key + old_capabilities_quantizer = dpa._qkv_capabilities_quantizer + + second_recipe = recipe.CustomRecipe( + qfactory=current_scaling_factory, + qfactory_key=("dpa-direct-commit", 2), + fp8_dpa=True, + ) + update = dpa._plan_quantization_update( + recipe=second_recipe, + recipe_config=second_recipe.quantizer_config(), + recipe_config_revision=old_runtime.recipe_config_revision + 1, + num_gemms=dpa._get_quantization_runtime_num_gemms(), + ) + + assert dpa._quantization_runtime is old_runtime + assert dpa.fp8_meta["local_recipes"] is old_local_recipes + assert dpa._custom_dpa_local_recipes_cache_key is old_cache_key + assert dpa._qkv_capabilities_quantizer is old_capabilities_quantizer + + assert dpa._apply_quantization_update(update) + assert dpa._quantization_runtime is update.candidate + assert "local_recipes" not in dpa.fp8_meta + assert dpa._custom_dpa_local_recipes_cache is None + assert dpa._custom_dpa_local_recipes_cache_key != old_cache_key + assert dpa._qkv_capabilities_quantizer is None + assert dpa._qkv_capabilities_cache is None + def test_candidate_failure_preserves_active_runtime_and_caches(self): """A failed candidate validation does not partially update DPA.""" from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( @@ -2144,31 +2190,11 @@ def test_inference_runs_once_per_quantizer_state_and_clears_stale_labels(self, m from transformer_engine.pytorch.attention.dot_product_attention import ( dot_product_attention as dpa_module, ) - from transformer_engine.pytorch.module.base import TransformerEngineBaseModule - from transformer_engine.pytorch.quantization import FP8GlobalStateManager custom_recipe = recipe.CustomRecipe( qfactory=lambda _role: IdentityQuantizer(), qfactory_key=("test_dpa_identity", 1), ) - monkeypatch.setattr( - FP8GlobalStateManager, - "get_fp8_recipe", - classmethod(lambda _cls: custom_recipe), - ) - - state = [object()] - quantizer = [object()] - - def fake_base_init(module, num_gemms=1): # pylint: disable=unused-argument - module.fp8_meta["scaling_fwd"] = state[0] - module.quantizers["scaling_fwd"] = [quantizer[0]] - - monkeypatch.setattr( - TransformerEngineBaseModule, - "init_fp8_metadata", - fake_base_init, - ) inferred_labels = [recipe.MXFP8BlockScaling()] mutated_recipe_labels = [recipe.MXFP8BlockScaling(fp8_mha=True)] @@ -2187,32 +2213,34 @@ def fake_infer(*_args, **_kwargs): kv_channels=16, attention_dropout=0.0, ) - dpa.init_fp8_metadata() - assert dpa.fp8_meta["local_recipes"] is inferred_labels - dpa.init_fp8_metadata() - assert inference_calls == 1 - assert dpa.fp8_meta["local_recipes"] is inferred_labels - - # Native labels also copy these mutable fields from CustomRecipe. They - # must refresh even when the quantizer generation itself is unchanged. + with autocast(enabled=True, recipe=custom_recipe): + dpa.init_fp8_metadata(num_gemms=3) + assert dpa.fp8_meta["local_recipes"] is inferred_labels + dpa.init_fp8_metadata(num_gemms=3) + assert inference_calls == 1 + assert dpa.fp8_meta["local_recipes"] is inferred_labels + + # Native labels also copy these mutable fields from CustomRecipe, so a + # semantic recipe update must refresh them with the candidate runtime. custom_recipe.fp8_mha = True - dpa.init_fp8_metadata() - assert inference_calls == 2 - assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels - dpa.init_fp8_metadata() - assert inference_calls == 2 - assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels + with autocast(enabled=True, recipe=custom_recipe): + dpa.init_fp8_metadata(num_gemms=3) + assert inference_calls == 2 + assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels + dpa.init_fp8_metadata(num_gemms=3) + assert inference_calls == 2 + assert dpa.fp8_meta["local_recipes"] is mutated_recipe_labels # A rebuilt recipe state/quantizer list invalidates the cache. If the # new family has no native label, the old label must not survive. - state[0] = object() - quantizer[0] = object() - dpa.init_fp8_metadata() - assert inference_calls == 3 - assert "local_recipes" not in dpa.fp8_meta - dpa.init_fp8_metadata() - assert inference_calls == 3 - assert "local_recipes" not in dpa.fp8_meta + custom_recipe.qfactory_key = ("test_dpa_identity", 2) + with autocast(enabled=True, recipe=custom_recipe): + dpa.init_fp8_metadata(num_gemms=3) + assert inference_calls == 3 + assert "local_recipes" not in dpa.fp8_meta + dpa.init_fp8_metadata(num_gemms=3) + assert inference_calls == 3 + assert "local_recipes" not in dpa.fp8_meta @pytest.mark.parametrize( "factory_name,expected", diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index add0ec3f135..65bfb5f82d2 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -6,18 +6,25 @@ import pytest import torch +import transformer_engine.pytorch.ops as te_ops from transformer_engine.common.recipe import CustomRecipe, DelayedScaling, Float8CurrentScaling from transformer_engine.pytorch import ( + DotProductAttention, GroupedLinear, LayerNormLinear, LayerNormMLP, Linear, + MultiheadAttention, + TransformerLayer, + apply_recipe, autocast, is_fp8_available, ) +from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import ( DelayedScalingRequest, + FP8GlobalStateManager, QuantizerRole, ) from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer @@ -38,7 +45,16 @@ def qfactory(role): def _ensure_runtime(module, recipe, revision, *, num_gemms=1): - return module._ensure_active_quantization_runtime( # pylint: disable=protected-access + return module._ensure_quantization_runtime( # pylint: disable=protected-access + recipe=recipe, + recipe_config=recipe.quantizer_config(), + recipe_config_revision=revision, + num_gemms=num_gemms, + ) + + +def _prepare_runtime_update(module, recipe, revision, *, num_gemms=1): + return module._plan_quantization_update( # pylint: disable=protected-access recipe=recipe, recipe_config=recipe.quantizer_config(), recipe_config_revision=revision, @@ -73,6 +89,39 @@ def _runtime_views(module): ) +def _run_update_step(module, recipe, inp, *args): + """Run a real forward/backward step and return the output.""" + module.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=recipe): + output = module(inp, *args) + if isinstance(output, tuple): + output = output[0] + output.float().sum().backward() + assert inp.grad is not None + assert all(param.grad is not None for param in module.parameters() if param.requires_grad) + return output + + +def _active_runtime_owners(module): + """Return executed TE runtime owners in module traversal order.""" + return [ + submodule + for submodule in module.modules() + if isinstance(submodule, TransformerEngineBaseModule) + and submodule._quantization_runtime is not None # pylint: disable=protected-access + ] + + +def _global_recipe_state(): + """Return the identity-bearing global recipe state.""" + state = FP8GlobalStateManager.quantization_state + return ( + state.fp8_recipe, + state.quantizer_config, + state.quantizer_config_revision, + ) + + @pytest.mark.parametrize( "replacement", ( @@ -268,6 +317,137 @@ def unexpected_role_resolution(**_kwargs): assert len(calls) == 5 +def test_revision_only_update_is_not_published_during_planning(): + """An equal semantic update synchronizes revisions only when committed.""" + calls = [] + recipe = _make_counting_recipe(("revision-plan", 1), calls) + equal_recipe = _make_counting_recipe(("revision-plan", 1), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + _ensure_runtime(module, recipe, revision=1) + active = module._quantization_runtime # pylint: disable=protected-access + factory_call_count = len(calls) + + update = _prepare_runtime_update(module, equal_recipe, revision=3) + + assert update.candidate is None + assert module._quantization_runtime is active # pylint: disable=protected-access + assert active.recipe_config_revision == 1 + assert len(calls) == factory_call_count + + assert not module._apply_quantization_update( # pylint: disable=protected-access + update + ) + assert active.recipe_config_revision == 3 + assert len(calls) == factory_call_count + + +def test_candidate_update_is_not_published_during_planning(): + """Candidate state, views, and workspaces change only during commit.""" + calls = [] + active_recipe = _make_counting_recipe(("candidate-plan", 1), calls) + replacement_recipe = _make_counting_recipe(("candidate-plan", 2), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + _ensure_runtime(module, active_recipe, revision=1) + old_views = _runtime_views(module) + workspace = object() + module._fp8_workspaces["weight"] = workspace # pylint: disable=protected-access + + update = _prepare_runtime_update(module, replacement_recipe, revision=2) + + assert update.candidate is not None + assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + assert old_views[0].recipe_config_revision == 1 + assert module._fp8_workspaces["weight"] is workspace # pylint: disable=protected-access + + assert module._apply_quantization_update( # pylint: disable=protected-access + update + ) + assert module._quantization_runtime is update.candidate # pylint: disable=protected-access + assert not module._fp8_workspaces # pylint: disable=protected-access + + +def test_uninitialized_runtime_can_be_prepared_without_publication(): + """Planning supports modules that have not yet executed a quantized forward.""" + calls = [] + recipe = _make_counting_recipe(("initial-plan", 1), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + + update = _prepare_runtime_update(module, recipe, revision=1) + + assert update.candidate is not None + assert module._quantization_runtime is None # pylint: disable=protected-access + assert "scaling_fwd" not in module.fp8_meta + assert "scaling_bwd" not in module.fp8_meta + assert module.quantizers == {"scaling_fwd": [], "scaling_bwd": []} + + assert module._apply_quantization_update( # pylint: disable=protected-access + update + ) + assert module._quantization_runtime is update.candidate # pylint: disable=protected-access + + +def test_later_planning_failure_leaves_earlier_module_unchanged(monkeypatch): + """Prepared candidates can be discarded without changing any module.""" + calls = [] + active_recipe = _make_counting_recipe(("multi-plan", 1), calls) + replacement_recipe = _make_counting_recipe(("multi-plan", 2), calls) + first = Linear(16, 16, bias=False, device="cuda", name="first") + second = Linear(16, 16, bias=False, device="cuda", name="second") + _ensure_runtime(first, active_recipe, revision=1) + _ensure_runtime(second, active_recipe, revision=1) + old_first_views = _runtime_views(first) + old_second_views = _runtime_views(second) + + first_update = _prepare_runtime_update(first, replacement_recipe, revision=2) + assert first_update.candidate is not None + + def reject_candidate(_candidate): + raise RuntimeError("later module validation failure") + + monkeypatch.setattr(second, "_validate_quantization_runtime", reject_candidate) + with pytest.raises(RuntimeError, match="later module validation failure"): + _prepare_runtime_update(second, replacement_recipe, revision=2) + + assert all(current is old for current, old in zip(_runtime_views(first), old_first_views)) + assert all(current is old for current, old in zip(_runtime_views(second), old_second_views)) + assert old_first_views[0].recipe_config_revision == 1 + assert old_second_views[0].recipe_config_revision == 1 + + +@pytest.mark.parametrize( + ("module_factory", "expected_num_gemms"), + ( + pytest.param( + lambda: Linear(16, 16, bias=False, device="cuda"), + 1, + id="linear", + ), + pytest.param( + lambda: LayerNormLinear(16, 16, bias=False, device="cuda"), + 1, + id="layernorm-linear", + ), + pytest.param( + lambda: LayerNormMLP(16, 32, bias=False, device="cuda"), + 2, + id="layernorm-mlp", + ), + pytest.param( + lambda: GroupedLinear(2, 16, 16, bias=False, device="cuda"), + 2, + id="grouped-linear", + ), + ), +) +def test_module_owns_runtime_slot_layout(module_factory, expected_num_gemms): + """Model-wide orchestration does not need a module-type dispatch table.""" + module = module_factory() + assert ( + module._get_quantization_runtime_num_gemms() # pylint: disable=protected-access + == expected_num_gemms + ) + + def test_unchanged_forward_uses_revision_hot_path(monkeypatch): """Repeated forwards do not enter any runtime-construction cold path.""" calls = [] @@ -286,9 +466,10 @@ def unexpected_cold_path(*_args, **_kwargs): raise AssertionError("unchanged forward entered the runtime cold path") monkeypatch.setattr(module, "get_quantizer_roles", unexpected_cold_path) - monkeypatch.setattr(module, "_prepare_quantization_runtime", unexpected_cold_path) + monkeypatch.setattr(module, "_plan_quantization_update", unexpected_cold_path) + monkeypatch.setattr(module, "_build_quantization_runtime", unexpected_cold_path) monkeypatch.setattr(module, "_validate_quantization_runtime", unexpected_cold_path) - monkeypatch.setattr(module, "_commit_quantization_runtime", unexpected_cold_path) + monkeypatch.setattr(module, "_activate_quantization_runtime", unexpected_cold_path) workspace_sentinel = object() module._fp8_workspaces["sentinel"] = workspace_sentinel # pylint: disable=protected-access @@ -489,6 +670,426 @@ def test_runtime_update_workspace_lifecycle( assert not module._fp8_workspaces # pylint: disable=protected-access +@pytest.mark.parametrize( + ("module_factory", "input_factory", "forward_args"), + ( + pytest.param( + lambda: LayerNormLinear(16, 16, bias=False, device="cuda"), + lambda: torch.randn(8, 16, device="cuda", requires_grad=True), + (), + id="layernorm-linear", + ), + pytest.param( + lambda: LayerNormMLP(16, 32, bias=False, device="cuda"), + lambda: torch.randn(8, 16, device="cuda", requires_grad=True), + (), + id="layernorm-mlp", + ), + pytest.param( + lambda: GroupedLinear(2, 16, 16, bias=False, device="cuda"), + lambda: torch.randn(8, 16, device="cuda", requires_grad=True), + ([4, 4],), + id="grouped-linear", + ), + ), +) +def test_module_family_executes_after_recipe_update( + module_factory, + input_factory, + forward_args, +): + """Every base module family executes forward/backward across a runtime replacement.""" + calls = [] + first_recipe = _make_counting_recipe(("module-execution", 1), calls) + second_recipe = _make_counting_recipe(("module-execution", 2), calls) + module = module_factory() + + _run_update_step(module, first_recipe, input_factory(), *forward_args) + first_runtime = module._quantization_runtime # pylint: disable=protected-access + assert first_runtime is not None + + _run_update_step(module, second_recipe, input_factory(), *forward_args) + second_runtime = module._quantization_runtime # pylint: disable=protected-access + assert second_runtime is not first_runtime + assert second_runtime.key.recipe_config == second_recipe.quantizer_config() + + +@pytest.mark.parametrize( + "module_factory", + ( + pytest.param( + lambda: MultiheadAttention( + hidden_size=32, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + bias=False, + device="cuda", + name="mha", + ), + id="multihead-attention", + ), + pytest.param( + lambda: TransformerLayer( + hidden_size=32, + ffn_hidden_size=64, + num_attention_heads=2, + hidden_dropout=0.0, + attention_dropout=0.0, + self_attn_mask_type="no_mask", + bias=False, + device="cuda", + ), + id="transformer-layer", + ), + ), +) +def test_composed_module_executes_after_recipe_update(module_factory): + """A composed module updates every runtime owner reached by real execution.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + calls = [] + first_recipe = _make_counting_recipe(("composed-execution", 1), calls) + second_recipe = _make_counting_recipe(("composed-execution", 2), calls) + module = module_factory() + + first_inp = torch.randn(8, 2, 32, device="cuda", requires_grad=True) + _run_update_step(module, first_recipe, first_inp) + owners = _active_runtime_owners(module) + assert len(owners) >= 2 + first_runtimes = { + id(owner): owner._quantization_runtime # pylint: disable=protected-access + for owner in owners + } + + second_inp = torch.randn(8, 2, 32, device="cuda", requires_grad=True) + _run_update_step(module, second_recipe, second_inp) + assert _active_runtime_owners(module) == owners + for owner in owners: + runtime = owner._quantization_runtime # pylint: disable=protected-access + assert runtime is not first_runtimes[id(owner)] + assert runtime.key.recipe_config == second_recipe.quantizer_config() + + +def test_apply_recipe_success_noop_mutation_and_forward_fast_path(): + """Explicit application publishes once and matching forwards remain lazy-path no-ops.""" + FP8GlobalStateManager.reset() + calls = [] + first_recipe = _make_counting_recipe(("apply-success", 1), calls) + equal_recipe = _make_counting_recipe(("apply-success", 1), calls) + model = torch.nn.Sequential( + Linear(16, 16, bias=False, device="cuda", name="first"), + Linear(16, 16, bias=False, device="cuda", name="second"), + ) + + try: + apply_recipe(model, first_recipe) + first_runtimes = [module._quantization_runtime for module in model] + first_revision = FP8GlobalStateManager.get_quantizer_config_revision() + first_factory_calls = len(calls) + assert first_factory_calls == 10 + assert FP8GlobalStateManager.get_fp8_recipe() is first_recipe + + # An independent equal recipe updates only the manager's requested + # recipe object. Runtime identity, revision, and factories stay fixed. + apply_recipe(model, equal_recipe) + assert [module._quantization_runtime for module in model] == first_runtimes + assert FP8GlobalStateManager.get_quantizer_config_revision() == first_revision + assert FP8GlobalStateManager.get_fp8_recipe() is equal_recipe + assert len(calls) == first_factory_calls + + # The first matching autocast forward after explicit application uses + # each module's revision fast path and performs no factory work. + inp = torch.randn(8, 16, device="cuda", requires_grad=True) + _run_update_step(model, equal_recipe, inp) + assert [module._quantization_runtime for module in model] == first_runtimes + assert len(calls) == first_factory_calls + + # Mutating and reusing the same recipe object produces a real model-wide + # replacement and one new global revision. + equal_recipe.qfactory_key = ("apply-success", 2) + apply_recipe(model, equal_recipe) + assert all( + module._quantization_runtime is not old_runtime + for module, old_runtime in zip(model, first_runtimes) + ) + assert FP8GlobalStateManager.get_quantizer_config_revision() == first_revision + 1 + assert len(calls) == first_factory_calls + 10 + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_computes_config_once(monkeypatch): + """All participant plans receive one shared semantic configuration.""" + FP8GlobalStateManager.reset() + config_calls = [] + factory_calls = [] + make_config = CustomRecipe._make_quantizer_config + + def counted_make_config(recipe): + config_calls.append(recipe) + return make_config(recipe) + + monkeypatch.setattr(CustomRecipe, "_make_quantizer_config", counted_make_config) + recipe = _make_counting_recipe(("apply-config-once", 1), factory_calls) + model = torch.nn.Sequential( + Linear(16, 16, bias=False, device="cuda", name="first"), + Linear(16, 16, bias=False, device="cuda", name="second"), + ) + + try: + apply_recipe(model, recipe) + assert config_calls == [recipe] + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_reconstructs_stateless_runtime_after_checkpoint_restore(): + """A restored model can rebuild its committed recipe before its first forward.""" + FP8GlobalStateManager.reset() + source_calls = [] + source = Linear(16, 16, bias=False, device="cuda", name="linear") + initial_recipe = _make_counting_recipe(("checkpoint-resume", 1), source_calls) + committed_recipe = _make_counting_recipe(("checkpoint-resume", 2), source_calls) + inp = torch.randn(8, 16, device="cuda") + + try: + apply_recipe(source, initial_recipe) + apply_recipe(source, committed_recipe) + checkpoint = source.state_dict() + with torch.no_grad(), autocast(enabled=True, recipe=committed_recipe): + expected = source(inp) + + # Model/checkpoint restoration happens in a fresh runtime context. The + # intended recipe is reconstructed and explicitly applied before the + # restored model executes its first forward. + FP8GlobalStateManager.reset() + restored = Linear(16, 16, bias=False, device="cuda", name="linear") + restored.load_state_dict(checkpoint) + assert restored._quantization_runtime is None # pylint: disable=protected-access + + resumed_calls = [] + resumed_recipe = _make_counting_recipe(("checkpoint-resume", 2), resumed_calls) + apply_recipe(restored, resumed_recipe) + runtime = restored._quantization_runtime # pylint: disable=protected-access + assert runtime is not None + assert runtime.key.recipe_config == committed_recipe.quantizer_config() + + with torch.no_grad(), autocast(enabled=True, recipe=resumed_recipe): + actual = restored(inp) + torch.testing.assert_close(actual, expected) + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_includes_unexecuted_conditional_branch(): + """Traversal updates owners independently of the branch executed by forward.""" + + class ConditionalModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.executed = Linear(16, 16, bias=False, device="cuda", name="executed") + self.unexecuted = Linear(16, 16, bias=False, device="cuda", name="unexecuted") + + def forward(self, inp, *, use_unexecuted=False): + module = self.unexecuted if use_unexecuted else self.executed + return module(inp) + + FP8GlobalStateManager.reset() + calls = [] + first_recipe = _make_counting_recipe(("apply-conditional", 1), calls) + replacement_recipe = _make_counting_recipe(("apply-conditional", 2), calls) + model = ConditionalModel() + + try: + # Lazy execution initializes only one branch. + inp = torch.randn(8, 16, device="cuda", requires_grad=True) + with autocast(enabled=True, recipe=first_recipe): + model(inp).sum().backward() + old_runtime = model.executed._quantization_runtime + assert old_runtime is not None + assert model.unexecuted._quantization_runtime is None + + apply_recipe(model, replacement_recipe) + assert model.executed._quantization_runtime is not old_runtime + assert model.unexecuted._quantization_runtime is not None + assert model.executed._quantization_runtime.key.recipe_config == ( + replacement_recipe.quantizer_config() + ) + assert model.unexecuted._quantization_runtime.key.recipe_config == ( + replacement_recipe.quantizer_config() + ) + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_supports_custom_recipe_dpa(): + """Runtime-managed CustomRecipe DPA participates in model-wide application.""" + FP8GlobalStateManager.reset() + calls = [] + recipe = _make_counting_recipe(("apply-custom-dpa", 1), calls) + dpa = DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + name="dpa", + ).cuda() + + try: + apply_recipe(dpa, recipe) + assert dpa._quantization_runtime is not None + assert dpa._quantization_runtime.key.recipe_config == recipe.quantizer_config() + assert len(calls) == 15 + finally: + FP8GlobalStateManager.reset() + + +@pytest.mark.parametrize("failure_index", (0, 1, 2)) +def test_apply_recipe_planning_failure_is_model_wide_atomic(failure_index): + """A factory failure in any participant leaves modules and manager unchanged.""" + FP8GlobalStateManager.reset() + calls = [] + active_recipe = _make_counting_recipe(("apply-failure", 1), calls) + modules = [ + Linear(16, 16, bias=False, device="cuda", name=f"line{index}") + for index in range(3) + ] + model = torch.nn.Sequential(*modules) + + try: + apply_recipe(model, active_recipe) + old_global_state = _global_recipe_state() + old_views = [_runtime_views(module) for module in modules] + workspaces = [] + for module in modules: + workspace = object() + module._fp8_workspaces["weight"] = workspace + workspaces.append(workspace) + + def failing_factory(role): + if role is not None and role.name == f"line{failure_index}": + raise RuntimeError("model-wide factory failure") + return IdentityQuantizer() + + replacement_recipe = CustomRecipe( + qfactory=failing_factory, + qfactory_key=("apply-failure", 2, failure_index), + ) + with pytest.raises( + RuntimeError, + match=rf"planning module '{failure_index}': model-wide factory failure", + ): + apply_recipe(model, replacement_recipe) + + assert _global_recipe_state() == old_global_state + for module, expected_views, workspace in zip(modules, old_views, workspaces): + assert all( + current is expected + for current, expected in zip(_runtime_views(module), expected_views) + ) + assert module._fp8_workspaces["weight"] is workspace + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_deduplicates_shared_runtime_owner(): + """A shared module is planned and applied exactly once.""" + FP8GlobalStateManager.reset() + calls = [] + recipe = _make_counting_recipe(("apply-shared", 1), calls) + shared = Linear(16, 16, bias=False, device="cuda", name="shared") + model = torch.nn.Module() + model.add_module("first", shared) + model.add_module("second", shared) + + try: + apply_recipe(model, recipe) + assert shared._quantization_runtime is not None + assert len(calls) == 5 + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_rejects_fusible_owner_before_factory(): + """Excluded fusible owners are discovered before any participant planning.""" + FP8GlobalStateManager.reset() + + def unexpected_factory(_role): + raise AssertionError("excluded model invoked qfactory") + + recipe = CustomRecipe( + qfactory=unexpected_factory, + qfactory_key=("apply-fusible", 1), + ) + model = torch.nn.ModuleList( + [ + Linear(16, 16, bias=False, device="cuda", name="linear"), + te_ops.Quantize(), + ] + ) + old_global_state = _global_recipe_state() + + try: + with pytest.raises(RuntimeError, match="does not support fusible operations"): + apply_recipe(model, recipe) + assert model[0]._quantization_runtime is None + assert _global_recipe_state() == old_global_state + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_rejects_legacy_dpa_path(): + """Model-wide application does not alter the built-in NVTE_DPA_* mechanism.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + dpa = DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + name="dpa", + ).cuda() + old_global_state = _global_recipe_state() + try: + with pytest.raises(RuntimeError, match="only through CustomRecipe"): + apply_recipe(dpa, Float8CurrentScaling()) + assert dpa._quantization_runtime is None + assert _global_recipe_state() == old_global_state + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_rejects_active_autocast_and_graph_capture(monkeypatch): + """Explicit model-wide application is limited to the documented safe boundary.""" + FP8GlobalStateManager.reset() + calls = [] + recipe = _make_counting_recipe(("apply-boundary", 1), calls) + model = Linear(16, 16, bias=False, device="cuda", name="linear") + + try: + with autocast(enabled=False): + with pytest.raises(RuntimeError, match="outside te.autocast"): + apply_recipe(model, recipe) + assert not calls + assert model._quantization_runtime is None + + monkeypatch.setattr( + FP8GlobalStateManager, + "fp8_graph_capturing", + classmethod(lambda _cls: True), + ) + with pytest.raises(RuntimeError, match="outside CUDA graph capture"): + apply_recipe(model, recipe) + assert not calls + assert model._quantization_runtime is None + finally: + FP8GlobalStateManager.reset() + + def test_candidate_validation_failure_keeps_complete_active_runtime(monkeypatch): """Validation runs before any candidate state is published.""" calls = [] @@ -584,7 +1185,13 @@ def qfactory(role): dtype=torch.float16, unsafe_inputs=True, ) - assert _ensure_runtime(module, replacement_recipe, revision=3, num_gemms=2) + update = _prepare_runtime_update(module, replacement_recipe, revision=3, num_gemms=2) + assert module._validated_quantizer_generations is old_validated_generations + assert module._delayed_scaling_input_quantizer is old_delayed_quantizer + assert module._unsafe_requantization_input_quantizer is old_unsafe_quantizer + assert module._apply_quantization_update( # pylint: disable=protected-access + update + ) replacement_runtime = module._quantization_runtime # pylint: disable=protected-access assert replacement_runtime is not old_runtime assert ( diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index fa481e5e8b4..8784b0a8dd4 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -42,9 +42,10 @@ moe_sort_chunks_by_index, moe_sort_chunks_by_index_with_probs, ) +from transformer_engine.pytorch.quantization import apply_recipe +from transformer_engine.pytorch.quantization import autocast from transformer_engine.pytorch.quantization import fp8_autocast from transformer_engine.pytorch.quantization import fp8_model_init -from transformer_engine.pytorch.quantization import autocast from transformer_engine.pytorch.quantization import quantized_model_init from transformer_engine.pytorch.quantization import is_fp8_available from transformer_engine.pytorch.quantization import is_mxfp8_available diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index edf0571d7d4..aa670643ccb 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -629,6 +629,10 @@ def nvfp4_linear_mxfp8_dpa_factory(role): p2p between sub-groups (e.g., via IBLink). """ + def _get_quantization_runtime_num_gemms(self) -> int: + """Return the fixed CustomRecipe DPA quantizer slot layout.""" + return 3 + def __init__( self, num_attention_heads: int, @@ -883,30 +887,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # global recipe set in autocast() fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): - old_runtime = getattr(self, "_quantization_runtime", None) super().init_fp8_metadata(num_gemms=num_gemms) - if self._quantization_runtime is not old_runtime: - self._invalidate_dpa_runtime_caches() - fwd_quantizers = self.quantizers.get("scaling_fwd", ()) - cache_key = ( - id(self.fp8_meta.get("scaling_fwd")), - tuple(id(quantizer) for quantizer in fwd_quantizers), - fp8_recipe.fp8_format, - fp8_recipe.fp8_dpa, - fp8_recipe.fp8_mha, - ) - if cache_key != self._custom_dpa_local_recipes_cache_key: - self._custom_dpa_local_recipes_cache = _infer_custom_dpa_local_recipes( - fp8_recipe, self.fp8_meta, self.quantizers - ) - self._custom_dpa_local_recipes_cache_key = cache_key - - if self._custom_dpa_local_recipes_cache is None: - # Do not leave labels from an earlier supported quantizer - # family attached after a rebuild to an unsupported family. - self.fp8_meta.pop("local_recipes", None) - else: - self.fp8_meta["local_recipes"] = self._custom_dpa_local_recipes_cache return # switch/append recipe: fp8_recipe stays unchanged, but DPA.fp8_meta["recipe"] may be set to @@ -1121,23 +1102,63 @@ def _invalidate_dpa_runtime_caches(self) -> None: self._qkv_capabilities_cache = None _attention_backends["backend_selection_requires_update"] = True - def _validate_quantization_runtime(self, candidate: _QuantizationRuntime) -> Any: + def _activate_quantization_runtime( + self, + candidate: _QuantizationRuntime, + *, + validation_result: Optional[List[Recipe]] = None, + ) -> None: + """Publish a CustomRecipe DPA runtime and its derived cache state.""" + super()._activate_quantization_runtime( + candidate, + validation_result=validation_result, + ) + self._invalidate_dpa_runtime_caches() + + cache_key = ( + id(self.fp8_meta.get("scaling_fwd")), + tuple(id(quantizer) for quantizer in candidate.forward_quantizers), + candidate.recipe.fp8_format, + candidate.recipe.fp8_dpa, + candidate.recipe.fp8_mha, + ) + local_recipes = validation_result + self._custom_dpa_local_recipes_cache_key = cache_key + self._custom_dpa_local_recipes_cache = local_recipes + if local_recipes is None: + # Do not leave labels from an earlier supported quantizer family + # attached after a rebuild to an unsupported family. + self.fp8_meta.pop("local_recipes", None) + else: + self.fp8_meta["local_recipes"] = local_recipes + + def _validate_quantization_runtime( + self, + candidate: _QuantizationRuntime, + ) -> Optional[List[Recipe]]: """Validate an FP8-attention candidate before publishing it.""" super()._validate_quantization_runtime(candidate) - if not candidate.recipe.fp8_dpa: - return None - if len(candidate.forward_quantizers) != 9 or len(candidate.backward_quantizers) != 6: - raise ValueError( - "FP8 DotProductAttention requires 9 forward and 6 backward quantizer slots." + if candidate.recipe.fp8_dpa: + if len(candidate.forward_quantizers) != 9 or len(candidate.backward_quantizers) != 6: + raise ValueError( + "FP8 DotProductAttention requires 9 forward and 6 backward quantizer slots." + ) + dpa_utils.get_attention_quantizers( + True, + { + "scaling_fwd": candidate.forward_quantizers, + "scaling_bwd": candidate.backward_quantizers, + }, ) - dpa_utils.get_attention_quantizers( - True, + + return _infer_custom_dpa_local_recipes( + candidate.recipe, + {"scaling_fwd": candidate.forward_states[0]}, { "scaling_fwd": candidate.forward_quantizers, "scaling_bwd": candidate.backward_quantizers, }, ) - return None def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: """Return MHA boundary capabilities from the canonical QKV quantizer. diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index de16379b60d..7eecaa57abb 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -10,6 +10,7 @@ import pickle import warnings from collections.abc import Hashable +from dataclasses import dataclass from enum import Enum from abc import ABC, abstractmethod from typing import Any, Dict, Generator, List, Optional, Tuple, Union @@ -79,6 +80,17 @@ "UserBufferQuantizationMode", ] + +@dataclass(frozen=True) +class _QuantizationUpdate: + """A validated runtime update that is ready to commit.""" + + candidate: Optional[_QuantizationRuntime] + validation_result: Any + recipe_config_revision: int + role_revision: int + + _2X_ACC_FPROP = False _2X_ACC_DGRAD = True _2X_ACC_WGRAD = True @@ -887,7 +899,30 @@ def quantize_weight( class TransformerEngineBaseModule(torch.nn.Module, ABC): - """Base TE module.""" + """Base TE module. + + Private quantization runtime lifecycle + -------------------------------------- + ``init_fp8_metadata`` calls ``_ensure_quantization_runtime``. An unchanged + forward returns from its constant-time revision fast path. A cold-path + mismatch follows this staged workflow:: + + _ensure_quantization_runtime + -> _plan_quantization_update + -> resolve roles and requested runtime key + -> recognize an equal semantic runtime + -> reject unsupported transitions + -> _build_quantization_runtime + -> _validate_quantization_runtime + -> _apply_quantization_update + -> synchronize revisions for an equal runtime, or + -> _activate_quantization_runtime for a replacement + + Planning constructs and validates both forward and backward state without + publishing it. Applying a validated update is limited to revision updates + or no-fail publication of precomputed module state. Subclasses customize + candidate validation and activation through their corresponding hooks. + """ def __init__(self, name: Optional[str] = None) -> None: super().__init__() @@ -1136,7 +1171,7 @@ def _resolve_quantizer_roles( state_roles = list(roles) return tuple(state_roles), state_roles - def _prepare_quantization_runtime( + def _build_quantization_runtime( self, *, recipe: Recipe, @@ -1147,7 +1182,7 @@ def _prepare_quantization_runtime( forward_state_roles: Optional[List[Optional[QuantizerRole]]], backward_state_roles: Optional[List[Optional[QuantizerRole]]], ) -> _QuantizationRuntime: - """Build both quantization directions without changing live module state.""" + """Construct both candidate directions without validating or publishing them.""" num_forward_quantizers = num_gemms * 3 num_backward_quantizers = num_gemms * 2 @@ -1207,16 +1242,16 @@ def _prepare_quantization_runtime( ) def _validate_quantization_runtime(self, candidate: _QuantizationRuntime) -> Any: - """Validate a candidate and return optional transient commit data.""" + """Validate a built candidate and return optional precomputed activation data.""" del candidate - def _commit_quantization_runtime( + def _activate_quantization_runtime( self, candidate: _QuantizationRuntime, *, validation_result: Any = None, ) -> None: - """Publish a fully prepared and validated runtime as one unit.""" + """Publish a validated replacement runtime and its compatibility views.""" del validation_result recipe = candidate.recipe forward_state = candidate.forward_states[0] @@ -1242,31 +1277,76 @@ def _commit_quantization_runtime( if workspaces is not None: workspaces.clear() - def _ensure_active_quantization_runtime( + def _get_quantization_runtime_num_gemms(self) -> int: + """Return the module's fixed quantizer slot layout in GEMM units.""" + return 1 + + @staticmethod + def _runtime_has_delayed_scaling(runtime: _QuantizationRuntime) -> bool: + """Whether either direction of a runtime owns delayed-scaling state.""" + return any( + _is_delayed_scaling_state(state) + for state in runtime.forward_states + runtime.backward_states + ) + + @staticmethod + def _reject_delayed_scaling_update() -> None: + """Reject a delayed-state transition and protect any active reduction.""" + FP8GlobalStateManager.abort_current_amax_reduction() + raise RuntimeError( + "Mid-training recipe updates do not support delayed scaling. " + "A delayed-scaling runtime is frozen after initialization; only " + "unchanged execution is supported." + ) + + def _check_quantization_update_supported( + self, + *, + recipe: Recipe, + requested_key: _QuantizationRuntimeKey, + num_gemms: int, + ) -> None: + """Reject unsupported transitions known before candidate construction.""" + active = getattr(self, "_quantization_runtime", None) + if active is None: + return + + active_has_delayed_scaling = self._runtime_has_delayed_scaling(active) + + # Until per-weight semantic keys are available, a recipe change may + # not retarget already-quantized primary storage. Keep this more + # specific diagnostic ahead of the general delayed-scaling boundary. + if getattr(self, "primary_weights_in_fp8", False) and ( + active.key.recipe_config != requested_key.recipe_config + or active.num_gemms != num_gemms + ): + if active_has_delayed_scaling or recipe.delayed(): + FP8GlobalStateManager.abort_current_amax_reduction() + raise RuntimeError( + "Recipe mismatch for quantized primary weights: mid-training recipe " + "configuration changes require per-weight semantic compatibility and " + "are not supported yet." + ) + + # Delayed state participates in global reduction registration and is + # intentionally outside this round's migration contract. Reject any + # effective update before invoking a CustomRecipe factory or preparing + # replacement tensors. + if active_has_delayed_scaling or recipe.delayed(): + self._reject_delayed_scaling_update() + + def _plan_quantization_update( self, *, recipe: Recipe, recipe_config: Hashable, recipe_config_revision: int, num_gemms: int, - ) -> bool: - """Ensure the active runtime matches the requested recipe, roles, and slot layout. - - Return whether a new runtime was committed. - """ + ) -> _QuantizationUpdate: + """Plan and validate a complete cold-path update without publishing it.""" active = getattr(self, "_quantization_runtime", None) role_revision = getattr(self, "_role_revision", 0) - # Constant-time steady-state path: no role construction, factory - # calls, recipe traversal, or candidate validation. - if ( - active is not None - and active.recipe_config_revision == recipe_config_revision - and active.role_revision == role_revision - and active.num_gemms == num_gemms - ): - return False - num_forward_quantizers = num_gemms * 3 num_backward_quantizers = num_gemms * 2 forward_key_roles, forward_state_roles = self._resolve_quantizer_roles( @@ -1283,58 +1363,23 @@ def _ensure_active_quantization_runtime( backward_roles=backward_key_roles, ) - # A module can miss an A -> B -> A sequence. Synchronize its cheap - # revision markers without rebuilding an equal semantic runtime. + # A module can miss an A -> B -> A sequence. Defer synchronizing its + # revision markers until commit so model-wide preparation is mutation-free. if active is not None and active.key == requested_key and active.num_gemms == num_gemms: - active.recipe_config_revision = recipe_config_revision - active.role_revision = role_revision - return False - - def has_delayed_scaling(runtime: _QuantizationRuntime) -> bool: - return any( - _is_delayed_scaling_state(state) - for state in runtime.forward_states + runtime.backward_states - ) - - active_has_delayed_scaling = active is not None and has_delayed_scaling(active) - - # Until per-weight semantic keys are available, a recipe change may - # not retarget already-quantized primary storage. Keep this more - # specific diagnostic ahead of the general delayed-scaling boundary. - if ( - active is not None - and getattr(self, "primary_weights_in_fp8", False) - and ( - active.key.recipe_config != requested_key.recipe_config - or active.num_gemms != num_gemms - ) - ): - if active_has_delayed_scaling or recipe.delayed(): - FP8GlobalStateManager.abort_current_amax_reduction() - raise RuntimeError( - "Recipe mismatch for quantized primary weights: mid-training recipe " - "configuration changes require per-weight semantic compatibility and " - "are not supported yet." + return _QuantizationUpdate( + candidate=None, + validation_result=None, + recipe_config_revision=recipe_config_revision, + role_revision=role_revision, ) - def reject_delayed_scaling_update() -> None: - FP8GlobalStateManager.abort_current_amax_reduction() - raise RuntimeError( - "Mid-training recipe updates do not support delayed scaling. " - "A delayed-scaling runtime is frozen after initialization; only " - "unchanged execution is supported." - ) - - # Delayed state participates in global reduction registration and is - # intentionally outside this round's migration contract. Reject any - # effective update before invoking a CustomRecipe factory or preparing - # replacement tensors. - if active_has_delayed_scaling: - reject_delayed_scaling_update() - if active is not None and not active_has_delayed_scaling and recipe.delayed(): - reject_delayed_scaling_update() + self._check_quantization_update_supported( + recipe=recipe, + requested_key=requested_key, + num_gemms=num_gemms, + ) - candidate = self._prepare_quantization_runtime( + candidate = self._build_quantization_runtime( recipe=recipe, key=requested_key, num_gemms=num_gemms, @@ -1343,18 +1388,71 @@ def reject_delayed_scaling_update() -> None: forward_state_roles=forward_state_roles, backward_state_roles=backward_state_roles, ) - if active is not None and has_delayed_scaling(candidate): - reject_delayed_scaling_update() + if active is not None and self._runtime_has_delayed_scaling(candidate): + self._reject_delayed_scaling_update() # Subclasses may return transient commit data even though the base implementation # returns None, which pylint cannot infer through dynamic dispatch. # pylint: disable-next=assignment-from-no-return validation_result = self._validate_quantization_runtime(candidate) - self._commit_quantization_runtime( - candidate, + return _QuantizationUpdate( + candidate=candidate, validation_result=validation_result, + recipe_config_revision=recipe_config_revision, + role_revision=role_revision, + ) + + def _apply_quantization_update( + self, + update: _QuantizationUpdate, + ) -> bool: + """Apply a validated update by synchronizing revisions or activating its candidate.""" + candidate = update.candidate + if candidate is None: + active = self._quantization_runtime + assert active is not None + active.recipe_config_revision = update.recipe_config_revision + active.role_revision = update.role_revision + return False + + self._activate_quantization_runtime( + candidate, + validation_result=update.validation_result, ) return True + def _ensure_quantization_runtime( + self, + *, + recipe: Recipe, + recipe_config: Hashable, + recipe_config_revision: int, + num_gemms: int, + ) -> bool: + """Use the revision hot path or plan and immediately apply a module-local update. + + Return whether a new runtime was committed. + """ + active = getattr(self, "_quantization_runtime", None) + role_revision = getattr(self, "_role_revision", 0) + + # Constant-time steady-state path: no role construction, factory + # calls, recipe traversal, or candidate validation. + if ( + active is not None + and active.recipe_config_revision == recipe_config_revision + and active.role_revision == role_revision + and active.num_gemms == num_gemms + ): + return False + + update = self._plan_quantization_update( + recipe=recipe, + recipe_config=recipe_config, + recipe_config_revision=recipe_config_revision, + num_gemms=num_gemms, + ) + return self._apply_quantization_update(update) + def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Initialize both quantization directions atomically. @@ -1364,7 +1462,7 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: # TODO(negvet): Remove set_meta_tensor after checkpoint, GroupedLinear, and DPA migration. del fwd num_gemms = self.fp8_meta.get("num_gemms", 1) - self._ensure_active_quantization_runtime( + self._ensure_quantization_runtime( recipe=recipe, recipe_config=recipe.quantizer_config(), recipe_config_revision=-1, @@ -1747,7 +1845,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: if fp8_enabled: meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() - runtime_changed = self._ensure_active_quantization_runtime( + runtime_changed = self._ensure_quantization_runtime( recipe=recipe, recipe_config=recipe_config, recipe_config_revision=recipe_config_revision, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 0dcf7de0777..12d1498194a 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1765,6 +1765,10 @@ class GroupedLinear(TransformerEngineBaseModule): The TP communication should be handled in the dispatch and combine stages of MoE models. """ + def _get_quantization_runtime_num_gemms(self) -> int: + """Return the configured grouped-GEMM quantizer slot layout.""" + return self.num_gemms + def __init__( self, num_gemms: int, @@ -1995,7 +1999,7 @@ def _validate_quantization_runtime( ) return validation_result - def _commit_quantization_runtime( + def _activate_quantization_runtime( self, candidate: _QuantizationRuntime, *, @@ -2008,7 +2012,7 @@ def _commit_quantization_runtime( "scaling_bwd": candidate.backward_quantizers, } - super()._commit_quantization_runtime( + super()._activate_quantization_runtime( candidate, validation_result=validation_result, ) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 83080eb14f4..3c833929946 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -1946,6 +1946,10 @@ class LayerNormMLP(TransformerEngineBaseModule): for memory. default is false, in which activations are saved in fwd. not supported for onnx forward """ + def _get_quantization_runtime_num_gemms(self) -> int: + """Return the fixed FC1/FC2 quantizer slot layout.""" + return 2 + def __init__( self, hidden_size: int, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 79498cb9094..e550af7c3a8 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -38,6 +38,7 @@ from .quantized_tensor import Quantizer __all__ = [ + "apply_recipe", "autocast", "quantized_model_init", "is_fp8_available", @@ -996,6 +997,123 @@ def restore_fp8_meta_tensors(fp8_meta: Dict[str, Any]) -> None: fp8_meta["scaling_fwd"].scale.copy_(fp8_meta["updated_scale_fwd"]) +def apply_recipe(model: torch.nn.Module, recipe: Recipe) -> None: + """Prepare, validate, and apply a recipe across one model atomically. + + This is an optional synchronous cold-path API for controllers that need a + model-wide update instead of lazy per-module migration. Call it outside + :class:`autocast`, CUDA graph capture, and compiled regions, after all + outstanding forward, backward, recompute, optimizer, and communication + work has completed. + + Every participating Transformer Engine module is planned and validated + before any module or global recipe state is changed. If planning fails, + all active runtimes and the global recipe remain unchanged. Fusible + operations and the legacy built-in DPA recipe path are not yet supported by + this model-wide API. + + This operation is atomic only within the calling process. In distributed + training, the framework or controller must distribute an identical recipe, + invoke this function consistently on the appropriate ranks at a synchronized + boundary, and coordinate failures before training resumes. Transformer + Engine does not infer data-, tensor-, pipeline-, or expert-parallel process + groups or perform cross-rank agreement in this API. + + When resuming from a checkpoint, reconstruct the intended recipe and call + this function after restoring the model and before its first forward pass. + + Parameters + ---------- + model : torch.nn.Module + Model or composed module whose participating Transformer Engine + runtime owners should receive the recipe. + recipe : transformer_engine.common.recipe.Recipe + Complete precision recipe to apply. + """ + if not isinstance(model, torch.nn.Module): + raise TypeError(f"model must be a torch.nn.Module, got {type(model).__name__}") + if not isinstance(recipe, Recipe): + raise TypeError(f"recipe must be a Recipe, got {type(recipe).__name__}") + if torch.compiler.is_compiling(): + raise RuntimeError("te.apply_recipe() must be called outside torch.compile regions.") + + qstate = FP8GlobalStateManager.quantization_state + if qstate.autocast_depth != 0: + raise RuntimeError("te.apply_recipe() must be called outside te.autocast regions.") + if FP8GlobalStateManager.fp8_graph_capturing(): + raise RuntimeError("te.apply_recipe() must be called outside CUDA graph capture.") + + check_recipe_support(recipe) + recipe_config = recipe.quantizer_config() + recipe_config_revision = qstate.quantizer_config_revision + int( + recipe_config != qstate.quantizer_config + ) + + # Import locally to keep quantization.py independent from module/base.py + # during package initialization. + from .attention.dot_product_attention.dot_product_attention import DotProductAttention + from .module.base import TransformerEngineBaseModule + from .ops.op import FusibleOperation + + participants = [] + fusible_owners = [] + legacy_dpa_owners = [] + seen = set() + for fqn, module in model.named_modules(): + module_id = id(module) + if module_id in seen: + continue + seen.add(module_id) + diagnostic_name = fqn or "" + if isinstance(module, FusibleOperation): + fusible_owners.append(diagnostic_name) + if isinstance(module, DotProductAttention) and not recipe.custom(): + legacy_dpa_owners.append(diagnostic_name) + if isinstance(module, TransformerEngineBaseModule): + participants.append((diagnostic_name, module)) + + # Reject excluded owners after one complete discovery pass and before any + # qfactory invocation or candidate construction. + if fusible_owners: + owners = ", ".join(repr(name) for name in fusible_owners) + raise RuntimeError( + "te.apply_recipe() does not support fusible operations yet; " + f"recreate or update these owners separately: {owners}." + ) + if legacy_dpa_owners: + owners = ", ".join(repr(name) for name in legacy_dpa_owners) + raise RuntimeError( + "te.apply_recipe() supports DotProductAttention only through CustomRecipe; " + f"the built-in NVTE_DPA_* path remains lazy and unchanged for: {owners}." + ) + if not participants: + raise ValueError("te.apply_recipe() found no Transformer Engine runtime owners in model.") + + updates = [] + for fqn, module in participants: + try: + # pylint: disable-next=protected-access + num_gemms = module._get_quantization_runtime_num_gemms() + update = module._plan_quantization_update( # pylint: disable=protected-access + recipe=recipe, + recipe_config=recipe_config, + recipe_config_revision=recipe_config_revision, + num_gemms=num_gemms, + ) + except Exception as exc: + raise RuntimeError( + f"te.apply_recipe() failed while planning module {fqn!r}: {exc}" + ) from exc + updates.append((module, update)) + + # Applying an update only publishes state that was fully constructed and + # validated above. Publish the manager state last so planning failures + # cannot expose a requested recipe globally. + for module, update in updates: + module._apply_quantization_update(update) # pylint: disable=protected-access + FP8GlobalStateManager._set_recipe_config(recipe, recipe_config) + + @contextmanager def fp8_model_init( enabled: bool = True, From 364abe49aa27486575b6a7497ef9c562424a9052 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 13 Aug 2026 13:30:13 +0000 Subject: [PATCH 74/97] Minor fixes Signed-off-by: Evgeny --- tests/pytorch/hybrid_quantization_utils.py | 13 +++++++ .../pytorch/test_hybrid_quantization_fsdp2.py | 35 ++++++++++++++++--- tests/pytorch/test_recipe_update.py | 33 ++++++++--------- transformer_engine/pytorch/module/base.py | 6 ++-- 4 files changed, 65 insertions(+), 22 deletions(-) diff --git a/tests/pytorch/hybrid_quantization_utils.py b/tests/pytorch/hybrid_quantization_utils.py index adcbee592f8..163a6bf28ab 100644 --- a/tests/pytorch/hybrid_quantization_utils.py +++ b/tests/pytorch/hybrid_quantization_utils.py @@ -231,6 +231,7 @@ def qfactory(role): return recipe.CustomRecipe(qfactory=qfactory, qfactory_key=qfactory_key) +@recipe.quantizer_factory(key=("test_hybrid_fp8_current", 1)) def hybrid_fp8_current_qfactory(role): """FP8 current scaling in both hybrid directions for forward tensor roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -241,6 +242,7 @@ def hybrid_fp8_current_qfactory(role): return current_scaling_factory(role) +@recipe.quantizer_factory(key=("test_hybrid_fp8_current_e5m2_grads", 1)) def hybrid_fp8_current_e5m2_grads_qfactory(role): """FP8 current-scaling hybrid with E5M2 for both grad boundary roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -253,6 +255,7 @@ def hybrid_fp8_current_e5m2_grads_qfactory(role): return _make_fp8_current() +@recipe.quantizer_factory(key=("test_hybrid_mxfp8", 1)) def hybrid_mxfp8_qfactory(role): """MXFP8 in both hybrid directions for forward tensor roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -263,6 +266,7 @@ def hybrid_mxfp8_qfactory(role): return mxfp8_factory(role) +@recipe.quantizer_factory(key=("test_hybrid_float8_block", 1)) def hybrid_float8_block_qfactory(role): """Float8 block scaling in both hybrid directions for forward roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -273,6 +277,7 @@ def hybrid_float8_block_qfactory(role): return float8_block_scaling_factory(role) +@recipe.quantizer_factory(key=("test_hybrid_block_fp8_e4m3", 1)) def hybrid_block_fp8_e4m3_qfactory(role): """Hybrid E4M3 Block-FP8 with role-aware 1D/2D block geometry.""" is_linear = _is_linear_role(role) @@ -295,6 +300,7 @@ def make_quantizer(): ) +@recipe.quantizer_factory(key=("test_hybrid_mixed_mxfp8_fp8", 1)) def hybrid_mixed_mxfp8_fp8_qfactory(role): """MXFP8 rowwise plus FP8 current-scaling columnwise.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -305,6 +311,7 @@ def hybrid_mixed_mxfp8_fp8_qfactory(role): return current_scaling_factory(role) +@recipe.quantizer_factory(key=("test_hybrid_fp8_current_identity", 1)) def hybrid_fp8_current_identity_qfactory(role): """FP8 current-scaling forward plus Identity backward.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -317,6 +324,7 @@ def hybrid_fp8_current_identity_qfactory(role): return current_scaling_factory(role) +@recipe.quantizer_factory(key=("test_hybrid_mxfp8_identity", 1)) def hybrid_mxfp8_identity_qfactory(role): """MXFP8 forward plus Identity backward.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -329,11 +337,13 @@ def hybrid_mxfp8_identity_qfactory(role): return mxfp8_factory(role) +@recipe.quantizer_factory(key=("test_identity", 1)) def identity_qfactory(role): # pylint: disable=unused-argument """High-precision passthrough for every quantizer slot.""" return te.IdentityQuantizer() +@recipe.quantizer_factory(key=("test_hybrid_nvfp4", 1)) def hybrid_nvfp4_qfactory(role): """NVFP4 in both hybrid directions for forward tensor roles.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -344,6 +354,7 @@ def hybrid_nvfp4_qfactory(role): return nvfp4_factory(role) +@recipe.quantizer_factory(key=("test_hybrid_tp_mxfp8_nvfp4", 1)) def hybrid_tp_mxfp8_nvfp4_qfactory(role): """TP/SP MXFP8 rowwise plus NVFP4 columnwise, including boundary roles.""" if _is_linear_role(role) and role.tensor_type in _GRAD_TENSOR_TYPES: @@ -354,6 +365,7 @@ def hybrid_tp_mxfp8_nvfp4_qfactory(role): ) +@recipe.quantizer_factory(key=("test_hybrid_fp8_mxfp8", 1)) def hybrid_fp8_mxfp8_qfactory(role): """FP8 current-scaling rowwise plus MXFP8 columnwise for CPU-offload tests.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: @@ -366,6 +378,7 @@ def hybrid_fp8_mxfp8_qfactory(role): return _make_fp8_current() +@recipe.quantizer_factory(key=("test_hybrid_mxfp8_nvfp4", 1)) def hybrid_mxfp8_nvfp4_qfactory(role): """MXFP8 rowwise plus NVFP4 columnwise for CPU-offload tests.""" if _is_linear_role(role) and role.tensor_type in _FORWARD_TENSOR_TYPES: diff --git a/tests/pytorch/test_hybrid_quantization_fsdp2.py b/tests/pytorch/test_hybrid_quantization_fsdp2.py index cda3e7184af..c148ba55cc9 100644 --- a/tests/pytorch/test_hybrid_quantization_fsdp2.py +++ b/tests/pytorch/test_hybrid_quantization_fsdp2.py @@ -89,7 +89,17 @@ def _make_hybrid_param_for_dispatch( row_factory, col_factory, grad_factory=None, in_features=256, out_features=256 ): """Create a HybridQuantizedTensor weight via quantized_model_init for dispatch tests.""" - hybrid_recipe = _hybrid_custom_recipe(row_factory, col_factory, grad_factory) + factory_names = tuple( + getattr(factory, "__name__", type(factory).__name__) + for factory in (row_factory, col_factory, grad_factory) + if factory is not None + ) + hybrid_recipe = _hybrid_custom_recipe( + row_factory, + col_factory, + grad_factory, + qfactory_key=("fsdp2-dispatch", factory_names), + ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(in_features, out_features, params_dtype=torch.bfloat16).cuda() return model.weight @@ -772,11 +782,24 @@ def test_clone_returns_hybrid(self, hybrid_param): def _make_fsdp_protocol_param(config_name): """Create a HybridQuantizedTensor weight for FSDP protocol tests.""" if config_name == "fp8_fp8": - r = _hybrid_custom_recipe(_fp8_row_factory, _fp8_col_factory, _fp8_grad_factory) + r = _hybrid_custom_recipe( + _fp8_row_factory, + _fp8_col_factory, + _fp8_grad_factory, + qfactory_key=("fsdp2-protocol", config_name), + ) elif config_name == "mxfp8_fp8": - r = _hybrid_custom_recipe(_mxfp8_factory, _fp8_col_factory, _fp8_grad_factory) + r = _hybrid_custom_recipe( + _mxfp8_factory, + _fp8_col_factory, + _fp8_grad_factory, + qfactory_key=("fsdp2-protocol", config_name), + ) elif config_name == "block_fp8": - r = recipe.CustomRecipe(qfactory=_hybrid_block_fp8_qfactory) + r = recipe.CustomRecipe( + qfactory=_hybrid_block_fp8_qfactory, + qfactory_key=("fsdp2-protocol", config_name), + ) else: raise ValueError(f"Unknown config: {config_name}") with quantized_model_init(enabled=True, recipe=r): @@ -1080,6 +1103,7 @@ def test_scale_refresh_across_iterations(self): _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + qfactory_key=("fsdp2-scale-refresh", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -1160,6 +1184,7 @@ def test_nvfp4_sub_storage_raises_on_pre_all_gather(self): hybrid_recipe = _hybrid_custom_recipe( row_factory=lambda: NVFP4Quantizer(), col_factory=lambda: NVFP4Quantizer(), + qfactory_key=("fsdp2-nvfp4-rejection", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -1199,6 +1224,7 @@ def _make_hybrid_param(self): _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + qfactory_key=("fsdp2-make-like", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(256, 256, params_dtype=torch.bfloat16).cuda() @@ -1471,6 +1497,7 @@ def _make_param(self): _fp8_row_factory, _fp8_col_factory, _fp8_grad_factory, + qfactory_key=("fsdp2-post-gather-update-usage", 1), ) with quantized_model_init(enabled=True, recipe=hybrid_recipe): model = Linear(64, 64, params_dtype=torch.bfloat16).cuda() diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 65bfb5f82d2..8e2d5635eea 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -334,9 +334,7 @@ def test_revision_only_update_is_not_published_during_planning(): assert active.recipe_config_revision == 1 assert len(calls) == factory_call_count - assert not module._apply_quantization_update( # pylint: disable=protected-access - update - ) + assert not module._apply_quantization_update(update) # pylint: disable=protected-access assert active.recipe_config_revision == 3 assert len(calls) == factory_call_count @@ -359,13 +357,23 @@ def test_candidate_update_is_not_published_during_planning(): assert old_views[0].recipe_config_revision == 1 assert module._fp8_workspaces["weight"] is workspace # pylint: disable=protected-access - assert module._apply_quantization_update( # pylint: disable=protected-access - update - ) + assert module._apply_quantization_update(update) # pylint: disable=protected-access assert module._quantization_runtime is update.candidate # pylint: disable=protected-access assert not module._fp8_workspaces # pylint: disable=protected-access +def test_runtime_snapshot_does_not_copy_caller_repr_cache(): + """Caller display history must not affect a runtime snapshot or its checkpoint bytes.""" + recipe = _make_counting_recipe(("runtime-repr-cache", 1), []) + str(recipe) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + + update = _prepare_runtime_update(module, recipe, revision=1) + + assert recipe.__dict__["_cached_repr"] is not None + assert update.candidate.recipe.__dict__["_cached_repr"] is None + + def test_uninitialized_runtime_can_be_prepared_without_publication(): """Planning supports modules that have not yet executed a quantized forward.""" calls = [] @@ -380,9 +388,7 @@ def test_uninitialized_runtime_can_be_prepared_without_publication(): assert "scaling_bwd" not in module.fp8_meta assert module.quantizers == {"scaling_fwd": [], "scaling_bwd": []} - assert module._apply_quantization_update( # pylint: disable=protected-access - update - ) + assert module._apply_quantization_update(update) # pylint: disable=protected-access assert module._quantization_runtime is update.candidate # pylint: disable=protected-access @@ -952,10 +958,7 @@ def test_apply_recipe_planning_failure_is_model_wide_atomic(failure_index): FP8GlobalStateManager.reset() calls = [] active_recipe = _make_counting_recipe(("apply-failure", 1), calls) - modules = [ - Linear(16, 16, bias=False, device="cuda", name=f"line{index}") - for index in range(3) - ] + modules = [Linear(16, 16, bias=False, device="cuda", name=f"line{index}") for index in range(3)] model = torch.nn.Sequential(*modules) try: @@ -1189,9 +1192,7 @@ def qfactory(role): assert module._validated_quantizer_generations is old_validated_generations assert module._delayed_scaling_input_quantizer is old_delayed_quantizer assert module._unsafe_requantization_input_quantizer is old_unsafe_quantizer - assert module._apply_quantization_update( # pylint: disable=protected-access - update - ) + assert module._apply_quantization_update(update) # pylint: disable=protected-access replacement_runtime = module._quantization_runtime # pylint: disable=protected-access assert replacement_runtime is not old_runtime assert ( diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 7eecaa57abb..e704256a891 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1191,6 +1191,9 @@ def _build_quantization_runtime( # candidate validated, even if the caller later mutates and reuses the # same recipe instance. runtime_recipe = copy.copy(recipe) + # Do not snapshot a caller-populated display cache. Whether ``repr`` was + # requested before runtime construction must not change checkpoint bytes. + object.__setattr__(runtime_recipe, "_cached_repr", None) forward_state = RecipeState.create( # pylint: disable=assignment-from-none runtime_recipe, @@ -1317,8 +1320,7 @@ def _check_quantization_update_supported( # not retarget already-quantized primary storage. Keep this more # specific diagnostic ahead of the general delayed-scaling boundary. if getattr(self, "primary_weights_in_fp8", False) and ( - active.key.recipe_config != requested_key.recipe_config - or active.num_gemms != num_gemms + active.key.recipe_config != requested_key.recipe_config or active.num_gemms != num_gemms ): if active_has_delayed_scaling or recipe.delayed(): FP8GlobalStateManager.abort_current_amax_reduction() From 8e7873fa092de65a042d2bddd7aea0853e6c0576 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 19 Aug 2026 10:21:56 +0000 Subject: [PATCH 75/97] Handle boundaries before runtime Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 134 +++-- tests/pytorch/test_identity_quantizer.py | 12 +- tests/pytorch/test_recipe.py | 44 +- tests/pytorch/test_recipe_update.py | 516 +++++++++++++++++- tests/pytorch/test_torch_compile.py | 3 +- .../dot_product_attention.py | 45 +- .../attention/dot_product_attention/utils.py | 4 +- .../pytorch/attention/multi_head_attention.py | 102 ++-- transformer_engine/pytorch/module/base.py | 60 +- .../pytorch/module/grouped_linear.py | 12 +- .../pytorch/module/layernorm_linear.py | 10 +- .../pytorch/module/layernorm_mlp.py | 15 +- transformer_engine/pytorch/module/linear.py | 10 +- 13 files changed, 792 insertions(+), 175 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index fb1ac1ccf1b..5f900c844ea 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -2616,19 +2616,17 @@ def _run_mha_model(self, model, inp, grad, fp8_recipe, seed): ) @pytest.mark.parametrize( - "case_name,native_dpa_recipe,qfactory,expected_flags", + "case_name,native_dpa_recipe,qfactory", [ ( "fp8_dpa", "Float8CurrentScaling", nvfp4_linear_fp8_dpa_factory, - (False, False, False), ), ( "mxfp8_dpa", "MXFP8BlockScaling", _nvfp4_linear_mxfp8_dpa_factory, - (False, False, False), ), ], ) @@ -2638,31 +2636,16 @@ def test_multihead_attention_matches_native_env_recipe_bitwise( case_name, native_dpa_recipe, qfactory, - expected_flags, ): if case_name == "mxfp8_dpa" and not mxfp8_available: pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") - from transformer_engine.pytorch.attention import multi_head_attention as mha_module from transformer_engine.pytorch.utils import get_device_compute_capability cc = get_device_compute_capability() if cc < (9, 0) or cc >= (12, 0): pytest.skip(f"FP8 attention not supported on sm{cc[0] * 10 + cc[1]}") - recorded_flags = [] - orig_update_roles = mha_module.MultiheadAttention._update_output_quantizer_roles - - def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output): - recorded_flags.append((qkv_fp8_output, dpa_fp8_output, proj_fp8_grad)) - return orig_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) - - monkeypatch.setattr( - mha_module.MultiheadAttention, - "_update_output_quantizer_roles", - _recording_update_roles, - ) - self._set_native_dpa_recipe(monkeypatch, native_dpa_recipe) torch.manual_seed(2301) @@ -2710,7 +2693,6 @@ def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) native_recipe, seed=2303, ) - native_flags = recorded_flags[-1] self._clear_native_dpa_recipe(monkeypatch) qfactory_out, qfactory_dx, qfactory_grads = self._run_mha_model( model_qfactory, @@ -2719,10 +2701,7 @@ def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) qfactory_recipe, seed=2303, ) - qfactory_flags = recorded_flags[-1] - assert native_flags == expected_flags - assert qfactory_flags == expected_flags self._assert_equal(qfactory_out, native_out, f"{case_name} MHA output") self._assert_equal(qfactory_dx, native_dx, f"{case_name} MHA input grad") assert qfactory_grads.keys() == native_grads.keys() @@ -2733,7 +2712,8 @@ def _recording_update_roles(self, qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) f"{case_name} MHA param grad {name}", ) - def test_update_output_quantizer_roles_wires_independent_boundaries(self): + def test_mha_declared_boundaries_resolve_for_prospective_recipe(self): + """Topology is recipe-independent; DPA alone adapts its legacy inactive slots.""" from transformer_engine.pytorch.quantization import QuantizerRole model = te.MultiheadAttention( @@ -2770,25 +2750,105 @@ def test_update_output_quantizer_roles_wires_independent_boundaries(self): name=qkv.name or "", ) - def boundary_roles(): - return ( - qkv.output_quantizer_role, - model.proj.grad_input_quantizer_role, - model.core_attention.output_quantizer_role, - model.core_attention.grad_input_quantizer_role, + inactive_recipe = recipe.CustomRecipe( + qfactory=lambda _role: IdentityQuantizer(), + qfactory_key=("mha-declared-boundaries", 0), + fp8_mha=False, + ) + active_recipe = recipe.CustomRecipe( + qfactory=lambda _role: IdentityQuantizer(), + qfactory_key=("mha-declared-boundaries", 1), + fp8_mha=True, + ) + + assert qkv.output_quantizer_role is None + assert model.proj.grad_input_quantizer_role is None + assert model.core_attention.output_quantizer_role is None + assert model.core_attention.grad_input_quantizer_role is None + assert all(owner._role_revision == 0 for owner in (qkv, model.proj, model.core_attention)) + + def resolved_roles(owner, target_recipe, *, fwd, num_quantizers): + _, roles = owner._resolve_quantizer_roles( # pylint: disable=protected-access + recipe=target_recipe, + fwd=fwd, + num_quantizers=num_quantizers, ) + assert roles is not None + return roles - model._update_output_quantizer_roles(True, False, False) - assert boundary_roles() == (expected_qkv, None, None, None) + assert resolved_roles(qkv, inactive_recipe, fwd=True, num_quantizers=3)[-1] == expected_qkv + assert ( + resolved_roles(model.proj, inactive_recipe, fwd=False, num_quantizers=2)[-1] + == expected_do + ) - model._update_output_quantizer_roles(False, True, False) - assert boundary_roles() == (None, expected_do, None, None) + # O/dQKV are also internal fused-attention descriptor slots. The + # DPA-local compatibility adapter retains its hint roles while the + # legacy CustomRecipe fp8_mha switch keeps those boundaries in BF16. + dpa_name = model.core_attention.name or "" + expected_o_hint = QuantizerRole(name=f"{dpa_name}.dpa_output" if dpa_name else "dpa_output") + expected_dqkv_hint = QuantizerRole( + name=f"{dpa_name}.dpa_grad_input" if dpa_name else "dpa_grad_input" + ) + assert ( + resolved_roles( + model.core_attention, + inactive_recipe, + fwd=True, + num_quantizers=9, + )[3:6] + == [expected_o_hint] * 3 + ) + assert ( + resolved_roles( + model.core_attention, + inactive_recipe, + fwd=False, + num_quantizers=6, + )[:2] + == [expected_dqkv_hint] * 2 + ) - model._update_output_quantizer_roles(False, False, True) - assert boundary_roles() == (None, None, expected_o, expected_dqkv) + explicit_o = QuantizerRole( + module_type="linear", + tensor_type="input", + name="explicit-dpa-consumer", + ) + model.core_attention.output_quantizer_role = explicit_o + assert ( + resolved_roles( + model.core_attention, + inactive_recipe, + fwd=True, + num_quantizers=9, + )[3:6] + == [explicit_o] * 3 + ) + model.core_attention.output_quantizer_role = None - model._update_output_quantizer_roles(False, False, False) - assert boundary_roles() == (None, None, None, None) + assert resolved_roles(qkv, active_recipe, fwd=True, num_quantizers=3)[-1] == expected_qkv + assert ( + resolved_roles(model.proj, active_recipe, fwd=False, num_quantizers=2)[-1] + == expected_do + ) + assert ( + resolved_roles( + model.core_attention, + active_recipe, + fwd=True, + num_quantizers=9, + )[3:6] + == [expected_o] * 3 + ) + assert ( + resolved_roles( + model.core_attention, + active_recipe, + fwd=False, + num_quantizers=6, + )[:2] + == [expected_dqkv] * 2 + ) def test_mxfp8_qfactory_uses_plain_bf16_mha_boundaries(self, monkeypatch): """MXFP8 DPA stays internal; MHA boundary tensors remain plain BF16.""" diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 3cbfab0aa2e..1de95e0db4b 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -1498,7 +1498,11 @@ def _assert_zoo_layernorm_mlp_role_semantics(case_name): module = _make_zoo_dequantized_module("LayerNormMLP") qfactory = _zoo_dequantized_qfactory(case_name) - fwd_roles = module.get_quantizer_roles(fwd=True, num_quantizers=6) + fwd_roles = module.get_quantizer_roles( + fwd=True, + num_quantizers=6, + boundary_role=None, + ) fwd_gemm_roles = [ role for role in fwd_roles if role is not None and role.tensor_type in ("input", "weight") ] @@ -1514,7 +1518,11 @@ def _assert_zoo_layernorm_mlp_role_semantics(case_name): assert isinstance(quantizer.rowwise_quantizer, NVFP4Quantizer) assert quantizer.rowwise_quantizer.row_scaled_nvfp4 == (role.tensor_type == "input") - bwd_roles = module.get_quantizer_roles(fwd=False, num_quantizers=4) + bwd_roles = module.get_quantizer_roles( + fwd=False, + num_quantizers=4, + boundary_role=None, + ) grad_output_roles = [ role for role in bwd_roles if role is not None and role.tensor_type == "grad_output" ] diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index bdfecb3983a..1bf988e7018 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -359,19 +359,35 @@ def settings(mode, roles): ).make_quantizers() return [(q.force_pow_2_scales, q.amax_epsilon) for q in quantizers] - module_owner = SimpleNamespace( - name="test", - _output_quantizer_role=None, - _grad_input_quantizer_role=None, - ) + module_owner = SimpleNamespace(name="test") for module_type in (Linear, LayerNormLinear): - forward_roles = module_type.get_quantizer_roles(module_owner, fwd=True, num_quantizers=3) - backward_roles = module_type.get_quantizer_roles(module_owner, fwd=False, num_quantizers=2) + forward_roles = module_type.get_quantizer_roles( + module_owner, + fwd=True, + num_quantizers=3, + boundary_role=None, + ) + backward_roles = module_type.get_quantizer_roles( + module_owner, + fwd=False, + num_quantizers=2, + boundary_role=None, + ) assert settings("forward", forward_roles) == [inp, weight, boundary] assert settings("backward", backward_roles) == [grad, boundary] - mlp_forward_roles = LayerNormMLP.get_quantizer_roles(module_owner, fwd=True, num_quantizers=6) - mlp_backward_roles = LayerNormMLP.get_quantizer_roles(module_owner, fwd=False, num_quantizers=4) + mlp_forward_roles = LayerNormMLP.get_quantizer_roles( + module_owner, + fwd=True, + num_quantizers=6, + boundary_role=None, + ) + mlp_backward_roles = LayerNormMLP.get_quantizer_roles( + module_owner, + fwd=False, + num_quantizers=4, + boundary_role=None, + ) assert settings("forward", mlp_forward_roles) == [ inp, weight, @@ -383,10 +399,16 @@ def settings(mode, roles): assert settings("backward", mlp_backward_roles) == [grad, boundary, grad, grad] grouped_forward_roles = GroupedLinear.get_quantizer_roles( - module_owner, fwd=True, num_quantizers=6 + module_owner, + fwd=True, + num_quantizers=6, + boundary_role=None, ) grouped_backward_roles = GroupedLinear.get_quantizer_roles( - module_owner, fwd=False, num_quantizers=4 + module_owner, + fwd=False, + num_quantizers=4, + boundary_role=None, ) assert settings("forward", grouped_forward_roles) == [ inp, diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 8e2d5635eea..f12c641f716 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -20,6 +20,8 @@ apply_recipe, autocast, is_fp8_available, + is_nvfp4_available, + quantized_model_init, ) from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import ( @@ -89,11 +91,11 @@ def _runtime_views(module): ) -def _run_update_step(module, recipe, inp, *args): +def _run_update_step(module, recipe, inp, *args, **kwargs): """Run a real forward/backward step and return the output.""" module.zero_grad(set_to_none=True) with autocast(enabled=True, recipe=recipe): - output = module(inp, *args) + output = module(inp, *args, **kwargs) if isinstance(output, tuple): output = output[0] output.float().sum().backward() @@ -122,6 +124,37 @@ def _global_recipe_state(): ) +def _make_compatible_fp8_mha_recipe(key): + """Build an executable FP8-MHA recipe with compatible boundary formats.""" + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + delayed_scaling_factory, + ) + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + from transformer_engine.pytorch.utils import get_device_compute_capability + + compute_capability = get_device_compute_capability() + + def qfactory(role): + # Hopper supports the coherent all-delayed configuration. On Blackwell, + # use the shipped native FP8-attention mix for DPA-owned slots and + # current scaling for linears and the two DPA/linear boundaries. + if compute_capability < (10, 0): + return delayed_scaling_factory(role) + if role is not None and role.module_type == "dpa": + return nvfp4_linear_fp8_dpa_factory(role) + return current_scaling_factory(role) + + return CustomRecipe( + qfactory=qfactory, + qfactory_key=key, + fp8_dpa=True, + fp8_mha=True, + ) + + @pytest.mark.parametrize( "replacement", ( @@ -779,6 +812,485 @@ def test_composed_module_executes_after_recipe_update(module_factory): assert runtime.key.recipe_config == second_recipe.quantizer_config() +@pytest.mark.parametrize( + "module_factory", + ( + pytest.param( + lambda: MultiheadAttention( + hidden_size=128, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="mha", + ), + id="multihead-attention", + ), + pytest.param( + lambda: TransformerLayer( + hidden_size=128, + ffn_hidden_size=256, + num_attention_heads=2, + hidden_dropout=0.0, + attention_dropout=0.0, + self_attn_mask_type="no_mask", + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + ), + id="transformer-layer", + ), + ), +) +def test_apply_recipe_before_first_mha_forward_has_stable_boundary_topology( + module_factory, +): + """E2E TE-1 regression: first forward must not migrate unwarmed MHA roles.""" + from transformer_engine.pytorch.utils import get_device_compute_capability + + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + compute_capability = get_device_compute_capability() + if compute_capability < (9, 0) or compute_capability >= (12, 0): + pytest.skip("FP8 attention is not supported on this compute capability") + + FP8GlobalStateManager.reset() + module = module_factory() + recipe = _make_compatible_fp8_mha_recipe(("apply-unwarmed-mha", type(module).__name__, 1)) + + try: + assert not _active_runtime_owners(module) + apply_recipe(module, recipe) + owners = _active_runtime_owners(module) + assert owners + before = [ + ( + owner._quantization_runtime, # pylint: disable=protected-access + owner._quantization_runtime.key, # pylint: disable=protected-access + owner._role_revision, # pylint: disable=protected-access + owner._quantization_runtime.role_revision, # pylint: disable=protected-access + ) + for owner in owners + ] + assert all(requested == active for _, _, requested, active in before) + + mha = next(child for child in module.modules() if isinstance(child, MultiheadAttention)) + qkv = mha.layernorm_qkv if mha.input_layernorm else mha.qkv + assert qkv._quantization_runtime.key.forward_roles[ + -1 + ] == QuantizerRole( # pylint: disable=protected-access + module_type="dpa", + tensor_type="qkv", + name=mha.core_attention.name or "", + ) + assert mha.proj._quantization_runtime.key.backward_roles[ + -1 + ] == QuantizerRole( # pylint: disable=protected-access + module_type="dpa", + tensor_type="do", + name=mha.core_attention.name or "", + ) + + for _ in range(2): + inp = torch.randn( + 128, + 2, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + _run_update_step(module, recipe, inp) + after = [ + ( + owner._quantization_runtime, # pylint: disable=protected-access + owner._quantization_runtime.key, # pylint: disable=protected-access + owner._role_revision, # pylint: disable=protected-access + owner._quantization_runtime.role_revision, # pylint: disable=protected-access + ) + for owner in owners + ] + assert after == before + finally: + FP8GlobalStateManager.reset() + + +def test_mha_warmed_and_unwarmed_topology_matches_and_rope_does_not_migrate_runtime(): + """Cover the TE-1 invariants not asserted by existing attention numerics tests.""" + from transformer_engine.pytorch import RotaryPositionEmbedding + from transformer_engine.pytorch.utils import get_device_compute_capability + + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + compute_capability = get_device_compute_capability() + if compute_capability < (9, 0) or compute_capability >= (12, 0): + pytest.skip("FP8 attention is not supported on this compute capability") + + FP8GlobalStateManager.reset() + recipe = _make_compatible_fp8_mha_recipe(("mha-warm-rope-topology", 1)) + + def make_mha(): + return MultiheadAttention( + hidden_size=128, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="mha", + ) + + def runtime_signature(module): + signature = [] + for fqn, owner in module.named_modules(): + if not isinstance(owner, TransformerEngineBaseModule): + continue + runtime = owner._quantization_runtime # pylint: disable=protected-access + assert runtime is not None + signature.append( + ( + fqn, + runtime.key.forward_roles, + runtime.key.backward_roles, + owner._role_revision, # pylint: disable=protected-access + runtime.role_revision, + tuple(type(quantizer) for quantizer in runtime.forward_quantizers), + tuple(type(quantizer) for quantizer in runtime.backward_quantizers), + ) + ) + return signature + + try: + warmed = make_mha() + unwarmed = make_mha() + + warm_inp = torch.randn( + 128, + 2, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + _run_update_step(warmed, recipe, warm_inp) + assert not _active_runtime_owners(unwarmed) + + apply_recipe(unwarmed, recipe) + apply_recipe(warmed, recipe) + assert runtime_signature(unwarmed) == runtime_signature(warmed) + + # Toggle RoPE on the already-warmed instance. The old forward-time + # wiring changed QKV boundary roles between these two call shapes. + owners = _active_runtime_owners(warmed) + before_rope = [ + ( + owner._quantization_runtime, # pylint: disable=protected-access + owner._quantization_runtime.key, # pylint: disable=protected-access + owner._role_revision, # pylint: disable=protected-access + owner._quantization_runtime.role_revision, # pylint: disable=protected-access + ) + for owner in owners + ] + rotary_pos_emb = RotaryPositionEmbedding(dim=64)(128).to(device="cuda") + rope_inp = torch.randn( + 128, + 2, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + _run_update_step( + warmed, + recipe, + rope_inp, + rotary_pos_emb=rotary_pos_emb, + ) + after_rope = [ + ( + owner._quantization_runtime, # pylint: disable=protected-access + owner._quantization_runtime.key, # pylint: disable=protected-access + owner._role_revision, # pylint: disable=protected-access + owner._quantization_runtime.role_revision, # pylint: disable=protected-access + ) + for owner in owners + ] + assert after_rope == before_rope + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_rejects_e2e_mha_reproducer_during_planning(): + """The shipped factory's incompatible FP8-MHA boundaries fail before commit.""" + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + if not is_nvfp4_available(): + pytest.skip("NVFP4 is required for the E2E TE-1 reproducer") + + FP8GlobalStateManager.reset() + module = MultiheadAttention( + hidden_size=32, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + bias=False, + device="cuda", + name="mha", + ) + recipe = CustomRecipe( + qfactory=nvfp4_linear_fp8_dpa_factory, + fp8_dpa=True, + fp8_mha=True, + ) + owners = [child for child in module.modules() if isinstance(child, TransformerEngineBaseModule)] + old_global_state = _global_recipe_state() + + try: + with pytest.raises(RuntimeError, match="O quantizer is NVFP4Quantizer"): + apply_recipe(module, recipe) + + assert _global_recipe_state() == old_global_state + assert all( + owner._quantization_runtime is None for owner in owners + ) # pylint: disable=protected-access + assert all( + owner._role_revision == 0 for owner in owners + ) # pylint: disable=protected-access + assert all("scaling_fwd" not in owner.fp8_meta for owner in owners) + finally: + FP8GlobalStateManager.reset() + + +@pytest.mark.parametrize("attention_type", ("self", "cross")) +@pytest.mark.parametrize("input_layernorm", (False, True)) +def test_mha_declares_all_boundary_topology_variants(attention_type, input_layernorm): + """Self/cross-attention QKV producers expose stable recipe-resolved roles.""" + module = MultiheadAttention( + hidden_size=32, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + attention_type=attention_type, + input_layernorm=input_layernorm, + bias=False, + device="cuda", + name="mha", + ) + recipe = CustomRecipe( + qfactory=lambda _role: IdentityQuantizer(), + qfactory_key=("mha-topology-variants", attention_type, input_layernorm), + fp8_mha=True, + ) + expected_qkv = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=module.core_attention.name or "", + ) + if attention_type == "self": + qkv_producers = [module.layernorm_qkv if input_layernorm else module.qkv] + else: + qkv_producers = [ + module.layernorm_query if input_layernorm else module.query_layer, + module.key_value, + ] + + def resolved_roles(owner, *, fwd, num_quantizers): + _, roles = owner._resolve_quantizer_roles( # pylint: disable=protected-access + recipe=recipe, + fwd=fwd, + num_quantizers=num_quantizers, + ) + assert roles is not None + return roles + + for producer in qkv_producers: + assert producer.output_quantizer_role is None + assert ( + producer._declared_output_quantizer_role # pylint: disable=protected-access + == expected_qkv + ) + assert resolved_roles(producer, fwd=True, num_quantizers=3)[-1] == expected_qkv + assert producer._role_revision == 0 # pylint: disable=protected-access + + assert resolved_roles(module.proj, fwd=False, num_quantizers=2)[-1] == QuantizerRole( + module_type="dpa", + tensor_type="do", + name=module.core_attention.name or "", + ) + assert resolved_roles(module.core_attention, fwd=True, num_quantizers=9)[3] == QuantizerRole( + module_type="linear", + tensor_type="input", + name=module.proj.name or "", + ) + assert resolved_roles(module.core_attention, fwd=False, num_quantizers=6)[0] == QuantizerRole( + module_type="linear", + tensor_type="grad_output", + name=qkv_producers[0].name or "", + ) + + +def test_explicit_mha_boundary_override_wins_and_can_restore_declared_role(): + """Public role overrides retain revision semantics above declared topology.""" + FP8GlobalStateManager.reset() + module = MultiheadAttention( + hidden_size=32, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + bias=False, + device="cuda", + name="mha", + ) + recipe = CustomRecipe( + qfactory=lambda _role: IdentityQuantizer(), + qfactory_key=("mha-explicit-boundary-override", 1), + fp8_mha=True, + ) + qkv = module.qkv + + try: + apply_recipe(module, recipe) + declared_runtime = qkv._quantization_runtime # pylint: disable=protected-access + declared_role = declared_runtime.key.forward_roles[-1] + + override = QuantizerRole( + module_type="linear", + tensor_type="input", + name="explicit-consumer", + ) + qkv.output_quantizer_role = override + assert qkv._role_revision == 1 # pylint: disable=protected-access + apply_recipe(module, recipe) + override_runtime = qkv._quantization_runtime # pylint: disable=protected-access + assert override_runtime is not declared_runtime + assert override_runtime.key.forward_roles[-1] == override + + qkv.output_quantizer_role = None + assert qkv._role_revision == 2 # pylint: disable=protected-access + apply_recipe(module, recipe) + restored_runtime = qkv._quantization_runtime # pylint: disable=protected-access + assert restored_runtime is not override_runtime + assert restored_runtime.key.forward_roles[-1] == declared_role + finally: + FP8GlobalStateManager.reset() + + +def test_shipped_dpa_factory_unwarmed_mha_keeps_bf16_boundaries_stable(): + """The documented fp8_dpa=True/fp8_mha=False mode remains lazy-forward compatible.""" + from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( + nvfp4_linear_fp8_dpa_factory, + ) + from transformer_engine.pytorch.utils import get_device_compute_capability + + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + if not is_nvfp4_available(): + pytest.skip("NVFP4 is required for the shipped DPA factory") + compute_capability = get_device_compute_capability() + if compute_capability < (9, 0) or compute_capability >= (12, 0): + pytest.skip("FP8 attention is not supported on this compute capability") + + FP8GlobalStateManager.reset() + module = MultiheadAttention( + hidden_size=128, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="mha", + ) + recipe = CustomRecipe( + qfactory=nvfp4_linear_fp8_dpa_factory, + fp8_dpa=True, + fp8_mha=False, + ) + + try: + apply_recipe(module, recipe) + owners = _active_runtime_owners(module) + before = [ + ( + owner._quantization_runtime, # pylint: disable=protected-access + owner._quantization_runtime.key, # pylint: disable=protected-access + owner._role_revision, # pylint: disable=protected-access + ) + for owner in owners + ] + inp = torch.randn( + 128, + 2, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + _run_update_step(module, recipe, inp) + after = [ + ( + owner._quantization_runtime, # pylint: disable=protected-access + owner._quantization_runtime.key, # pylint: disable=protected-access + owner._role_revision, # pylint: disable=protected-access + ) + for owner in owners + ] + assert after == before + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_rejects_quantized_primary_update_before_candidate_construction(): + """The model-wide API preserves the quantized-model-init support boundary.""" + FP8GlobalStateManager.reset() + initial_recipe = CustomRecipe( + qfactory=lambda _role: IdentityQuantizer(), + qfactory_key=("quantized-primary-apply", 1), + ) + replacement_calls = [] + replacement_recipe = _make_counting_recipe( + ("quantized-primary-apply", 2), + replacement_calls, + ) + + try: + with quantized_model_init(recipe=initial_recipe): + module = Linear( + 16, + 16, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="linear", + ) + assert module.primary_weights_in_fp8 + old_views = _runtime_views(module) + old_global_state = _global_recipe_state() + + with pytest.raises( + RuntimeError, + match="Recipe mismatch for quantized primary weights", + ): + apply_recipe(module, replacement_recipe) + + assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + assert _global_recipe_state() == old_global_state + assert not replacement_calls + finally: + FP8GlobalStateManager.reset() + + def test_apply_recipe_success_noop_mutation_and_forward_fast_path(): """Explicit application publishes once and matching forwards remain lazy-path no-ops.""" FP8GlobalStateManager.reset() diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index e5d1666f8fc..27145567314 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -181,11 +181,12 @@ def __init__( ) torch.nn.init.normal_(self.weight) - def get_quantizer_roles(self, *, fwd: bool, num_quantizers: int): + def get_quantizer_roles(self, *, fwd: bool, num_quantizers: int, boundary_role): # Supplying explicit roles keeps CustomRecipeState from emitting a # warning (which would graph-break under fullgraph=True) and lets the # qfactory dispatch per tensor slot. Order must match the module's # quantizer array (FP8FwdTensorIdx / FP8BwdTensorIdx). + del boundary_role if fwd: return [ QuantizerRole(module_type="linear", tensor_type="input"), diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index aa670643ccb..eeab6865674 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1255,15 +1255,6 @@ def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> Non for i in range(len(recipe)) ] - # Reached the rebuild path because ``fp8_meta_tensors_initialized`` - # was flipped to False after first init — most commonly because the - # base-class ``output_quantizer_role`` / ``grad_input_quantizer_role`` - # setter invalidated state when MHA wired boundary roles. That - # setter is recipe-agnostic, so this code fires for built-in - # recipes too even though they don't consume role information here - # (e.g. ``test_dpa_fp8_extra_state`` reaches this path with pure - # DelayedScaling). - # # Rebuilding the recipe state must preserve persistent training # buffers (delayed-scaling ``scale`` / ``amax_history``) so the new # quantizer instances and the ``FP8GlobalStateManager`` reduction @@ -1286,11 +1277,37 @@ def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> Non for recipe_state in recipe_states: self.quantizers[fp8_meta_tensor_key].extend(recipe_state.make_quantizers()) + def _resolve_boundary_quantizer_role( + self, + *, + recipe: Recipe, + fwd: bool, + ) -> Optional[QuantizerRole]: + """Keep the legacy inactive-boundary policy local to DPA. + + O and dQKV are both composed-module boundaries and internal + fused-attention descriptor slots. With a ``CustomRecipe`` and the + legacy ``fp8_mha`` switch disabled, the external boundaries stay in + BF16 but those internal slots still require DPA-compatible placeholder + quantizers. Ignore only the declared default in that mode so + ``get_quantizer_roles`` emits its existing DPA hints. Explicit public + overrides continue to win. + """ + explicit_role = ( + self._output_quantizer_role if fwd else self._grad_input_quantizer_role + ) + if explicit_role is not None: + return explicit_role + if recipe.custom() and not recipe.fp8_mha: + return None + return super()._resolve_boundary_quantizer_role(recipe=recipe, fwd=fwd) + def get_quantizer_roles( self, *, fwd: bool, num_quantizers: int, + boundary_role: Optional[QuantizerRole], ) -> Optional[List[QuantizerRole]]: """QuantizerRole list for quantizers used by ``DotProductAttention``. @@ -1329,10 +1346,8 @@ def get_quantizer_roles( **Boundary slots** — O (fwd) and dQKV (bwd) leave DPA and enter the next module (e.g. proj linear). DPA does not know that - consumer, so these default to ``None``. The parent module - (e.g. ``MultiheadAttention``) can set - :attr:`output_quantizer_role` / :attr:`grad_input_quantizer_role` - to fill in the consumer identity. + consumer, so the planner supplies ``boundary_role`` from composed MHA + topology or an explicit caller override. When not set, a hint-only ``QuantizerRole`` with empty ``module_type`` / ``tensor_type`` is emitted, with ``name`` @@ -1343,7 +1358,7 @@ def get_quantizer_roles( name = self.name or "" if fwd: qkv_role = QuantizerRole(module_type="dpa", tensor_type="qkv", name=name) - o_role = self._output_quantizer_role + o_role = boundary_role if o_role is None: o_role = QuantizerRole(name=f"{name}.dpa_output" if name else "dpa_output") s_role = QuantizerRole(module_type="dpa", tensor_type="s", name=name) @@ -1359,7 +1374,7 @@ def get_quantizer_roles( s_role, # Group 3: S (post-softmax, input to S·V) ] else: - dqkv_role = self._grad_input_quantizer_role + dqkv_role = boundary_role if dqkv_role is None: dqkv_role = QuantizerRole( name=f"{name}.dpa_grad_input" if name else "dpa_grad_input" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6eb3ce54f18..2d2a5aef70e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2685,8 +2685,8 @@ def get_attention_quantizers(fp8, quantizers): f"but {_name} quantizer is {type(_q).__name__}. " "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an " "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or " - "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles " - "(boundary slots like O output and dQKV grad-input)." + "MXFP8Quantizer) for every role assigned to a DPA tensor slot. In composed " + "attention, O and dQKV carry their downstream linear consumer roles." ) return QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index f87365bf7cf..d5db6f938f8 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -480,82 +480,60 @@ def __init__( **common_gemm_kwargs, ) - def _update_output_quantizer_roles( - self, - qkv_fp8_output: bool, - proj_fp8_grad: bool, - dpa_fp8_output: bool, - ) -> None: - """Set quantizer roles at the boundaries between QKV, DPA, and proj. + self._declare_quantizer_roles() + + def _declare_quantizer_roles(self) -> None: + """Declare stable topology at the boundaries between QKV, DPA, and proj. MHA contains three submodules connected as follows:: Forward: QKV linear ──(QKV tensor)──> DPA ──(O tensor)──> Proj linear Backward: QKV linear <──(dQKV tensor)── DPA <──(dO tensor)── Proj linear - Each submodule owns quantizers for its internal tensors, but the - *boundary* tensors (the arrows above) need to know which module - will *consume* them so the quantizer factory can pick the right - format. This method sets those boundary roles on all four edges: - - 1. ``qkv_fp8_output`` — **QKV linear → DPA (fwd)**: the QKV - linear's ``output_quantizer_role`` is told its consumer is DPA. - 2. ``proj_fp8_grad`` — **Proj linear → DPA (bwd)**: proj's - ``grad_input_quantizer_role`` is told its consumer is DPA. - 3. ``dpa_fp8_output`` — **DPA → Proj linear (fwd)**: DPA's - ``output_quantizer_role`` is told its consumer is the proj linear. - 4. ``dpa_fp8_output`` — **DPA → QKV linear (bwd)**: DPA's - ``grad_input_quantizer_role`` is told its consumer is QKV linear. - - When a flag is ``False`` the corresponding role is reset to ``None`` - so the module falls back to its own default. + Each child owns its quantizers, while these declarations describe the + consumer of each boundary tensor. The base runtime planner combines + them with explicit caller overrides independently of recipe policy. + Whether a particular forward uses a prepared boundary quantizer remains + a call-local execution decision. DPA alone retains a local compatibility + adapter for its dual-purpose O/dQKV fused-kernel slots while the legacy + ``fp8_mha`` switch remains supported. """ dpa_name = self.core_attention.name or "" # ── Boundary 1 (fwd): QKV linear output → consumed by DPA ──────── - qkv_output_role = ( - QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) - if qkv_fp8_output - else None - ) + qkv_output_role = QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) if self.attention_type == "self": - if self.input_layernorm: - self.layernorm_qkv.output_quantizer_role = qkv_output_role - else: - self.qkv.output_quantizer_role = qkv_output_role - elif self.attention_type == "cross": - if self.input_layernorm: - self.layernorm_query.output_quantizer_role = qkv_output_role - else: - self.query_layer.output_quantizer_role = qkv_output_role - self.key_value.output_quantizer_role = qkv_output_role + qkv_producers = [self.layernorm_qkv if self.input_layernorm else self.qkv] + else: + qkv_producers = [ + self.layernorm_query if self.input_layernorm else self.query_layer, + self.key_value, + ] + for qkv_producer in qkv_producers: + qkv_producer.fast_setattr( + "_declared_output_quantizer_role", + qkv_output_role, + ) # ── Boundary 2 (bwd): Proj grad-input (dO) → consumed by DPA ───── - proj_grad_input_role = ( - QuantizerRole(module_type="dpa", tensor_type="do", name=dpa_name) - if proj_fp8_grad - else None + self.proj.fast_setattr( + "_declared_grad_input_quantizer_role", + QuantizerRole(module_type="dpa", tensor_type="do", name=dpa_name), ) - self.proj.grad_input_quantizer_role = proj_grad_input_role # ── Boundary 3 (fwd): DPA output (O) → consumed by Proj linear ─── proj_name = self.proj.name or "" - self.core_attention.output_quantizer_role = ( - QuantizerRole(module_type="linear", tensor_type="input", name=proj_name) - if dpa_fp8_output - else None + self.core_attention.fast_setattr( + "_declared_output_quantizer_role", + QuantizerRole(module_type="linear", tensor_type="input", name=proj_name), ) # ── Boundary 4 (bwd): DPA grad-input (dQKV) → consumed by QKV linear - if self.attention_type == "self": - qkv_linear = self.layernorm_qkv if self.input_layernorm else self.qkv - else: - qkv_linear = self.layernorm_query if self.input_layernorm else self.query_layer + qkv_linear = qkv_producers[0] qkv_name = qkv_linear.name or "" - self.core_attention.grad_input_quantizer_role = ( - QuantizerRole(module_type="linear", tensor_type="grad_output", name=qkv_name) - if dpa_fp8_output - else None + self.core_attention.fast_setattr( + "_declared_grad_input_quantizer_role", + QuantizerRole(module_type="linear", tensor_type="grad_output", name=qkv_name), ) def fast_setattr(self, name: str, value: Any) -> None: @@ -879,15 +857,6 @@ def forward( float8_current_scaling = fp8_recipe.float8_current_scaling() mxfp8_scaling = fp8_recipe.mxfp8() if fp8 and custom_recipe and fp8_mha: - # Wire every boundary this CustomRecipe may quantize before - # DPA materializes its recipe state. Some quantizer families - # disable the corresponding output below, but pre-wiring avoids - # rebuilding that state after inspecting its canonical QKV slot. - self._update_output_quantizer_roles( - rotary_pos_emb is None, - True, - True, - ) float8_current_scaling, mxfp8_scaling = ( self.core_attention.get_qkv_quantization_capabilities() ) @@ -920,11 +889,6 @@ def forward( # 1. FP8CS recipe: produce F16 grads; again, due to cuBLAS limitation proj_fp8_grad = dpa_fp8_output and not float8_current_scaling - # Custom fp8_mha boundaries were wired before DPA recipe-state setup so - # querying its canonical QKV quantizer cannot trigger a second build. - if not (fp8 and custom_recipe and fp8_mha and _dpa_fp8_recipe == ""): - self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) - # Packed pass-through to DotProductAttention: the fused QKV/KV projection # already produces one packed buffer, which DPA accepts directly via its # declarative qkv_layer/kv_layer arguments (deriving q/k/v as zero-copy diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index e704256a891..62f72f00ddb 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -952,6 +952,9 @@ def __init__(self, name: Optional[str] = None) -> None: self.wgrad_store = None self._output_quantizer_role: Optional[QuantizerRole] = None self._grad_input_quantizer_role: Optional[QuantizerRole] = None + # Stable defaults declared by a composed parent before runtime planning. + self._declared_output_quantizer_role: Optional[QuantizerRole] = None + self._declared_grad_input_quantizer_role: Optional[QuantizerRole] = None self._role_revision = 0 self._quantization_runtime: Optional[_QuantizationRuntime] = None @@ -1007,10 +1010,11 @@ def _apply(self, *args, **kwargs): def output_quantizer_role(self) -> Optional[QuantizerRole]: """Caller-configurable :class:`QuantizerRole` for the forward output quantizer. - When set, overrides the default role used by :meth:`get_quantizer_roles` - for the forward-pass output quantizer slot. Setting this after - quantizers have been created forces their recreation on the next - forward pass. + When set, overrides the boundary role resolved for the forward-pass + output quantizer slot, including a role declared by a composed module + such as ``MultiheadAttention``. Clearing it restores the declared or + default role. Setting it after quantizers have been created forces + their recreation during the next explicit or lazy planning pass. See also :attr:`grad_input_quantizer_role` for the backward-pass counterpart. @@ -1059,14 +1063,20 @@ def _warn_missing_output_quantizer_role( recipe = FP8GlobalStateManager.get_fp8_recipe() if not recipe.custom(): return - if fp8_output and self._output_quantizer_role is None: + output_role = self._output_quantizer_role + grad_input_role = self._grad_input_quantizer_role + if output_role is None: + output_role = getattr(self, "_declared_output_quantizer_role", None) + if grad_input_role is None: + grad_input_role = getattr(self, "_declared_grad_input_quantizer_role", None) + if fp8_output and output_role is None: warnings.warn( f"{type(self).__name__}: fp8_output=True but " "output_quantizer_role is not set. The CustomRecipe qfactory " "will receive None for the output quantizer role.", stacklevel=3, ) - if fp8_grad and self._grad_input_quantizer_role is None: + if fp8_grad and grad_input_role is None: warnings.warn( f"{type(self).__name__}: fp8_grad=True but " "grad_input_quantizer_role is not set. The CustomRecipe " @@ -1144,9 +1154,33 @@ def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> meta_key ].amax_history[0] + def _resolve_boundary_quantizer_role( + self, + *, + recipe: Recipe, + fwd: bool, + ) -> Optional[QuantizerRole]: + """Resolve the caller override or composed-module boundary declaration. + + ``recipe`` is part of this hook's contract so specialized runtime + owners can keep recipe-dependent inventory compatibility local. + Generic modules resolve topology independently of recipe policy. + """ + del recipe + if fwd: + explicit_role = self._output_quantizer_role + declared_role = getattr(self, "_declared_output_quantizer_role", None) + else: + explicit_role = self._grad_input_quantizer_role + declared_role = getattr(self, "_declared_grad_input_quantizer_role", None) + if explicit_role is not None: + return explicit_role + return declared_role + def _resolve_quantizer_roles( self, *, + recipe: Recipe, fwd: bool, num_quantizers: int, ) -> Tuple[ @@ -1154,9 +1188,12 @@ def _resolve_quantizer_roles( Optional[List[Optional[QuantizerRole]]], ]: """Resolve semantic key roles and the role list passed to ``RecipeState``.""" + boundary_role = self._resolve_boundary_quantizer_role(recipe=recipe, fwd=fwd) + roles = self.get_quantizer_roles( # pylint: disable=assignment-from-none fwd=fwd, num_quantizers=num_quantizers, + boundary_role=boundary_role, ) if roles is None: # ``RecipeState`` treats a missing role list as one bare role per @@ -1352,10 +1389,12 @@ def _plan_quantization_update( num_forward_quantizers = num_gemms * 3 num_backward_quantizers = num_gemms * 2 forward_key_roles, forward_state_roles = self._resolve_quantizer_roles( + recipe=recipe, fwd=True, num_quantizers=num_forward_quantizers, ) backward_key_roles, backward_state_roles = self._resolve_quantizer_roles( + recipe=recipe, fwd=False, num_quantizers=num_backward_quantizers, ) @@ -1476,6 +1515,7 @@ def get_quantizer_roles( *, fwd: bool, # pylint: disable=unused-argument num_quantizers: int, # pylint: disable=unused-argument + boundary_role: Optional[QuantizerRole], # pylint: disable=unused-argument ) -> Optional[List[QuantizerRole]]: """Return an ordered list of :class:`QuantizerRole` for quantizers. @@ -1519,12 +1559,8 @@ def get_quantizer_roles( The last slot of a forward GEMM group (output) and the last slot of a backward group (grad_input) are **boundary** slots — the tensor leaves this module and enters an unknown consumer. For - these slots, use ``self._output_quantizer_role`` (fwd) and - ``self._grad_input_quantizer_role`` (bwd), which default to - ``None``. A parent module (e.g. ``MultiheadAttention``) can set - these attributes to fill in the consumer identity; see - ``MultiheadAttention._update_output_quantizer_roles`` for an - example. + these slots, use ``boundary_role``, which the planner resolves from an + explicit caller override or composed-module topology. Return value ------------ diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 12d1498194a..dfa21bfb11a 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -2113,25 +2113,25 @@ def get_quantizer_roles( *, fwd: bool, num_quantizers: int, + boundary_role: Optional[QuantizerRole], ) -> Optional[List[QuantizerRole]]: """QuantizerRole list for quantizers used by ``GroupedLinear``. - For grouped GEMMs we repeat the same pattern for each GEMM in - order. The output (fwd) and grad-input (bwd) slots default to - ``None`` (unknown consumer). Set :attr:`output_quantizer_role` / - :attr:`grad_input_quantizer_role` to provide consumer identity. + For grouped GEMMs we repeat the same pattern for each GEMM in order. + ``boundary_role`` is the planner-resolved output (fwd) or grad-input + (bwd) consumer role. """ name = self.name or "" if fwd: base = [ QuantizerRole(module_type="grouped_linear", tensor_type="input", name=name), QuantizerRole(module_type="grouped_linear", tensor_type="weight", name=name), - self._output_quantizer_role, + boundary_role, ] else: base = [ QuantizerRole(module_type="grouped_linear", tensor_type="grad_output", name=name), - self._grad_input_quantizer_role, + boundary_role, ] return [base[i % len(base)] for i in range(num_quantizers)] diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 427a00e2b70..ad9d4e4782d 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1599,24 +1599,24 @@ def get_quantizer_roles( *, fwd: bool, num_quantizers: int, + boundary_role: Optional[QuantizerRole], ) -> Optional[List[QuantizerRole]]: """QuantizerRole list for quantizers used by ``LayerNormLinear``. - The output (fwd) and grad-input (bwd) slots default to ``None`` - (unknown consumer). Set :attr:`output_quantizer_role` / - :attr:`grad_input_quantizer_role` to provide consumer identity. + ``boundary_role`` is the planner-resolved output (fwd) or grad-input + (bwd) consumer role. """ name = self.name or "" if fwd: base = [ QuantizerRole(module_type="linear", tensor_type="input", name=name), QuantizerRole(module_type="linear", tensor_type="weight", name=name), - self._output_quantizer_role, + boundary_role, ] else: base = [ QuantizerRole(module_type="linear", tensor_type="grad_output", name=name), - self._grad_input_quantizer_role, + boundary_role, ] return [base[i % len(base)] for i in range(num_quantizers)] diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 3c833929946..97b7f4cbe0d 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2182,16 +2182,15 @@ def get_quantizer_roles( *, fwd: bool, num_quantizers: int, + boundary_role: Optional[QuantizerRole], ) -> Optional[List[QuantizerRole]]: """QuantizerRole list for quantizers used by ``LayerNormMLP``. Each internal GEMM (fc1, fc2) gets a distinct name suffix so that custom-recipe factories can target them individually. - The module's final output (fc2 fwd) and final grad (fc1 bwd) - slots default to ``None`` (unknown consumer). Set - :attr:`output_quantizer_role` / :attr:`grad_input_quantizer_role` - to provide consumer identity. Internal boundaries use fixed + ``boundary_role`` is the planner-resolved final output (fc2 fwd) or + final grad-input (fc1 bwd) consumer role. Internal boundaries use fixed roles with known consumer identity. """ base_name = self.name or "" @@ -2210,14 +2209,14 @@ def get_quantizer_roles( QuantizerRole(module_type="linear", tensor_type="input", name=fc2_name), QuantizerRole(module_type="linear", tensor_type="input", name=fc2_name), QuantizerRole(module_type="linear", tensor_type="weight", name=fc2_name), - # fc2 output — boundary, consumer unknown - self._output_quantizer_role, + # fc2 output — resolved boundary consumer + boundary_role, ] else: base = [ QuantizerRole(module_type="linear", tensor_type="grad_output", name=fc1_name), - # fc1 grad_input — boundary, consumer unknown - self._grad_input_quantizer_role, + # fc1 grad_input — resolved boundary consumer + boundary_role, QuantizerRole(module_type="linear", tensor_type="grad_output", name=fc2_name), # fc2 grad_input — consumed by fc1 (via activation'), so labeled as fc1 grad_output QuantizerRole(module_type="linear", tensor_type="grad_output", name=fc1_name), diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 5246c9b3c22..3749676673f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1817,24 +1817,24 @@ def get_quantizer_roles( *, fwd: bool, num_quantizers: int, + boundary_role: Optional[QuantizerRole], ) -> Optional[List[QuantizerRole]]: """QuantizerRole list for quantizers used by ``Linear``. - The output (fwd) and grad-input (bwd) slots default to ``None`` - (unknown consumer). Set :attr:`output_quantizer_role` / - :attr:`grad_input_quantizer_role` to provide consumer identity. + ``boundary_role`` is the planner-resolved output (fwd) or grad-input + (bwd) consumer role. """ name = self.name or "" if fwd: base = [ QuantizerRole(module_type="linear", tensor_type="input", name=name), QuantizerRole(module_type="linear", tensor_type="weight", name=name), - self._output_quantizer_role, + boundary_role, ] else: base = [ QuantizerRole(module_type="linear", tensor_type="grad_output", name=name), - self._grad_input_quantizer_role, + boundary_role, ] return [base[i % len(base)] for i in range(num_quantizers)] From 328efbed49988d69ee8e38eea95f9d0ad8a7cc65 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 19 Aug 2026 12:10:00 +0000 Subject: [PATCH 76/97] Propagate recipe flags to per-tensor QParams Signed-off-by: Evgeny --- tests/pytorch/test_recipe.py | 93 +++++++++++++ tests/pytorch/test_recipe_update.py | 136 ++++++++++++++++++- transformer_engine/common/recipe/__init__.py | 120 ++++++++++++++-- 3 files changed, 339 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 1bf988e7018..dfd58a34f7f 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -29,6 +29,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.quantization import ( FP8GlobalStateManager, + Float8BlockScalingRecipeState, Float8CurrentScalingRecipeState, NVFP4BlockScalingRecipeState, QuantizerRole, @@ -264,6 +265,98 @@ def test_same_recipe_mutation_invalidates_config_for_direct_and_nested_parameter ) +def test_high_level_recipe_flags_configure_concrete_quantizers_at_construction(): + """Constructor flags must reach derived QParams and concrete quantizers.""" + current = Float8CurrentScaling(use_power_2_scales=True) + current_quantizers = Float8CurrentScalingRecipeState( + current, + mode="forward", + num_quantizers=2, + device=torch.device("cpu"), + roles=[QuantizerRole(tensor_type="input"), QuantizerRole(tensor_type="weight")], + ).make_quantizers() + assert all(quantizer.force_pow_2_scales for quantizer in current_quantizers) + + block = Float8BlockScaling(use_f32_scales=True) + block_quantizers = Float8BlockScalingRecipeState( + block, + mode="forward", + num_quantizers=2, + device=torch.device("cpu"), + roles=[QuantizerRole(tensor_type="input"), QuantizerRole(tensor_type="weight")], + ).make_quantizers() + assert all(not quantizer.force_pow_2_scales for quantizer in block_quantizers) + + nvfp4 = NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + ) + nvfp4_forward = NVFP4BlockScalingRecipeState( + nvfp4, + mode="forward", + num_quantizers=2, + device=torch.device("cpu"), + roles=[QuantizerRole(tensor_type="input"), QuantizerRole(tensor_type="weight")], + ).make_quantizers() + nvfp4_backward = NVFP4BlockScalingRecipeState( + nvfp4, + mode="backward", + num_quantizers=1, + device=torch.device("cpu"), + roles=[QuantizerRole(tensor_type="grad_output")], + ).make_quantizers() + assert not nvfp4_forward[0].with_rht + assert not nvfp4_forward[1].with_2d_quantization + assert not nvfp4_backward[0].with_rht + assert not nvfp4_backward[0].stochastic_rounding + + +def test_high_level_recipe_mutation_preserves_unowned_qparams_fields(): + """Convenience fields update their traits without discarding nested customization.""" + current = Float8CurrentScaling() + current.fp8_quant_fwd_inp = replace(current.fp8_quant_fwd_inp, amax_epsilon=0.125) + current.use_power_2_scales = True + assert current.fp8_quant_fwd_inp.power_2_scale + assert current.fp8_quant_fwd_inp.amax_epsilon == 0.125 + assert current.fp8_quant_fwd_weight.power_2_scale + assert current.fp8_quant_bwd_grad.power_2_scale + + block = Float8BlockScaling() + block.fp8_quant_fwd_weight = replace(block.fp8_quant_fwd_weight, amax_epsilon=0.25) + block.use_f32_scales = True + assert not block.fp8_quant_fwd_inp.power_2_scale + assert not block.fp8_quant_fwd_weight.power_2_scale + assert block.fp8_quant_fwd_weight.amax_epsilon == 0.25 + assert not block.fp8_quant_bwd_grad.power_2_scale + block_quantizer = Float8BlockScalingRecipeState( + block, + mode="forward", + num_quantizers=1, + device=torch.device("cpu"), + roles=[QuantizerRole(tensor_type="weight")], + ).make_quantizers()[0] + assert not block_quantizer.force_pow_2_scales + + nvfp4 = NVFP4BlockScaling() + nvfp4.fp4_quant_fwd_inp = replace(nvfp4.fp4_quant_fwd_inp, amax_epsilon=0.375) + nvfp4.fp4_quant_fwd_weight = replace( + nvfp4.fp4_quant_fwd_weight, + stochastic_rounding=True, + ) + nvfp4.fp4_quant_bwd_grad = replace(nvfp4.fp4_quant_bwd_grad, amax_epsilon=0.5) + nvfp4.disable_rht = True + nvfp4.disable_stochastic_rounding = True + nvfp4.disable_2d_quantization = True + assert not nvfp4.fp4_quant_fwd_inp.random_hadamard_transform + assert nvfp4.fp4_quant_fwd_inp.amax_epsilon == 0.375 + assert not nvfp4.fp4_quant_fwd_weight.fp4_2d_quantization + assert nvfp4.fp4_quant_fwd_weight.stochastic_rounding + assert not nvfp4.fp4_quant_bwd_grad.random_hadamard_transform + assert not nvfp4.fp4_quant_bwd_grad.stochastic_rounding + assert nvfp4.fp4_quant_bwd_grad.amax_epsilon == 0.5 + + def test_current_scaling_recipe_state_configures_roles_and_preserves_boundary_defaults(): """Current-scaling construction follows role dispatch and legacy boundary defaults.""" recipe = Float8CurrentScaling() diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index f12c641f716..9241d2b0afc 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -8,7 +8,13 @@ import torch import transformer_engine.pytorch.ops as te_ops -from transformer_engine.common.recipe import CustomRecipe, DelayedScaling, Float8CurrentScaling +from transformer_engine.common.recipe import ( + CustomRecipe, + DelayedScaling, + Float8BlockScaling, + Float8CurrentScaling, + NVFP4BlockScaling, +) from transformer_engine.pytorch import ( DotProductAttention, GroupedLinear, @@ -20,6 +26,7 @@ apply_recipe, autocast, is_fp8_available, + is_fp8_block_scaling_available, is_nvfp4_available, quantized_model_init, ) @@ -536,6 +543,133 @@ def test_same_recipe_object_semantic_mutation_rebuilds_runtime(): assert len(calls) == 10 +def test_current_scaling_flag_mutation_rebuilds_concrete_quantizers(): + """A same-object power-of-two update must change the active quantizers.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + recipe = Float8CurrentScaling(use_power_2_scales=False) + module = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="current", + ) + try: + apply_recipe(module, recipe) + old_runtime = module._quantization_runtime # pylint: disable=protected-access + assert not old_runtime.forward_quantizers[0].force_pow_2_scales + assert not old_runtime.forward_quantizers[1].force_pow_2_scales + assert not old_runtime.backward_quantizers[0].force_pow_2_scales + + recipe.use_power_2_scales = True + apply_recipe(module, recipe) + runtime = module._quantization_runtime # pylint: disable=protected-access + assert runtime is not old_runtime + assert runtime.forward_quantizers[0].force_pow_2_scales + assert runtime.forward_quantizers[1].force_pow_2_scales + assert runtime.backward_quantizers[0].force_pow_2_scales + + inp = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + _run_update_step(module, recipe, inp) + assert module._quantization_runtime is runtime # pylint: disable=protected-access + finally: + FP8GlobalStateManager.reset() + + +def test_float8_block_scaling_flag_mutation_rebuilds_concrete_quantizers(): + """A same-object FP32-scale update must change the active quantizers.""" + available, reason = is_fp8_block_scaling_available(return_reason=True) + if not available: + pytest.skip(reason) + if torch.cuda.get_device_capability() >= (10, 0): + pytest.skip("Blackwell FP8 block-scaling emulation requires power-of-two scales") + + FP8GlobalStateManager.reset() + recipe = Float8BlockScaling(use_f32_scales=False) + module = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="block", + ) + try: + apply_recipe(module, recipe) + old_runtime = module._quantization_runtime # pylint: disable=protected-access + assert old_runtime.forward_quantizers[0].force_pow_2_scales + assert old_runtime.forward_quantizers[1].force_pow_2_scales + assert old_runtime.backward_quantizers[0].force_pow_2_scales + + recipe.use_f32_scales = True + apply_recipe(module, recipe) + runtime = module._quantization_runtime # pylint: disable=protected-access + assert runtime is not old_runtime + assert not runtime.forward_quantizers[0].force_pow_2_scales + assert not runtime.forward_quantizers[1].force_pow_2_scales + assert not runtime.backward_quantizers[0].force_pow_2_scales + finally: + FP8GlobalStateManager.reset() + + +def test_nvfp4_flag_mutations_rebuild_concrete_quantizers(): + """Every mutable NVFP4 convenience flag must change its concrete trait.""" + if not is_nvfp4_available(): + pytest.skip("NVFP4 is not available") + + FP8GlobalStateManager.reset() + recipe = NVFP4BlockScaling( + disable_rht=False, + disable_stochastic_rounding=False, + disable_2d_quantization=False, + ) + module = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="nvfp4", + ) + try: + apply_recipe(module, recipe) + runtime = module._quantization_runtime # pylint: disable=protected-access + assert runtime.forward_quantizers[0].with_rht + assert runtime.forward_quantizers[1].with_2d_quantization + assert runtime.backward_quantizers[0].with_rht + assert runtime.backward_quantizers[0].stochastic_rounding + + recipe.disable_rht = True + apply_recipe(module, recipe) + rht_runtime = module._quantization_runtime # pylint: disable=protected-access + assert rht_runtime is not runtime + assert not rht_runtime.forward_quantizers[0].with_rht + assert not rht_runtime.backward_quantizers[0].with_rht + + recipe.disable_stochastic_rounding = True + apply_recipe(module, recipe) + rounding_runtime = module._quantization_runtime # pylint: disable=protected-access + assert rounding_runtime is not rht_runtime + assert not rounding_runtime.backward_quantizers[0].stochastic_rounding + + recipe.disable_2d_quantization = True + apply_recipe(module, recipe) + block_runtime = module._quantization_runtime # pylint: disable=protected-access + assert block_runtime is not rounding_runtime + assert not block_runtime.forward_quantizers[1].with_2d_quantization + + inp = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) + _run_update_step(module, recipe, inp) + assert module._quantization_runtime is block_runtime # pylint: disable=protected-access + finally: + FP8GlobalStateManager.reset() + + def test_role_revision_is_requested_until_atomic_runtime_commit(): """Changing a boundary role does not mutate live state before the next update.""" calls = [] diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 9a32d71d5b9..1d8f54c025e 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -9,7 +9,7 @@ from collections.abc import Hashable from enum import Enum from typing import Any, Callable, Literal, NamedTuple, Optional, Union -from dataclasses import field +from dataclasses import field, replace from pydantic.dataclasses import dataclass @@ -192,15 +192,40 @@ class Recipe: def __setattr__(self, name: str, value: Any) -> None: # Invalidate derived caches on any recipe mutation. Updating either # cache itself must not invalidate the other cache. - if name not in ( + cache_attribute = name in ( "_cached_repr", "_cached_quantizer_config", "_has_cached_quantizer_config", - ): + ) + if not cache_attribute: object.__setattr__(self, "_cached_repr", None) object.__setattr__(self, "_cached_quantizer_config", None) object.__setattr__(self, "_has_cached_quantizer_config", False) object.__setattr__(self, name, value) + if not cache_attribute: + self._synchronize_derived_qparams(name) + + def _synchronize_derived_qparams(self, name: str) -> None: + """Propagate a changed recipe option into per-tensor ``QParams``. + + Some public recipe options are shorthand for one field in several + derived ``QParams`` bundles. For example, ``use_power_2_scales`` sets + ``power_2_scale`` for the input, weight, and gradient bundles. A recipe + with such an option overrides this hook and updates only that owned + ``QParams`` field, preserving unrelated fields such as + ``amax_epsilon``. + + Assigning a complete ``QParams`` bundle directly is still supported. + That assignment controls every field until a later assignment to a + shorthand recipe option updates the specific field it owns. + """ + del name + + def _replace_derived_qparams(self, name: str, **changes: Any) -> None: + """Replace fields in an initialized derived ``QParams`` bundle.""" + qparams = self.__dict__.get(name) + if qparams is not None: + object.__setattr__(self, name, replace(qparams, **changes)) def quantizer_config(self) -> Hashable: """Return the immutable semantic configuration used to build quantizers. @@ -445,9 +470,9 @@ class Float8CurrentScaling(Recipe): use_power_2_scales: bool = os.getenv("NVTE_FP8_CURRENT_SCALING_POWER_2_SCALES", "0") == "1" fp8_format: Format = Format.HYBRID - fp8_quant_fwd_inp = QParams(power_2_scale=use_power_2_scales, amax_epsilon=0.0) - fp8_quant_fwd_weight = QParams(power_2_scale=use_power_2_scales, amax_epsilon=0.0) - fp8_quant_bwd_grad = QParams(power_2_scale=use_power_2_scales, amax_epsilon=0.0) + fp8_quant_fwd_inp: QParams = field(init=False) + fp8_quant_fwd_weight: QParams = field(init=False) + fp8_quant_bwd_grad: QParams = field(init=False) fp8_gemm_fprop: MMParams = MMParams(use_split_accumulator=False) fp8_gemm_dgrad: MMParams = MMParams(use_split_accumulator=True) fp8_gemm_wgrad: MMParams = MMParams(use_split_accumulator=True) @@ -460,6 +485,31 @@ def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + self.fp8_quant_fwd_inp = QParams( + power_2_scale=self.use_power_2_scales, + amax_epsilon=0.0, + ) + self.fp8_quant_fwd_weight = QParams( + power_2_scale=self.use_power_2_scales, + amax_epsilon=0.0, + ) + self.fp8_quant_bwd_grad = QParams( + power_2_scale=self.use_power_2_scales, + amax_epsilon=0.0, + ) + + def _synchronize_derived_qparams(self, name: str) -> None: + if name != "use_power_2_scales": + return + for qparams_name in ( + "fp8_quant_fwd_inp", + "fp8_quant_fwd_weight", + "fp8_quant_bwd_grad", + ): + self._replace_derived_qparams( + qparams_name, + power_2_scale=self.use_power_2_scales, + ) def _make_quantizer_config(self) -> Hashable: return ( @@ -595,9 +645,9 @@ class Float8BlockScaling(Recipe): use_f32_scales: bool = os.getenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") == "1" fp8_format: Format = Format.E4M3 - fp8_quant_fwd_inp = QParams(power_2_scale=not use_f32_scales, amax_epsilon=0.0) - fp8_quant_fwd_weight = QParams(power_2_scale=not use_f32_scales, amax_epsilon=0.0) - fp8_quant_bwd_grad = QParams(power_2_scale=not use_f32_scales, amax_epsilon=0.0) + fp8_quant_fwd_inp: QParams = field(init=False) + fp8_quant_fwd_weight: QParams = field(init=False) + fp8_quant_bwd_grad: QParams = field(init=False) x_block_scaling_dim: int = 1 w_block_scaling_dim: int = 2 grad_block_scaling_dim: int = 1 @@ -631,6 +681,33 @@ def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + power_2_scale = not self.use_f32_scales + self.fp8_quant_fwd_inp = QParams( + power_2_scale=power_2_scale, + amax_epsilon=0.0, + ) + self.fp8_quant_fwd_weight = QParams( + power_2_scale=power_2_scale, + amax_epsilon=0.0, + ) + self.fp8_quant_bwd_grad = QParams( + power_2_scale=power_2_scale, + amax_epsilon=0.0, + ) + + def _synchronize_derived_qparams(self, name: str) -> None: + if name != "use_f32_scales": + return + power_2_scale = not self.use_f32_scales + for qparams_name in ( + "fp8_quant_fwd_inp", + "fp8_quant_fwd_weight", + "fp8_quant_bwd_grad", + ): + self._replace_derived_qparams( + qparams_name, + power_2_scale=power_2_scale, + ) def _make_quantizer_config(self) -> Hashable: return ( @@ -752,6 +829,9 @@ class NVFP4BlockScaling(Recipe): fp4_format: Format = Format.E2M1 fp8_format: Format = Format.E4M3 + fp4_quant_fwd_inp: QParams = field(init=False) + fp4_quant_fwd_weight: QParams = field(init=False) + fp4_quant_bwd_grad: QParams = field(init=False) # Not applying quantization to attention for now fp8_dpa: bool = False @@ -793,6 +873,28 @@ def __post_init__(self) -> None: fp4_2d_quantization=False, ) + def _synchronize_derived_qparams(self, name: str) -> None: + updates = () + if name == "disable_rht": + updates = ( + ("fp4_quant_fwd_inp", "random_hadamard_transform", not self.disable_rht), + ("fp4_quant_bwd_grad", "random_hadamard_transform", not self.disable_rht), + ) + elif name == "disable_stochastic_rounding": + updates = ( + ( + "fp4_quant_bwd_grad", + "stochastic_rounding", + not self.disable_stochastic_rounding, + ), + ) + elif name == "disable_2d_quantization": + updates = ( + ("fp4_quant_fwd_weight", "fp4_2d_quantization", not self.disable_2d_quantization), + ) + for qparams_name, field_name, value in updates: + self._replace_derived_qparams(qparams_name, **{field_name: value}) + def _make_quantizer_config(self) -> Hashable: return ( ("recipe_type", "NVFP4BlockScaling"), From f469a1abce5002112581546812c12a7ab61f7a7e Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 19 Aug 2026 12:25:23 +0000 Subject: [PATCH 77/97] Handle failed autocast activation Signed-off-by: Evgeny --- tests/pytorch/test_recipe.py | 94 ++++++++++++++++++++++ transformer_engine/pytorch/quantization.py | 67 ++++++++++----- 2 files changed, 140 insertions(+), 21 deletions(-) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index dfd58a34f7f..bf9b5a26abb 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -820,6 +820,100 @@ def unkeyed_factory(role): FP8GlobalStateManager.reset() +def _autocast_activation_state(): + """Capture identity-bearing state that autocast entry may publish.""" + qstate = FP8GlobalStateManager.quantization_state + return ( + qstate.fp8_enabled, + qstate.fp8_calibration, + id(qstate.fp8_recipe), + id(qstate.quantizer_config), + qstate.quantizer_config_revision, + id(qstate.fp8_distributed_group), + qstate.is_first_fp8_module, + qstate.fp8_graph_capturing, + qstate.autocast_depth, + qstate.abort_amax_reduction, + tuple( + (key, id(recipe), id(group)) + for key, (recipe, group) in qstate.autocast_arguments.items() + ), + ) + + +def test_failed_autocast_recipe_activation_is_atomic_and_context_is_reusable(): + """A semantic-config error must not publish state or poison the context instance.""" + FP8GlobalStateManager.reset() + active_recipe = Float8CurrentScaling() + FP8GlobalStateManager.activate_recipe(active_recipe) + original_state = _autocast_activation_state() + + def unkeyed_factory(role): + del role + + context = te.autocast( + enabled=True, + calibrating=True, + recipe=CustomRecipe(qfactory=unkeyed_factory), + ) + for _ in range(2): + with pytest.raises(ValueError, match="requires a semantic qfactory key"): + context.__enter__() + assert context._fp8_state is None # pylint: disable=protected-access + assert _autocast_activation_state() == original_state + + FP8GlobalStateManager.reset() + + +def test_failed_autocast_support_check_is_atomic(monkeypatch): + """Platform validation must complete before autocast state is published.""" + FP8GlobalStateManager.reset() + active_recipe = Float8CurrentScaling() + FP8GlobalStateManager.activate_recipe(active_recipe) + original_state = _autocast_activation_state() + + @quantizer_factory(key=("autocast_support_failure", 1)) + def keyed_factory(role): + del role + + monkeypatch.setattr( + FP8GlobalStateManager, + "is_fp8_available", + classmethod(lambda _cls: (False, "injected FP8 support failure")), + ) + context = te.autocast(enabled=True, recipe=CustomRecipe(qfactory=keyed_factory)) + with pytest.raises(AssertionError, match="injected FP8 support failure"): + context.__enter__() + + assert context._fp8_state is None # pylint: disable=protected-access + assert _autocast_activation_state() == original_state + FP8GlobalStateManager.reset() + + +def test_failed_nested_autocast_activation_preserves_outer_state(): + """A rejected inner recipe must leave its active outer autocast unchanged.""" + FP8GlobalStateManager.reset() + + def unkeyed_factory(role): + del role + + invalid_context = te.autocast( + enabled=True, + calibrating=True, + recipe=CustomRecipe(qfactory=unkeyed_factory), + ) + outer_recipe = Float8CurrentScaling() + with te.autocast(enabled=False, calibrating=True, recipe=outer_recipe): + outer_state = _autocast_activation_state() + with pytest.raises(ValueError, match="requires a semantic qfactory key"): + invalid_context.__enter__() + assert invalid_context._fp8_state is None # pylint: disable=protected-access + assert _autocast_activation_state() == outer_state + + assert FP8GlobalStateManager.quantization_state.autocast_depth == 0 + FP8GlobalStateManager.reset() + + def test_autocast_restores_recipe_and_quantizer_config_together(): """Leaving autocast restores the previously active recipe/configuration pair.""" FP8GlobalStateManager.reset() diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index e550af7c3a8..b5750eb4b15 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -859,6 +859,38 @@ def get_unique_autocast_key( group_id = id(group) if group is not None else None return f"recipe={recipe_repr},group={group_id}" + @classmethod + def _prepare_autocast_enter( + cls, + enabled: bool, + fp8_recipe: Optional[Recipe], + fp8_group: Optional[dist_group_type], + ) -> Tuple[Recipe, Hashable, str]: + """Resolve and validate an autocast activation without publishing state.""" + fp8_recipe = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe + if enabled: + check_recipe_support(fp8_recipe) + + quantizer_config = fp8_recipe.quantizer_config() + autocast_key = cls.get_unique_autocast_key(fp8_recipe, fp8_group) + + if enabled: + fp8_available, reason_for_no_fp8 = cls.is_fp8_available() + assert fp8_available, reason_for_no_fp8 + if isinstance(fp8_recipe, MXFP8BlockScaling): + mxfp8_available, reason_for_no_mxfp8 = cls.is_mxfp8_available() + assert mxfp8_available, reason_for_no_mxfp8 + if isinstance(fp8_recipe, Float8BlockScaling): + fp8_block_available, reason_for_no_fp8_block = ( + cls.is_fp8_block_scaling_available() + ) + assert fp8_block_available, reason_for_no_fp8_block + if isinstance(fp8_recipe, NVFP4BlockScaling): + nvfp4_available, reason_for_no_nvfp4 = cls.is_nvfp4_available() + assert nvfp4_available, reason_for_no_nvfp4 + + return fp8_recipe, quantizer_config, autocast_key + @classmethod def autocast_enter( cls, @@ -868,11 +900,18 @@ def autocast_enter( fp8_group: Optional[dist_group_type] = None, _graph: bool = False, ) -> None: - """Set state and tracking variables for entry into FP8 region.""" + """Prepare and publish state for entry into an FP8 region.""" - fp8_recipe = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe - autocast_key = cls.get_unique_autocast_key(fp8_recipe, fp8_group) + fp8_recipe, quantizer_config, autocast_key = cls._prepare_autocast_enter( + enabled, + fp8_recipe, + fp8_group, + ) qstate = cls.quantization_state + + # Preparation above contains every operation that can fail due to the + # requested recipe or platform. Publish only after it has succeeded. + cls._set_recipe_config(fp8_recipe, quantizer_config) # Once a delayed bucket is registered, its committed recipe snapshot # owns the reduction semantics for this key. Do not replace it with a # caller-owned recipe object merely by entering another autocast. @@ -883,7 +922,6 @@ def autocast_enter( qstate.fp8_enabled = enabled qstate.fp8_calibration = calibrating - cls.activate_recipe(fp8_recipe) qstate.fp8_distributed_group = fp8_group qstate.fp8_graph_capturing = _graph @@ -892,19 +930,6 @@ def autocast_enter( qstate.is_first_fp8_module = True qstate.autocast_depth += 1 - if enabled: - fp8_available, reason_for_no_fp8 = cls.is_fp8_available() - assert fp8_available, reason_for_no_fp8 - if isinstance(fp8_recipe, MXFP8BlockScaling): - mxfp8_available, reason_for_no_mxfp8 = cls.is_mxfp8_available() - assert mxfp8_available, reason_for_no_mxfp8 - if isinstance(fp8_recipe, Float8BlockScaling): - fp8_block_available, reason_for_no_fp8_block = cls.is_fp8_block_scaling_available() - assert fp8_block_available, reason_for_no_fp8_block - if isinstance(fp8_recipe, NVFP4BlockScaling): - nvfp4_available, reason_for_no_nvfp4 = cls.is_nvfp4_available() - assert nvfp4_available, reason_for_no_nvfp4 - @classmethod def abort_current_amax_reduction(cls) -> None: """Prevent delayed-state updates when the active autocast has failed.""" @@ -1318,10 +1343,9 @@ def __enter__(self) -> "autocast": raise RuntimeError( "autocast context manager cannot be entered more than once concurrently" ) - if self._enabled: - check_recipe_support(self._recipe) - # Save current state so we always restore it on exit. - self._fp8_state = FP8GlobalStateManager.get_autocast_state() + # Mark this instance active only after entry has succeeded. Failed + # preparation is mutation-free, so the same instance remains reusable. + fp8_state = FP8GlobalStateManager.get_autocast_state() FP8GlobalStateManager.autocast_enter( enabled=self._enabled, calibrating=self._calibrating, @@ -1329,6 +1353,7 @@ def __enter__(self) -> "autocast": fp8_group=self._amax_reduction_group, _graph=self._graph, ) + self._fp8_state = fp8_state return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: From b2beaa13691a02aaf673f934a1007ee00a039ba6 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 19 Aug 2026 15:34:29 +0000 Subject: [PATCH 78/97] Fix attention quant config validation Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 153 +++++++++++++++--- .../dot_product_attention.py | 32 +++- 2 files changed, 154 insertions(+), 31 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 5f900c844ea..8364b862a7b 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -32,6 +32,8 @@ ) from transformer_engine.common import recipe from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + delayed_scaling_factory, + mxfp8_factory, nvfp4_factory, ) from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( @@ -149,14 +151,13 @@ def _make_dpa(): def test_custom_recipe_update_replaces_both_directions(self): """A qfactory change commits one complete DPA runtime.""" - from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( - current_scaling_factory, - ) + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") dpa = self._make_dpa() def first_factory(role): - return current_scaling_factory(role) + return mxfp8_factory(role) first_recipe = recipe.CustomRecipe( qfactory=first_factory, @@ -181,7 +182,7 @@ def first_factory(role): dpa._fp8_workspaces["old"] = object() def second_factory(role): - return current_scaling_factory(role) + return mxfp8_factory(role) second_recipe = recipe.CustomRecipe( qfactory=second_factory, @@ -197,7 +198,7 @@ def second_factory(role): assert dpa.fp8_meta["scaling_bwd"] is second_runtime.backward_states[0] assert dpa.quantizers["scaling_fwd"] is second_runtime.forward_quantizers assert dpa.quantizers["scaling_bwd"] is second_runtime.backward_quantizers - assert second_capabilities == first_capabilities == (True, False) + assert second_capabilities == first_capabilities == (False, True) assert dpa._qkv_capabilities_quantizer is not first_capabilities_quantizer assert any( dpa._qkv_capabilities_quantizer is quantizer @@ -207,13 +208,12 @@ def second_factory(role): def test_direct_runtime_commit_updates_dpa_derived_state(self): """DPA cache and recipe labels are part of its runtime commit hook.""" - from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( - current_scaling_factory, - ) + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") dpa = self._make_dpa() first_recipe = recipe.CustomRecipe( - qfactory=current_scaling_factory, + qfactory=mxfp8_factory, qfactory_key=("dpa-direct-commit", 1), fp8_dpa=True, ) @@ -227,7 +227,7 @@ def test_direct_runtime_commit_updates_dpa_derived_state(self): old_capabilities_quantizer = dpa._qkv_capabilities_quantizer second_recipe = recipe.CustomRecipe( - qfactory=current_scaling_factory, + qfactory=mxfp8_factory, qfactory_key=("dpa-direct-commit", 2), fp8_dpa=True, ) @@ -245,21 +245,23 @@ def test_direct_runtime_commit_updates_dpa_derived_state(self): assert dpa._apply_quantization_update(update) assert dpa._quantization_runtime is update.candidate - assert "local_recipes" not in dpa.fp8_meta - assert dpa._custom_dpa_local_recipes_cache is None + assert dpa.fp8_meta["local_recipes"] is update.validation_result + assert dpa._custom_dpa_local_recipes_cache is update.validation_result + assert [type(item).__name__ for item in update.validation_result] == [ + "MXFP8BlockScaling" + ] assert dpa._custom_dpa_local_recipes_cache_key != old_cache_key assert dpa._qkv_capabilities_quantizer is None assert dpa._qkv_capabilities_cache is None def test_candidate_failure_preserves_active_runtime_and_caches(self): """A failed candidate validation does not partially update DPA.""" - from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( - current_scaling_factory, - ) + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") dpa = self._make_dpa() active_recipe = recipe.CustomRecipe( - qfactory=current_scaling_factory, + qfactory=mxfp8_factory, qfactory_key=("dpa-runtime-valid", 1), fp8_dpa=True, ) @@ -309,6 +311,89 @@ def test_candidate_failure_preserves_active_runtime_and_caches(self): ) assert dpa._fp8_workspaces["old"] is workspace + def test_all_current_scaling_runtime_is_rejected(self): + """Current-scaling QKV requires delayed-scaling S/dP kernel slots.""" + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + ) + + dpa = self._make_dpa() + invalid_recipe = recipe.CustomRecipe( + qfactory=current_scaling_factory, + qfactory_key=("dpa-all-current", 1), + fp8_dpa=True, + ) + + with autocast(enabled=True, recipe=invalid_recipe): + with pytest.raises( + TypeError, + match="Float8CurrentScaling DPA requires delayed scaling for S and dP", + ): + dpa.get_qkv_quantization_capabilities() + assert dpa._quantization_runtime is None + + def test_all_current_scaling_update_is_rejected_before_commit(self): + """An invalid DPA update leaves the active MXFP8 runtime intact.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + ) + + dpa = self._make_dpa() + active_recipe = recipe.CustomRecipe( + qfactory=mxfp8_factory, + qfactory_key=("dpa-supported-mxfp8", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=active_recipe): + assert dpa.get_qkv_quantization_capabilities() == (False, True) + + old_runtime = dpa._quantization_runtime + old_views = ( + dpa.fp8_meta["recipe"], + dpa.fp8_meta["scaling_fwd"], + dpa.fp8_meta["scaling_bwd"], + dpa.quantizers["scaling_fwd"], + dpa.quantizers["scaling_bwd"], + dpa._custom_dpa_local_recipes_cache, + dpa._qkv_capabilities_quantizer, + dpa._qkv_capabilities_cache, + ) + invalid_recipe = recipe.CustomRecipe( + qfactory=current_scaling_factory, + qfactory_key=("dpa-all-current-update", 1), + fp8_dpa=True, + ) + + with pytest.raises( + RuntimeError, + match="Float8CurrentScaling DPA requires delayed scaling for S and dP", + ): + te.apply_recipe(dpa, invalid_recipe) + + assert dpa._quantization_runtime is old_runtime + assert all( + current is old + for current, old in zip( + ( + dpa.fp8_meta["recipe"], + dpa.fp8_meta["scaling_fwd"], + dpa.fp8_meta["scaling_bwd"], + dpa.quantizers["scaling_fwd"], + dpa.quantizers["scaling_bwd"], + dpa._custom_dpa_local_recipes_cache, + dpa._qkv_capabilities_quantizer, + dpa._qkv_capabilities_cache, + ), + old_views, + ) + ) + with autocast(enabled=True, recipe=active_recipe): + assert dpa.get_qkv_quantization_capabilities() == (False, True) + assert dpa._quantization_runtime is old_runtime + def test_hybrid_storage_snapshots_parent_quantizer(): """Caller mutations must not change an existing Hybrid tensor's behavior.""" @@ -2243,17 +2328,35 @@ def fake_infer(*_args, **_kwargs): assert "local_recipes" not in dpa.fp8_meta @pytest.mark.parametrize( - "factory_name,expected", + "factory_name,base_qfactory,expected,update_supported", [ - ("current_scaling_factory", (True, False)), - ("delayed_scaling_factory", (False, False)), + ( + "current_scaling_dpa", + nvfp4_linear_fp8_dpa_factory, + (True, False), + False, + ), + ("delayed_scaling", delayed_scaling_factory, (False, False), False), + pytest.param( + "mxfp8", + mxfp8_factory, + (False, True), + True, + marks=pytest.mark.skipif( + not mxfp8_available, + reason=f"MXFP8: {reason_for_no_mxfp8}", + ), + ), ], ) - def test_qkv_capabilities_reuse_canonical_quantizer(self, factory_name, expected): + def test_qkv_capabilities_reuse_canonical_quantizer( + self, + factory_name, + base_qfactory, + expected, + update_supported, + ): """Capability queries must not call qfactory outside recipe-state setup.""" - from transformer_engine.pytorch.custom_recipes import quantizer_factories - - base_qfactory = getattr(quantizer_factories, factory_name) calls = [] def counting_qfactory(role): @@ -2299,7 +2402,7 @@ def rebuilt_qfactory(role): fp8_mha=True, qfactory_key=("test_dpa_capability_rebuilt", factory_name, 1), ) - if factory_name == "delayed_scaling_factory": + if not update_supported: with autocast(enabled=True, recipe=rebuilt_recipe): with pytest.raises(RuntimeError, match="do not support delayed scaling"): dpa.get_qkv_quantization_capabilities() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index eeab6865674..887514ea41f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1143,14 +1143,34 @@ def _validate_quantization_runtime( raise ValueError( "FP8 DotProductAttention requires 9 forward and 6 backward quantizer slots." ) - dpa_utils.get_attention_quantizers( - True, - { - "scaling_fwd": candidate.forward_quantizers, - "scaling_bwd": candidate.backward_quantizers, - }, + qkv_quantizer, _, s_quantizer, _, _, dp_quantizer = ( + dpa_utils.get_attention_quantizers( + True, + { + "scaling_fwd": candidate.forward_quantizers, + "scaling_bwd": candidate.backward_quantizers, + }, + ) + ) + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, ) + # Native current-scaling DPA uses delayed scaling for the two + # softmax-intermediate kernel slots. The fused kernel does not + # support an all-current-scaling slot topology. + if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer) and not ( + isinstance(s_quantizer, Float8Quantizer) + and isinstance(dp_quantizer, Float8Quantizer) + ): + raise TypeError( + "Float8CurrentScaling DPA requires delayed scaling for S and dP, " + f"but got S={type(s_quantizer).__name__} and " + f"dP={type(dp_quantizer).__name__}. Return DelayedScalingRequest " + "for DPA tensor types 's' and 'dp'." + ) + return _infer_custom_dpa_local_recipes( candidate.recipe, {"scaling_fwd": candidate.forward_states[0]}, From b56c9a7c74f107e8220e5c184e0c7a73fe3409a8 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 19 Aug 2026 15:50:04 +0000 Subject: [PATCH 79/97] Canonicalize FP8 MHA recipe flags Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 9 ++-- tests/pytorch/test_recipe.py | 44 +++++++++++++++++++ tests/pytorch/test_recipe_update.py | 18 +++++--- transformer_engine/common/recipe/__init__.py | 18 ++++++++ .../dot_product_attention.py | 8 ---- 5 files changed, 81 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 8364b862a7b..6b86e854618 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -2272,13 +2272,16 @@ class TestCustomDPALocalRecipeCache: """Custom-DPA native recipe labels track the quantizer rebuild.""" def test_inference_runs_once_per_quantizer_state_and_clears_stale_labels(self, monkeypatch): + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + from transformer_engine.pytorch.attention.dot_product_attention import ( dot_product_attention as dpa_module, ) custom_recipe = recipe.CustomRecipe( - qfactory=lambda _role: IdentityQuantizer(), - qfactory_key=("test_dpa_identity", 1), + qfactory=mxfp8_factory, + qfactory_key=("test_dpa_local_recipe_cache", 1), ) inferred_labels = [recipe.MXFP8BlockScaling()] @@ -2318,7 +2321,7 @@ def fake_infer(*_args, **_kwargs): # A rebuilt recipe state/quantizer list invalidates the cache. If the # new family has no native label, the old label must not survive. - custom_recipe.qfactory_key = ("test_dpa_identity", 2) + custom_recipe.qfactory_key = ("test_dpa_local_recipe_cache", 2) with autocast(enabled=True, recipe=custom_recipe): dpa.init_fp8_metadata(num_gemms=3) assert inference_calls == 3 diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index bf9b5a26abb..6263c9987a3 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -265,6 +265,50 @@ def test_same_recipe_mutation_invalidates_config_for_direct_and_nested_parameter ) +@pytest.mark.parametrize( + "make_recipe", + ( + pytest.param(DelayedScaling, id="delayed-scaling"), + pytest.param(Float8CurrentScaling, id="current-scaling"), + pytest.param(MXFP8BlockScaling, id="mxfp8"), + pytest.param(NVFP4BlockScaling, id="nvfp4"), + pytest.param( + lambda **kwargs: CustomRecipe( + qfactory=lambda _role: None, + qfactory_key=("canonical-attention-flags", 1), + **kwargs, + ), + id="custom", + ), + ), +) +def test_fp8_mha_canonicalizes_fp8_dpa_during_construction_and_mutation(make_recipe): + """FP8 MHA must never produce a semantic configuration with DPA disabled.""" + constructed = make_recipe(fp8_mha=True) + assert constructed.fp8_mha is True + assert constructed.fp8_dpa is True + assert dict(constructed.quantizer_config())["fp8_dpa"] is True + + mutated = make_recipe() + inactive_config = mutated.quantizer_config() + mutated.fp8_mha = True + canonical_config = mutated.quantizer_config() + assert mutated.fp8_dpa is True + assert canonical_config != inactive_config + assert dict(canonical_config)["fp8_dpa"] is True + assert dict(canonical_config)["fp8_mha"] is True + + # DPA cannot be disabled while MHA still depends on it. + mutated.fp8_dpa = False + assert mutated.fp8_dpa is True + assert mutated.quantizer_config() == canonical_config + + # Disable MHA first when disabling both attention modes. + mutated.fp8_mha = False + mutated.fp8_dpa = False + assert mutated.quantizer_config() == inactive_config + + def test_high_level_recipe_flags_configure_concrete_quantizers_at_construction(): """Constructor flags must reach derived QParams and concrete quantizers.""" current = Float8CurrentScaling(use_power_2_scales=True) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 9241d2b0afc..42275d10d00 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -154,12 +154,13 @@ def qfactory(role): return nvfp4_linear_fp8_dpa_factory(role) return current_scaling_factory(role) - return CustomRecipe( + recipe = CustomRecipe( qfactory=qfactory, qfactory_key=key, - fp8_dpa=True, fp8_mha=True, ) + assert recipe.fp8_dpa + return recipe @pytest.mark.parametrize( @@ -997,6 +998,7 @@ def test_apply_recipe_before_first_mha_forward_has_stable_boundary_topology( try: assert not _active_runtime_owners(module) + expected_recipe_config = recipe.quantizer_config() apply_recipe(module, recipe) owners = _active_runtime_owners(module) assert owners @@ -1010,6 +1012,11 @@ def test_apply_recipe_before_first_mha_forward_has_stable_boundary_topology( for owner in owners ] assert all(requested == active for _, _, requested, active in before) + assert all( + owner._quantization_runtime.key.recipe_config == expected_recipe_config + and owner._quantization_runtime.recipe.fp8_dpa + for owner in owners + ) mha = next(child for child in module.modules() if isinstance(child, MultiheadAttention)) qkv = mha.layernorm_qkv if mha.input_layernorm else mha.qkv @@ -1048,6 +1055,7 @@ def test_apply_recipe_before_first_mha_forward_has_stable_boundary_topology( for owner in owners ] assert after == before + assert recipe.quantizer_config() == expected_recipe_config finally: FP8GlobalStateManager.reset() @@ -1293,7 +1301,7 @@ def test_explicit_mha_boundary_override_wins_and_can_restore_declared_role(): qkv = module.qkv try: - apply_recipe(module, recipe) + apply_recipe(qkv, recipe) declared_runtime = qkv._quantization_runtime # pylint: disable=protected-access declared_role = declared_runtime.key.forward_roles[-1] @@ -1304,14 +1312,14 @@ def test_explicit_mha_boundary_override_wins_and_can_restore_declared_role(): ) qkv.output_quantizer_role = override assert qkv._role_revision == 1 # pylint: disable=protected-access - apply_recipe(module, recipe) + apply_recipe(qkv, recipe) override_runtime = qkv._quantization_runtime # pylint: disable=protected-access assert override_runtime is not declared_runtime assert override_runtime.key.forward_roles[-1] == override qkv.output_quantizer_role = None assert qkv._role_revision == 2 # pylint: disable=protected-access - apply_recipe(module, recipe) + apply_recipe(qkv, recipe) restored_runtime = qkv._quantization_runtime # pylint: disable=protected-access assert restored_runtime is not override_runtime assert restored_runtime.key.forward_roles[-1] == declared_role diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 1d8f54c025e..0fc0e001661 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -190,6 +190,10 @@ class Recipe: _has_cached_quantizer_config: bool = False def __setattr__(self, name: str, value: Any) -> None: + # FP8 MHA requires FP8 DPA; keep both flags canonical across mutations. + if name == "fp8_dpa" and not value and self.__dict__.get("fp8_mha", False): + value = True + # Invalidate derived caches on any recipe mutation. Updating either # cache itself must not invalidate the other cache. cache_attribute = name in ( @@ -202,9 +206,16 @@ def __setattr__(self, name: str, value: Any) -> None: object.__setattr__(self, "_cached_quantizer_config", None) object.__setattr__(self, "_has_cached_quantizer_config", False) object.__setattr__(self, name, value) + if name == "fp8_mha" and value: + object.__setattr__(self, "fp8_dpa", True) if not cache_attribute: self._synchronize_derived_qparams(name) + def __post_init__(self) -> None: + """Canonicalize fields after Pydantic dataclass construction.""" + if getattr(self, "fp8_mha", False): + self.fp8_dpa = True + def _synchronize_derived_qparams(self, name: str) -> None: """Propagate a changed recipe option into per-tensor ``QParams``. @@ -363,6 +374,7 @@ def scaling_factor_compute(amax: Tensor, Whether to enable FP8 multi-head attention (MHA). When `True`, it removes the casting operations mentioned above at the DPA boundaries. Currently only standard MHA modules i.e. `LayerNormLinear/Linear + DPA + Linear`, are supported for this feature. When + enabled, `fp8_dpa` is also enabled because FP8 MHA requires FP8 DPA. When `fp8_mha = False, fp8_dpa = True`, a typical MHA module works as `LayerNormLinear (BF16 output) -> (cast to FP8 ) FP8 DPA (cast to BF16) -> Linear`. When `fp8_mha = True, fp8_dpa = True`, it becomes @@ -404,6 +416,7 @@ def scaling_factor_compute(amax: Tensor, scaling_factor_compute_algo_key: Optional[Hashable] = None def __post_init__(self) -> None: + super().__post_init__() assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." assert ( self.backward_override in _BACKWARD_OVERRIDES @@ -481,6 +494,7 @@ class Float8CurrentScaling(Recipe): backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: + super().__post_init__() assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." assert ( self.backward_override in _BACKWARD_OVERRIDES @@ -580,6 +594,7 @@ class MXFP8BlockScaling(Recipe): backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: + super().__post_init__() assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." assert ( self.backward_override in _BACKWARD_OVERRIDES @@ -659,6 +674,7 @@ class Float8BlockScaling(Recipe): backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: + super().__post_init__() assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x" assert self.w_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for w" assert self.grad_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for grad" @@ -839,6 +855,7 @@ class NVFP4BlockScaling(Recipe): backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: + super().__post_init__() assert self.fp4_format == Format.E2M1, "Only E2M1 is supported for NVFP4 scaling" assert self.fp8_format == Format.E4M3, "Only E4M3 is supported for NVFP4 scaling" assert ( @@ -1006,6 +1023,7 @@ class CustomRecipe(Recipe): qfactory_key: Optional[Hashable] = None def __post_init__(self) -> None: + super().__post_init__() assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 887514ea41f..ff97d18df22 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1713,14 +1713,6 @@ def forward( ), "Upgrade PyTorch version to get RNG manipulation support for cuda graph capture." # checks for FP8 - if self.fp8: - if self.fp8_meta["recipe"].fp8_mha: - if not self.fp8_meta["recipe"].fp8_dpa: - self.fp8_meta["recipe"].fp8_dpa = True - self.logger.warning( - """Forcing fp8_meta["recipe"].fp8_dpa=True due to """ - """fp8_meta["recipe"].fp8_mha=True""" - ) if self.fp8 and self.fp8_meta["recipe"].fp8_dpa: forward_dtype = get_fp8_te_dtype(self.fp8_meta["recipe"], fprop_tensor=True) backward_dtype = get_fp8_te_dtype(self.fp8_meta["recipe"], fprop_tensor=False) From ee614b5ecf5f5c4389370323548aaa3ae62d7021 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 19 Aug 2026 16:03:24 +0000 Subject: [PATCH 80/97] make sure forward does not change the recipe Signed-off-by: Evgeny --- tests/pytorch/test_recipe_update.py | 78 ++++++++++++++++++++++++++++- 1 file changed, 76 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 42275d10d00..7ff034d941a 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -98,15 +98,38 @@ def _runtime_views(module): ) +def _assert_runtime_recipes_match_keys(module): + """Check every committed recipe snapshot in a module hierarchy.""" + for module_name, owner in module.named_modules(): + if not isinstance(owner, TransformerEngineBaseModule): + continue + runtime = owner._quantization_runtime # pylint: disable=protected-access + if runtime is None: + continue + active_recipe = owner.fp8_meta.get("recipe") + assert ( + active_recipe is runtime.recipe + ), f"Runtime recipe view changed after commit for {module_name or ''}" + recipe_config = active_recipe.quantizer_config() + assert recipe_config == runtime.key.recipe_config, ( + f"Runtime recipe mutated after commit for {module_name or ''}: " + f"{recipe_config} != {runtime.key.recipe_config}" + ) + + def _run_update_step(module, recipe, inp, *args, **kwargs): """Run a real forward/backward step and return the output.""" + inputs = inp if isinstance(inp, tuple) else (inp,) + _assert_runtime_recipes_match_keys(module) module.zero_grad(set_to_none=True) with autocast(enabled=True, recipe=recipe): - output = module(inp, *args, **kwargs) + output = module(*inputs, *args, **kwargs) + _assert_runtime_recipes_match_keys(module) if isinstance(output, tuple): output = output[0] output.float().sum().backward() - assert inp.grad is not None + _assert_runtime_recipes_match_keys(module) + assert all(input_tensor.grad is not None for input_tensor in inputs) assert all(param.grad is not None for param in module.parameters() if param.requires_grad) return output @@ -844,9 +867,37 @@ def test_runtime_update_workspace_lifecycle( assert not module._fp8_workspaces # pylint: disable=protected-access +def test_forward_recipe_mutation_breaks_committed_key_invariant(): + """A forward path must not change its committed recipe snapshot.""" + + class MutatingLinear(Linear): + def forward(self, inp, *args, **kwargs): + output = super().forward(inp, *args, **kwargs) + self.fp8_meta["recipe"].fp8_dpa = True + return output + + FP8GlobalStateManager.reset() + calls = [] + recipe = _make_counting_recipe(("forward-recipe-mutation", 1), calls) + module = MutatingLinear(16, 16, bias=False, device="cuda") + inp = torch.randn(8, 16, device="cuda", requires_grad=True) + + try: + with pytest.raises(AssertionError, match="Runtime recipe mutated after commit"): + _run_update_step(module, recipe, inp) + finally: + FP8GlobalStateManager.reset() + + @pytest.mark.parametrize( ("module_factory", "input_factory", "forward_args"), ( + pytest.param( + lambda: Linear(16, 16, bias=False, device="cuda"), + lambda: torch.randn(8, 16, device="cuda", requires_grad=True), + (), + id="linear", + ), pytest.param( lambda: LayerNormLinear(16, 16, bias=False, device="cuda"), lambda: torch.randn(8, 16, device="cuda", requires_grad=True), @@ -865,6 +916,28 @@ def test_runtime_update_workspace_lifecycle( ([4, 4],), id="grouped-linear", ), + pytest.param( + lambda: DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + qkv_format="bshd", + name="dpa", + ).cuda(), + lambda: tuple( + torch.randn( + 2, + 8, + 2, + 16, + device="cuda", + requires_grad=True, + ) + for _ in range(3) + ), + (), + id="dot-product-attention", + ), ), ) def test_module_family_executes_after_recipe_update( @@ -933,6 +1006,7 @@ def test_composed_module_executes_after_recipe_update(module_factory): _run_update_step(module, first_recipe, first_inp) owners = _active_runtime_owners(module) assert len(owners) >= 2 + assert any(isinstance(owner, DotProductAttention) for owner in owners) first_runtimes = { id(owner): owner._quantization_runtime # pylint: disable=protected-access for owner in owners From 5ab662bf6afeda88577ae8ee36328a66fe1e9fee Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:35:59 +0000 Subject: [PATCH 81/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_hybrid_quantization.py | 4 +--- .../dot_product_attention.py | 18 +++++++----------- transformer_engine/pytorch/quantization.py | 4 +--- 3 files changed, 9 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index dd8deadcfd4..2631a19c5d6 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -248,9 +248,7 @@ def test_direct_runtime_commit_updates_dpa_derived_state(self): assert dpa._quantization_runtime is update.candidate assert dpa.fp8_meta["local_recipes"] is update.validation_result assert dpa._custom_dpa_local_recipes_cache is update.validation_result - assert [type(item).__name__ for item in update.validation_result] == [ - "MXFP8BlockScaling" - ] + assert [type(item).__name__ for item in update.validation_result] == ["MXFP8BlockScaling"] assert dpa._custom_dpa_local_recipes_cache_key != old_cache_key assert dpa._qkv_capabilities_quantizer is None assert dpa._qkv_capabilities_cache is None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index ebba652de78..14b5f2ce88e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1196,14 +1196,12 @@ def _validate_quantization_runtime( raise ValueError( "FP8 DotProductAttention requires 9 forward and 6 backward quantizer slots." ) - qkv_quantizer, _, s_quantizer, _, _, dp_quantizer = ( - dpa_utils.get_attention_quantizers( - True, - { - "scaling_fwd": candidate.forward_quantizers, - "scaling_bwd": candidate.backward_quantizers, - }, - ) + qkv_quantizer, _, s_quantizer, _, _, dp_quantizer = dpa_utils.get_attention_quantizers( + True, + { + "scaling_fwd": candidate.forward_quantizers, + "scaling_bwd": candidate.backward_quantizers, + }, ) from transformer_engine.pytorch.tensor.float8_tensor import ( Float8CurrentScalingQuantizer, @@ -1366,9 +1364,7 @@ def _resolve_boundary_quantizer_role( ``get_quantizer_roles`` emits its existing DPA hints. Explicit public overrides continue to win. """ - explicit_role = ( - self._output_quantizer_role if fwd else self._grad_input_quantizer_role - ) + explicit_role = self._output_quantizer_role if fwd else self._grad_input_quantizer_role if explicit_role is not None: return explicit_role if recipe.custom() and not recipe.fp8_mha: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 0e3198571d0..ebeaabd8a7f 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -881,9 +881,7 @@ def _prepare_autocast_enter( mxfp8_available, reason_for_no_mxfp8 = cls.is_mxfp8_available() assert mxfp8_available, reason_for_no_mxfp8 if isinstance(fp8_recipe, Float8BlockScaling): - fp8_block_available, reason_for_no_fp8_block = ( - cls.is_fp8_block_scaling_available() - ) + fp8_block_available, reason_for_no_fp8_block = cls.is_fp8_block_scaling_available() assert fp8_block_available, reason_for_no_fp8_block if isinstance(fp8_recipe, NVFP4BlockScaling): nvfp4_available, reason_for_no_nvfp4 = cls.is_nvfp4_available() From 5becfd960312d4564d63d00bf90ab87d8f97ae6f Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 19 Aug 2026 17:07:05 +0000 Subject: [PATCH 82/97] improve numerics tests Signed-off-by: Evgeny --- tests/pytorch/test_recipe_update.py | 442 ++++++++++++++++++++++++++-- 1 file changed, 425 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 7ff034d941a..6c6f31c6bbb 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -13,6 +13,7 @@ DelayedScaling, Float8BlockScaling, Float8CurrentScaling, + MXFP8BlockScaling, NVFP4BlockScaling, ) from transformer_engine.pytorch import ( @@ -27,9 +28,11 @@ autocast, is_fp8_available, is_fp8_block_scaling_available, + is_mxfp8_available, is_nvfp4_available, quantized_model_init, ) +from transformer_engine.pytorch.custom_recipes.quantizer_factories import current_scaling_factory from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import ( DelayedScalingRequest, @@ -134,6 +137,193 @@ def _run_update_step(module, recipe, inp, *args, **kwargs): return output +def _clone_inputs(inp): + """Clone one tensor or a tuple of tensors while preserving grad requirements.""" + inputs = inp if isinstance(inp, tuple) else (inp,) + clones = tuple( + input_tensor.detach().clone().requires_grad_(input_tensor.requires_grad) + for input_tensor in inputs + ) + return clones if isinstance(inp, tuple) else clones[0] + + +def _run_numerical_step(module, recipe, inp, *args, output_grad=None, **kwargs): + """Run forward/backward and capture values used by exact numerical oracles.""" + inputs = inp if isinstance(inp, tuple) else (inp,) + _assert_runtime_recipes_match_keys(module) + module.zero_grad(set_to_none=True) + with autocast(enabled=True, recipe=recipe): + output = module(*inputs, *args, **kwargs) + _assert_runtime_recipes_match_keys(module) + if isinstance(output, tuple): + output = output[0] + if output_grad is None: + output_grad = torch.linspace( + -1, + 1, + output.numel(), + device=output.device, + dtype=output.dtype, + ).reshape(output.shape) + output.backward(output_grad) + _assert_runtime_recipes_match_keys(module) + return ( + ( + output.detach().clone(), + tuple(input_tensor.grad.detach().clone() for input_tensor in inputs), + { + name: param.grad.detach().clone() + for name, param in module.named_parameters() + if param.requires_grad + }, + ), + output_grad, + ) + + +def _assert_numerical_traces_match(actual, expected): + """Require bitwise-identical output, input gradients, and parameter gradients.""" + actual_output, actual_input_grads, actual_param_grads = actual + expected_output, expected_input_grads, expected_param_grads = expected + torch.testing.assert_close(actual_output, expected_output, rtol=0, atol=0) + assert len(actual_input_grads) == len(expected_input_grads) + for actual_grad, expected_grad in zip(actual_input_grads, expected_input_grads): + torch.testing.assert_close(actual_grad, expected_grad, rtol=0, atol=0) + assert actual_param_grads.keys() == expected_param_grads.keys() + for name, actual_grad in actual_param_grads.items(): + torch.testing.assert_close(actual_grad, expected_param_grads[name], rtol=0, atol=0) + + +def _make_identical_linears(): + """Construct identical BF16 Linear modules for numerical comparisons.""" + torch.manual_seed(2026) + first = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="linear", + ) + second = Linear( + 128, + 128, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="linear", + ) + second.load_state_dict(first.state_dict()) + return first, second + + +def _check_same_object_mutation_against_fresh_recipe( + active_recipe, + fresh_recipe, + mutations, + *, + seed, + prepare_models=None, +): + """Compare a warmed, mutated runtime with a fresh target runtime.""" + switching, oracle = _make_identical_linears() + if prepare_models is not None: + prepare_models(switching, oracle) + torch.manual_seed(seed) + warmup_inp = torch.randn( + 32, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + _run_numerical_step(switching, active_recipe, warmup_inp) + old_runtime = switching._quantization_runtime # pylint: disable=protected-access + + for name, value in mutations.items(): + setattr(active_recipe, name, value) + apply_recipe(switching, active_recipe) + runtime = switching._quantization_runtime # pylint: disable=protected-access + apply_recipe(oracle, fresh_recipe) + + base_inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + switching_trace, output_grad = _run_numerical_step( + switching, + active_recipe, + base_inp.detach().clone().requires_grad_(True), + ) + oracle_trace, _ = _run_numerical_step( + oracle, + fresh_recipe, + base_inp.detach().clone().requires_grad_(True), + output_grad=output_grad, + ) + _assert_numerical_traces_match(switching_trace, oracle_trace) + assert runtime is not old_runtime + + +def _check_recipe_transition_sequence(recipes, *, seed): + """Compare each runtime in a warmed recipe sequence with a fresh target.""" + switching, _ = _make_identical_linears() + previous_runtime = None + previous_config = None + for step, recipe in enumerate(recipes): + _, oracle = _make_identical_linears() + if previous_runtime is not None: + switching._fp8_workspaces["sentinel"] = object() # pylint: disable=protected-access + apply_recipe(switching, recipe) + runtime = switching._quantization_runtime # pylint: disable=protected-access + recipe_config = recipe.quantizer_config() + if previous_config is not None and recipe_config != previous_config: + assert runtime is not previous_runtime + assert not switching._fp8_workspaces # pylint: disable=protected-access + + apply_recipe(oracle, recipe) + oracle_runtime = oracle._quantization_runtime # pylint: disable=protected-access + assert tuple(type(quantizer) for quantizer in runtime.forward_quantizers) == tuple( + type(quantizer) for quantizer in oracle_runtime.forward_quantizers + ) + assert tuple(type(quantizer) for quantizer in runtime.backward_quantizers) == tuple( + type(quantizer) for quantizer in oracle_runtime.backward_quantizers + ) + + torch.manual_seed(seed + step) + base_inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + switching_trace, output_grad = _run_numerical_step( + switching, + recipe, + base_inp.detach().clone().requires_grad_(True), + ) + oracle_trace, _ = _run_numerical_step( + oracle, + recipe, + base_inp.detach().clone().requires_grad_(True), + output_grad=output_grad, + ) + _assert_numerical_traces_match(switching_trace, oracle_trace) + previous_runtime = runtime + previous_config = recipe_config + + +def _assert_modules_and_sgd_state_match(first, second, first_optimizer, second_optimizer): + """Require exact parameters and per-parameter SGD state.""" + first_params = dict(first.named_parameters()) + second_params = dict(second.named_parameters()) + assert first_params.keys() == second_params.keys() + for name, first_param in first_params.items(): + second_param = second_params[name] + torch.testing.assert_close(first_param, second_param, rtol=0, atol=0) + first_state = first_optimizer.state[first_param] + second_state = second_optimizer.state[second_param] + assert first_state.keys() == second_state.keys() + for state_name, first_value in first_state.items(): + second_value = second_state[state_name] + if isinstance(first_value, torch.Tensor): + torch.testing.assert_close(first_value, second_value, rtol=0, atol=0) + else: + assert first_value == second_value + + def _active_runtime_owners(module): """Return executed TE runtime owners in module traversal order.""" return [ @@ -156,10 +346,7 @@ def _global_recipe_state(): def _make_compatible_fp8_mha_recipe(key): """Build an executable FP8-MHA recipe with compatible boundary formats.""" - from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( - current_scaling_factory, - delayed_scaling_factory, - ) + from transformer_engine.pytorch.custom_recipes.quantizer_factories import delayed_scaling_factory from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) @@ -605,6 +792,118 @@ def test_current_scaling_flag_mutation_rebuilds_concrete_quantizers(): FP8GlobalStateManager.reset() +@pytest.mark.parametrize(("initial", "target"), ((False, True), (True, False))) +def test_current_scaling_same_object_mutation_matches_fresh_recipe(initial, target): + """Same-object flag updates match fresh target recipes in both directions.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + try: + _check_same_object_mutation_against_fresh_recipe( + Float8CurrentScaling(use_power_2_scales=initial), + Float8CurrentScaling(use_power_2_scales=target), + {"use_power_2_scales": target}, + seed=6200 + int(target), + ) + finally: + FP8GlobalStateManager.reset() + + +def test_float8_block_same_object_layout_mutation_matches_fresh_recipe(): + """A warmed Float8BlockScaling layout update matches a fresh target.""" + available, reason = is_fp8_block_scaling_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + try: + _check_same_object_mutation_against_fresh_recipe( + Float8BlockScaling(x_block_scaling_dim=1, w_block_scaling_dim=2), + Float8BlockScaling(x_block_scaling_dim=2, w_block_scaling_dim=1), + {"x_block_scaling_dim": 2, "w_block_scaling_dim": 1}, + seed=6250, + ) + finally: + FP8GlobalStateManager.reset() + + +def test_mxfp8_same_object_2d_mutation_matches_fresh_recipe(): + """A warmed MXFP8 2D-weight update matches a fresh target.""" + available, reason = is_mxfp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + def prepare_models(switching, oracle): + torch.manual_seed(6261) + weight = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16) + row_scales = torch.pow( + 2.0, + torch.arange(128, device="cuda", dtype=torch.float32).remainder(16) - 8, + ).to(torch.bfloat16) + weight *= row_scales[:, None] + with torch.no_grad(): + switching.weight.copy_(weight) + oracle.weight.copy_(weight) + + FP8GlobalStateManager.reset() + try: + _check_same_object_mutation_against_fresh_recipe( + MXFP8BlockScaling(enable_2d_quantization=False), + MXFP8BlockScaling(enable_2d_quantization=True), + {"enable_2d_quantization": True}, + seed=6260, + prepare_models=prepare_models, + ) + finally: + FP8GlobalStateManager.reset() + + +def test_nvfp4_same_object_2d_mutation_matches_fresh_recipe(): + """A warmed deterministic NVFP4 weight-layout update matches a fresh target.""" + if not is_nvfp4_available(): + pytest.skip("NVFP4 is not available") + + FP8GlobalStateManager.reset() + try: + _check_same_object_mutation_against_fresh_recipe( + NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=False, + ), + NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + ), + {"disable_2d_quantization": True}, + seed=6270, + ) + finally: + FP8GlobalStateManager.reset() + + +def test_native_recipe_transition_matches_fresh_target(): + """A warmed stateless runtime matches a fresh target across A -> B -> A.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + available, reason = is_mxfp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + try: + _check_recipe_transition_sequence( + (Float8CurrentScaling(), MXFP8BlockScaling(), Float8CurrentScaling()), + seed=6300, + ) + finally: + FP8GlobalStateManager.reset() + + def test_float8_block_scaling_flag_mutation_rebuilds_concrete_quantizers(): """A same-object FP32-scale update must change the active quantizers.""" available, reason = is_fp8_block_scaling_available(return_reason=True) @@ -945,20 +1244,88 @@ def test_module_family_executes_after_recipe_update( input_factory, forward_args, ): - """Every base module family executes forward/backward across a runtime replacement.""" + """A semantic-only runtime replacement is numerically transparent.""" calls = [] first_recipe = _make_counting_recipe(("module-execution", 1), calls) second_recipe = _make_counting_recipe(("module-execution", 2), calls) module = module_factory() - _run_update_step(module, first_recipe, input_factory(), *forward_args) + first_inp = input_factory() + second_inp = _clone_inputs(first_inp) + first_trace, output_grad = _run_numerical_step( + module, + first_recipe, + first_inp, + *forward_args, + ) first_runtime = module._quantization_runtime # pylint: disable=protected-access assert first_runtime is not None - _run_update_step(module, second_recipe, input_factory(), *forward_args) + second_trace, _ = _run_numerical_step( + module, + second_recipe, + second_inp, + *forward_args, + output_grad=output_grad, + ) second_runtime = module._quantization_runtime # pylint: disable=protected-access assert second_runtime is not first_runtime assert second_runtime.key.recipe_config == second_recipe.quantizer_config() + _assert_numerical_traces_match(second_trace, first_trace) + + +def test_equivalent_recipe_round_trip_matches_uninterrupted_training(): + """A -> equivalent B -> A training is bitwise identical to uninterrupted A.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + reference, switching = _make_identical_linears() + reference_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01, momentum=0.9) + switching_optimizer = torch.optim.SGD(switching.parameters(), lr=0.01, momentum=0.9) + reference_recipe = Float8CurrentScaling() + native_recipe = Float8CurrentScaling() + equivalent_recipe = CustomRecipe(qfactory=current_scaling_factory) + schedule = (native_recipe, equivalent_recipe, native_recipe, native_recipe) + previous_runtime = None + torch.manual_seed(6100) + + try: + for step, switching_recipe in enumerate(schedule): + base_inp = torch.randn(32, 128, device="cuda", dtype=torch.bfloat16) + reference_inp = base_inp.detach().clone().requires_grad_(True) + switching_inp = base_inp.detach().clone().requires_grad_(True) + reference_trace, output_grad = _run_numerical_step( + reference, + reference_recipe, + reference_inp, + ) + switching_trace, _ = _run_numerical_step( + switching, + switching_recipe, + switching_inp, + output_grad=output_grad, + ) + _assert_numerical_traces_match(switching_trace, reference_trace) + + runtime = switching._quantization_runtime # pylint: disable=protected-access + if step in (1, 2): + assert runtime is not previous_runtime + if step == 3: + assert runtime is previous_runtime + previous_runtime = runtime + + reference_optimizer.step() + switching_optimizer.step() + _assert_modules_and_sgd_state_match( + reference, + switching, + reference_optimizer, + switching_optimizer, + ) + finally: + FP8GlobalStateManager.reset() @pytest.mark.parametrize( @@ -992,7 +1359,7 @@ def test_module_family_executes_after_recipe_update( ), ) def test_composed_module_executes_after_recipe_update(module_factory): - """A composed module updates every runtime owner reached by real execution.""" + """A composed module's semantic-only runtime replacements are bitwise exact.""" available, reason = is_fp8_available(return_reason=True) if not available: pytest.skip(reason) @@ -1003,7 +1370,8 @@ def test_composed_module_executes_after_recipe_update(module_factory): module = module_factory() first_inp = torch.randn(8, 2, 32, device="cuda", requires_grad=True) - _run_update_step(module, first_recipe, first_inp) + second_inp = _clone_inputs(first_inp) + first_trace, output_grad = _run_numerical_step(module, first_recipe, first_inp) owners = _active_runtime_owners(module) assert len(owners) >= 2 assert any(isinstance(owner, DotProductAttention) for owner in owners) @@ -1012,13 +1380,18 @@ def test_composed_module_executes_after_recipe_update(module_factory): for owner in owners } - second_inp = torch.randn(8, 2, 32, device="cuda", requires_grad=True) - _run_update_step(module, second_recipe, second_inp) + second_trace, _ = _run_numerical_step( + module, + second_recipe, + second_inp, + output_grad=output_grad, + ) assert _active_runtime_owners(module) == owners for owner in owners: runtime = owner._quantization_runtime # pylint: disable=protected-access assert runtime is not first_runtimes[id(owner)] assert runtime.key.recipe_config == second_recipe.quantizer_config() + _assert_numerical_traces_match(second_trace, first_trace) @pytest.mark.parametrize( @@ -1587,14 +1960,13 @@ def test_apply_recipe_reconstructs_stateless_runtime_after_checkpoint_restore(): source = Linear(16, 16, bias=False, device="cuda", name="linear") initial_recipe = _make_counting_recipe(("checkpoint-resume", 1), source_calls) committed_recipe = _make_counting_recipe(("checkpoint-resume", 2), source_calls) - inp = torch.randn(8, 16, device="cuda") + inp = torch.randn(8, 16, device="cuda", requires_grad=True) try: apply_recipe(source, initial_recipe) apply_recipe(source, committed_recipe) checkpoint = source.state_dict() - with torch.no_grad(), autocast(enabled=True, recipe=committed_recipe): - expected = source(inp) + expected, output_grad = _run_numerical_step(source, committed_recipe, inp) # Model/checkpoint restoration happens in a fresh runtime context. The # intended recipe is reconstructed and explicitly applied before the @@ -1611,9 +1983,14 @@ def test_apply_recipe_reconstructs_stateless_runtime_after_checkpoint_restore(): assert runtime is not None assert runtime.key.recipe_config == committed_recipe.quantizer_config() - with torch.no_grad(), autocast(enabled=True, recipe=resumed_recipe): - actual = restored(inp) - torch.testing.assert_close(actual, expected) + restored_inp = _clone_inputs(inp) + actual, _ = _run_numerical_step( + restored, + resumed_recipe, + restored_inp, + output_grad=output_grad, + ) + _assert_numerical_traces_match(actual, expected) finally: FP8GlobalStateManager.reset() @@ -1688,9 +2065,26 @@ def test_apply_recipe_planning_failure_is_model_wide_atomic(failure_index): active_recipe = _make_counting_recipe(("apply-failure", 1), calls) modules = [Linear(16, 16, bias=False, device="cuda", name=f"line{index}") for index in range(3)] model = torch.nn.Sequential(*modules) + control = torch.nn.Sequential( + *[ + Linear(16, 16, bias=False, device="cuda", name=f"line{index}") + for index in range(3) + ] + ) + control.load_state_dict(model.state_dict()) try: apply_recipe(model, active_recipe) + apply_recipe(control, active_recipe) + warmup_inp = torch.randn(8, 16, device="cuda", requires_grad=True) + expected, output_grad = _run_numerical_step(control, active_recipe, warmup_inp) + actual, _ = _run_numerical_step( + model, + active_recipe, + _clone_inputs(warmup_inp), + output_grad=output_grad, + ) + _assert_numerical_traces_match(actual, expected) old_global_state = _global_recipe_state() old_views = [_runtime_views(module) for module in modules] workspaces = [] @@ -1721,6 +2115,20 @@ def failing_factory(role): for current, expected in zip(_runtime_views(module), expected_views) ) assert module._fp8_workspaces["weight"] is workspace + + continuation_inp = torch.randn(8, 16, device="cuda", requires_grad=True) + expected, output_grad = _run_numerical_step( + control, + active_recipe, + continuation_inp, + ) + actual, _ = _run_numerical_step( + model, + active_recipe, + _clone_inputs(continuation_inp), + output_grad=output_grad, + ) + _assert_numerical_traces_match(actual, expected) finally: FP8GlobalStateManager.reset() From 1355ff241ccd0997537a9ce43155f568ad9b366b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 19 Aug 2026 17:09:31 +0000 Subject: [PATCH 83/97] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_recipe_update.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 6c6f31c6bbb..b3c02684430 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -346,7 +346,9 @@ def _global_recipe_state(): def _make_compatible_fp8_mha_recipe(key): """Build an executable FP8-MHA recipe with compatible boundary formats.""" - from transformer_engine.pytorch.custom_recipes.quantizer_factories import delayed_scaling_factory + from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + delayed_scaling_factory, + ) from transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo import ( nvfp4_linear_fp8_dpa_factory, ) @@ -2066,10 +2068,7 @@ def test_apply_recipe_planning_failure_is_model_wide_atomic(failure_index): modules = [Linear(16, 16, bias=False, device="cuda", name=f"line{index}") for index in range(3)] model = torch.nn.Sequential(*modules) control = torch.nn.Sequential( - *[ - Linear(16, 16, bias=False, device="cuda", name=f"line{index}") - for index in range(3) - ] + *[Linear(16, 16, bias=False, device="cuda", name=f"line{index}") for index in range(3)] ) control.load_state_dict(model.state_dict()) From c6d3ed5ee85d7599fd9a9df48605c30dac38d6f7 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 19 Aug 2026 17:28:24 +0000 Subject: [PATCH 84/97] add mxfp8 2d config Signed-off-by: Evgeny --- tests/pytorch/test_recipe.py | 13 +++++++++++++ transformer_engine/common/recipe/__init__.py | 1 + 2 files changed, 14 insertions(+) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 6263c9987a3..d5d868c357c 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -243,6 +243,19 @@ def test_equal_builtin_recipes_have_equal_quantizer_configs(recipe_type): assert first == second +def test_mxfp8_2d_quantization_changes_quantizer_config(): + """MXFP8 weight layout participates in construction and mutation semantics.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=False) + original_config = recipe.quantizer_config() + + recipe.enable_2d_quantization = True + updated_config = recipe.quantizer_config() + + assert updated_config != original_config + assert dict(updated_config)["enable_2d_quantization"] is True + assert updated_config == MXFP8BlockScaling(enable_2d_quantization=True).quantizer_config() + + def test_same_recipe_mutation_invalidates_config_for_direct_and_nested_parameters(): """Built-in recipe config changes for same-object and nested-qparam updates.""" recipe = Float8CurrentScaling() diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 18339e3a41d..db91c13e9a1 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -611,6 +611,7 @@ def _make_quantizer_config(self) -> Hashable: ("fp8_dpa", self.fp8_dpa), ("fp8_mha", self.fp8_mha), ("backward_override", self.backward_override), + ("enable_2d_quantization", self.enable_2d_quantization), ) def _make_repr(self) -> str: From b455cd95fc810e70880d80bcb82308c5b4cdbd95 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 20 Aug 2026 13:09:13 +0000 Subject: [PATCH 85/97] CI fixes Signed-off-by: Evgeny --- .../pytorch/distributed/run_numerics_exact.py | 2 + .../pytorch/nvfp4/test_nvfp4_module_exact.py | 32 ++++++ tests/pytorch/test_backward_override.py | 9 ++ tests/pytorch/test_fusible_ops.py | 60 +++++++++- tests/pytorch/test_recipe_update.py | 104 ++++++++++++++++++ .../dot_product_attention.py | 9 +- .../pytorch/attention/multi_head_attention.py | 100 +++++++---------- transformer_engine/pytorch/module/base.py | 18 ++- .../pytorch/module/layernorm_linear.py | 9 +- transformer_engine/pytorch/module/linear.py | 9 +- transformer_engine/pytorch/ops/fuser.py | 9 +- transformer_engine/pytorch/ops/op.py | 50 ++++++++- 12 files changed, 331 insertions(+), 80 deletions(-) diff --git a/tests/pytorch/distributed/run_numerics_exact.py b/tests/pytorch/distributed/run_numerics_exact.py index 46782122b29..6e5cc44d32b 100644 --- a/tests/pytorch/distributed/run_numerics_exact.py +++ b/tests/pytorch/distributed/run_numerics_exact.py @@ -19,6 +19,7 @@ Recipe, QParams, CustomRecipe, + quantizer_factory, ) from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.constants import NVFP4_BLOCK_SCALING_SIZE @@ -66,6 +67,7 @@ def get_nvfp4_quantizer_factory(): A factory function that takes a QuantizerRole and returns a quantizer instance """ + @quantizer_factory(key=("distributed_numerics_exact_nvfp4_reference", 1)) def factory(role): is_weight = ( role is not None diff --git a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py index 9fdabf68dd0..2923c794548 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py @@ -110,6 +110,15 @@ def _default_quantizer(): with_rht=with_rht, ) + @recipe.quantizer_factory( + key=( + "nvfp4_module_exact_reference", + 1, + ("with_rht", with_rht), + ("with_2d_quantization", with_2d_quantization), + ("row_scaled", row_scaled), + ) + ) def factory(role): if role is None: return _default_quantizer() @@ -137,6 +146,29 @@ def factory(role): return factory +def test_nvfp4_reference_factory_semantic_keys(monkeypatch): + """Factory keys are stable, capture all behavior, and require no factory probe.""" + + def unexpected_factory_invocation(*_args, **_kwargs): + pytest.fail("Creating semantic recipe configuration invoked the quantizer factory") + + monkeypatch.setattr(reference_nvfp4, "NVFP4QuantizerRef", unexpected_factory_invocation) + flags = ( + (False, False, False), + (False, False, False), + (True, False, False), + (False, True, False), + (False, False, True), + ) + factories = [get_nvfp4_quantizer_factory(*values) for values in flags] + configs = [recipe.CustomRecipe(qfactory=factory).quantizer_config() for factory in factories] + + assert factories[0].qfactory_key == factories[1].qfactory_key + assert configs[0] == configs[1] + assert len({factory.qfactory_key for factory in factories[1:]}) == 4 + assert len(set(configs[1:])) == 4 + + def reset_rng_states(): seed = 1234 torch.manual_seed(seed) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index c0acf2e6b3b..ce9e13864e8 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -1905,6 +1905,15 @@ def test_operation_fuser_rebuilds_userbuffers_fusion_on_backward_override_switch ) assert not any(isinstance(op, UserbuffersForwardLinear) for op, _ in fuser._forward_ops) + current_recipe["value"] = quant_recipe + fuser.maybe_fuse_ops( + is_grad_enabled=True, + recipe=quant_recipe, + input_=x, + extra_inputs=extra_inputs, + ) + assert any(isinstance(op, UserbuffersForwardLinear) for op, _ in fuser._forward_ops) + @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) @pytest.mark.parametrize("dtype", _core_dtypes, ids=str) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 9f59ad30dc5..d7bebe1394b 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -153,7 +153,7 @@ def maybe_skip_quantization( def test_fusible_operation_rejects_unsupported_same_class_recipe_update( persistent_fuser: bool, ) -> None: - """Fusible ops must not silently retain stale same-class quantizers.""" + """Unsupported and mixed same-class deltas reject without partial publication.""" op = te_ops.Quantize() model = te_ops.Sequential(op) if persistent_fuser else op x = torch.randn(16, 16, dtype=torch.bfloat16, device="cuda") @@ -161,15 +161,69 @@ def test_fusible_operation_rejects_unsupported_same_class_recipe_update( equal_recipe = copy.deepcopy(initial_recipe) changed_recipe = copy.deepcopy(initial_recipe) changed_recipe.fp8_quant_fwd_inp = transformer_engine.common.recipe.QParams(amax_epsilon=0.25) + mixed_recipe = copy.deepcopy(changed_recipe) + mixed_recipe.backward_override = "high_precision" try: with torch.no_grad(), te.autocast(recipe=initial_recipe): model(x) with torch.no_grad(), te.autocast(recipe=equal_recipe): model(x) - with pytest.raises(RuntimeError, match="not supported for fusible operations"): - with torch.no_grad(), te.autocast(recipe=changed_recipe): + + op_recipe_config = op._recipe_config + op_fp8_metas = op._fp8_metas + op_quantizers = op._quantizers + if persistent_fuser: + fuser = model._module_groups[0] + fuser_recipe_config = fuser.recipe_config + forward_ops = fuser._forward_ops + backward_ops = fuser._backward_ops + + for rejected_recipe in (changed_recipe, mixed_recipe): + with pytest.raises(RuntimeError, match="not supported for fusible operations"): + with torch.no_grad(), te.autocast(recipe=rejected_recipe): + model(x) + + assert op._recipe_config == op_recipe_config + assert op._fp8_metas is op_fp8_metas + assert op._quantizers is op_quantizers + if persistent_fuser: + assert fuser.recipe_config == fuser_recipe_config + assert fuser.backward_override is None + assert fuser._forward_ops is forward_ops + assert fuser._backward_ops is backward_ops + finally: + FP8GlobalStateManager.reset() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("persistent_fuser", (False, True)) +def test_fusible_operation_preserves_backward_override_transition( + persistent_fuser: bool, +) -> None: + """The existing fusion-control field can change without rebuilding quantizers.""" + op = te_ops.Quantize() + model = te_ops.Sequential(op) if persistent_fuser else op + x = torch.randn(16, 16, dtype=torch.bfloat16, device="cuda") + recipes = [ + transformer_engine.common.recipe.Float8CurrentScaling(backward_override=mode) + for mode in (None, "high_precision", "dequantized", None) + ] + + try: + quantizers = None + for active_recipe in recipes: + with torch.no_grad(), te.autocast(recipe=active_recipe): model(x) + assert op._recipe_config == active_recipe.quantizer_config() + if quantizers is None: + quantizers = op._quantizers + else: + assert op._quantizers is quantizers + if persistent_fuser: + fuser = model._module_groups[0] + assert fuser.recipe_config == active_recipe.quantizer_config() + assert fuser.backward_override == active_recipe.backward_override finally: FP8GlobalStateManager.reset() diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index b3c02684430..6900b81fc59 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -1730,6 +1730,110 @@ def resolved_roles(owner, *, fwd, num_quantizers): ) +@pytest.mark.parametrize( + "name", + (pytest.param(None, id="unnamed"), pytest.param("mha", id="named")), +) +@pytest.mark.parametrize("input_layernorm", (False, True)) +def test_delayed_qmi_mha_first_forward_preserves_initialized_runtimes( + name, + input_layernorm, +): + """Initial composed roles are not mistaken for a delayed runtime update.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + def runtime_object_ids(owner): + runtime = owner._quantization_runtime # pylint: disable=protected-access + assert runtime is not None + forward_state = runtime.forward_states[0] + backward_state = runtime.backward_states[0] + return tuple( + id(obj) + for obj in ( + runtime, + forward_state, + backward_state, + forward_state.scale, + forward_state.amax_history, + backward_state.scale, + backward_state.amax_history, + *runtime.forward_quantizers, + *runtime.backward_quantizers, + ) + ) + + FP8GlobalStateManager.reset() + recipe = DelayedScaling(amax_history_len=4) + + try: + with quantized_model_init(recipe=recipe): + module = MultiheadAttention( + hidden_size=32, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + input_layernorm=input_layernorm, + fuse_qkv_params=True, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name=name, + ) + + qkv = module.layernorm_qkv if input_layernorm else module.qkv + initialized_owners = _active_runtime_owners(module) + assert initialized_owners == [qkv, module.proj] + assert all(owner.primary_weights_in_fp8 for owner in initialized_owners) + semantic_dpa_name = "" if name is None else module.core_attention.name + semantic_qkv_name = "" if name is None else qkv.name + semantic_proj_name = "" if name is None else module.proj.name + assert all( + owner._role_revision == owner._quantization_runtime.role_revision == 0 + for owner in initialized_owners + ) # pylint: disable=protected-access + assert qkv._quantization_runtime.key.forward_roles[-1] == QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=semantic_dpa_name, + ) # pylint: disable=protected-access + assert module.proj._quantization_runtime.key.backward_roles[-1] == QuantizerRole( + module_type="dpa", + tensor_type="do", + name=semantic_dpa_name, + ) # pylint: disable=protected-access + assert module.core_attention._declared_output_quantizer_role == QuantizerRole( + module_type="linear", + tensor_type="input", + name=semantic_proj_name, + ) # pylint: disable=protected-access + assert module.core_attention._declared_grad_input_quantizer_role == QuantizerRole( + module_type="linear", + tensor_type="grad_output", + name=semantic_qkv_name, + ) # pylint: disable=protected-access + + before = {owner: runtime_object_ids(owner) for owner in initialized_owners} + inp = torch.randn( + 8, + 2, + 32, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + _run_update_step(module, recipe, inp) + + assert all(runtime_object_ids(owner) == before[owner] for owner in initialized_owners) + assert all( + owner._role_revision == owner._quantization_runtime.role_revision == 0 + for owner in initialized_owners + ) # pylint: disable=protected-access + finally: + FP8GlobalStateManager.reset() + + def test_explicit_mha_boundary_override_wins_and_can_restore_declared_role(): """Public role overrides retain revision semantics above declared topology.""" FP8GlobalStateManager.reset() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 14b5f2ce88e..1719caa8487 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -698,8 +698,15 @@ def __init__( softmax_type: str = "vanilla", return_max_logit: Optional[bool] = False, name: Optional[str] = None, + *, + _declared_output_quantizer_role: Optional[QuantizerRole] = None, + _declared_grad_input_quantizer_role: Optional[QuantizerRole] = None, ) -> None: - super().__init__(name=name) + super().__init__( + name=name, + _declared_output_quantizer_role=_declared_output_quantizer_role, + _declared_grad_input_quantizer_role=_declared_grad_input_quantizer_role, + ) # Cache the native recipe labels inferred from custom DPA quantizers. # ``init_fp8_metadata`` runs on every forward, while the quantizers only diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index d5db6f938f8..1d12db28ac1 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -5,7 +5,7 @@ """Multi-head Attention.""" import os import collections -from typing import Any, Callable, List, Optional, Tuple, Union +from typing import Callable, List, Optional, Tuple, Union import torch from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole @@ -361,6 +361,15 @@ def __init__( qkv_parallel_mode = "column" if set_parallel_mode else None + # quantized_model_init eagerly creates runtimes for weight-owning children. + # Pass the final composed topology into their constructors so each initial + # runtime is created with the same roles it will use during execution. + qkv_output_role = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name=dpa_name, + ) + if self.attention_type == "self": parameters_split = None if not fuse_qkv_params: @@ -390,8 +399,10 @@ def __init__( normalization=normalization, ub_name="qkv", name=name + ".layernorm_linear_qkv" if name is not None else None, + _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) + qkv_producers = [self.layernorm_qkv] else: self.qkv = Linear( hidden_size, @@ -402,8 +413,10 @@ def __init__( parallel_mode=qkv_parallel_mode, parameters_split=parameters_split, name=name + ".linear_qkv" if name is not None else None, + _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) + qkv_producers = [self.qkv] elif self.attention_type == "cross": if self.input_layernorm: self.layernorm_query = LayerNormLinear( @@ -424,6 +437,7 @@ def __init__( normalization=normalization, ub_name="qkv", name=name + ".layernorm_linear_q" if name is not None else None, + _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) else: @@ -434,6 +448,7 @@ def __init__( bias=bias, return_bias=False, parallel_mode=qkv_parallel_mode, + _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) self.key_value = Linear( @@ -445,8 +460,16 @@ def __init__( parallel_mode=qkv_parallel_mode, parameters_split=("key", "value") if not fuse_qkv_params else None, name=name + ".linear_kv" if name is not None else None, + _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) + qkv_producers = [ + self.layernorm_query if self.input_layernorm else self.query_layer, + self.key_value, + ] + + qkv_name = qkv_producers[0].name if name is not None else "" + proj_name = name + ".proj" if name is not None else "" # Attention. self.core_attention = DotProductAttention( @@ -463,6 +486,16 @@ def __init__( attention_type=self.attention_type, softmax_type=self.softmax_type, name=name + ".core_attention" if name is not None else None, + _declared_output_quantizer_role=QuantizerRole( + module_type="linear", + tensor_type="input", + name=proj_name, + ), + _declared_grad_input_quantizer_role=QuantizerRole( + module_type="linear", + tensor_type="grad_output", + name=qkv_name, + ), ) # Linear @@ -477,69 +510,14 @@ def __init__( ub_overlap_ag=ub_overlap_ag, ub_name="proj", name=name + ".proj" if name is not None else None, + _declared_grad_input_quantizer_role=QuantizerRole( + module_type="dpa", + tensor_type="do", + name=dpa_name, + ), **common_gemm_kwargs, ) - self._declare_quantizer_roles() - - def _declare_quantizer_roles(self) -> None: - """Declare stable topology at the boundaries between QKV, DPA, and proj. - - MHA contains three submodules connected as follows:: - - Forward: QKV linear ──(QKV tensor)──> DPA ──(O tensor)──> Proj linear - Backward: QKV linear <──(dQKV tensor)── DPA <──(dO tensor)── Proj linear - - Each child owns its quantizers, while these declarations describe the - consumer of each boundary tensor. The base runtime planner combines - them with explicit caller overrides independently of recipe policy. - Whether a particular forward uses a prepared boundary quantizer remains - a call-local execution decision. DPA alone retains a local compatibility - adapter for its dual-purpose O/dQKV fused-kernel slots while the legacy - ``fp8_mha`` switch remains supported. - """ - dpa_name = self.core_attention.name or "" - - # ── Boundary 1 (fwd): QKV linear output → consumed by DPA ──────── - qkv_output_role = QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) - if self.attention_type == "self": - qkv_producers = [self.layernorm_qkv if self.input_layernorm else self.qkv] - else: - qkv_producers = [ - self.layernorm_query if self.input_layernorm else self.query_layer, - self.key_value, - ] - for qkv_producer in qkv_producers: - qkv_producer.fast_setattr( - "_declared_output_quantizer_role", - qkv_output_role, - ) - - # ── Boundary 2 (bwd): Proj grad-input (dO) → consumed by DPA ───── - self.proj.fast_setattr( - "_declared_grad_input_quantizer_role", - QuantizerRole(module_type="dpa", tensor_type="do", name=dpa_name), - ) - - # ── Boundary 3 (fwd): DPA output (O) → consumed by Proj linear ─── - proj_name = self.proj.name or "" - self.core_attention.fast_setattr( - "_declared_output_quantizer_role", - QuantizerRole(module_type="linear", tensor_type="input", name=proj_name), - ) - - # ── Boundary 4 (bwd): DPA grad-input (dQKV) → consumed by QKV linear - qkv_linear = qkv_producers[0] - qkv_name = qkv_linear.name or "" - self.core_attention.fast_setattr( - "_declared_grad_input_quantizer_role", - QuantizerRole(module_type="linear", tensor_type="grad_output", name=qkv_name), - ) - - def fast_setattr(self, name: str, value: Any) -> None: - """Fast attribute set for non-parameter fields.""" - self.__dict__[name] = value - def _create_qk_norm_modules( self, qk_norm_type: Optional[str], diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 62f72f00ddb..d9e8ff7e092 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -924,10 +924,22 @@ class TransformerEngineBaseModule(torch.nn.Module, ABC): candidate validation and activation through their corresponding hooks. """ - def __init__(self, name: Optional[str] = None) -> None: + def __init__( + self, + name: Optional[str] = None, + *, + _declared_output_quantizer_role: Optional[QuantizerRole] = None, + _declared_grad_input_quantizer_role: Optional[QuantizerRole] = None, + ) -> None: super().__init__() if not torch.cuda.is_available(): raise RuntimeError("TransformerEngine needs CUDA.") + for attribute, role in ( + ("_declared_output_quantizer_role", _declared_output_quantizer_role), + ("_declared_grad_input_quantizer_role", _declared_grad_input_quantizer_role), + ): + if role is not None and not isinstance(role, QuantizerRole): + raise TypeError(f"{attribute} must be a QuantizerRole or None") self.name = name self.next_iter_when_debug_should_be_run = 0 self.fp8_initialized = False @@ -953,8 +965,8 @@ def __init__(self, name: Optional[str] = None) -> None: self._output_quantizer_role: Optional[QuantizerRole] = None self._grad_input_quantizer_role: Optional[QuantizerRole] = None # Stable defaults declared by a composed parent before runtime planning. - self._declared_output_quantizer_role: Optional[QuantizerRole] = None - self._declared_grad_input_quantizer_role: Optional[QuantizerRole] = None + self._declared_output_quantizer_role = _declared_output_quantizer_role + self._declared_grad_input_quantizer_role = _declared_grad_input_quantizer_role self._role_revision = 0 self._quantization_runtime: Optional[_QuantizationRuntime] = None diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index ad9d4e4782d..69a2760fbac 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1332,8 +1332,15 @@ def __init__( delay_wgrad_compute: bool = False, symmetric_ar_type: Optional[str] = None, name: Optional[str] = None, + *, + _declared_output_quantizer_role: Optional[QuantizerRole] = None, + _declared_grad_input_quantizer_role: Optional[QuantizerRole] = None, ) -> None: - super().__init__(name) + super().__init__( + name, + _declared_output_quantizer_role=_declared_output_quantizer_role, + _declared_grad_input_quantizer_role=_declared_grad_input_quantizer_role, + ) params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype self.in_features = in_features diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 3749676673f..dd95204a2ac 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1585,8 +1585,15 @@ def __init__( symmetric_ar_type: Optional[str] = None, save_original_input: bool = False, name: Optional[str] = None, + *, + _declared_output_quantizer_role: Optional[QuantizerRole] = None, + _declared_grad_input_quantizer_role: Optional[QuantizerRole] = None, ) -> None: - super().__init__(name) + super().__init__( + name, + _declared_output_quantizer_role=_declared_output_quantizer_role, + _declared_grad_input_quantizer_role=_declared_grad_input_quantizer_role, + ) params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype self.in_features = in_features diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index a1f953fdc06..53b4304c175 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -18,7 +18,7 @@ FusibleOperation, FusedOperation, OperationContext, - _only_delayed_history_length_changed, + _is_preserved_fusible_recipe_transition, ) @@ -635,9 +635,10 @@ def maybe_fuse_ops( if ( recipe_type is self.recipe_type and self.recipe_config != recipe_config - and not ( - isinstance(recipe, DelayedScaling) - and _only_delayed_history_length_changed(self.recipe_config, recipe_config) + and not _is_preserved_fusible_recipe_transition( + recipe, + self.recipe_config, + recipe_config, ) ): FP8GlobalStateManager.abort_current_amax_reduction() diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 8d7d9aaeab3..40408b711ea 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -24,16 +24,53 @@ from ..tensor import Quantizer +def _only_recipe_fields_changed( + old_config: Hashable, + new_config: Hashable, + allowed_fields: Iterable[str], +) -> bool: + """Whether unequal recipe configs differ only in explicitly allowed top-level fields.""" + if old_config == new_config: + return False + allowed_field_names = frozenset(allowed_fields) + + def without_allowed_fields(config: Hashable) -> tuple: + return tuple(item for item in config if item[0] not in allowed_field_names) + + return without_allowed_fields(old_config) == without_allowed_fields(new_config) + + def _only_delayed_history_length_changed( old_config: Hashable, new_config: Hashable, ) -> bool: """Whether delayed-scaling configs differ only in history length.""" + return _only_recipe_fields_changed( + old_config, + new_config, + ("amax_history_len",), + ) + - def without_history_length(config: Hashable) -> tuple: - return tuple(item for item in config if item[0] != "amax_history_len") +def _is_preserved_fusible_recipe_transition( + recipe: Recipe, + old_config: Hashable, + new_config: Hashable, +) -> bool: + """Whether this config delta is a transition already preserved by fusible ops. - return without_history_length(old_config) == without_history_length(new_config) + Delayed recipes retain only their legacy history-resize path. For other recipes, + ``backward_override`` may change because the fuser already rebuilds the affected + topology and the field does not alter quantizer construction. Every other + same-class semantic delta remains unsupported. + """ + if recipe.delayed(): + return _only_delayed_history_length_changed(old_config, new_config) + return _only_recipe_fields_changed( + old_config, + new_config, + ("backward_override",), + ) @dataclasses.dataclass @@ -350,9 +387,10 @@ def reset_recipe_state( self._recipe_type is recipe_type and self._recipe_config is not None and self._recipe_config != recipe_config - and not ( - recipe.delayed() - and _only_delayed_history_length_changed(self._recipe_config, recipe_config) + and not _is_preserved_fusible_recipe_transition( + recipe, + self._recipe_config, + recipe_config, ) ): FP8GlobalStateManager.abort_current_amax_reduction() From 8ed5fa2c68e62330a632f4db052e0da9672bc889 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 10 Sep 2026 09:40:24 +0000 Subject: [PATCH 86/97] Minor fixes Signed-off-by: Evgeny --- qa/L0_pytorch_unittest/test.sh | 2 + tests/pytorch/test_fusible_ops.py | 59 +++++++++++ tests/pytorch/test_hybrid_quantization.py | 13 ++- tests/pytorch/test_recipe_update.py | 100 ++++++++++++++++++ .../pytorch/attention/multi_head_attention.py | 1 + 5 files changed, 172 insertions(+), 3 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a78a99d7f95..b0d16a9dbfe 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -32,6 +32,8 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe_update.xml $TE_PATH/tests/pytorch/test_recipe_update.py || test_fail "test_recipe_update.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization_fsdp2.xml $TE_PATH/tests/pytorch/test_hybrid_quantization_fsdp2.py || test_fail "test_hybrid_quantization_fsdp2.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py || test_fail "test_deferred_init.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/test_numerics.py || test_fail "test_numerics.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cuda_graphs.xml $TE_PATH/tests/pytorch/test_cuda_graphs.py || test_fail "test_cuda_graphs.py" diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3cd48a0173a..5b1aef1c911 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -414,6 +414,65 @@ def track_first_forward(): assert fusion_calls == 3 +def test_operation_fuser_tracks_recipe_config_and_caches_plans(monkeypatch) -> None: + """The cached recipe config follows the active recipe and plans are built once. + + ``recipe_config`` must be published whenever the fuser resets recipe state, + so a cache hit is decided on recipe semantics rather than recipe identity, + and the plan cache must stop growing once the no-grad and grad-enabled plans + exist for a stable recipe. + """ + + op = te_ops.Identity() + + # Identity has no quantizers; neutralize its state hooks to keep the test + # CPU-only and isolate the fuser's own bookkeeping. + monkeypatch.setattr(op, "reset_recipe_state", lambda *, recipe: None) + monkeypatch.setattr(op, "pre_first_fuser_forward", lambda: None) + + fuser = OperationFuser([op]) + plan_cache = fuser._fused_ops_cache # pylint: disable=protected-access + x = torch.ones(1) + x_requiring_grad = torch.ones(1, requires_grad=True) + extra_inputs = [()] + + assert fuser.recipe_config is None + assert not plan_cache + + recipe = transformer_engine.common.recipe.Float8CurrentScaling(backward_override=None) + fuser.maybe_fuse_ops(False, recipe, x, extra_inputs) + assert fuser.recipe_config == recipe.quantizer_config() + assert len(plan_cache) == 1 + + # A fresh but equal recipe object hits both the config and the plan cache. + equivalent_recipe = transformer_engine.common.recipe.Float8CurrentScaling( + backward_override=None + ) + fuser.maybe_fuse_ops(False, equivalent_recipe, x, extra_inputs) + assert fuser.recipe_config == equivalent_recipe.quantizer_config() + assert len(plan_cache) == 1 + + # Training needs a distinct backward boundary, so it builds the second plan. + fuser.maybe_fuse_ops(True, equivalent_recipe, x_requiring_grad, extra_inputs) + assert fuser.recipe_config == equivalent_recipe.quantizer_config() + assert len(plan_cache) == 2 + + # With both configurations built, the registry stops growing. + for is_grad_enabled, tensor in ( + (False, x), + (True, x_requiring_grad), + (False, x), + (True, x_requiring_grad), + ): + fuser.maybe_fuse_ops(is_grad_enabled, equivalent_recipe, tensor, extra_inputs) + assert fuser.recipe_config == equivalent_recipe.quantizer_config() + assert len(plan_cache) == 2 + + # An unquantized region is its own configuration and publishes no config. + fuser.maybe_fuse_ops(False, None, x, extra_inputs) + assert fuser.recipe_config is None + + @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index c1a5697eef3..572213921aa 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -5274,19 +5274,26 @@ def tracked_validate(quantizers, *, operand_name="operand"): first_call_count = len(validation_calls) first_generation = model._validated_quantizer_generations["scaling_fwd"] assert first_call_count > 0 + # Candidate validation covers both directions before the commit, so the + # first forward validates the backward operand too, whatever the grad mode. + assert {name for name, _ in validation_calls} == {"input", "weight", "grad_output"} with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): model(tensor, m_splits) assert len(validation_calls) == first_call_count assert model._validated_quantizer_generations["scaling_fwd"] is first_generation - # Backward quantizers are validated the first time quantizers are selected - # with gradients enabled. Do not run a full forward here: this validation + # Selecting quantizers with gradients enabled hits the O(1) generation + # guard: the committed runtime already recorded both generations, so + # nothing is revalidated. Do not run a full forward here: this validation # test intentionally uses a columnwise-only configuration that is not # supported by the split-quantization kernel. with torch.enable_grad(), autocast(enabled=True, recipe=original_recipe): model._get_quantizers() - assert len(validation_calls) == first_call_count + 1 + assert len(validation_calls) == first_call_count + assert ( + model._validated_quantizer_generations["scaling_bwd"] is model.quantizers["scaling_bwd"] + ) def unexpected_grad_mode_query(): pytest.fail("cached validation must not query grad mode") diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 6900b81fc59..c5df9762e4c 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -1662,6 +1662,106 @@ def test_apply_recipe_rejects_e2e_mha_reproducer_during_planning(): FP8GlobalStateManager.reset() +@pytest.mark.parametrize("attention_type", ("self", "cross")) +@pytest.mark.parametrize("input_layernorm", (False, True)) +@pytest.mark.parametrize( + "name", + (pytest.param(None, id="unnamed"), pytest.param("mha", id="named")), +) +def test_mha_constructor_declares_literal_dpa_qkv_role(attention_type, input_layernorm, name): + """Every MHA variant constructs and declares the QKV boundary role by literal name. + + Construction alone is the regression guard: ``dpa_name`` used to be + referenced before assignment, so every one of these variants raised + ``NameError`` from ``__init__``. + """ + module = MultiheadAttention( + hidden_size=32, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + attention_type=attention_type, + input_layernorm=input_layernorm, + bias=False, + device="cuda", + name=name, + ) + + # Literal, not derived from the DPA child: an unnamed MHA declares "" like + # qkv_name/proj_name do, a named one declares the child's composed name. + expected_qkv = QuantizerRole( + module_type="dpa", + tensor_type="qkv", + name="mha.core_attention" if name is not None else "", + ) + if attention_type == "self": + qkv_producers = [module.layernorm_qkv if input_layernorm else module.qkv] + else: + qkv_producers = [ + module.layernorm_query if input_layernorm else module.query_layer, + module.key_value, + ] + for producer in qkv_producers: + assert ( + producer._declared_output_quantizer_role # pylint: disable=protected-access + == expected_qkv + ) + + dpa_own_qkv_role = module.core_attention.get_quantizer_roles( + fwd=True, + num_quantizers=9, + boundary_role=None, + )[0] + if name is not None: + # Producer-declared and DPA own-slot roles agree on both sides. + assert dpa_own_qkv_role == expected_qkv + else: + # WP9(b) will give unnamed modules "" on both sides. Today the DPA child + # still auto-names itself in __init__, so the two sides disagree; pin that + # so the equality above can be extended deliberately. + assert module.core_attention.name.startswith("Layer_") + assert dpa_own_qkv_role.name == module.core_attention.name + assert dpa_own_qkv_role != expected_qkv + + +@pytest.mark.parametrize( + "name", + (pytest.param(None, id="unnamed"), pytest.param("layer", id="named")), +) +def test_transformer_layer_constructor_declares_dpa_qkv_role(name): + """TransformerLayer builds its MHA and both boundary sides agree. + + A TransformerLayer always has a name (its own auto name when unnamed) and + passes it down, so the unnamed case composes a name here rather than "". + """ + layer = TransformerLayer( + hidden_size=32, + ffn_hidden_size=64, + num_attention_heads=2, + hidden_dropout=0.0, + attention_dropout=0.0, + self_attn_mask_type="no_mask", + bias=False, + device="cuda", + name=name, + ) + + mha = layer.self_attention + dpa = mha.core_attention + assert dpa.name == f"{mha.name}.core_attention" + if name is not None: + assert dpa.name == "layer.self_attention.core_attention" + + expected_qkv = QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa.name) + assert ( + mha.layernorm_qkv._declared_output_quantizer_role # pylint: disable=protected-access + == expected_qkv + ) + assert ( + dpa.get_quantizer_roles(fwd=True, num_quantizers=9, boundary_role=None)[0] == expected_qkv + ) + + @pytest.mark.parametrize("attention_type", ("self", "cross")) @pytest.mark.parametrize("input_layernorm", (False, True)) def test_mha_declares_all_boundary_topology_variants(attention_type, input_layernorm): diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 94d8a0f311c..78f6140a839 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -365,6 +365,7 @@ def __init__( # quantized_model_init eagerly creates runtimes for weight-owning children. # Pass the final composed topology into their constructors so each initial # runtime is created with the same roles it will use during execution. + dpa_name = name + ".core_attention" if name is not None else "" qkv_output_role = QuantizerRole( module_type="dpa", tensor_type="qkv", From 28592f19a05ccd2ccfbaaa0b3b527f0d144818b9 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 10 Sep 2026 14:10:15 +0000 Subject: [PATCH 87/97] Centralize recipe configuration and validation Signed-off-by: Evgeny --- ...torch_fine_grained_quantization_example.py | 10 +- .../fine_grained_quantization.rst | 12 +- tests/pytorch/test_hybrid_quantization.py | 3 + tests/pytorch/test_recipe.py | 213 +++++++- tests/pytorch/test_recipe_update.py | 18 +- transformer_engine/common/recipe/__init__.py | 515 +++++++++++------- .../dot_product_attention.py | 5 +- transformer_engine/pytorch/module/base.py | 6 +- transformer_engine/pytorch/quantization.py | 26 +- 9 files changed, 564 insertions(+), 244 deletions(-) diff --git a/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py index b511fc41dc4..902fa9ba58a 100644 --- a/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py +++ b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py @@ -53,7 +53,7 @@ def require_supported_hardware() -> None: import torch import transformer_engine.pytorch as te -from transformer_engine.common.recipe import CustomRecipe +from transformer_engine.common.recipe import CustomRecipe, quantizer_factory from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( mxfp8_factory, nvfp4_factory, @@ -65,9 +65,13 @@ def require_supported_hardware() -> None: BASE_FACTORY = mxfp8_factory -def quantizer_factory(role: Optional[te.QuantizerRole]): +@quantizer_factory(key=("fine_grained_quantization_example", 1)) +def demo_quantizer_factory(role: Optional[te.QuantizerRole]): """Return a fresh quantizer for every role, including ``None``. + The key stands in for this factory's behavior: bump its revision whenever the + body changes, so quantizers built from the previous behavior are rebuilt. + ``BASE_FACTORY`` makes the factory total: unknown roles, future role values, and untargeted modules all retain valid MXFP8 behavior. """ @@ -111,7 +115,7 @@ def quantizer_factory(role: Optional[te.QuantizerRole]): te.Linear(256, 128, name="demo.output", **linear_options), ) inputs = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) -recipe = CustomRecipe(qfactory=quantizer_factory) +recipe = CustomRecipe(qfactory=demo_quantizer_factory) with te.autocast(enabled=True, recipe=recipe): outputs = model(inputs) diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst index 301b655b3f1..23a911cbf85 100644 --- a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst +++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst @@ -53,14 +53,15 @@ TE autocast path: .. code-block:: python import transformer_engine.pytorch as te - from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.common.recipe import CustomRecipe, quantizer_factory from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( mxfp8_factory, nvfp4_factory, ) - def quantizer_factory(role): + @quantizer_factory(key=("demo_factory", 1)) + def demo_factory(role): if role is not None and role.name == "demo.fc1": if role.tensor_type == "input": # wgrad keeps the original BF16 input @@ -88,7 +89,7 @@ TE autocast path: return mxfp8_factory(role) # every other TE module in MXFP8 - recipe = CustomRecipe(qfactory=quantizer_factory) + recipe = CustomRecipe(qfactory=demo_factory) with te.autocast(enabled=True, recipe=recipe): output = model(inputs) @@ -130,7 +131,8 @@ It does not call the factory on every unchanged forward. # name: str = "" - def quantizer_factory(role: Optional[te.QuantizerRole]): + @quantizer_factory(key=("demo_factory", 1)) + def demo_factory(role: Optional[te.QuantizerRole]): # construct a fresh quantizer on every call ... # Boundary slots may pass role=None or a role with empty fields, so @@ -139,7 +141,7 @@ It does not call the factory on every unchanged forward. # The factory plugs into the standard TE autocast path: - recipe = CustomRecipe(qfactory=quantizer_factory) + recipe = CustomRecipe(qfactory=demo_factory) with te.autocast(enabled=True, recipe=recipe): output = model(inputs) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 572213921aa..413dd466f30 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -2364,6 +2364,9 @@ def fake_infer(*_args, **_kwargs): # Native labels also copy these mutable fields from CustomRecipe, so a # semantic recipe update must refresh them with the candidate runtime. + # Enabling FP8 MHA after construction requires enabling FP8 DPA with it: + # only construction canonicalizes the pair. + custom_recipe.fp8_dpa = True custom_recipe.fp8_mha = True with autocast(enabled=True, recipe=custom_recipe): dpa.init_fp8_metadata(num_gemms=3) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index d5d868c357c..5ef613af559 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. +import copy from dataclasses import fields, replace from types import SimpleNamespace from typing import Optional @@ -44,8 +45,10 @@ DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + Format, MXFP8BlockScaling, NVFP4BlockScaling, + QParams, Recipe, quantizer_factory, ) @@ -113,16 +116,20 @@ def _make_repr(self) -> str: UnhashableConfigRecipe().quantizer_config() -def test_quantization_runtime_key_is_semantic_and_normalizes_slot_roles(): - """Runtime-key equality includes recipe config and each ordered role slot.""" +def test_quantization_runtime_key_is_semantic_over_ordered_slot_roles(): + """Runtime-key equality includes recipe config and each ordered role slot. + + Role containers are tuples by construction (``_resolve_quantizer_roles`` + builds them), so the key normalizes nothing and stays a plain frozen record. + """ input_role = QuantizerRole(module_type="linear", tensor_type="input", name="qkv") weight_role = QuantizerRole(module_type="linear", tensor_type="weight", name="qkv") grad_role = QuantizerRole(module_type="linear", tensor_type="grad_output", name="qkv") key = _QuantizationRuntimeKey( recipe_config=("recipe", "fp8"), - forward_roles=[input_role, None, weight_role], - backward_roles=[grad_role], + forward_roles=(input_role, None, weight_role), + backward_roles=(grad_role,), ) same_request = _QuantizationRuntimeKey( recipe_config=("recipe", "fp8"), @@ -295,33 +302,87 @@ def test_same_recipe_mutation_invalidates_config_for_direct_and_nested_parameter ), ), ) -def test_fp8_mha_canonicalizes_fp8_dpa_during_construction_and_mutation(make_recipe): +def test_fp8_mha_requires_fp8_dpa_at_construction_and_after_mutation(make_recipe): """FP8 MHA must never produce a semantic configuration with DPA disabled.""" constructed = make_recipe(fp8_mha=True) assert constructed.fp8_mha is True assert constructed.fp8_dpa is True assert dict(constructed.quantizer_config())["fp8_dpa"] is True + # Construction canonicalizes; a later assignment does not. Enabling MHA alone + # is rejected at the next activation instead of silently enabling DPA. mutated = make_recipe() inactive_config = mutated.quantizer_config() mutated.fp8_mha = True + assert mutated.fp8_dpa is False + with pytest.raises(ValueError, match="fp8_mha=True requires fp8_dpa=True"): + mutated.quantizer_config() + + mutated.fp8_dpa = True canonical_config = mutated.quantizer_config() - assert mutated.fp8_dpa is True assert canonical_config != inactive_config assert dict(canonical_config)["fp8_dpa"] is True assert dict(canonical_config)["fp8_mha"] is True - # DPA cannot be disabled while MHA still depends on it. + # DPA cannot be disabled while MHA still depends on it: the assignment is + # kept as written and rejected when the recipe is next activated. mutated.fp8_dpa = False - assert mutated.fp8_dpa is True - assert mutated.quantizer_config() == canonical_config + assert mutated.fp8_dpa is False + with pytest.raises(ValueError, match="fp8_mha=True requires fp8_dpa=True"): + mutated.quantizer_config() # Disable MHA first when disabling both attention modes. mutated.fp8_mha = False - mutated.fp8_dpa = False assert mutated.quantizer_config() == inactive_config +@pytest.mark.parametrize( + ("make_recipe", "field_name", "invalid_value", "match"), + ( + pytest.param( + DelayedScaling, + "fp8_format", + transformer_engine.common.recipe.Format.E5M2, + "Pure E5M2", + id="delayed-format", + ), + pytest.param( + Float8BlockScaling, + "w_block_scaling_dim", + 99, + "Only 1D or 2D blocks supported for w", + id="block-scaling-dim", + ), + pytest.param( + NVFP4BlockScaling, + "nvfp4_4over6_err_mode", + "BOGUS", + "NVTE_NVFP4_4OVER6_ERR_MODE", + id="nvfp4-err-mode", + ), + ), +) +def test_invalid_field_assignment_rejected_before_config( + make_recipe, field_name, invalid_value, match +): + """A field assigned after construction is validated at the next activation. + + The constructor's own message is reused, and the recipe never produces a + semantic configuration built from the invalid value. + """ + recipe_instance = make_recipe() + valid_config = recipe_instance.quantizer_config() + valid_value = getattr(recipe_instance, field_name) + + setattr(recipe_instance, field_name, invalid_value) + with pytest.raises((AssertionError, ValueError), match=match): + recipe_instance.quantizer_config() + + # Restoring the field makes the recipe usable again with its original config. + setattr(recipe_instance, field_name, valid_value) + assert recipe_instance.quantizer_config() == valid_config + + def test_high_level_recipe_flags_configure_concrete_quantizers_at_construction(): """Constructor flags must reach derived QParams and concrete quantizers.""" current = Float8CurrentScaling(use_power_2_scales=True) @@ -795,12 +856,142 @@ def unkeyed_factory(role): with pytest.raises( ValueError, - match=r"Pass qfactory_key=.*@quantizer_factory", + match=r"(?s)requires a semantic qfactory key.*@quantizer_factory", ): CustomRecipe(qfactory=unkeyed_factory).quantizer_config() assert calls == [] +def test_recipe_pickle_round_trip_drops_caches_and_restores_fields(): + """Checkpoint bytes must not depend on whether caches were populated.""" + warm = Float8CurrentScaling(use_power_2_scales=True) + warm.quantizer_config() + repr(warm) + cold = Float8CurrentScaling(use_power_2_scales=True) + + assert pickle.dumps(warm) == pickle.dumps(cold) + + restored = pickle.loads(pickle.dumps(warm)) + assert "_cached_quantizer_config" not in pickle.loads(pickle.dumps(warm)).__dict__ + assert restored.quantizer_config() == warm.quantizer_config() + assert restored.fp8_quant_fwd_inp == warm.fp8_quant_fwd_inp + + +def test_legacy_pickled_recipe_gains_defaults_and_canonical_attention_flags(): + """A payload written before a field existed still loads and validates.""" + current = Float8CurrentScaling() + + # A payload that predates ``backward_override`` and the derived bundles, and + # that stores the pre-canonicalization (fp8_dpa=False, fp8_mha=True) pair. + legacy_state = { + name: value + for name, value in current.__dict__.items() + if name + not in ( + "backward_override", + "fp8_quant_fwd_inp", + "fp8_quant_fwd_weight", + "fp8_quant_bwd_grad", + ) + } + legacy_state["use_power_2_scales"] = True + legacy_state["fp8_dpa"] = False + legacy_state["fp8_mha"] = True + + restored = object.__new__(Float8CurrentScaling) + restored.__setstate__(legacy_state) + + assert restored.backward_override is None + assert restored.fp8_dpa is True + assert restored.fp8_quant_fwd_inp == QParams(power_2_scale=True, amax_epsilon=0.0) + assert ( + restored.quantizer_config() + == Float8CurrentScaling(use_power_2_scales=True, fp8_mha=True).quantizer_config() + ) + + +def test_setstate_preserves_customized_qparams(): + """Restoring must not re-run construction, which would reset customized bundles.""" + customized = Float8CurrentScaling() + customized.fp8_quant_fwd_inp = QParams(power_2_scale=False, amax_epsilon=1e-5) + + restored = pickle.loads(pickle.dumps(customized)) + + assert restored.fp8_quant_fwd_inp == QParams(power_2_scale=False, amax_epsilon=1e-5) + assert restored.quantizer_config() == customized.quantizer_config() + + +def test_copy_keeps_cached_configuration_for_committed_snapshots(): + """Runtime owners snapshot recipes per commit; the copy keeps the cached config.""" + original = Float8CurrentScaling() + config = original.quantizer_config() + + snapshot = copy.copy(original) + + assert snapshot.quantizer_config() is config + assert snapshot is not original + + +def test_reassigning_an_unchanged_value_keeps_the_cached_configuration(): + """A framework that rewrites its recipe fields every step must not force rebuilds.""" + recipe_instance = Float8CurrentScaling() + config = recipe_instance.quantizer_config() + + recipe_instance.fp8_dpa = recipe_instance.fp8_dpa + recipe_instance.fp8_format = Format.HYBRID + recipe_instance.fp8_quant_fwd_inp = QParams(power_2_scale=False, amax_epsilon=0.0) + + assert recipe_instance.quantizer_config() is config + + +def test_qfactory_replacement_does_not_inherit_the_previous_key(): + """A key describes one factory, so replacing the factory re-derives it.""" + + @quantizer_factory(key=("first_factory", 1)) + def first_factory(role): + raise AssertionError("factory must not be called during key resolution") + + @quantizer_factory(key=("second_factory", 1)) + def second_factory(role): + raise AssertionError("factory must not be called during key resolution") + + def unkeyed_factory(role): + raise AssertionError("factory must not be called during key resolution") + + # An attached key is picked up at construction and re-derived on replacement. + recipe_instance = CustomRecipe(qfactory=first_factory) + assert recipe_instance.qfactory_key == ("first_factory", 1) + first_config = recipe_instance.quantizer_config() + + recipe_instance.qfactory = second_factory + assert recipe_instance.qfactory_key == ("second_factory", 1) + assert recipe_instance.quantizer_config() != first_config + + # An explicit key describes the factory it was passed with, so it does not + # survive a replacement either: the new factory must bring its own. + explicit = CustomRecipe(qfactory=first_factory, qfactory_key=("explicit", 7)) + assert explicit.qfactory_key == ("explicit", 7) + explicit.qfactory = unkeyed_factory + assert explicit.qfactory_key is None + with pytest.raises(ValueError, match="requires a semantic qfactory key"): + explicit.quantizer_config() + + +def test_unkeyed_qfactory_error_names_both_ways_to_supply_a_key(): + """The rejection is the whole migration path, so it must be actionable.""" + + def unkeyed_factory(role): + raise AssertionError("factory must not be called during key resolution") + + with pytest.raises(ValueError) as excinfo: + CustomRecipe(qfactory=unkeyed_factory).quantizer_config() + + message = str(excinfo.value) + assert "unkeyed_factory" in message + assert "@quantizer_factory(key=" in message + assert "qfactory_key=" in message + + def test_custom_recipe_qfactory_key_mutation_changes_semantic_configuration(): """Changing a CustomRecipe policy key invalidates its cached config without probing.""" calls = [] diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index c5df9762e4c..8bd2c59344a 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -615,16 +615,28 @@ def test_candidate_update_is_not_published_during_planning(): assert not module._fp8_workspaces # pylint: disable=protected-access -def test_runtime_snapshot_does_not_copy_caller_repr_cache(): - """Caller display history must not affect a runtime snapshot or its checkpoint bytes.""" +def test_runtime_snapshot_keeps_caches_but_not_in_checkpoint_bytes(): + """Caller display history must not affect a snapshot's checkpoint bytes. + + The snapshot keeps the caller's derived caches, because comparing a committed + runtime against later configurations must not rebuild one every time. Cache + state is excluded where it would actually matter: the serialized payload. + """ recipe = _make_counting_recipe(("runtime-repr-cache", 1), []) str(recipe) module = Linear(16, 16, bias=False, device="cuda", name="linear") update = _prepare_runtime_update(module, recipe, revision=1) + snapshot = update.candidate.recipe assert recipe.__dict__["_cached_repr"] is not None - assert update.candidate.recipe.__dict__["_cached_repr"] is None + assert snapshot is not recipe + assert snapshot.quantizer_config() is recipe.quantizer_config() + + # The serialized payload is where display history would otherwise leak. + payload = snapshot.__getstate__() + assert "_cached_repr" not in payload + assert "_cached_quantizer_config" not in payload def test_uninitialized_runtime_can_be_prepared_without_publication(): diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 28a6e2f0085..c3e694d1460 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -3,16 +3,16 @@ # See LICENSE for license information. """This module provides predefined FP8 recipes.""" + from __future__ import annotations import abc import os from collections.abc import Hashable from enum import Enum -from typing import Any, Callable, Literal, NamedTuple, Optional, Union -from dataclasses import field, replace +from typing import Any, Callable, ClassVar, Dict, Literal, NamedTuple, Optional, Tuple, Union +from dataclasses import MISSING, field, fields, replace from pydantic.dataclasses import dataclass - _BACKWARD_OVERRIDES = (None, "high_precision", "dequantized") _NVFP4_4OVER6_SCOPES = ("none", "weights", "activations", "all") _NVFP4_4OVER6_ERR_MODES = ("MAE", "MSE") @@ -137,6 +137,144 @@ def _algorithm_config( return (("callable",), semantic_key) +# Sentinel for "this attribute has no stored value yet", so that assigning ``None`` +# to a field for the first time is not mistaken for an unchanged value. +_UNSET = object() + +# Cached values derived from a recipe's fields; never part of its semantic state. +_CACHE_ATTRIBUTES = ("_cached_repr", "_cached_quantizer_config", "_has_cached_quantizer_config") + + +# Rows of a derived-``QParams`` table: {recipe option: ((bundle, field, transform), ...)}. +_DerivedQParams = Dict[str, Tuple[Tuple[str, str, Callable[[Any], Any]], ...]] + + +def _unchanged(value: Any) -> Any: + """Derived ``QParams`` field takes the recipe option's value unchanged.""" + return value + + +def _negated(value: Any) -> bool: + """Derived ``QParams`` field takes the negation of the recipe option's value.""" + return not value + + +def _attached_qfactory_key(qfactory: Any) -> Optional[Hashable]: + """Return the semantic key ``@quantizer_factory`` attached to a factory, if any.""" + attached_key = getattr(qfactory, "qfactory_key", None) + if attached_key is None: + return None + return _validate_qfactory_key(attached_key, source="qfactory.qfactory_key") + + +def _config_factory_key(recipe: Any, value: Any) -> Hashable: + """Configuration entry for the quantizer factory: its semantic key alone.""" + if value is None: + qfactory = getattr(recipe, "qfactory", None) + name = getattr(qfactory, "__qualname__", None) or repr(qfactory) + raise ValueError( + f"CustomRecipe requires a semantic qfactory key for {name}. The key stands in for" + " the factory's behavior, because callable identity is not stable across processes," + " ranks or checkpoints. Attach one where the factory is defined:\n" + ' @quantizer_factory(key=("my_factory", 1))\n' + " def my_factory(role): ...\n" + "or pass it with the recipe: CustomRecipe(qfactory=my_factory," + ' qfactory_key=("my_factory", 1)).\n' + "Bump the revision whenever the factory's behavior changes, so quantizers built from" + " the previous behavior are rebuilt instead of silently reused." + ) + return value + + +def _config_algorithm(field_name: str, *, half: int) -> Callable[[Any, Any], Any]: + """Converter for one half of a literal-or-callable algorithm field. + + ``half=0`` is the algorithm entry, ``half=1`` its semantic key; each keeps its + own declaration position in the configuration. + """ + key_name = f"{field_name}_key" + + def convert(recipe: Any, value: Any) -> Any: + algorithm = value if half == 0 else getattr(recipe, field_name) + return _algorithm_config(algorithm, getattr(recipe, key_name), field_name=field_name)[half] + + return convert + + +# Semantic-configuration fields per recipe class, populated on first use. Kept in a +# module dict rather than on the class: writing to a class ``__dict__`` is not +# traceable by torch.compile, and a first activation can happen inside a compiled +# region. +_QUANTIZER_CONFIG_FIELDS: Dict[type, tuple] = {} + + +# Converters for field values that are not their own configuration, by value type. +_VALUE_CONFIG_CONVERTERS = ( + (Enum, lambda recipe, value: value.name), + (QParams, lambda recipe, value: _qparams_config(value)), + (MMParams, lambda recipe, value: _mmparams_config(value)), +) + + +def _build_quantizer_config_fields(recipe: Any) -> tuple: + """Resolve how each dataclass field contributes to the semantic configuration. + + The result depends only on the recipe class, so it is computed once per class: + per instance the work is then one attribute read and one ``is None`` test per + field. Entries are ``(entry_name, source_field, converter)`` in declaration + order; a ``None`` converter means the value is already its own configuration. + """ + # ``fields`` is given the instance, not the class: torch.compile can trace the + # instance form, and a recipe's first activation may happen inside a compiled + # region. + recipe_fields = fields(recipe) + field_names = {f.name for f in recipe_fields} + config_fields = [] + for recipe_field in recipe_fields: + name = recipe_field.name + + # The quantizer factory is represented by its semantic key alone: the + # callable itself is deliberately not part of the configuration. + if name == "qfactory": + config_fields.append(("qfactory_key", "qfactory_key", _config_factory_key)) + continue + if name == "qfactory_key": + continue + + # A field with a sibling ``_key`` is a literal-or-callable algorithm. + if f"{name}_key" in field_names: + config_fields.append((name, name, _config_algorithm(name, half=0))) + continue + if name.endswith("_key") and name[: -len("_key")] in field_names: + config_fields.append((name, name, _config_algorithm(name[: -len("_key")], half=1))) + continue + + value = getattr(recipe, name) + converter = next( + ( + convert + for value_type, convert in _VALUE_CONFIG_CONVERTERS + if isinstance(value, value_type) + ), + None, + ) + config_fields.append((name, name, converter)) + return tuple(config_fields) + + +def _recipe_type_entry(cls: type) -> tuple: + """Semantic recipe-type name: the nearest recipe class defined by this module. + + A framework subclass (e.g. Megatron's ``TEDelayedScaling``) keeps the semantic + identity of the recipe it specializes, so it neither invalidates quantizers nor + rejects updates against an equivalent stock recipe. + """ + for klass in cls.__mro__: + if klass.__module__ == __name__: + return ("recipe_type", klass.__name__) + return ("recipe_type", cls.__name__) + + def _validate_qfactory_key(key: Any, *, source: str) -> Hashable: """Validate and return an immutable semantic quantizer-factory key.""" if key is None: @@ -190,31 +328,50 @@ class Recipe: _has_cached_quantizer_config: bool = False def __setattr__(self, name: str, value: Any) -> None: - # FP8 MHA requires FP8 DPA; keep both flags canonical across mutations. - if name == "fp8_dpa" and not value and self.__dict__.get("fp8_mha", False): - value = True - # Invalidate derived caches on any recipe mutation. Updating either # cache itself must not invalidate the other cache. - cache_attribute = name in ( - "_cached_repr", - "_cached_quantizer_config", - "_has_cached_quantizer_config", - ) + cache_attribute = name in _CACHE_ATTRIBUTES if not cache_attribute: + # Re-assigning the stored value is not a mutation: keep the caches, so + # a framework that rewrites its recipe fields every step does not force + # a rebuild of the semantic configuration. + current = self.__dict__.get(name, _UNSET) + if current is not _UNSET and (current is value or current == value): + return object.__setattr__(self, "_cached_repr", None) object.__setattr__(self, "_cached_quantizer_config", None) object.__setattr__(self, "_has_cached_quantizer_config", False) object.__setattr__(self, name, value) - if name == "fp8_mha" and value: - object.__setattr__(self, "fp8_dpa", True) if not cache_attribute: self._synchronize_derived_qparams(name) def __post_init__(self) -> None: - """Canonicalize fields after Pydantic dataclass construction.""" + """Canonicalize and validate after Pydantic dataclass construction. + + Subclasses that materialize derived bundles call this last, once their own + fields are final, so validation always sees the finished recipe. + """ if getattr(self, "fp8_mha", False): self.fp8_dpa = True + self.validate() + + def validate(self) -> None: + """Check that the recipe's field values are a supported combination. + + Called at construction and again at every activation whose semantic + configuration is not already cached, so a field assigned after construction + is checked before it can reach quantizer construction. Subclasses add their + own rules and call ``super().validate()``. + """ + if getattr(self, "fp8_mha", False) and not getattr(self, "fp8_dpa", False): + raise ValueError( + "fp8_mha=True requires fp8_dpa=True. Disable fp8_mha before disabling fp8_dpa." + ) + + # Shorthand recipe options that own one field in derived ``QParams`` bundles, + # as ``{option: ((bundle, qparams_field, transform), ...)}``. Subclasses with + # such options declare their rows here instead of overriding the hook below. + _DERIVED_QPARAMS: ClassVar[_DerivedQParams] = {} def _synchronize_derived_qparams(self, name: str) -> None: """Propagate a changed recipe option into per-tensor ``QParams``. @@ -222,15 +379,76 @@ def _synchronize_derived_qparams(self, name: str) -> None: Some public recipe options are shorthand for one field in several derived ``QParams`` bundles. For example, ``use_power_2_scales`` sets ``power_2_scale`` for the input, weight, and gradient bundles. A recipe - with such an option overrides this hook and updates only that owned - ``QParams`` field, preserving unrelated fields such as + with such an option declares it in ``_DERIVED_QPARAMS`` and only that + owned ``QParams`` field is updated, preserving unrelated fields such as ``amax_epsilon``. Assigning a complete ``QParams`` bundle directly is still supported. That assignment controls every field until a later assignment to a shorthand recipe option updates the specific field it owns. """ - del name + rows = self._DERIVED_QPARAMS.get(name) + if rows is None: + return + value = getattr(self, name) + for bundle, qparams_field, transform in rows: + self._replace_derived_qparams(bundle, **{qparams_field: transform(value)}) + + def __getstate__(self) -> Dict[str, Any]: + """Serialize the recipe's fields only. + + Derived caches are rebuilt on demand, so keeping them out of the payload + makes checkpoint bytes independent of whether ``repr`` or a semantic + configuration was requested before saving. + """ + state = self.__dict__.copy() + for cache_attribute in _CACHE_ATTRIBUTES: + state.pop(cache_attribute, None) + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + """Restore a recipe, including payloads written by earlier versions. + + A payload can predate fields added since it was written, so missing fields + take their declared defaults and missing derived bundles are rebuilt from + the shorthand options that own them. ``__post_init__`` is deliberately not + re-run: it would overwrite bundles the saved recipe customized. + """ + self.__dict__.update(state) + + for recipe_field in fields(self): + if recipe_field.name in self.__dict__: + continue + if recipe_field.default is not MISSING: + self.__dict__[recipe_field.name] = recipe_field.default + elif recipe_field.default_factory is not MISSING: + self.__dict__[recipe_field.name] = recipe_field.default_factory() + + rebuilt: Dict[str, Dict[str, Any]] = {} + for option, rows in self._DERIVED_QPARAMS.items(): + value = self.__dict__.get(option) + for bundle, qparams_field, transform in rows: + if bundle not in self.__dict__: + rebuilt.setdefault(bundle, {})[qparams_field] = transform(value) + for bundle, qparams_fields in rebuilt.items(): + self.__dict__[bundle] = QParams(**qparams_fields) + + # Construction canonicalizes this pair, so a payload saved before the rule + # was enforced loads as the recipe it would be constructed as today. + if self.__dict__.get("fp8_mha", False): + self.__dict__["fp8_dpa"] = True + self.validate() + + def __copy__(self) -> "Recipe": + """Shallow copy that keeps the derived caches, unlike pickling. + + Runtime owners snapshot the caller's recipe on every commit, and the + snapshot is compared against later configurations; dropping its cached + configuration would make each comparison rebuild one. + """ + clone = object.__new__(type(self)) + clone.__dict__.update(self.__dict__) + return clone def _replace_derived_qparams(self, name: str, **changes: Any) -> None: """Replace fields in an initialized derived ``QParams`` bundle.""" @@ -247,6 +465,7 @@ def quantizer_config(self) -> Hashable: :meth:`_make_quantizer_config` rather than overriding this cached accessor. """ if not self._has_cached_quantizer_config: + self.validate() config = self._make_quantizer_config() try: hash(config) @@ -260,10 +479,29 @@ def quantizer_config(self) -> Hashable: return self._cached_quantizer_config def _make_quantizer_config(self) -> Hashable: - """Build the immutable semantic configuration returned by :meth:`quantizer_config`.""" - raise NotImplementedError( - f"{self.__class__.__name__} must implement Recipe.quantizer_config()" - ) + """Build the immutable semantic configuration returned by :meth:`quantizer_config`. + + One entry per dataclass field, in declaration order, behind a leading + ``("recipe_type", ...)`` entry. ``Enum`` fields contribute their name, + ``QParams``/``MMParams`` bundles their own configuration, and a + literal-or-callable algorithm its identity-independent form. Subclasses + get this for free; adding a field to a recipe therefore cannot leave it + out of the semantic configuration. + """ + cls = type(self) + cached_fields = _QUANTIZER_CONFIG_FIELDS.get(cls) + if cached_fields is None: + cached_fields = (_recipe_type_entry(cls), _build_quantizer_config_fields(self)) + _QUANTIZER_CONFIG_FIELDS[cls] = cached_fields + recipe_type_entry, config_fields = cached_fields + + values = self.__dict__ + items = [recipe_type_entry] + append = items.append + for entry_name, source, convert in config_fields: + value = values[source] + append((entry_name, value if convert is None else convert(self, value))) + return tuple(items) @abc.abstractmethod def _make_repr(self) -> str: @@ -415,8 +653,8 @@ def scaling_factor_compute(amax: Tensor, amax_compute_algo_key: Optional[Hashable] = None scaling_factor_compute_algo_key: Optional[Hashable] = None - def __post_init__(self) -> None: - super().__post_init__() + def validate(self) -> None: + super().validate() assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." assert ( self.backward_override in _BACKWARD_OVERRIDES @@ -425,32 +663,6 @@ def __post_init__(self) -> None: self.backward_override is None ), "Delayed scaling only supports backward_override=None." - def _make_quantizer_config(self) -> Hashable: - amax_compute_algo, amax_compute_algo_key = _algorithm_config( - self.amax_compute_algo, - self.amax_compute_algo_key, - field_name="amax_compute_algo", - ) - scaling_factor_compute_algo, scaling_factor_compute_algo_key = _algorithm_config( - self.scaling_factor_compute_algo, - self.scaling_factor_compute_algo_key, - field_name="scaling_factor_compute_algo", - ) - return ( - ("recipe_type", "DelayedScaling"), - ("margin", self.margin), - ("fp8_format", self.fp8_format.name), - ("amax_history_len", self.amax_history_len), - ("amax_compute_algo", amax_compute_algo), - ("scaling_factor_compute_algo", scaling_factor_compute_algo), - ("reduce_amax", self.reduce_amax), - ("fp8_dpa", self.fp8_dpa), - ("fp8_mha", self.fp8_mha), - ("backward_override", self.backward_override), - ("amax_compute_algo_key", amax_compute_algo_key), - ("scaling_factor_compute_algo_key", scaling_factor_compute_algo_key), - ) - def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -493,12 +705,22 @@ class Float8CurrentScaling(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) - def __post_init__(self) -> None: - super().__post_init__() + _DERIVED_QPARAMS: ClassVar[_DerivedQParams] = { + "use_power_2_scales": ( + ("fp8_quant_fwd_inp", "power_2_scale", _unchanged), + ("fp8_quant_fwd_weight", "power_2_scale", _unchanged), + ("fp8_quant_bwd_grad", "power_2_scale", _unchanged), + ), + } + + def validate(self) -> None: + super().validate() assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + + def __post_init__(self) -> None: self.fp8_quant_fwd_inp = QParams( power_2_scale=self.use_power_2_scales, amax_epsilon=0.0, @@ -511,35 +733,7 @@ def __post_init__(self) -> None: power_2_scale=self.use_power_2_scales, amax_epsilon=0.0, ) - - def _synchronize_derived_qparams(self, name: str) -> None: - if name != "use_power_2_scales": - return - for qparams_name in ( - "fp8_quant_fwd_inp", - "fp8_quant_fwd_weight", - "fp8_quant_bwd_grad", - ): - self._replace_derived_qparams( - qparams_name, - power_2_scale=self.use_power_2_scales, - ) - - def _make_quantizer_config(self) -> Hashable: - return ( - ("recipe_type", "Float8CurrentScaling"), - ("use_power_2_scales", self.use_power_2_scales), - ("fp8_format", self.fp8_format.name), - ("fp8_quant_fwd_inp", _qparams_config(self.fp8_quant_fwd_inp)), - ("fp8_quant_fwd_weight", _qparams_config(self.fp8_quant_fwd_weight)), - ("fp8_quant_bwd_grad", _qparams_config(self.fp8_quant_bwd_grad)), - ("fp8_gemm_fprop", _mmparams_config(self.fp8_gemm_fprop)), - ("fp8_gemm_dgrad", _mmparams_config(self.fp8_gemm_dgrad)), - ("fp8_gemm_wgrad", _mmparams_config(self.fp8_gemm_wgrad)), - ("fp8_dpa", self.fp8_dpa), - ("fp8_mha", self.fp8_mha), - ("backward_override", self.backward_override), - ) + super().__post_init__() def _make_repr(self) -> str: return ( @@ -596,24 +790,13 @@ class MXFP8BlockScaling(Recipe): backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) enable_2d_quantization: bool = False - def __post_init__(self) -> None: - super().__post_init__() + def validate(self) -> None: + super().validate() assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." - def _make_quantizer_config(self) -> Hashable: - return ( - ("recipe_type", "MXFP8BlockScaling"), - ("margin", self.margin), - ("fp8_format", self.fp8_format.name), - ("fp8_dpa", self.fp8_dpa), - ("fp8_mha", self.fp8_mha), - ("backward_override", self.backward_override), - ("enable_2d_quantization", self.enable_2d_quantization), - ) - def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -678,8 +861,16 @@ class Float8BlockScaling(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) - def __post_init__(self) -> None: - super().__post_init__() + _DERIVED_QPARAMS: ClassVar[_DerivedQParams] = { + "use_f32_scales": ( + ("fp8_quant_fwd_inp", "power_2_scale", _negated), + ("fp8_quant_fwd_weight", "power_2_scale", _negated), + ("fp8_quant_bwd_grad", "power_2_scale", _negated), + ), + } + + def validate(self) -> None: + super().validate() assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x" assert self.w_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for w" assert self.grad_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for grad" @@ -702,6 +893,8 @@ def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + + def __post_init__(self) -> None: power_2_scale = not self.use_f32_scales self.fp8_quant_fwd_inp = QParams( power_2_scale=power_2_scale, @@ -715,39 +908,7 @@ def __post_init__(self) -> None: power_2_scale=power_2_scale, amax_epsilon=0.0, ) - - def _synchronize_derived_qparams(self, name: str) -> None: - if name != "use_f32_scales": - return - power_2_scale = not self.use_f32_scales - for qparams_name in ( - "fp8_quant_fwd_inp", - "fp8_quant_fwd_weight", - "fp8_quant_bwd_grad", - ): - self._replace_derived_qparams( - qparams_name, - power_2_scale=power_2_scale, - ) - - def _make_quantizer_config(self) -> Hashable: - return ( - ("recipe_type", "Float8BlockScaling"), - ("use_f32_scales", self.use_f32_scales), - ("fp8_format", self.fp8_format.name), - ("fp8_quant_fwd_inp", _qparams_config(self.fp8_quant_fwd_inp)), - ("fp8_quant_fwd_weight", _qparams_config(self.fp8_quant_fwd_weight)), - ("fp8_quant_bwd_grad", _qparams_config(self.fp8_quant_bwd_grad)), - ("x_block_scaling_dim", self.x_block_scaling_dim), - ("w_block_scaling_dim", self.w_block_scaling_dim), - ("grad_block_scaling_dim", self.grad_block_scaling_dim), - ("fp8_gemm_fprop", _mmparams_config(self.fp8_gemm_fprop)), - ("fp8_gemm_dgrad", _mmparams_config(self.fp8_gemm_dgrad)), - ("fp8_gemm_wgrad", _mmparams_config(self.fp8_gemm_wgrad)), - ("fp8_dpa", self.fp8_dpa), - ("fp8_mha", self.fp8_mha), - ("backward_override", self.backward_override), - ) + super().__post_init__() def _make_repr(self) -> str: return ( @@ -859,8 +1020,17 @@ class NVFP4BlockScaling(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) - def __post_init__(self) -> None: - super().__post_init__() + _DERIVED_QPARAMS: ClassVar[_DerivedQParams] = { + "disable_rht": ( + ("fp4_quant_fwd_inp", "random_hadamard_transform", _negated), + ("fp4_quant_bwd_grad", "random_hadamard_transform", _negated), + ), + "disable_stochastic_rounding": (("fp4_quant_bwd_grad", "stochastic_rounding", _negated),), + "disable_2d_quantization": (("fp4_quant_fwd_weight", "fp4_2d_quantization", _negated),), + } + + def validate(self) -> None: + super().validate() assert self.fp4_format == Format.E2M1, "Only E2M1 is supported for NVFP4 scaling" assert self.fp8_format == Format.E4M3, "Only E4M3 is supported for NVFP4 scaling" assert ( @@ -876,6 +1046,7 @@ def __post_init__(self) -> None: self.nvfp4_4over6_err_mode in _NVFP4_4OVER6_ERR_MODES ), "NVTE_NVFP4_4OVER6_ERR_MODE must be one of: 'MAE', 'MSE'." + def __post_init__(self) -> None: # Quantization params # Note: RHT is currently only applied to column-wise usage so that # it can be used for wgrad GEMM. @@ -894,48 +1065,7 @@ def __post_init__(self) -> None: stochastic_rounding=not self.disable_stochastic_rounding, fp4_2d_quantization=False, ) - - def _synchronize_derived_qparams(self, name: str) -> None: - updates = () - if name == "disable_rht": - updates = ( - ("fp4_quant_fwd_inp", "random_hadamard_transform", not self.disable_rht), - ("fp4_quant_bwd_grad", "random_hadamard_transform", not self.disable_rht), - ) - elif name == "disable_stochastic_rounding": - updates = ( - ( - "fp4_quant_bwd_grad", - "stochastic_rounding", - not self.disable_stochastic_rounding, - ), - ) - elif name == "disable_2d_quantization": - updates = ( - ("fp4_quant_fwd_weight", "fp4_2d_quantization", not self.disable_2d_quantization), - ) - for qparams_name, field_name, value in updates: - self._replace_derived_qparams(qparams_name, **{field_name: value}) - - def _make_quantizer_config(self) -> Hashable: - return ( - ("recipe_type", "NVFP4BlockScaling"), - ("disable_rht", self.disable_rht), - ("disable_stochastic_rounding", self.disable_stochastic_rounding), - ("disable_2d_quantization", self.disable_2d_quantization), - ("row_scaled_activation", self.row_scaled_activation), - ("nvfp4_4over6", self.nvfp4_4over6), - ("nvfp4_4over6_e4m3_use_256", self.nvfp4_4over6_e4m3_use_256), - ("nvfp4_4over6_err_mode", self.nvfp4_4over6_err_mode), - ("fp4_format", self.fp4_format.name), - ("fp8_format", self.fp8_format.name), - ("fp4_quant_fwd_inp", _qparams_config(self.fp4_quant_fwd_inp)), - ("fp4_quant_fwd_weight", _qparams_config(self.fp4_quant_fwd_weight)), - ("fp4_quant_bwd_grad", _qparams_config(self.fp4_quant_bwd_grad)), - ("fp8_dpa", self.fp8_dpa), - ("fp8_mha", self.fp8_mha), - ("backward_override", self.backward_override), - ) + super().__post_init__() def _make_repr(self) -> str: return ( @@ -1033,41 +1163,32 @@ class CustomRecipe(Recipe): quantization_alignment: int = 128 qfactory_key: Optional[Hashable] = None - def __post_init__(self) -> None: - super().__post_init__() + def validate(self) -> None: + super().validate() assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." if self.quantization_alignment <= 0: raise ValueError("CustomRecipe quantization_alignment must be positive.") + + def __setattr__(self, name: str, value: Any) -> None: + super().__setattr__(name, value) + # A key stands in for one factory's behavior, so it must not outlive the + # factory it came from: replacing the factory re-derives it, and a + # replacement without a key is reported at the next activation. During + # construction the key is resolved by ``__post_init__`` instead. + if name == "qfactory" and "qfactory_key" in self.__dict__: + object.__setattr__(self, "qfactory_key", _attached_qfactory_key(value)) + + def __post_init__(self) -> None: if self.qfactory_key is None: - attached_key = getattr(self.qfactory, "qfactory_key", None) - if attached_key is not None: - self.qfactory_key = _validate_qfactory_key( - attached_key, - source="qfactory.qfactory_key", - ) + self.qfactory_key = _attached_qfactory_key(self.qfactory) else: self.qfactory_key = _validate_qfactory_key( self.qfactory_key, source="CustomRecipe qfactory_key", ) - - def _make_quantizer_config(self) -> Hashable: - if self.qfactory_key is None: - raise ValueError( - "CustomRecipe requires a semantic qfactory key. Pass qfactory_key=... or decorate " - "qfactory with @quantizer_factory(key=...)." - ) - return ( - ("recipe_type", "CustomRecipe"), - ("qfactory_key", self.qfactory_key), - ("fp8_format", self.fp8_format.name), - ("fp8_dpa", self.fp8_dpa), - ("fp8_mha", self.fp8_mha), - ("backward_override", self.backward_override), - ("quantization_alignment", self.quantization_alignment), - ) + super().__post_init__() def _make_repr(self) -> str: return ( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 1b5ab884952..dd893febc55 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -668,12 +668,15 @@ class DotProductAttention(TransformerEngineBaseModule): tensor remains in BF16. For MXFP8 attention, the fused kernel handles the S/dP slots internally, so their factory-provided quantizers are not consumed:: - from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.common.recipe import CustomRecipe, quantizer_factory from transformer_engine.pytorch.constants import DType from transformer_engine.pytorch.custom_recipes.quantizer_factories import nvfp4_factory from transformer_engine.pytorch.quantization import autocast from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + # The key stands in for the factory's behavior; bump its revision whenever + # the body changes, so quantizers built from the old behavior are rebuilt. + @quantizer_factory(key=("nvfp4_linear_mxfp8_dpa", 1)) def nvfp4_linear_mxfp8_dpa_factory(role): # NVFP4 for Linear roles and MXFP8 for supported DPA roles. is_dpa = role is not None and role.module_type == "dpa" diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 9705f1dcc85..c1b9c8fb4a8 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Base modules and utilities for TransformerEngine PyTorch API""" + import copy import io import math @@ -1239,10 +1240,9 @@ def _build_quantization_runtime( # reduction buckets must retain the exact configuration that this # candidate validated, even if the caller later mutates and reuses the # same recipe instance. + # ``Recipe.__copy__`` keeps the derived caches, and ``__getstate__`` keeps + # them out of the checkpoint, so the snapshot needs no cache surgery here. runtime_recipe = copy.copy(recipe) - # Do not snapshot a caller-populated display cache. Whether ``repr`` was - # requested before runtime construction must not change checkpoint bytes. - object.__setattr__(runtime_recipe, "_cached_repr", None) forward_state = RecipeState.create( # pylint: disable=assignment-from-none runtime_recipe, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index ebeaabd8a7f..bdff9239f35 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -122,21 +122,6 @@ class _QuantizationRuntimeKey: forward_roles: Tuple[Optional[QuantizerRole], ...] backward_roles: Tuple[Optional[QuantizerRole], ...] - def __post_init__(self) -> None: - """Canonicalize role containers and reject non-semantic slot values.""" - try: - hash(self.recipe_config) - except TypeError as exc: - raise TypeError("_QuantizationRuntimeKey.recipe_config must be hashable") from exc - - for attribute in ("forward_roles", "backward_roles"): - roles = tuple(getattr(self, attribute)) - if any(role is not None and not isinstance(role, QuantizerRole) for role in roles): - raise TypeError( - f"_QuantizationRuntimeKey.{attribute} entries must be QuantizerRole or None" - ) - object.__setattr__(self, attribute, roles) - @dataclass class _QuantizationRuntime: @@ -850,14 +835,13 @@ def get_unique_autocast_key( ): """ For FP8, each autocast can be uniquely identified by the recipe and fp8 group. - Object identity is sufficient since autocast contexts never outlive a single - training session. + Two recipes that build the same quantizers share a key, so the semantic + configuration identifies the recipe rather than its ``repr``, which carries + fields that do not affect quantization. """ - recipe_repr = recipe.__dict__.get("_cached_repr") if recipe is not None else None - if recipe_repr is None: - recipe_repr = str(recipe) + recipe_config = recipe.quantizer_config() if recipe is not None else None group_id = id(group) if group is not None else None - return f"recipe={recipe_repr},group={group_id}" + return f"recipe={recipe_config!r},group={group_id}" @classmethod def _prepare_autocast_enter( From fd0adb4454392755fbfc2f612a5cfc693c1a53d6 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Thu, 10 Sep 2026 14:50:31 +0000 Subject: [PATCH 88/97] Decide quantizer reuse from the recipe, not a revision counter Signed-off-by: Evgeny --- tests/pytorch/test_cuda_graphs.py | 55 ++---- tests/pytorch/test_hybrid_quantization.py | 6 +- tests/pytorch/test_recipe.py | 57 +++--- tests/pytorch/test_recipe_update.py | 165 +++++++++--------- .../dot_product_attention.py | 2 +- transformer_engine/pytorch/graph.py | 19 +- transformer_engine/pytorch/module/base.py | 108 ++++++------ .../pytorch/module/grouped_linear.py | 2 +- transformer_engine/pytorch/quantization.py | 61 ++----- 9 files changed, 197 insertions(+), 278 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 29749d7dc05..f0d618525f5 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -754,8 +754,8 @@ def test_make_graphed_callables_with_kwargs( @pytest.mark.skipif(not fp8_available, reason="FP8 is not available") -def test_graphed_callable_rejects_changed_quantization_recipe(monkeypatch) -> None: - """Validate the capture recipe once per revision and reject mismatches.""" +def test_graphed_callable_rejects_changed_quantization_recipe() -> None: + """Accept replays under an equal configuration and reject changed ones.""" model_config = model_configs["small"] dtype = torch.bfloat16 captured_recipe = recipe.Float8CurrentScaling() @@ -775,47 +775,30 @@ def test_graphed_callable_rejects_changed_quantization_recipe(monkeypatch) -> No recipe=captured_recipe, ) - get_config_calls = 0 - get_quantizer_config = FP8GlobalStateManager.get_quantizer_config - - def counted_get_quantizer_config(cls): - nonlocal get_config_calls - get_config_calls += 1 - return get_quantizer_config() - - monkeypatch.setattr( - FP8GlobalStateManager, - "get_quantizer_config", - classmethod(counted_get_quantizer_config), - ) + def replay(): + output = model(generate_data(model_config, dtype)) + output.backward(generate_data(model_config, dtype, requires_grad=False)) try: - with autocast(enabled=True, recipe=equal_recipe): - for _ in range(2): - output = model(generate_data(model_config, dtype)) - output.backward(generate_data(model_config, dtype, requires_grad=False)) - assert get_config_calls == 1 - - # Re-entering after autocast restores the previous recipe advances the - # revision. Validate the equal semantic configuration once again, then - # stay on the integer-only fast path. - with autocast(enabled=True, recipe=equal_recipe): - for _ in range(2): - output = model(generate_data(model_config, dtype)) - output.backward(generate_data(model_config, dtype, requires_grad=False)) - assert get_config_calls == 2 + # An equal but independent recipe replays, in this scope and a later one. + for _ in range(2): + with autocast(enabled=True, recipe=equal_recipe): + for _ in range(2): + replay() # Mutating and reusing the caller-owned recipe must not mutate the # immutable semantic configuration saved at capture time. captured_recipe.fp8_quant_fwd_inp = recipe.QParams(amax_epsilon=0.25) with autocast(enabled=True, recipe=captured_recipe): - for _ in range(2): - with pytest.raises( - RuntimeError, - match="Recapture the graph with the new recipe", - ): - model(generate_data(model_config, dtype)) - assert get_config_calls == 3 + with pytest.raises(RuntimeError, match="Recapture the graph with the new recipe"): + model(generate_data(model_config, dtype)) + + # A manager reset cannot make a different recipe look like the captured + # one: the check compares configurations, not a recyclable counter. + FP8GlobalStateManager.reset() + with autocast(enabled=True, recipe=recipe.Float8CurrentScaling(use_power_2_scales=True)): + with pytest.raises(RuntimeError, match="Recapture the graph with the new recipe"): + model(generate_data(model_config, dtype)) finally: reset_graphs(model) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 413dd466f30..f719158ef2e 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -195,8 +195,8 @@ def second_factory(role): second_runtime = dpa._quantization_runtime assert second_runtime is not first_runtime - assert dpa.fp8_meta["scaling_fwd"] is second_runtime.forward_states[0] - assert dpa.fp8_meta["scaling_bwd"] is second_runtime.backward_states[0] + assert dpa.fp8_meta["scaling_fwd"] is second_runtime.forward_state + assert dpa.fp8_meta["scaling_bwd"] is second_runtime.backward_state assert dpa.quantizers["scaling_fwd"] is second_runtime.forward_quantizers assert dpa.quantizers["scaling_bwd"] is second_runtime.backward_quantizers assert second_capabilities == first_capabilities == (False, True) @@ -234,8 +234,6 @@ def test_direct_runtime_commit_updates_dpa_derived_state(self): ) update = dpa._plan_quantization_update( recipe=second_recipe, - recipe_config=second_recipe.quantizer_config(), - recipe_config_revision=old_runtime.recipe_config_revision + 1, num_gemms=dpa._get_quantization_runtime_num_gemms(), ) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 5ef613af559..313908b8ea9 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -166,10 +166,9 @@ def test_quantization_runtime_bundles_both_quantization_directions(): key=runtime_key, recipe=runtime_recipe, num_gemms=1, - recipe_config_revision=3, role_revision=2, - forward_states=(), - backward_states=(), + forward_state=None, + backward_state=None, forward_quantizers=forward_quantizers, backward_quantizers=backward_quantizers, ) @@ -667,7 +666,9 @@ def settings(quantizers): def make_module_owner(module_type, num_gemms): owner = object.__new__(module_type) owner.__dict__.update( - name="test", + # ``name`` is a property backed by ``_name``: writing the public name + # into ``__dict__`` would be shadowed by the descriptor. + _name="test", num_gemms=num_gemms, fp8_meta={"num_gemms": num_gemms}, fp8_meta_tensors_initialized=False, @@ -1008,34 +1009,26 @@ def factory(role): assert calls == [] -def test_global_state_caches_active_quantizer_config_and_revision(): - """Recipe activation publishes semantic configuration with a manager-owned revision.""" +def test_global_state_publishes_the_active_recipe_and_its_configuration(): + """Activation publishes one value: the recipe, which owns its configuration.""" FP8GlobalStateManager.reset() active_recipe = Float8CurrentScaling() - assert FP8GlobalStateManager.get_quantizer_config_revision() == 0 FP8GlobalStateManager.activate_recipe(active_recipe) original_config = active_recipe.quantizer_config() assert FP8GlobalStateManager.get_fp8_recipe() is active_recipe assert FP8GlobalStateManager.get_quantizer_config() is original_config - assert FP8GlobalStateManager.get_quantizer_config_revision() == 1 - # Equivalent recipe objects update the requested recipe without advancing the revision. + # An equivalent recipe object replaces the active recipe and compares equal. equivalent_recipe = Float8CurrentScaling() FP8GlobalStateManager.activate_recipe(equivalent_recipe) assert FP8GlobalStateManager.get_fp8_recipe() is equivalent_recipe - assert FP8GlobalStateManager.get_quantizer_config() is original_config - assert FP8GlobalStateManager.get_quantizer_config_revision() == 1 - - # Recipe mutation is requested state until the recipe is activated again. - active_recipe.fp8_dpa = True - assert FP8GlobalStateManager.get_quantizer_config() is original_config - assert FP8GlobalStateManager.get_quantizer_config_revision() == 1 + assert FP8GlobalStateManager.get_quantizer_config() == original_config - FP8GlobalStateManager.activate_recipe(active_recipe) - assert FP8GlobalStateManager.get_quantizer_config() == active_recipe.quantizer_config() + # Mutating the active recipe is visible immediately. + equivalent_recipe.fp8_dpa = True assert FP8GlobalStateManager.get_quantizer_config() != original_config - assert FP8GlobalStateManager.get_quantizer_config_revision() == 2 + assert FP8GlobalStateManager.get_quantizer_config() == equivalent_recipe.quantizer_config() FP8GlobalStateManager.reset() @@ -1051,7 +1044,6 @@ def keyed_factory(role): keyed_recipe = CustomRecipe(qfactory=keyed_factory) FP8GlobalStateManager.activate_recipe(keyed_recipe) keyed_config = keyed_recipe.quantizer_config() - keyed_revision = FP8GlobalStateManager.get_quantizer_config_revision() assert FP8GlobalStateManager.get_fp8_recipe() is keyed_recipe assert FP8GlobalStateManager.get_quantizer_config() is keyed_config assert calls == [] @@ -1063,7 +1055,6 @@ def unkeyed_factory(role): FP8GlobalStateManager.activate_recipe(CustomRecipe(qfactory=unkeyed_factory)) assert FP8GlobalStateManager.get_fp8_recipe() is keyed_recipe assert FP8GlobalStateManager.get_quantizer_config() is keyed_config - assert FP8GlobalStateManager.get_quantizer_config_revision() == keyed_revision assert calls == [] FP8GlobalStateManager.reset() @@ -1075,8 +1066,6 @@ def _autocast_activation_state(): qstate.fp8_enabled, qstate.fp8_calibration, id(qstate.fp8_recipe), - id(qstate.quantizer_config), - qstate.quantizer_config_revision, id(qstate.fp8_distributed_group), qstate.is_first_fp8_module, qstate.fp8_graph_capturing, @@ -1162,47 +1151,43 @@ def unkeyed_factory(role): FP8GlobalStateManager.reset() -def test_autocast_restores_recipe_and_quantizer_config_together(): - """Leaving autocast restores the previously active recipe/configuration pair.""" +def test_autocast_restores_the_outer_recipe_pointer(): + """Leaving autocast restores the recipe object that was active on entry.""" FP8GlobalStateManager.reset() outer_recipe = Float8CurrentScaling() inner_recipe = MXFP8BlockScaling() FP8GlobalStateManager.activate_recipe(outer_recipe) outer_config = FP8GlobalStateManager.get_quantizer_config() - outer_revision = FP8GlobalStateManager.get_quantizer_config_revision() - # Equal independent configurations do not advance the revision on entry or restoration. + # An equal but independent recipe leaves the same configuration active. with te.autocast(enabled=False, recipe=Float8CurrentScaling()): - assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision - assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + assert FP8GlobalStateManager.get_quantizer_config() == outer_config + assert FP8GlobalStateManager.get_fp8_recipe() is outer_recipe with te.autocast(enabled=False, recipe=inner_recipe): assert FP8GlobalStateManager.get_fp8_recipe() is inner_recipe assert FP8GlobalStateManager.get_quantizer_config() == inner_recipe.quantizer_config() - assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + 1 + # Restoration republishes the recipe object, so its cached config comes back. assert FP8GlobalStateManager.get_fp8_recipe() is outer_recipe assert FP8GlobalStateManager.get_quantizer_config() is outer_config - assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + 2 FP8GlobalStateManager.reset() -def test_quantized_model_init_restores_recipe_config_with_monotonic_revision(): - """Model-init recipe restoration republishes config rather than reusing an old revision.""" +def test_quantized_model_init_restores_the_outer_recipe_pointer(): + """Model-init restores the recipe object that was active on entry.""" FP8GlobalStateManager.reset() outer_recipe = Float8CurrentScaling() inner_recipe = MXFP8BlockScaling() FP8GlobalStateManager.activate_recipe(outer_recipe) outer_config = FP8GlobalStateManager.get_quantizer_config() - outer_revision = FP8GlobalStateManager.get_quantizer_config_revision() with te.quantized_model_init(enabled=False, recipe=inner_recipe): assert FP8GlobalStateManager.get_fp8_recipe() is inner_recipe - assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + 1 + assert FP8GlobalStateManager.get_quantizer_config() == inner_recipe.quantizer_config() assert FP8GlobalStateManager.get_fp8_recipe() is outer_recipe assert FP8GlobalStateManager.get_quantizer_config() is outer_config - assert FP8GlobalStateManager.get_quantizer_config_revision() == outer_revision + 2 FP8GlobalStateManager.reset() diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 8bd2c59344a..f6a02be75b8 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -56,20 +56,16 @@ def qfactory(role): return CustomRecipe(qfactory=qfactory, qfactory_key=key) -def _ensure_runtime(module, recipe, revision, *, num_gemms=1): +def _ensure_runtime(module, recipe, *, num_gemms=1): return module._ensure_quantization_runtime( # pylint: disable=protected-access recipe=recipe, - recipe_config=recipe.quantizer_config(), - recipe_config_revision=revision, num_gemms=num_gemms, ) -def _prepare_runtime_update(module, recipe, revision, *, num_gemms=1): +def _prepare_runtime_update(module, recipe, *, num_gemms=1): return module._plan_quantization_update( # pylint: disable=protected-access recipe=recipe, - recipe_config=recipe.quantizer_config(), - recipe_config_revision=revision, num_gemms=num_gemms, ) @@ -337,11 +333,7 @@ def _active_runtime_owners(module): def _global_recipe_state(): """Return the identity-bearing global recipe state.""" state = FP8GlobalStateManager.quantization_state - return ( - state.fp8_recipe, - state.quantizer_config, - state.quantizer_config_revision, - ) + return (state.fp8_recipe,) def _make_compatible_fp8_mha_recipe(key): @@ -387,24 +379,23 @@ def test_delayed_runtime_rejects_effective_recipe_updates_atomically(replacement """Delayed state is frozen once a module runtime has been initialized.""" module = Linear(16, 16, bias=False, device="cuda", name="linear") recipe = DelayedScaling(amax_history_len=4) - assert _ensure_runtime(module, recipe, revision=1) + assert _ensure_runtime(module, recipe) old_views = _runtime_views(module) with pytest.raises( RuntimeError, match="Mid-training recipe updates do not support delayed scaling", ): - _ensure_runtime(module, replacement, revision=2) + _ensure_runtime(module, replacement) assert all(current is old for current, old in zip(_runtime_views(module), old_views)) - assert old_views[0].recipe_config_revision == 1 def test_delayed_runtime_rejects_role_and_slot_layout_updates(): """Role and GEMM-layout changes cannot rebuild a delayed runtime.""" module = Linear(16, 16, bias=False, device="cuda", name="linear") recipe = DelayedScaling(amax_history_len=4) - assert _ensure_runtime(module, recipe, revision=1) + assert _ensure_runtime(module, recipe) old_views = _runtime_views(module) module.output_quantizer_role = QuantizerRole( @@ -413,20 +404,20 @@ def test_delayed_runtime_rejects_role_and_slot_layout_updates(): name="consumer", ) with pytest.raises(RuntimeError, match="do not support delayed scaling"): - _ensure_runtime(module, recipe, revision=1) + _ensure_runtime(module, recipe) assert all(current is old for current, old in zip(_runtime_views(module), old_views)) with pytest.raises(RuntimeError, match="do not support delayed scaling"): - _ensure_runtime(module, recipe, revision=1, num_gemms=2) + _ensure_runtime(module, recipe, num_gemms=2) assert all(current is old for current, old in zip(_runtime_views(module), old_views)) """An active stateless runtime cannot acquire delayed state mid-training.""" module = Linear(16, 16, bias=False, device="cuda", name="linear") - assert _ensure_runtime(module, Float8CurrentScaling(), revision=1) + assert _ensure_runtime(module, Float8CurrentScaling()) old_views = _runtime_views(module) with pytest.raises(RuntimeError, match="do not support delayed scaling"): - _ensure_runtime(module, DelayedScaling(amax_history_len=4), revision=2) + _ensure_runtime(module, DelayedScaling(amax_history_len=4)) assert all(current is old for current, old in zip(_runtime_views(module), old_views)) @@ -439,11 +430,11 @@ def test_custom_recipe_cannot_introduce_delayed_state(): qfactory_key=("custom-enter-delayed", 2), ) module = Linear(16, 16, bias=False, device="cuda", name="linear") - assert _ensure_runtime(module, active_recipe, revision=1) + assert _ensure_runtime(module, active_recipe) old_views = _runtime_views(module) with pytest.raises(RuntimeError, match="do not support delayed scaling"): - _ensure_runtime(module, replacement, revision=2) + _ensure_runtime(module, replacement) assert all(current is old for current, old in zip(_runtime_views(module), old_views)) @@ -459,7 +450,7 @@ def test_asymmetric_custom_delayed_scaling_is_rejected_clearly(): RuntimeError, match="This hybrid quantization configuration with delayed scaling is not supported", ): - _ensure_runtime(module, recipe, revision=1) + _ensure_runtime(module, recipe) assert module._quantization_runtime is None # pylint: disable=protected-access @@ -473,7 +464,7 @@ def active_factory(role): active_recipe = CustomRecipe(qfactory=active_factory, qfactory_key=("mixed", 1)) module = Linear(16, 16, bias=False, device="cuda", name="linear") - assert _ensure_runtime(module, active_recipe, revision=1) + assert _ensure_runtime(module, active_recipe) old_views = _runtime_views(module) active_call_count = len(calls) @@ -482,7 +473,7 @@ def replacement_factory(_role): replacement = CustomRecipe(qfactory=replacement_factory, qfactory_key=("mixed", 2)) with pytest.raises(RuntimeError, match="do not support delayed scaling"): - _ensure_runtime(module, replacement, revision=2) + _ensure_runtime(module, replacement) assert len(calls) == active_call_count assert all(current is old for current, old in zip(_runtime_views(module), old_views)) @@ -492,13 +483,13 @@ def test_same_delayed_recipe_object_mutation_keeps_committed_snapshot(): """Rejecting caller mutation must not mutate the active reduction recipe.""" module = Linear(16, 16, bias=False, device="cuda", name="linear") recipe = DelayedScaling(amax_history_len=4, margin=0) - assert _ensure_runtime(module, recipe, revision=1) + assert _ensure_runtime(module, recipe) committed_recipe = module.fp8_meta["recipe"] assert committed_recipe is not recipe recipe.margin = 1 with pytest.raises(RuntimeError, match="do not support delayed scaling"): - _ensure_runtime(module, recipe, revision=2) + _ensure_runtime(module, recipe) assert module.fp8_meta["recipe"] is committed_recipe assert committed_recipe.margin == 0 @@ -543,52 +534,63 @@ def test_rejected_delayed_update_aborts_autocast_reduction(): equal_recipe = _make_counting_recipe(("runtime-reuse", 1), calls) module = Linear(16, 16, bias=False, device="cuda", name="linear") - assert _ensure_runtime(module, first_recipe, revision=1) + assert _ensure_runtime(module, first_recipe) active = module._quantization_runtime # pylint: disable=protected-access assert active is not None assert len(calls) == 5 - # Revision 3 models a module that did not run while revision 2 was active. - assert not _ensure_runtime(module, equal_recipe, revision=3) + # A module that missed an intervening region still recognizes an equal recipe. + assert not _ensure_runtime(module, equal_recipe) assert module._quantization_runtime is active # pylint: disable=protected-access - assert active.recipe_config_revision == 3 assert len(calls) == 5 -def test_unchanged_runtime_uses_revision_hot_path(monkeypatch): - """The steady-state check does not resolve roles or invoke the factory.""" +@pytest.mark.parametrize("reuse_recipe_object", (True, False), ids=("same-object", "equal-object")) +def test_unchanged_runtime_uses_the_steady_state_path(monkeypatch, reuse_recipe_object): + """The steady-state check resolves no roles and invokes no factory. + + A reused recipe object hits on identity; an equivalent one built per region + falls through to the configuration comparison and must agree. + """ calls = [] recipe = _make_counting_recipe(("runtime-hot-path", 1), calls) module = Linear(16, 16, bias=False, device="cuda", name="linear") - _ensure_runtime(module, recipe, revision=1) + _ensure_runtime(module, recipe) + active = module._quantization_runtime # pylint: disable=protected-access + factory_call_count = len(calls) def unexpected_role_resolution(**_kwargs): raise AssertionError("unchanged runtime resolved quantizer roles") monkeypatch.setattr(module, "get_quantizer_roles", unexpected_role_resolution) - assert not _ensure_runtime(module, recipe, revision=1) - assert len(calls) == 5 + + for _ in range(3): + requested = ( + recipe if reuse_recipe_object else _make_counting_recipe(("runtime-hot-path", 1), calls) + ) + assert not _ensure_runtime(module, requested) + assert module._quantization_runtime is active # pylint: disable=protected-access + assert len(calls) == factory_call_count -def test_revision_only_update_is_not_published_during_planning(): - """An equal semantic update synchronizes revisions only when committed.""" +def test_equal_semantic_update_is_not_published_during_planning(): + """An equal semantic update commits nothing and calls no factory.""" calls = [] - recipe = _make_counting_recipe(("revision-plan", 1), calls) - equal_recipe = _make_counting_recipe(("revision-plan", 1), calls) + recipe = _make_counting_recipe(("equal-plan", 1), calls) + equal_recipe = _make_counting_recipe(("equal-plan", 1), calls) module = Linear(16, 16, bias=False, device="cuda", name="linear") - _ensure_runtime(module, recipe, revision=1) + _ensure_runtime(module, recipe) active = module._quantization_runtime # pylint: disable=protected-access factory_call_count = len(calls) - update = _prepare_runtime_update(module, equal_recipe, revision=3) + update = _prepare_runtime_update(module, equal_recipe) assert update.candidate is None assert module._quantization_runtime is active # pylint: disable=protected-access - assert active.recipe_config_revision == 1 assert len(calls) == factory_call_count assert not module._apply_quantization_update(update) # pylint: disable=protected-access - assert active.recipe_config_revision == 3 + assert module._quantization_runtime is active # pylint: disable=protected-access assert len(calls) == factory_call_count @@ -598,16 +600,15 @@ def test_candidate_update_is_not_published_during_planning(): active_recipe = _make_counting_recipe(("candidate-plan", 1), calls) replacement_recipe = _make_counting_recipe(("candidate-plan", 2), calls) module = Linear(16, 16, bias=False, device="cuda", name="linear") - _ensure_runtime(module, active_recipe, revision=1) + _ensure_runtime(module, active_recipe) old_views = _runtime_views(module) workspace = object() module._fp8_workspaces["weight"] = workspace # pylint: disable=protected-access - update = _prepare_runtime_update(module, replacement_recipe, revision=2) + update = _prepare_runtime_update(module, replacement_recipe) assert update.candidate is not None assert all(current is old for current, old in zip(_runtime_views(module), old_views)) - assert old_views[0].recipe_config_revision == 1 assert module._fp8_workspaces["weight"] is workspace # pylint: disable=protected-access assert module._apply_quantization_update(update) # pylint: disable=protected-access @@ -616,17 +617,12 @@ def test_candidate_update_is_not_published_during_planning(): def test_runtime_snapshot_keeps_caches_but_not_in_checkpoint_bytes(): - """Caller display history must not affect a snapshot's checkpoint bytes. - - The snapshot keeps the caller's derived caches, because comparing a committed - runtime against later configurations must not rebuild one every time. Cache - state is excluded where it would actually matter: the serialized payload. - """ + """A snapshot keeps the caller's caches; only its payload excludes them.""" recipe = _make_counting_recipe(("runtime-repr-cache", 1), []) str(recipe) module = Linear(16, 16, bias=False, device="cuda", name="linear") - update = _prepare_runtime_update(module, recipe, revision=1) + update = _prepare_runtime_update(module, recipe) snapshot = update.candidate.recipe assert recipe.__dict__["_cached_repr"] is not None @@ -645,7 +641,7 @@ def test_uninitialized_runtime_can_be_prepared_without_publication(): recipe = _make_counting_recipe(("initial-plan", 1), calls) module = Linear(16, 16, bias=False, device="cuda", name="linear") - update = _prepare_runtime_update(module, recipe, revision=1) + update = _prepare_runtime_update(module, recipe) assert update.candidate is not None assert module._quantization_runtime is None # pylint: disable=protected-access @@ -664,12 +660,12 @@ def test_later_planning_failure_leaves_earlier_module_unchanged(monkeypatch): replacement_recipe = _make_counting_recipe(("multi-plan", 2), calls) first = Linear(16, 16, bias=False, device="cuda", name="first") second = Linear(16, 16, bias=False, device="cuda", name="second") - _ensure_runtime(first, active_recipe, revision=1) - _ensure_runtime(second, active_recipe, revision=1) + _ensure_runtime(first, active_recipe) + _ensure_runtime(second, active_recipe) old_first_views = _runtime_views(first) old_second_views = _runtime_views(second) - first_update = _prepare_runtime_update(first, replacement_recipe, revision=2) + first_update = _prepare_runtime_update(first, replacement_recipe) assert first_update.candidate is not None def reject_candidate(_candidate): @@ -677,12 +673,10 @@ def reject_candidate(_candidate): monkeypatch.setattr(second, "_validate_quantization_runtime", reject_candidate) with pytest.raises(RuntimeError, match="later module validation failure"): - _prepare_runtime_update(second, replacement_recipe, revision=2) + _prepare_runtime_update(second, replacement_recipe) assert all(current is old for current, old in zip(_runtime_views(first), old_first_views)) assert all(current is old for current, old in zip(_runtime_views(second), old_second_views)) - assert old_first_views[0].recipe_config_revision == 1 - assert old_second_views[0].recipe_config_revision == 1 @pytest.mark.parametrize( @@ -759,11 +753,11 @@ def test_same_recipe_object_semantic_mutation_rebuilds_runtime(): calls = [] recipe = _make_counting_recipe(("same-object-update", 1), calls) module = Linear(16, 16, bias=False, device="cuda", name="linear") - _ensure_runtime(module, recipe, revision=1) + _ensure_runtime(module, recipe) old_runtime = module._quantization_runtime # pylint: disable=protected-access recipe.qfactory_key = ("same-object-update", 2) - assert _ensure_runtime(module, recipe, revision=2) + assert _ensure_runtime(module, recipe) assert module._quantization_runtime is not old_runtime # pylint: disable=protected-access assert len(calls) == 10 @@ -1012,7 +1006,7 @@ def test_role_revision_is_requested_until_atomic_runtime_commit(): calls = [] recipe = _make_counting_recipe(("role-update", 1), calls) module = Linear(16, 16, bias=False, device="cuda", name="linear") - _ensure_runtime(module, recipe, revision=1) + _ensure_runtime(module, recipe) old_runtime = module._quantization_runtime # pylint: disable=protected-access old_forward_state = module.fp8_meta["scaling_fwd"] @@ -1026,11 +1020,11 @@ def test_role_revision_is_requested_until_atomic_runtime_commit(): assert module.fp8_meta["scaling_fwd"] is old_forward_state assert module.quantizers["scaling_fwd"] is old_forward_quantizers - assert _ensure_runtime(module, recipe, revision=1) + assert _ensure_runtime(module, recipe) new_runtime = module._quantization_runtime # pylint: disable=protected-access assert new_runtime is not old_runtime assert new_runtime.key.forward_roles[-1] == role - assert module.fp8_meta["scaling_fwd"] is new_runtime.forward_states[0] + assert module.fp8_meta["scaling_fwd"] is new_runtime.forward_state assert module.quantizers["scaling_fwd"] is new_runtime.forward_quantizers assert len(calls) == 10 @@ -1039,7 +1033,7 @@ def test_role_revision_is_requested_until_atomic_runtime_commit(): module_type="linear", tensor_type="input", name="consumer" ) assert module._role_revision == old_role_revision + 1 # pylint: disable=protected-access - assert not _ensure_runtime(module, recipe, revision=1) + assert not _ensure_runtime(module, recipe) assert len(calls) == 10 @@ -1053,7 +1047,7 @@ def test_backward_factory_failure_keeps_complete_active_runtime(): fail_on_grad_output=True, ) module = Linear(16, 16, bias=False, device="cuda", name="linear") - _ensure_runtime(module, active_recipe, revision=1) + _ensure_runtime(module, active_recipe) old_runtime = module._quantization_runtime # pylint: disable=protected-access old_recipe = module.fp8_meta["recipe"] @@ -1063,7 +1057,7 @@ def test_backward_factory_failure_keeps_complete_active_runtime(): old_backward_quantizers = module.quantizers["scaling_bwd"] with pytest.raises(RuntimeError, match="backward factory failure"): - _ensure_runtime(module, failing_recipe, revision=2) + _ensure_runtime(module, failing_recipe) assert module._quantization_runtime is old_runtime # pylint: disable=protected-access assert module.fp8_meta["recipe"] is old_recipe @@ -1071,7 +1065,6 @@ def test_backward_factory_failure_keeps_complete_active_runtime(): assert module.fp8_meta["scaling_bwd"] is old_backward_state assert module.quantizers["scaling_fwd"] is old_forward_quantizers assert module.quantizers["scaling_bwd"] is old_backward_quantizers - assert old_runtime.recipe_config_revision == 1 @pytest.mark.parametrize( @@ -1138,7 +1131,7 @@ def test_runtime_update_workspace_lifecycle( ) replacement_recipe = _make_counting_recipe(("workspace-update", 3), calls) module = module_factory() - assert _ensure_runtime(module, active_recipe, revision=1, num_gemms=num_gemms) + assert _ensure_runtime(module, active_recipe, num_gemms=num_gemms) active_runtime = module._quantization_runtime # pylint: disable=protected-access active_views = ( @@ -1153,14 +1146,14 @@ def test_runtime_update_workspace_lifecycle( module._fp8_workspaces[cache_name] = workspace # pylint: disable=protected-access # A global revision change with an equal semantic runtime preserves caches. - assert not _ensure_runtime(module, equal_recipe, revision=2, num_gemms=num_gemms) + assert not _ensure_runtime(module, equal_recipe, num_gemms=num_gemms) assert module._quantization_runtime is active_runtime # pylint: disable=protected-access for cache_name, workspace in workspaces.items(): assert module._fp8_workspaces[cache_name] is workspace # pylint: disable=protected-access # Candidate construction failure preserves both compatibility views and caches. with pytest.raises(RuntimeError, match="backward factory failure"): - _ensure_runtime(module, failing_recipe, revision=3, num_gemms=num_gemms) + _ensure_runtime(module, failing_recipe, num_gemms=num_gemms) assert module._quantization_runtime is active_runtime # pylint: disable=protected-access current_views = ( module.fp8_meta["recipe"], @@ -1174,7 +1167,7 @@ def test_runtime_update_workspace_lifecycle( assert module._fp8_workspaces[cache_name] is workspace # pylint: disable=protected-access # Only a fully committed replacement clears cached workspaces. - assert _ensure_runtime(module, replacement_recipe, revision=4, num_gemms=num_gemms) + assert _ensure_runtime(module, replacement_recipe, num_gemms=num_gemms) replacement_runtime = module._quantization_runtime # pylint: disable=protected-access assert replacement_runtime is not active_runtime assert not module._fp8_workspaces # pylint: disable=protected-access @@ -1859,8 +1852,8 @@ def test_delayed_qmi_mha_first_forward_preserves_initialized_runtimes( def runtime_object_ids(owner): runtime = owner._quantization_runtime # pylint: disable=protected-access assert runtime is not None - forward_state = runtime.forward_states[0] - backward_state = runtime.backward_states[0] + forward_state = runtime.forward_state + backward_state = runtime.backward_state return tuple( id(obj) for obj in ( @@ -2112,35 +2105,35 @@ def test_apply_recipe_success_noop_mutation_and_forward_fast_path(): try: apply_recipe(model, first_recipe) first_runtimes = [module._quantization_runtime for module in model] - first_revision = FP8GlobalStateManager.get_quantizer_config_revision() + first_config = FP8GlobalStateManager.get_quantizer_config() first_factory_calls = len(calls) assert first_factory_calls == 10 assert FP8GlobalStateManager.get_fp8_recipe() is first_recipe # An independent equal recipe updates only the manager's requested - # recipe object. Runtime identity, revision, and factories stay fixed. + # recipe object. Runtime identity, configuration, and factories stay fixed. apply_recipe(model, equal_recipe) assert [module._quantization_runtime for module in model] == first_runtimes - assert FP8GlobalStateManager.get_quantizer_config_revision() == first_revision + assert FP8GlobalStateManager.get_quantizer_config() == first_config assert FP8GlobalStateManager.get_fp8_recipe() is equal_recipe assert len(calls) == first_factory_calls # The first matching autocast forward after explicit application uses - # each module's revision fast path and performs no factory work. + # each module's steady-state path and performs no factory work. inp = torch.randn(8, 16, device="cuda", requires_grad=True) _run_update_step(model, equal_recipe, inp) assert [module._quantization_runtime for module in model] == first_runtimes assert len(calls) == first_factory_calls # Mutating and reusing the same recipe object produces a real model-wide - # replacement and one new global revision. + # replacement and a different active configuration. equal_recipe.qfactory_key = ("apply-success", 2) apply_recipe(model, equal_recipe) assert all( module._quantization_runtime is not old_runtime for module, old_runtime in zip(model, first_runtimes) ) - assert FP8GlobalStateManager.get_quantizer_config_revision() == first_revision + 1 + assert FP8GlobalStateManager.get_quantizer_config() != first_config assert len(calls) == first_factory_calls + 10 finally: FP8GlobalStateManager.reset() @@ -2450,7 +2443,7 @@ def test_candidate_validation_failure_keeps_complete_active_runtime(monkeypatch) active_recipe = _make_counting_recipe(("validation-update", 1), calls) candidate_recipe = _make_counting_recipe(("validation-update", 2), calls) module = Linear(16, 16, bias=False, device="cuda", name="linear") - _ensure_runtime(module, active_recipe, revision=1) + _ensure_runtime(module, active_recipe) old_runtime = module._quantization_runtime # pylint: disable=protected-access old_forward_state = module.fp8_meta["scaling_fwd"] @@ -2464,7 +2457,7 @@ def reject_candidate(_candidate): monkeypatch.setattr(module, "_validate_quantization_runtime", reject_candidate) with pytest.raises(RuntimeError, match="candidate validation failure"): - _ensure_runtime(module, candidate_recipe, revision=2) + _ensure_runtime(module, candidate_recipe) assert module._quantization_runtime is old_runtime # pylint: disable=protected-access assert module.fp8_meta["scaling_fwd"] is old_forward_state @@ -2507,7 +2500,7 @@ def qfactory(role): module = GroupedLinear(2, 16, 16, bias=False, device="cuda", name="grouped") active_recipe = make_recipe(("grouped-atomic", "active", mismatched_tensor_type)) - _ensure_runtime(module, active_recipe, revision=1, num_gemms=2) + _ensure_runtime(module, active_recipe, num_gemms=2) old_runtime = module._quantization_runtime # pylint: disable=protected-access old_forward_state = module.fp8_meta["scaling_fwd"] @@ -2523,7 +2516,7 @@ def qfactory(role): mismatched_role=mismatched_tensor_type, ) with pytest.raises(ValueError, match="incompatible plain backend configurations"): - _ensure_runtime(module, invalid_recipe, revision=2, num_gemms=2) + _ensure_runtime(module, invalid_recipe, num_gemms=2) assert module._quantization_runtime is old_runtime # pylint: disable=protected-access assert module.fp8_meta["scaling_fwd"] is old_forward_state @@ -2539,7 +2532,7 @@ def qfactory(role): dtype=torch.float16, unsafe_inputs=True, ) - update = _prepare_runtime_update(module, replacement_recipe, revision=3, num_gemms=2) + update = _prepare_runtime_update(module, replacement_recipe, num_gemms=2) assert module._validated_quantizer_generations is old_validated_generations assert module._delayed_scaling_input_quantizer is old_delayed_quantizer assert module._unsafe_requantization_input_quantizer is old_unsafe_quantizer diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index dd893febc55..c77debb9bc2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1422,7 +1422,7 @@ def _validate_quantization_runtime( return _infer_custom_dpa_local_recipes( candidate.recipe, - {"scaling_fwd": candidate.forward_states[0]}, + {"scaling_fwd": candidate.forward_state}, { "scaling_fwd": candidate.forward_quantizers, "scaling_bwd": candidate.backward_quantizers, diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 647b018884a..952d4d433e8 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Functions for CUDA Graphs support in FP8""" + from collections.abc import Iterable import contextlib import gc @@ -1267,22 +1268,10 @@ def reset(): return backward_dw, reset - # Cache CUDA-graph recipe validation by the manager-owned configuration - # revision. Semantic configuration equality is only evaluated when recipe - # activation advances the revision, while unchanged replays only compare - # integers. - checked_recipe_config_revision = None - checked_recipe_config_matches = False - def active_recipe_config_matches_capture() -> bool: - nonlocal checked_recipe_config_revision, checked_recipe_config_matches - recipe_config_revision = FP8GlobalStateManager.get_quantizer_config_revision() - if recipe_config_revision != checked_recipe_config_revision: - checked_recipe_config_matches = ( - FP8GlobalStateManager.get_quantizer_config() == captured_recipe_config - ) - checked_recipe_config_revision = recipe_config_revision - return checked_recipe_config_matches + """Whether the active recipe still builds the quantizers this graph captured.""" + recipe_config = FP8GlobalStateManager.get_quantizer_config() + return recipe_config is captured_recipe_config or recipe_config == captured_recipe_config # Put together the final graphed callables ret = [] diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index c1b9c8fb4a8..02246a93813 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -10,7 +10,6 @@ import os import pickle import warnings -from collections.abc import Hashable from dataclasses import dataclass from enum import Enum from abc import ABC, abstractmethod @@ -88,7 +87,6 @@ class _QuantizationUpdate: candidate: Optional[_QuantizationRuntime] validation_result: Any - recipe_config_revision: int role_revision: int @@ -941,7 +939,7 @@ def __init__( ): if role is not None and not isinstance(role, QuantizerRole): raise TypeError(f"{attribute} must be a QuantizerRole or None") - self.name = name + self._name = name self.next_iter_when_debug_should_be_run = 0 self.fp8_initialized = False self.fp8 = False @@ -975,6 +973,24 @@ def __init__( TEDebugState.initialize() self._validate_name() + @property + def name(self) -> Optional[str]: + """Semantic name of this module, used to key its quantizer roles.""" + return self._name + + @name.setter + def name(self, value: Optional[str]) -> None: + """Rename the module, invalidating a committed runtime's role-keyed key. + + Never route a name through ``fast_setattr``: it writes ``__dict__`` and + bypasses this setter. + """ + if getattr(self, "_name", None) == value: + return + self._name = value + if getattr(self, "_quantization_runtime", None) is not None: + self.fast_setattr("_role_revision", self._role_revision + 1) + def fast_setattr(self, name: str, value: Any) -> None: """ Fast version of the Module's set attribute function. @@ -1227,7 +1243,6 @@ def _build_quantization_runtime( recipe: Recipe, key: _QuantizationRuntimeKey, num_gemms: int, - recipe_config_revision: int, role_revision: int, forward_state_roles: Optional[List[Optional[QuantizerRole]]], backward_state_roles: Optional[List[Optional[QuantizerRole]]], @@ -1259,8 +1274,8 @@ def _build_quantization_runtime( active = getattr(self, "_quantization_runtime", None) if active is not None: - old_forward_state = active.forward_states[0] - old_backward_state = active.backward_states[0] + old_forward_state = active.forward_state + old_backward_state = active.backward_state else: # Checkpoint loading and legacy direct initialization can populate # compatibility metadata before the first runtime is active. @@ -1285,10 +1300,9 @@ def _build_quantization_runtime( key=key, recipe=runtime_recipe, num_gemms=num_gemms, - recipe_config_revision=recipe_config_revision, role_revision=role_revision, - forward_states=(forward_state,), - backward_states=(backward_state,), + forward_state=forward_state, + backward_state=backward_state, forward_quantizers=forward_quantizers, backward_quantizers=backward_quantizers, ) @@ -1306,8 +1320,8 @@ def _activate_quantization_runtime( """Publish a validated replacement runtime and its compatibility views.""" del validation_result recipe = candidate.recipe - forward_state = candidate.forward_states[0] - backward_state = candidate.backward_states[0] + forward_state = candidate.forward_state + backward_state = candidate.backward_state # ``fp8_meta`` and ``quantizers`` remain compatibility views. Only # update them after both directions have been built and validated. @@ -1336,9 +1350,8 @@ def _get_quantization_runtime_num_gemms(self) -> int: @staticmethod def _runtime_has_delayed_scaling(runtime: _QuantizationRuntime) -> bool: """Whether either direction of a runtime owns delayed-scaling state.""" - return any( - _is_delayed_scaling_state(state) - for state in runtime.forward_states + runtime.backward_states + return _is_delayed_scaling_state(runtime.forward_state) or _is_delayed_scaling_state( + runtime.backward_state ) @staticmethod @@ -1390,11 +1403,14 @@ def _plan_quantization_update( self, *, recipe: Recipe, - recipe_config: Hashable, - recipe_config_revision: int, num_gemms: int, ) -> _QuantizationUpdate: - """Plan and validate a complete cold-path update without publishing it.""" + """Plan and validate a complete cold-path update without publishing it. + + The configuration is read once here, so the key and the quantizers below + are derived from the same value. + """ + recipe_config = recipe.quantizer_config() active = getattr(self, "_quantization_runtime", None) role_revision = getattr(self, "_role_revision", 0) @@ -1422,7 +1438,6 @@ def _plan_quantization_update( return _QuantizationUpdate( candidate=None, validation_result=None, - recipe_config_revision=recipe_config_revision, role_revision=role_revision, ) @@ -1436,7 +1451,6 @@ def _plan_quantization_update( recipe=recipe, key=requested_key, num_gemms=num_gemms, - recipe_config_revision=recipe_config_revision, role_revision=role_revision, forward_state_roles=forward_state_roles, backward_state_roles=backward_state_roles, @@ -1450,7 +1464,6 @@ def _plan_quantization_update( return _QuantizationUpdate( candidate=candidate, validation_result=validation_result, - recipe_config_revision=recipe_config_revision, role_revision=role_revision, ) @@ -1463,7 +1476,6 @@ def _apply_quantization_update( if candidate is None: active = self._quantization_runtime assert active is not None - active.recipe_config_revision = update.recipe_config_revision active.role_revision = update.role_revision return False @@ -1477,33 +1489,30 @@ def _ensure_quantization_runtime( self, *, recipe: Recipe, - recipe_config: Hashable, - recipe_config_revision: int, num_gemms: int, ) -> bool: - """Use the revision hot path or plan and immediately apply a module-local update. + """Reuse the committed runtime, or plan and immediately apply a module-local update. Return whether a new runtime was committed. """ active = getattr(self, "_quantization_runtime", None) role_revision = getattr(self, "_role_revision", 0) - # Constant-time steady-state path: no role construction, factory - # calls, recipe traversal, or candidate validation. - if ( - active is not None - and active.recipe_config_revision == recipe_config_revision - and active.role_revision == role_revision - and active.num_gemms == num_gemms - ): - return False - - update = self._plan_quantization_update( - recipe=recipe, - recipe_config=recipe_config, - recipe_config_revision=recipe_config_revision, - num_gemms=num_gemms, - ) + # Constant-time steady-state path: no role construction, factory calls, + # recipe traversal, or candidate validation. Identity settles a reused + # recipe object; equality covers an equivalent recipe built per region. + if active is not None and active.role_revision == role_revision: + recipe_config = recipe.quantizer_config() + active_config = active.key.recipe_config + if ( + active_config is recipe_config or active_config == recipe_config + ) and active.num_gemms == num_gemms: + # Frameworks write this compatibility view from outside TE. + if self.fp8_meta.get("recipe") is not active.recipe: + self.fp8_meta["recipe"] = active.recipe + return False + + update = self._plan_quantization_update(recipe=recipe, num_gemms=num_gemms) return self._apply_quantization_update(update) def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: @@ -1514,12 +1523,9 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """ # TODO(negvet): Remove set_meta_tensor after checkpoint, GroupedLinear, and DPA migration. del fwd - num_gemms = self.fp8_meta.get("num_gemms", 1) - self._ensure_quantization_runtime( - recipe=recipe, - recipe_config=recipe.quantizer_config(), - recipe_config_revision=-1, - num_gemms=num_gemms, + num_gemms = self._get_quantization_runtime_num_gemms() + self._apply_quantization_update( + self._plan_quantization_update(recipe=recipe, num_gemms=num_gemms) ) def get_quantizer_roles( @@ -1890,17 +1896,13 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: original_recipe = meta.get("recipe") recipe = FP8GlobalStateManager.get_fp8_recipe() - recipe_config = FP8GlobalStateManager.get_quantizer_config() - recipe_config_revision = FP8GlobalStateManager.get_quantizer_config_revision() - if fp8_enabled: - meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() runtime_changed = self._ensure_quantization_runtime( recipe=recipe, - recipe_config=recipe_config, - recipe_config_revision=recipe_config_revision, num_gemms=num_gemms, ) + if fp8_enabled: + meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() if fp8_enabled: self.fast_setattr("fp8_initialized", True) @@ -2371,6 +2373,8 @@ def _validate_name(self): if self.name is not None: return + # MultiheadAttention and TransformerLayer borrow this helper without + # inheriting the name property, so assign the public attribute. self.name = f"Layer_{TEDebugState.get_layer_count()}" def _check_weight_tensor_recipe_correspondence(self) -> None: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 97aaadae0dc..d7bd1ef562d 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1877,7 +1877,7 @@ def _validate_quantization_runtime( """Validate both grouped directions without changing live module state.""" super()._validate_quantization_runtime(candidate) if ( - isinstance(candidate.forward_states[0], Float8CurrentScalingRecipeState) + isinstance(candidate.forward_state, Float8CurrentScalingRecipeState) and self.tp_size > 1 ): raise ValueError( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index bdff9239f35..9d26dff5d59 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -138,10 +138,9 @@ class _QuantizationRuntime: key: _QuantizationRuntimeKey recipe: Recipe num_gemms: int - recipe_config_revision: int role_revision: int - forward_states: Tuple["RecipeState", ...] - backward_states: Tuple["RecipeState", ...] + forward_state: "RecipeState" + backward_state: "RecipeState" forward_quantizers: List["Quantizer"] backward_quantizers: List["Quantizer"] @@ -446,8 +445,6 @@ class FP8GlobalState: fp8_enabled: bool = False fp8_calibration: bool = False fp8_recipe: Optional[Recipe] = None - quantizer_config: Optional[Hashable] = None - quantizer_config_revision: int = 0 fp8_distributed_group: Optional[dist_group_type] = None fp8_parameters: bool = False high_precision_init_val: bool = False @@ -684,39 +681,17 @@ def get_fp8_recipe(cls) -> Recipe: return cls.quantization_state.fp8_recipe return get_default_fp8_recipe() - @classmethod - def _set_recipe_config( - cls, - recipe: Optional[Recipe], - quantizer_config: Optional[Hashable], - ) -> None: - """Publish a recipe/configuration pair and advance its monotonic revision.""" - qstate = cls.quantization_state - config_changed = quantizer_config != qstate.quantizer_config - if config_changed: - qstate.quantizer_config_revision += 1 - qstate.quantizer_config = quantizer_config - qstate.fp8_recipe = recipe - @classmethod def activate_recipe(cls, recipe: Recipe) -> Hashable: - """Make a recipe and its semantic quantizer configuration active together.""" + """Make a recipe active and return its semantic quantizer configuration.""" quantizer_config = recipe.quantizer_config() - cls._set_recipe_config(recipe, quantizer_config) + cls.quantization_state.fp8_recipe = recipe return quantizer_config @classmethod def get_quantizer_config(cls) -> Hashable: """Return the active recipe's cached semantic quantizer configuration.""" - quantizer_config = cls.quantization_state.quantizer_config - if quantizer_config is None: - return cls.activate_recipe(cls.get_fp8_recipe()) - return quantizer_config - - @classmethod - def get_quantizer_config_revision(cls) -> int: - """Return the manager-owned active recipe-configuration revision.""" - return cls.quantization_state.quantizer_config_revision + return cls.get_fp8_recipe().quantizer_config() @classmethod def get_fp8_group(cls) -> Union[dist_group_type, None]: @@ -729,9 +704,8 @@ def get_autocast_state(cls) -> tuple: qstate = cls.quantization_state return ( qstate.fp8_enabled, - qstate.fp8_calibration, qstate.fp8_recipe, - qstate.quantizer_config, + qstate.fp8_calibration, qstate.fp8_distributed_group, qstate.is_first_fp8_module, qstate.fp8_graph_capturing, @@ -739,18 +713,17 @@ def get_autocast_state(cls) -> tuple: @classmethod def set_autocast_state(cls, state: tuple) -> None: - """Restore an autocast snapshot without restoring its old configuration revision.""" + """Restore an autocast snapshot.""" qstate = cls.quantization_state ( fp8_enabled, - fp8_calibration, fp8_recipe, - quantizer_config, + fp8_calibration, fp8_distributed_group, is_first_fp8_module, fp8_graph_capturing, ) = state - cls._set_recipe_config(fp8_recipe, quantizer_config) + qstate.fp8_recipe = fp8_recipe qstate.fp8_enabled = fp8_enabled qstate.fp8_calibration = fp8_calibration qstate.fp8_distributed_group = fp8_distributed_group @@ -884,7 +857,7 @@ def autocast_enter( ) -> None: """Prepare and publish state for entry into an FP8 region.""" - fp8_recipe, quantizer_config, autocast_key = cls._prepare_autocast_enter( + fp8_recipe, _, autocast_key = cls._prepare_autocast_enter( enabled, fp8_recipe, fp8_group, @@ -893,7 +866,7 @@ def autocast_enter( # Preparation above contains every operation that can fail due to the # requested recipe or platform. Publish only after it has succeeded. - cls._set_recipe_config(fp8_recipe, quantizer_config) + qstate.fp8_recipe = fp8_recipe # Once a delayed bucket is registered, its committed recipe snapshot # owns the reduction semantics for this key. Do not replace it with a # caller-owned recipe object merely by entering another autocast. @@ -1051,10 +1024,7 @@ def apply_recipe(model: torch.nn.Module, recipe: Recipe) -> None: raise RuntimeError("te.apply_recipe() must be called outside CUDA graph capture.") check_recipe_support(recipe) - recipe_config = recipe.quantizer_config() - recipe_config_revision = qstate.quantizer_config_revision + int( - recipe_config != qstate.quantizer_config - ) + recipe.quantizer_config() # Import locally to keep quantization.py independent from module/base.py # during package initialization. @@ -1103,8 +1073,6 @@ def apply_recipe(model: torch.nn.Module, recipe: Recipe) -> None: num_gemms = module._get_quantization_runtime_num_gemms() update = module._plan_quantization_update( # pylint: disable=protected-access recipe=recipe, - recipe_config=recipe_config, - recipe_config_revision=recipe_config_revision, num_gemms=num_gemms, ) except Exception as exc: @@ -1118,7 +1086,7 @@ def apply_recipe(model: torch.nn.Module, recipe: Recipe) -> None: # cannot expose a requested recipe globally. for module, update in updates: module._apply_quantization_update(update) # pylint: disable=protected-access - FP8GlobalStateManager._set_recipe_config(recipe, recipe_config) + FP8GlobalStateManager.quantization_state.fp8_recipe = recipe @contextmanager @@ -1204,7 +1172,6 @@ def quantized_model_init( qstate = FP8GlobalStateManager.quantization_state _fp8_parameters = qstate.fp8_parameters _fp8_recipe = qstate.fp8_recipe - _quantizer_config = qstate.quantizer_config _high_precision_init_val = qstate.high_precision_init_val qstate.fp8_parameters = enabled FP8GlobalStateManager.activate_recipe(get_default_fp8_recipe() if recipe is None else recipe) @@ -1213,7 +1180,7 @@ def quantized_model_init( yield finally: qstate.fp8_parameters = _fp8_parameters - FP8GlobalStateManager._set_recipe_config(_fp8_recipe, _quantizer_config) + qstate.fp8_recipe = _fp8_recipe qstate.high_precision_init_val = _high_precision_init_val From 092f00f07d387391baba452a7443a424f68d5a30 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 11 Sep 2026 10:36:26 +0000 Subject: [PATCH 89/97] Validate recipe entry points before they publish state Signed-off-by: Evgeny --- tests/pytorch/test_recipe.py | 137 +++++++++++++++++++- tests/pytorch/test_recipe_update.py | 141 +++++++++++++++++++++ transformer_engine/pytorch/module/base.py | 20 +-- transformer_engine/pytorch/ops/fuser.py | 10 +- transformer_engine/pytorch/ops/op.py | 1 - transformer_engine/pytorch/quantization.py | 81 +++++------- 6 files changed, 327 insertions(+), 63 deletions(-) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 313908b8ea9..6b839e52945 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -1102,6 +1102,135 @@ def unkeyed_factory(role): FP8GlobalStateManager.reset() +def test_quantized_model_init_failure_restores_fp8_parameters(): + """A recipe this context cannot use must leave the process as it found it.""" + FP8GlobalStateManager.reset() + + def unkeyed_factory(role): + raise AssertionError("factory must not be called") + + assert not FP8GlobalStateManager.with_fp8_parameters() + with pytest.raises(ValueError, match="requires a semantic qfactory key"): + with te.quantized_model_init(enabled=True, recipe=CustomRecipe(qfactory=unkeyed_factory)): + pass + + assert not FP8GlobalStateManager.with_fp8_parameters() + assert FP8GlobalStateManager.get_fp8_recipe() is not None + FP8GlobalStateManager.reset() + + +def test_disabled_autocast_does_not_validate_the_recipe(): + """A region that quantizes nothing must not reject a recipe it never uses.""" + FP8GlobalStateManager.reset() + + def unkeyed_factory(role): + raise AssertionError("factory must not be called") + + unusable = CustomRecipe(qfactory=unkeyed_factory) + + # Disabled: the recipe is published as a pointer and never asked for a config. + with te.autocast(enabled=False, recipe=unusable): + assert FP8GlobalStateManager.get_fp8_recipe() is unusable + + # Calibration builds quantizers, so it validates like an enabled region. + with pytest.raises(ValueError, match="requires a semantic qfactory key"): + with te.autocast(enabled=False, calibrating=True, recipe=unusable): + pass + FP8GlobalStateManager.reset() + + +def test_calibrating_autocast_runs_platform_preflight(monkeypatch): + """Calibration builds quantizers, so it needs the platform gate an fp8 region needs.""" + FP8GlobalStateManager.reset() + try: + # A plain disabled region publishes a pointer without consulting the platform. + monkeypatch.setattr( + transformer_engine.pytorch.quantization, + "check_fp8_support", + lambda: (False, "injected FP8 support failure"), + ) + with te.autocast(enabled=False, recipe=Float8CurrentScaling()): + pass + + original_state = _autocast_activation_state() + with pytest.raises(RuntimeError, match="injected FP8 support failure"): + with te.autocast(enabled=False, calibrating=True, recipe=Float8CurrentScaling()): + pass + assert _autocast_activation_state() == original_state + finally: + FP8GlobalStateManager.reset() + + +def test_stateless_regions_do_not_register_autocast_arguments(): + """Only delayed buckets own reduction arguments, so only they register any.""" + FP8GlobalStateManager.reset() + qstate = FP8GlobalStateManager.quantization_state + try: + for index in range(25): + + @quantizer_factory(key=("stateless_region", index)) + def stateless_factory(role): + del role + + with te.autocast(enabled=False, recipe=CustomRecipe(qfactory=stateless_factory)): + pass + assert not qstate.autocast_arguments + finally: + FP8GlobalStateManager.reset() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_delayed_bucket_registers_exactly_one_autocast_argument(): + """Reduction arguments are registered where the bucket is created, once.""" + FP8GlobalStateManager.reset() + qstate = FP8GlobalStateManager.quantization_state + try: + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + stateless = Linear(16, 16, bias=False, params_dtype=torch.bfloat16, device="cuda") + delayed = Linear(16, 16, bias=False, params_dtype=torch.bfloat16, device="cuda") + + # Stateless regions quantize without a bucket, so they register nothing. + for _ in range(25): + with te.autocast(enabled=True, recipe=Float8CurrentScaling()): + stateless(inp) + assert not qstate.autocast_arguments + + # A delayed bucket is created once and owns one set of reduction arguments. + for _ in range(3): + with te.autocast(enabled=True, recipe=DelayedScaling(amax_history_len=4)): + delayed(inp) + assert len(qstate.autocast_arguments) == 1 + finally: + FP8GlobalStateManager.reset() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_failed_buffer_registration_leaves_no_partial_entry(monkeypatch): + """A half-registered module would later unpack four buffer positions from two.""" + FP8GlobalStateManager.reset() + try: + module = Linear(16, 16, bias=False, params_dtype=torch.bfloat16, device="cuda") + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + + # The backward direction is registered second, so this fails mid-walk. + original_key = FP8GlobalStateManager.get_key_in_buffer.__func__ + + def failing_key(cls, forward, fp8_recipe, fp8_group): + if not forward: + raise RuntimeError("injected buffer registration failure") + return original_key(cls, forward, fp8_recipe, fp8_group) + + monkeypatch.setattr(FP8GlobalStateManager, "get_key_in_buffer", classmethod(failing_key)) + + with pytest.raises(RuntimeError, match="injected buffer registration failure"): + with te.autocast(enabled=True, recipe=DelayedScaling(amax_history_len=4)): + module(inp) + + assert FP8GlobalStateManager.get_buffer_info() not in module.fp8_meta + finally: + FP8GlobalStateManager.reset() + + def test_failed_autocast_support_check_is_atomic(monkeypatch): """Platform validation must complete before autocast state is published.""" FP8GlobalStateManager.reset() @@ -1114,12 +1243,12 @@ def keyed_factory(role): del role monkeypatch.setattr( - FP8GlobalStateManager, - "is_fp8_available", - classmethod(lambda _cls: (False, "injected FP8 support failure")), + transformer_engine.pytorch.quantization, + "check_fp8_support", + lambda: (False, "injected FP8 support failure"), ) context = te.autocast(enabled=True, recipe=CustomRecipe(qfactory=keyed_factory)) - with pytest.raises(AssertionError, match="injected FP8 support failure"): + with pytest.raises(RuntimeError, match="injected FP8 support failure"): context.__enter__() assert context._fp8_state is None # pylint: disable=protected-access diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index f6a02be75b8..343cb425ffe 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -528,6 +528,147 @@ def test_rejected_delayed_update_aborts_autocast_reduction(): assert torch.equal(state.scale, expected_scale) assert torch.equal(state.amax_history, expected_history) + +def test_failed_stateless_activation_aborts_a_registered_delayed_reduction(): + """A caught failure in one module must not update another module's delayed bucket.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + try: + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + delayed_module = Linear( + 16, + 16, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="delayed", + ) + with autocast(enabled=True, recipe=DelayedScaling(amax_history_len=4, margin=0)): + delayed_module(inp) + + state = delayed_module.fp8_meta["scaling_fwd"] + state.scale.fill_(3) + state.amax_history.fill_(7) + expected_scale = state.scale.clone() + expected_history = state.amax_history.clone() + + def failing_qfactory(role): + del role + raise RuntimeError("deliberate qfactory failure") + + stateless_module = Linear( + 16, + 16, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="stateless", + ) + recipe = CustomRecipe(qfactory=failing_qfactory, qfactory_key=("abort-stateless", 1)) + with autocast(enabled=True, recipe=recipe): + with pytest.raises(RuntimeError, match="deliberate qfactory failure"): + stateless_module(inp) + + assert torch.equal(state.scale, expected_scale) + assert torch.equal(state.amax_history, expected_history) + finally: + FP8GlobalStateManager.reset() + + +def test_uncaught_region_failure_aborts_the_autocast_reduction(): + """An exception that escapes the region leaves its delayed tensors untouched.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + try: + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + module = Linear( + 16, + 16, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="linear", + ) + recipe = DelayedScaling(amax_history_len=4, margin=0) + with autocast(enabled=True, recipe=recipe): + module(inp) + + state = module.fp8_meta["scaling_fwd"] + state.scale.fill_(3) + state.amax_history.fill_(7) + expected_scale = state.scale.clone() + expected_history = state.amax_history.clone() + + with pytest.raises(RuntimeError, match="deliberate region failure"): + with autocast(enabled=True, recipe=recipe): + raise RuntimeError("deliberate region failure") + + assert torch.equal(state.scale, expected_scale) + assert torch.equal(state.amax_history, expected_history) + finally: + FP8GlobalStateManager.reset() + + +def test_rejected_fusible_update_aborts_autocast_reduction(monkeypatch): + """The fusible rejection fires before the op reset loop, so it must still abort.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + try: + model = te_ops.Sequential(te_ops.Linear(16, 16, bias=False, device="cuda")) + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + with autocast(enabled=True, recipe=DelayedScaling(amax_history_len=4, margin=0)): + model(inp) + + reductions = [] + monkeypatch.setattr( + FP8GlobalStateManager, + "reduce_and_update_fp8_tensors", + classmethod(lambda _cls, forward=True: reductions.append(forward)), + ) + with autocast(enabled=True, recipe=DelayedScaling(amax_history_len=4, margin=1)): + with pytest.raises(RuntimeError, match="not supported for fusible operations"): + model(inp) + assert not reductions + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_preflight_rejects_unsupported_platform_before_commit(monkeypatch): + """A recipe this platform cannot run is rejected before any module is planned.""" + monkeypatch.setattr( + "transformer_engine.pytorch.quantization._NVFP4_SUPPORT", + (False, "injected NVFP4 support failure"), + ) + calls = [] + active_recipe = _make_counting_recipe(("apply-preflight", 1), calls) + module = Linear(16, 16, bias=False, device="cuda", name="linear") + assert _ensure_runtime(module, active_recipe) + active_runtime = module._quantization_runtime # pylint: disable=protected-access + active_recipe_pointer = FP8GlobalStateManager.quantization_state.fp8_recipe + call_count = len(calls) + + try: + with pytest.raises(RuntimeError, match="injected NVFP4 support failure"): + apply_recipe(module, NVFP4BlockScaling()) + + # pylint: disable-next=protected-access + assert module._quantization_runtime is active_runtime + assert FP8GlobalStateManager.quantization_state.fp8_recipe is active_recipe_pointer + assert len(calls) == call_count + finally: + FP8GlobalStateManager.reset() + + +def test_equal_recipe_objects_and_missed_updates_reuse_the_runtime(): """Equal independent recipes and missed A -> B -> A updates reuse the runtime.""" calls = [] first_recipe = _make_counting_recipe(("runtime-reuse", 1), calls) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 02246a93813..234384afe1f 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1291,7 +1291,6 @@ def _build_quantization_runtime( backward_quantizers = backward_state.make_quantizers() if _is_delayed_scaling_state(forward_state) != _is_delayed_scaling_state(backward_state): - FP8GlobalStateManager.abort_current_amax_reduction() raise RuntimeError( "This hybrid quantization configuration with delayed scaling is not supported." ) @@ -1356,8 +1355,7 @@ def _runtime_has_delayed_scaling(runtime: _QuantizationRuntime) -> bool: @staticmethod def _reject_delayed_scaling_update() -> None: - """Reject a delayed-state transition and protect any active reduction.""" - FP8GlobalStateManager.abort_current_amax_reduction() + """Reject a delayed-state transition.""" raise RuntimeError( "Mid-training recipe updates do not support delayed scaling. " "A delayed-scaling runtime is frozen after initialization; only " @@ -1384,8 +1382,6 @@ def _check_quantization_update_supported( if getattr(self, "primary_weights_in_fp8", False) and ( active.key.recipe_config != requested_key.recipe_config or active.num_gemms != num_gemms ): - if active_has_delayed_scaling or recipe.delayed(): - FP8GlobalStateManager.abort_current_amax_reduction() raise RuntimeError( "Recipe mismatch for quantized primary weights: mid-training recipe " "configuration changes require per-weight semantic compatibility and " @@ -1897,10 +1893,16 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: original_recipe = meta.get("recipe") recipe = FP8GlobalStateManager.get_fp8_recipe() - runtime_changed = self._ensure_quantization_runtime( - recipe=recipe, - num_gemms=num_gemms, - ) + try: + runtime_changed = self._ensure_quantization_runtime( + recipe=recipe, + num_gemms=num_gemms, + ) + except BaseException: + # Any failure here leaves this module's delayed state unusable for a + # reduction the region may still complete. + FP8GlobalStateManager.abort_current_amax_reduction() + raise if fp8_enabled: meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() if fp8_enabled: diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 0f8a378fd40..a22501b8886 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -654,7 +654,6 @@ def maybe_fuse_ops( recipe_config, ) ): - FP8GlobalStateManager.abort_current_amax_reduction() raise RuntimeError( "Mid-training recipe updates are not supported for fusible operations. " "Recreate the fusible operation or operation pipeline with the new recipe." @@ -756,8 +755,13 @@ def __call__( recipe = FP8GlobalStateManager.get_fp8_recipe() is_grad_enabled = torch.is_grad_enabled() - # Attempt to fuse operations if neccesary - self.maybe_fuse_ops(is_grad_enabled, recipe, input, basic_op_extra_inputs) + # Attempt to fuse operations if neccesary. The recipe transition inside + # can reject before the op reset loop runs, so guard the whole call. + try: + self.maybe_fuse_ops(is_grad_enabled, recipe, input, basic_op_extra_inputs) + except BaseException: + FP8GlobalStateManager.abort_current_amax_reduction() + raise # Initialization before forward for idx, op in enumerate(self._basic_ops): diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 40408b711ea..b7218f696c8 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -393,7 +393,6 @@ def reset_recipe_state( recipe_config, ) ): - FP8GlobalStateManager.abort_current_amax_reduction() raise RuntimeError( "Mid-training recipe updates are not supported for fusible operations. " "Recreate the fusible operation or operation pipeline with the new recipe." diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 9d26dff5d59..1855a20f01f 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -276,14 +276,16 @@ def check_recipe_support(recipe: Recipe) -> None: raise RuntimeError( "DelayedScaling is not supported under torch.compile. Please use other recipes instead." ) - recipe_supported = True - unsupported_reason = "" - if isinstance(recipe, (DelayedScaling, Float8CurrentScaling)): - recipe_supported, unsupported_reason = check_fp8_support() - elif isinstance(recipe, Float8BlockScaling): + if isinstance(recipe, Float8BlockScaling): recipe_supported, unsupported_reason = check_fp8_block_scaling_support() elif isinstance(recipe, MXFP8BlockScaling): recipe_supported, unsupported_reason = check_mxfp8_support() + elif isinstance(recipe, NVFP4BlockScaling): + recipe_supported, unsupported_reason = check_nvfp4_support() + else: + # Covers DelayedScaling, Float8CurrentScaling and CustomRecipe, whose + # quantizers are chosen by a factory this check cannot inspect. + recipe_supported, unsupported_reason = check_fp8_support() if not recipe_supported: raise RuntimeError(unsupported_reason) @@ -578,7 +580,6 @@ def add_fp8_tensors_to_global_buffer( forward_delayed = _is_delayed_scaling_state(fp8_meta[forward_key]) backward_delayed = _is_delayed_scaling_state(fp8_meta[backward_key]) if forward_delayed != backward_delayed: - cls.abort_current_amax_reduction() raise RuntimeError( "This hybrid quantization configuration with delayed scaling is not supported." ) @@ -590,8 +591,8 @@ def add_fp8_tensors_to_global_buffer( if index_in_buffer in fp8_meta: return - fp8_meta[index_in_buffer] = [] qstate = cls.quantization_state + buffer_positions = [] for forward in (True, False): fp8_meta_tensor_key = cls.get_meta_tensor_key(forward=forward) if fp8_meta_tensor_key not in fp8_meta: @@ -632,8 +633,11 @@ def add_fp8_tensors_to_global_buffer( fp8_meta[fp8_meta_tensor_key].amax_history ) qstate.global_scale_buffer[key].append(fp8_meta[fp8_meta_tensor_key].scale) - fp8_meta[index_in_buffer].append(len(qstate.global_amax_buffer[key]) - 1) - fp8_meta[index_in_buffer].append(key) + buffer_positions.append(len(qstate.global_amax_buffer[key]) - 1) + buffer_positions.append(key) + + # Publish last: a module is registered only once the whole walk succeeded. + fp8_meta[index_in_buffer] = buffer_positions @classmethod def is_fp8_enabled(cls) -> bool: @@ -820,31 +824,18 @@ def get_unique_autocast_key( def _prepare_autocast_enter( cls, enabled: bool, + calibrating: bool, fp8_recipe: Optional[Recipe], - fp8_group: Optional[dist_group_type], - ) -> Tuple[Recipe, Hashable, str]: + ) -> Recipe: """Resolve and validate an autocast activation without publishing state.""" fp8_recipe = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe - if enabled: + if enabled or calibrating: check_recipe_support(fp8_recipe) - - quantizer_config = fp8_recipe.quantizer_config() - autocast_key = cls.get_unique_autocast_key(fp8_recipe, fp8_group) - - if enabled: - fp8_available, reason_for_no_fp8 = cls.is_fp8_available() - assert fp8_available, reason_for_no_fp8 - if isinstance(fp8_recipe, MXFP8BlockScaling): - mxfp8_available, reason_for_no_mxfp8 = cls.is_mxfp8_available() - assert mxfp8_available, reason_for_no_mxfp8 - if isinstance(fp8_recipe, Float8BlockScaling): - fp8_block_available, reason_for_no_fp8_block = cls.is_fp8_block_scaling_available() - assert fp8_block_available, reason_for_no_fp8_block - if isinstance(fp8_recipe, NVFP4BlockScaling): - nvfp4_available, reason_for_no_nvfp4 = cls.is_nvfp4_available() - assert nvfp4_available, reason_for_no_nvfp4 - - return fp8_recipe, quantizer_config, autocast_key + if enabled or calibrating or cls.quantization_state.fp8_parameters: + # Building the configuration validates the recipe. A region that + # quantizes nothing must not reject a recipe it will never use. + fp8_recipe.quantizer_config() + return fp8_recipe @classmethod def autocast_enter( @@ -857,24 +848,12 @@ def autocast_enter( ) -> None: """Prepare and publish state for entry into an FP8 region.""" - fp8_recipe, _, autocast_key = cls._prepare_autocast_enter( - enabled, - fp8_recipe, - fp8_group, - ) + fp8_recipe = cls._prepare_autocast_enter(enabled, calibrating, fp8_recipe) qstate = cls.quantization_state # Preparation above contains every operation that can fail due to the # requested recipe or platform. Publish only after it has succeeded. qstate.fp8_recipe = fp8_recipe - # Once a delayed bucket is registered, its committed recipe snapshot - # owns the reduction semantics for this key. Do not replace it with a - # caller-owned recipe object merely by entering another autocast. - qstate.autocast_arguments.setdefault( - autocast_key, - (fp8_recipe, fp8_group), - ) - qstate.fp8_enabled = enabled qstate.fp8_calibration = calibrating qstate.fp8_distributed_group = fp8_group @@ -1170,13 +1149,21 @@ def quantized_model_init( """ qstate = FP8GlobalStateManager.quantization_state + + # Resolve and validate before publishing anything: a recipe this context + # cannot use must leave the process as it found it. + recipe = get_default_fp8_recipe() if recipe is None else recipe + if enabled: + check_recipe_support(recipe) + recipe.quantizer_config() + _fp8_parameters = qstate.fp8_parameters _fp8_recipe = qstate.fp8_recipe _high_precision_init_val = qstate.high_precision_init_val - qstate.fp8_parameters = enabled - FP8GlobalStateManager.activate_recipe(get_default_fp8_recipe() if recipe is None else recipe) - qstate.high_precision_init_val = preserve_high_precision_init_val try: + qstate.fp8_parameters = enabled + qstate.fp8_recipe = recipe + qstate.high_precision_init_val = preserve_high_precision_init_val yield finally: qstate.fp8_parameters = _fp8_parameters @@ -1306,6 +1293,8 @@ def __enter__(self) -> "autocast": return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: + if exc_type is not None: + FP8GlobalStateManager.abort_current_amax_reduction() try: FP8GlobalStateManager.set_autocast_state(self._fp8_state) FP8GlobalStateManager.autocast_exit(self._enabled, _graph=self._graph) From b61e2f4b36a216a2808bdf331b21a733ee08e6c9 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 11 Sep 2026 12:07:50 +0000 Subject: [PATCH 90/97] Freeze the built-in attention recipe path Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 311 ++++++++++++++++-- transformer_engine/common/recipe/__init__.py | 6 + .../dot_product_attention.py | 161 +++++---- .../attention/dot_product_attention/utils.py | 80 +++-- 4 files changed, 439 insertions(+), 119 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index f719158ef2e..e83136e3a80 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -171,7 +171,7 @@ def first_factory(role): first_runtime = dpa._quantization_runtime first_forward_quantizers = dpa.quantizers["scaling_fwd"] first_backward_quantizers = dpa.quantizers["scaling_bwd"] - first_capabilities_quantizer = dpa._qkv_capabilities_quantizer + first_qkv_quantizer = _canonical_qkv_quantizer(dpa) assert dpa.get_qkv_quantization_capabilities() == first_capabilities assert dpa._quantization_runtime is first_runtime @@ -200,13 +200,255 @@ def second_factory(role): assert dpa.quantizers["scaling_fwd"] is second_runtime.forward_quantizers assert dpa.quantizers["scaling_bwd"] is second_runtime.backward_quantizers assert second_capabilities == first_capabilities == (False, True) - assert dpa._qkv_capabilities_quantizer is not first_capabilities_quantizer + assert _canonical_qkv_quantizer(dpa) is not first_qkv_quantizer assert any( - dpa._qkv_capabilities_quantizer is quantizer + _canonical_qkv_quantizer(dpa) is quantizer for quantizer in second_runtime.forward_quantizers ) assert not dpa._fp8_workspaces + def test_custom_to_builtin_crossing_is_rejected(self): + """The built-in path never takes over a module a CustomRecipe runtime owns.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + for builtin in ( + recipe.MXFP8BlockScaling(fp8_dpa=True), # stateless + recipe.DelayedScaling(fp8_dpa=True), # owns reduction buckets + ): + dpa = self._make_dpa() + custom = recipe.CustomRecipe( + qfactory=mxfp8_factory, + qfactory_key=("dpa-crossing", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=custom): + dpa.get_qkv_quantization_capabilities() + runtime = dpa._quantization_runtime + assert runtime is not None + old_views = ( + dpa.fp8_meta["recipe"], + dpa.fp8_meta["scaling_fwd"], + dpa.quantizers["scaling_fwd"], + ) + + with autocast(enabled=True, recipe=builtin): + with pytest.raises( + RuntimeError, match="built-in DotProductAttention recipe path is frozen" + ): + dpa.init_fp8_metadata(num_gemms=3) + + assert dpa._quantization_runtime is runtime + assert ( + dpa.fp8_meta["recipe"], + dpa.fp8_meta["scaling_fwd"], + dpa.quantizers["scaling_fwd"], + ) == old_views + + def test_unlabelled_custom_runtime_still_blocks_the_builtin_path(self): + """A factory TE cannot label leaves no local recipes for the change guard.""" + + def prebuilt_factory(role): + del role + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + + dpa = self._make_dpa() + custom = recipe.CustomRecipe( + qfactory=prebuilt_factory, + qfactory_key=("dpa-unlabelled-crossing", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=custom): + dpa.get_qkv_quantization_capabilities() + runtime = dpa._quantization_runtime + assert runtime is not None + assert "local_recipes" not in dpa.fp8_meta + + with autocast(enabled=True, recipe=recipe.DelayedScaling(fp8_dpa=True)): + with pytest.raises( + RuntimeError, match="built-in DotProductAttention recipe path is frozen" + ): + dpa.init_fp8_metadata(num_gemms=3) + assert dpa._quantization_runtime is runtime + + def test_builtin_to_custom_crossing_is_rejected_before_mutation(self): + """An initialized built-in DPA cannot be retargeted by a CustomRecipe.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + for builtin in ( + recipe.MXFP8BlockScaling(fp8_dpa=True), # stateless + recipe.Float8CurrentScaling(fp8_dpa=True), # delayed S/dP half + recipe.DelayedScaling(fp8_dpa=True), + ): + dpa = self._make_dpa() + with autocast(enabled=True, recipe=builtin): + dpa.init_fp8_metadata(num_gemms=3) + assert dpa._quantization_runtime is None + old_state = dpa.fp8_meta["scaling_fwd"] + + calls = [] + + def counting_factory(role): + calls.append(role) + return mxfp8_factory(role) + + custom = recipe.CustomRecipe( + qfactory=counting_factory, + qfactory_key=("dpa-builtin-crossing", 1), + fp8_dpa=True, + ) + + # Both the lazy forward and the model-wide planner reject before the + # factory is invoked or a candidate is built. + with autocast(enabled=True, recipe=custom): + with pytest.raises( + RuntimeError, match="built-in DotProductAttention recipe path is frozen" + ): + dpa.get_qkv_quantization_capabilities() + with pytest.raises( + RuntimeError, match="built-in DotProductAttention recipe path is frozen" + ): + te.apply_recipe(dpa, custom) + + assert not calls + assert dpa._quantization_runtime is None + assert dpa.fp8_meta["scaling_fwd"] is old_state + + def test_builtin_dpa_recipe_change_is_rejected(self): + """Any built-in recipe change is rejected instead of silently ignored.""" + # CurrentScaling: the recipe object compared on re-entry is the fake + # DelayedScaling rebuilt from env values, so this flip is invisible to it. + changed = recipe.Float8CurrentScaling(fp8_dpa=True) + changed.fp8_quant_fwd_inp = recipe.QParams(power_2_scale=True) + cases = [ + (recipe.Float8CurrentScaling(fp8_dpa=True), changed), + # DelayedScaling reaches set_meta_tensor, whose RecipeState-type early + # return dropped the change. + ( + recipe.DelayedScaling(fp8_dpa=True, margin=0), + recipe.DelayedScaling(fp8_dpa=True, margin=1), + ), + ( + recipe.DelayedScaling(fp8_dpa=True, amax_history_len=4), + recipe.DelayedScaling(fp8_dpa=True, amax_history_len=8), + ), + ] + for initial, updated in cases: + dpa = self._make_dpa() + with autocast(enabled=True, recipe=initial): + dpa.init_fp8_metadata(num_gemms=3) + committed = dpa.fp8_meta["scaling_fwd"] + + with autocast(enabled=True, recipe=updated): + with pytest.raises( + RuntimeError, match="built-in DotProductAttention recipe path is frozen" + ): + dpa.init_fp8_metadata(num_gemms=3) + assert dpa.fp8_meta["scaling_fwd"] is committed + + # An unchanged recipe still re-enters freely. + dpa = self._make_dpa() + for _ in range(3): + with autocast(enabled=True, recipe=recipe.DelayedScaling(fp8_dpa=True)): + dpa.init_fp8_metadata(num_gemms=3) + + def test_failed_dpa_candidate_does_not_mutate_shared_quantizer(self): + """Validating a rejected candidate must not write to an aliased live quantizer.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + dpa = self._make_dpa() + active = recipe.CustomRecipe( + qfactory=mxfp8_factory, + qfactory_key=("dpa-aliased-active", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=active): + dpa.get_qkv_quantization_capabilities() + + live = _canonical_qkv_quantizer(dpa) + # Give the live quantizer a state the preparing wrapper would overwrite: + # it sets internal=False on the QKV slot before any slot is type-checked. + live.internal = True + before = (live.internal, live.rowwise_usage, live.columnwise_usage) + + # A factory that hands back the live quantizer for the QKV slot while + # making the O slot invalid, so validation rejects the candidate. + def aliasing_factory(role): + if role is not None and role.name.endswith("dpa_output"): + return IdentityQuantizer() + return live + + invalid = recipe.CustomRecipe( + qfactory=aliasing_factory, + qfactory_key=("dpa-aliased-invalid", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=invalid): + with pytest.raises((TypeError, RuntimeError)): + dpa.get_qkv_quantization_capabilities() + + assert (live.internal, live.rowwise_usage, live.columnwise_usage) == before + assert _canonical_qkv_quantizer(dpa) is live + + def test_builtin_dpa_same_class_change_is_rejected(self): + """A built-in recipe change that would otherwise be silently ignored must raise.""" + # CurrentScaling: the recipe object compared on re-entry is the fake + # DelayedScaling rebuilt from env values, so this flip is invisible to it. + dpa = self._make_dpa() + with autocast(enabled=True, recipe=recipe.Float8CurrentScaling(fp8_dpa=True)): + dpa.init_fp8_metadata(num_gemms=3) + changed = recipe.Float8CurrentScaling(fp8_dpa=True) + changed.fp8_quant_fwd_inp = recipe.QParams(power_2_scale=True) + with autocast(enabled=True, recipe=changed): + with pytest.raises( + RuntimeError, match="built-in DotProductAttention recipe path is frozen" + ): + dpa.init_fp8_metadata(num_gemms=3) + + # DelayedScaling: reaches set_meta_tensor, whose RecipeState-type early + # return dropped the change. + delayed_dpa = self._make_dpa() + with autocast(enabled=True, recipe=recipe.DelayedScaling(fp8_dpa=True, margin=0)): + delayed_dpa.init_fp8_metadata(num_gemms=3) + with autocast(enabled=True, recipe=recipe.DelayedScaling(fp8_dpa=True, margin=1)): + with pytest.raises( + RuntimeError, match="built-in DotProductAttention recipe path is frozen" + ): + delayed_dpa.init_fp8_metadata(num_gemms=3) + + def test_qkv_capabilities_fall_back_to_the_quantizer_when_unlabelled(self): + """A factory returning pre-built quantizers leaves no local recipes to read.""" + + def prebuilt_factory(role): + del role + return Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + + dpa = self._make_dpa() + custom = recipe.CustomRecipe( + qfactory=prebuilt_factory, + qfactory_key=("dpa-prebuilt-float8", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=custom): + capabilities = dpa.get_qkv_quantization_capabilities() + + assert "local_recipes" not in dpa.fp8_meta + assert capabilities == (False, False) + def test_direct_runtime_commit_updates_dpa_derived_state(self): """DPA cache and recipe labels are part of its runtime commit hook.""" if not mxfp8_available: @@ -224,8 +466,7 @@ def test_direct_runtime_commit_updates_dpa_derived_state(self): old_runtime = dpa._quantization_runtime old_local_recipes = object() dpa.fp8_meta["local_recipes"] = old_local_recipes - old_cache_key = dpa._custom_dpa_local_recipes_cache_key - old_capabilities_quantizer = dpa._qkv_capabilities_quantizer + old_qkv_quantizer = _canonical_qkv_quantizer(dpa) second_recipe = recipe.CustomRecipe( qfactory=mxfp8_factory, @@ -239,17 +480,13 @@ def test_direct_runtime_commit_updates_dpa_derived_state(self): assert dpa._quantization_runtime is old_runtime assert dpa.fp8_meta["local_recipes"] is old_local_recipes - assert dpa._custom_dpa_local_recipes_cache_key is old_cache_key - assert dpa._qkv_capabilities_quantizer is old_capabilities_quantizer + assert _canonical_qkv_quantizer(dpa) is old_qkv_quantizer assert dpa._apply_quantization_update(update) assert dpa._quantization_runtime is update.candidate assert dpa.fp8_meta["local_recipes"] is update.validation_result - assert dpa._custom_dpa_local_recipes_cache is update.validation_result assert [type(item).__name__ for item in update.validation_result] == ["MXFP8BlockScaling"] - assert dpa._custom_dpa_local_recipes_cache_key != old_cache_key - assert dpa._qkv_capabilities_quantizer is None - assert dpa._qkv_capabilities_cache is None + assert _canonical_qkv_quantizer(dpa) is not old_qkv_quantizer def test_candidate_failure_preserves_active_runtime_and_caches(self): """A failed candidate validation does not partially update DPA.""" @@ -272,8 +509,7 @@ def test_candidate_failure_preserves_active_runtime_and_caches(self): dpa.fp8_meta["scaling_bwd"], dpa.quantizers["scaling_fwd"], dpa.quantizers["scaling_bwd"], - dpa._qkv_capabilities_quantizer, - dpa._qkv_capabilities_cache, + _canonical_qkv_quantizer(dpa), ) workspace = object() dpa._fp8_workspaces["old"] = workspace @@ -300,8 +536,7 @@ def test_candidate_failure_preserves_active_runtime_and_caches(self): dpa.fp8_meta["scaling_bwd"], dpa.quantizers["scaling_fwd"], dpa.quantizers["scaling_bwd"], - dpa._qkv_capabilities_quantizer, - dpa._qkv_capabilities_cache, + _canonical_qkv_quantizer(dpa), ), old_views, ) @@ -354,9 +589,8 @@ def test_all_current_scaling_update_is_rejected_before_commit(self): dpa.fp8_meta["scaling_bwd"], dpa.quantizers["scaling_fwd"], dpa.quantizers["scaling_bwd"], - dpa._custom_dpa_local_recipes_cache, - dpa._qkv_capabilities_quantizer, - dpa._qkv_capabilities_cache, + dpa.fp8_meta.get("local_recipes"), + _canonical_qkv_quantizer(dpa), ) invalid_recipe = recipe.CustomRecipe( qfactory=current_scaling_factory, @@ -380,9 +614,8 @@ def test_all_current_scaling_update_is_rejected_before_commit(self): dpa.fp8_meta["scaling_bwd"], dpa.quantizers["scaling_fwd"], dpa.quantizers["scaling_bwd"], - dpa._custom_dpa_local_recipes_cache, - dpa._qkv_capabilities_quantizer, - dpa._qkv_capabilities_cache, + dpa.fp8_meta.get("local_recipes"), + _canonical_qkv_quantizer(dpa), ), old_views, ) @@ -392,6 +625,34 @@ def test_all_current_scaling_update_is_rejected_before_commit(self): assert dpa._quantization_runtime is old_runtime +def test_attention_params_distinguish_local_recipe_families(): + """Backend selection keys on the DPA-local families, not only fp8_meta['recipe'].""" + from transformer_engine.pytorch.attention.dot_product_attention.utils import AttentionParams + + shared = recipe.DelayedScaling() + + def params(local_recipes): + return AttentionParams( + fp8=True, + fp8_meta={"recipe": shared, "local_recipes": local_recipes}, + ) + + current_scaling = params([recipe.Float8CurrentScaling()]) + mxfp8 = params([recipe.MXFP8BlockScaling()]) + equivalent = params([recipe.Float8CurrentScaling()]) + + assert current_scaling != mxfp8 + assert current_scaling == equivalent + assert current_scaling != params(None) + + +def _canonical_qkv_quantizer(dpa): + """The slot the capability query reads, formerly mirrored in a private cache.""" + from transformer_engine.pytorch.attention.dot_product_attention.utils import META_QKV + + return dpa.quantizers["scaling_fwd"][META_QKV] + + def test_hybrid_storage_snapshots_parent_quantizer(): """Caller mutations must not change an existing Hybrid tensor's behavior.""" quantizer = HybridQuantizer( @@ -2439,13 +2700,13 @@ def counting_qfactory(role): with autocast(enabled=True, recipe=custom_recipe): first = dpa.get_qkv_quantization_capabilities() - canonical_qkv = dpa._qkv_capabilities_quantizer + canonical_qkv = _canonical_qkv_quantizer(dpa) calls_after_first = len(calls) second = dpa.get_qkv_quantization_capabilities() assert first == expected assert second == first - assert dpa._qkv_capabilities_quantizer is canonical_qkv + assert _canonical_qkv_quantizer(dpa) is canonical_qkv assert len(calls) == calls_after_first # A supported recipe-state rebuild creates a new canonical slot and @@ -2464,13 +2725,13 @@ def rebuilt_qfactory(role): with autocast(enabled=True, recipe=rebuilt_recipe): with pytest.raises(RuntimeError, match="do not support delayed scaling"): dpa.get_qkv_quantization_capabilities() - assert dpa._qkv_capabilities_quantizer is canonical_qkv + assert _canonical_qkv_quantizer(dpa) is canonical_qkv assert len(calls) == calls_after_first return with autocast(enabled=True, recipe=rebuilt_recipe): rebuilt = dpa.get_qkv_quantization_capabilities() - rebuilt_qkv = dpa._qkv_capabilities_quantizer + rebuilt_qkv = _canonical_qkv_quantizer(dpa) assert rebuilt == first assert rebuilt_qkv is not canonical_qkv diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index c3e694d1460..e4f4385955d 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -1110,6 +1110,12 @@ class CustomRecipe(Recipe): ``IdentityQuantizer`` for an intentional high-precision slot instead of returning ``None``. + Prefer returning a fresh quantizer per call. Transformer Engine does + not write to a candidate's quantizers before that candidate is + validated and committed, so returning shared instances is supported, + but a quantizer aliased across roles carries any usage or ``internal`` + flags a committed runtime sets on it. + ``QuantizerRole`` is a frozen dataclass with the following fields: - ``module_type`` (str): module type (empty string when not set), e.g. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index c77debb9bc2..647e6da4222 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -32,6 +32,7 @@ Float8CurrentScalingRecipeState, Float8BlockScalingRecipeState, _QuantizationRuntime, + _QuantizationRuntimeKey, ) from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -890,14 +891,6 @@ def __init__( _declared_grad_input_quantizer_role=_declared_grad_input_quantizer_role, ) - # Cache the native recipe labels inferred from custom DPA quantizers. - # ``init_fp8_metadata`` runs on every forward, while the quantizers only - # change when their recipe state is rebuilt. - self._custom_dpa_local_recipes_cache_key: Optional[Tuple[Any, ...]] = None - self._custom_dpa_local_recipes_cache: Optional[List[Recipe]] = None - self._qkv_capabilities_quantizer: Optional[Any] = None - self._qkv_capabilities_cache: Optional[Tuple[bool, bool]] = None - self.logger = logging.getLogger("DotProductAttention") self.logger.setLevel(attn_log._log_level) if not self.logger.hasHandlers(): @@ -1141,6 +1134,22 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: super().init_fp8_metadata(num_gemms=num_gemms) return + try: + self._init_builtin_fp8_metadata(num_gemms, fp8_recipe, _original_recipe) + except BaseException: + # The built-in path bypasses the base init_fp8_metadata abort wrapper, + # so its rejections must still protect a reduction this region may + # complete with delayed tensors this module can no longer update. + FP8GlobalStateManager.abort_current_amax_reduction() + raise + + def _init_builtin_fp8_metadata( + self, + num_gemms: int, + fp8_recipe: Recipe, + _original_recipe: Optional[Recipe], + ) -> None: + """Resolve and initialize metadata for the built-in ``NVTE_DPA_*`` path.""" # switch/append recipe: fp8_recipe stays unchanged, but DPA.fp8_meta["recipe"] may be set to # a different recipe than fp8_recipe. DPA.quantizers may be a mix of different quantizers as well. # @@ -1152,7 +1161,10 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: fp8_recipe_dpa = fp8_recipe fp8_recipes = fp8_recipe if _dpa_fp8_recipe == "F16": - # ignore the recipe from autocast, set fp8_dpa = False, fp8_mha = False + # ignore the recipe from autocast, set fp8_dpa = False, fp8_mha = False. + # Keep these two assignments adjacent and in this order: the recipe is + # invalid in between (fp8_mha=True requires fp8_dpa=True), so anything + # that reads its quantizer configuration here would raise. fp8_recipe.fp8_dpa = False fp8_recipe.fp8_mha = False elif ( @@ -1279,26 +1291,32 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: self.fast_setattr("fp8_calibration", FP8GlobalStateManager.is_fp8_calibration()) fp8_enabled = self.fp8 or self.fp8_calibration self.fp8_meta["fp8_checkpoint"] = self.fp8 or self.fp8_calibration + local_recipes = fp8_recipes if isinstance(fp8_recipes, List) else [fp8_recipes] + previous_local_recipes = self.fp8_meta.get("local_recipes") if self.fp8_parameters or fp8_enabled: self.fp8_meta["global_recipe"] = fp8_recipe - self.fp8_meta["local_recipes"] = ( - fp8_recipes if isinstance(fp8_recipes, List) else [fp8_recipes] - ) + self.fp8_meta["local_recipes"] = local_recipes if self.fp8_parameters or fp8_enabled: + # Taking ownership here would strand a CustomRecipe runtime still + # pointing at this module's state. The two paths do not hand state + # back and forth; with quantization off nothing is published, so this + # only guards the case where the built-in path would actually write. + if getattr(self, "_quantization_runtime", None) is not None: + self._reject_builtin_dpa_update() + # ``fp8_recipe_dpa`` below is often a fake rebuilt from env values, so it + # hides a change in the recipes that actually build the quantizers. + if ( + self.fp8_initialized + and previous_local_recipes + and previous_local_recipes != local_recipes + ): + self._reject_builtin_dpa_update() if self.fp8_initialized and fp8_recipe_dpa == self.fp8_meta["recipe"]: # FP8 init has already been run and recipe is the same, don't do anything. return self.fp8_meta["recipe"] = fp8_recipe_dpa - if fp8_recipe != fp8_recipe_dpa: - # fp8_recipe has changed, rehash the key. - autocast_key = FP8GlobalStateManager.get_unique_autocast_key( - fp8_recipe_dpa, fp8_group - ) - FP8GlobalStateManager.quantization_state.autocast_arguments[autocast_key] = ( - fp8_recipe_dpa, - fp8_group, - ) + self._rehash_autocast_arguments(fp8_recipe, fp8_recipe_dpa, fp8_group) else: # If fp8 isn't enabled, turn off and return. self.fast_setattr("fp8_initialized", False) @@ -1322,15 +1340,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: self.fast_setattr("fp8_initialized", True) self.fp8_meta["recipe"] = fp8_recipe_dpa - if fp8_recipe != fp8_recipe_dpa: - # fp8_recipe has changed, rehash the key. - autocast_key = FP8GlobalStateManager.get_unique_autocast_key( - fp8_recipe_dpa, fp8_group - ) - FP8GlobalStateManager.quantization_state.autocast_arguments[autocast_key] = ( - fp8_recipe_dpa, - fp8_group, - ) + self._rehash_autocast_arguments(fp8_recipe, fp8_recipe_dpa, fp8_group) _current_recipe = self.fp8_meta["recipe"] if _original_recipe is not None and not ( @@ -1345,12 +1355,49 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # Clear cached workspaces as they were created with the old recipe/quantizer type self._fp8_workspaces.clear() - def _invalidate_dpa_runtime_caches(self) -> None: - """Invalidate DPA caches after a committed quantizer replacement.""" - self._custom_dpa_local_recipes_cache_key = None - self._custom_dpa_local_recipes_cache = None - self._qkv_capabilities_quantizer = None - self._qkv_capabilities_cache = None + @staticmethod + def _reject_builtin_dpa_update() -> None: + """Reject a recipe change involving the built-in ``NVTE_DPA_*`` path.""" + raise RuntimeError( + "The built-in DotProductAttention recipe path is frozen after " + "initialization and does not support mid-training recipe updates. " + "Recreate the module, or use a CustomRecipe throughout." + ) + + @staticmethod + def _rehash_autocast_arguments(fp8_recipe, fp8_recipe_dpa, fp8_group) -> None: + """Re-register reduction arguments when DPA resolved its own recipe.""" + if fp8_recipe == fp8_recipe_dpa: + return + autocast_key = FP8GlobalStateManager.get_unique_autocast_key(fp8_recipe_dpa, fp8_group) + FP8GlobalStateManager.quantization_state.autocast_arguments[autocast_key] = ( + fp8_recipe_dpa, + fp8_group, + ) + + def _check_quantization_update_supported( + self, + *, + recipe: Recipe, + requested_key: _QuantizationRuntimeKey, + num_gemms: int, + ) -> None: + """Also reject a CustomRecipe update onto an initialized built-in DPA.""" + if getattr(self, "_quantization_runtime", None) is None and self.fp8_initialized: + # The built-in path owns this module's state, and its delayed half is + # registered in the global reduction buckets. Running before super() + # means both the lazy forward and apply_recipe planning reject here, + # before any candidate is built. + self._reject_builtin_dpa_update() + super()._check_quantization_update_supported( + recipe=recipe, + requested_key=requested_key, + num_gemms=num_gemms, + ) + + @staticmethod + def _invalidate_dpa_runtime_caches() -> None: + """Force backend reselection after a committed quantizer replacement.""" _attention_backends["backend_selection_requires_update"] = True def _activate_quantization_runtime( @@ -1366,16 +1413,7 @@ def _activate_quantization_runtime( ) self._invalidate_dpa_runtime_caches() - cache_key = ( - id(self.fp8_meta.get("scaling_fwd")), - tuple(id(quantizer) for quantizer in candidate.forward_quantizers), - candidate.recipe.fp8_format, - candidate.recipe.fp8_dpa, - candidate.recipe.fp8_mha, - ) local_recipes = validation_result - self._custom_dpa_local_recipes_cache_key = cache_key - self._custom_dpa_local_recipes_cache = local_recipes if local_recipes is None: # Do not leave labels from an earlier supported quantizer family # attached after a rebuild to an unsupported family. @@ -1394,8 +1432,16 @@ def _validate_quantization_runtime( raise ValueError( "FP8 DotProductAttention requires 9 forward and 6 backward quantizer slots." ) - qkv_quantizer, _, s_quantizer, _, _, dp_quantizer = dpa_utils.get_attention_quantizers( - True, + # Selection only: a qfactory may alias quantizers the active runtime + # still uses, so validating a candidate must not write to them. + ( + qkv_quantizer, + _, + s_quantizer, + _, + _, + dp_quantizer, + ) = dpa_utils._select_attention_quantizers( { "scaling_fwd": candidate.forward_quantizers, "scaling_bwd": candidate.backward_quantizers, @@ -1447,10 +1493,6 @@ def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: f"DotProductAttention did not materialize the canonical QKV quantizer for {role}." ) from exc - if qkv_quantizer is self._qkv_capabilities_quantizer: - assert self._qkv_capabilities_cache is not None - return self._qkv_capabilities_cache - from transformer_engine.pytorch.tensor.float8_tensor import ( Float8CurrentScalingQuantizer, Float8Quantizer, @@ -1458,18 +1500,11 @@ def get_qkv_quantization_capabilities(self) -> Tuple[bool, bool]: from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): - capabilities = (True, False) - elif isinstance(qkv_quantizer, MXFP8Quantizer): - capabilities = (False, True) - elif isinstance(qkv_quantizer, Float8Quantizer): - capabilities = (False, False) - else: - capabilities = None - - if capabilities is not None: - self._qkv_capabilities_quantizer = qkv_quantizer - self._qkv_capabilities_cache = capabilities - return capabilities + return (True, False) + if isinstance(qkv_quantizer, MXFP8Quantizer): + return (False, True) + if isinstance(qkv_quantizer, Float8Quantizer): + return (False, False) role = QuantizerRole( module_type="dpa", diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 71265958402..1ac1c1f50b1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -194,6 +194,11 @@ def set_flash_attention_4_params(): FlashAttentionUtils.v4_is_installed = True +def _local_recipe_families(fp8_meta): + """Quantizer families the DPA resolved locally, used for backend-cache equality.""" + return tuple(type(recipe) for recipe in (fp8_meta or {}).get("local_recipes") or ()) + + @dataclass(eq=True) class AttentionParams: """ @@ -326,7 +331,13 @@ def __eq__(self, other): if fname != "fp8_meta": if sf != of: return False - elif (sf or {}).get("recipe", None) != (of or {}).get("recipe", None): + continue + if (sf or {}).get("recipe", None) != (of or {}).get("recipe", None): + return False + # Backend selection reads fp8_meta["local_recipes"], so two DPAs under + # one role-sensitive CustomRecipe must not share a backend decision + # when their resolved quantizer families differ. + if _local_recipe_families(sf) != _local_recipe_families(of): return False return True @@ -2720,54 +2731,33 @@ def check_set_window_size( return window_size -def get_attention_quantizers(fp8, quantizers): - """Get the list of quantizers used in attention from the quantizers list.""" - if not fp8: - return [None] * 6 +def _select_attention_quantizers(quantizers): + """Select the six DPA slot quantizers and check their families. + Pure: this inspects ``quantizers`` and writes no attributes, so validating a + candidate runtime cannot mutate a quantizer the active runtime still uses + (a ``qfactory`` is free to return aliased instances). + """ QKV_quantizer = quantizers["scaling_fwd"][META_QKV] - QKV_quantizer.internal = False - QKV_quantizer.set_usage(rowwise=True, columnwise=False) - S_quantizer = quantizers["scaling_fwd"][META_S] - S_quantizer.internal = True - S_quantizer.set_usage(rowwise=True, columnwise=False) - O_quantizer = quantizers["scaling_fwd"][META_O] - O_quantizer.internal = False - O_quantizer.set_usage(rowwise=True, columnwise=False) - dO_quantizer = quantizers["scaling_bwd"][META_DO] - dO_quantizer.internal = False - dO_quantizer.set_usage(rowwise=True, columnwise=False) - dP_quantizer = quantizers["scaling_bwd"][META_DP] - dP_quantizer.internal = True - dP_quantizer.set_usage(rowwise=True, columnwise=False) - dQKV_quantizer = quantizers["scaling_bwd"][META_DQKV] - dQKV_quantizer.internal = False - dQKV_quantizer.set_usage(rowwise=True, columnwise=False) # MXFP8 attention: detect from the QKV quantizer instance rather than the # recipe predicate so that CustomRecipe (whose `mxfp8()` predicate returns # False) gets the same treatment as the built-in MXFP8 recipe. The kernel # handles S/dP internally for MXFP8, hence S/dP are nulled out. - if isinstance(QKV_quantizer, MXFP8Quantizer): - QKV_quantizer.columnwise_usage = True - QKV_quantizer.optimize_for_gemm = True + is_mxfp8 = isinstance(QKV_quantizer, MXFP8Quantizer) + if is_mxfp8: S_quantizer = None - O_quantizer.columnwise_usage = True - - dO_quantizer.columnwise_usage = True - dO_quantizer.optimize_for_gemm = True dP_quantizer = None - dQKV_quantizer.columnwise_usage = True _fp8_types = (Float8Quantizer, Float8CurrentScalingQuantizer, MXFP8Quantizer) # S/dP are intentionally None under MXFP8 attention; skip the type check # for those slots in that case. - _allow_none = {"S", "dP"} if isinstance(QKV_quantizer, MXFP8Quantizer) else set() + _allow_none = {"S", "dP"} if is_mxfp8 else set() for _name, _q in [ ("QKV", QKV_quantizer), ("O", O_quantizer), @@ -2791,6 +2781,34 @@ def get_attention_quantizers(fp8, quantizers): return QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer +def get_attention_quantizers(fp8, quantizers): + """Get the list of quantizers used in attention, prepared for the fused kernel.""" + if not fp8: + return [None] * 6 + + for _key, _index, _internal in ( + ("scaling_fwd", META_QKV, False), + ("scaling_fwd", META_S, True), + ("scaling_fwd", META_O, False), + ("scaling_bwd", META_DO, False), + ("scaling_bwd", META_DP, True), + ("scaling_bwd", META_DQKV, False), + ): + quantizer = quantizers[_key][_index] + quantizer.internal = _internal + quantizer.set_usage(rowwise=True, columnwise=False) + + if isinstance(quantizers["scaling_fwd"][META_QKV], MXFP8Quantizer): + quantizers["scaling_fwd"][META_QKV].columnwise_usage = True + quantizers["scaling_fwd"][META_QKV].optimize_for_gemm = True + quantizers["scaling_fwd"][META_O].columnwise_usage = True + quantizers["scaling_bwd"][META_DO].columnwise_usage = True + quantizers["scaling_bwd"][META_DO].optimize_for_gemm = True + quantizers["scaling_bwd"][META_DQKV].columnwise_usage = True + + return _select_attention_quantizers(quantizers) + + def print_quantizers( label, layer_number, From f3879936efbb20f0c95bb87f940c4d1a30cf6fa8 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 11 Sep 2026 12:46:18 +0000 Subject: [PATCH 91/97] Dispatch apply_recipe on an owner protocol Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 27 ++- tests/pytorch/test_recipe_update.py | 195 ++++++++++++++++-- .../dot_product_attention.py | 42 +++- transformer_engine/pytorch/module/base.py | 23 +++ transformer_engine/pytorch/ops/fuser.py | 4 + transformer_engine/pytorch/ops/op.py | 34 +++ transformer_engine/pytorch/quantization.py | 71 ++----- 7 files changed, 324 insertions(+), 72 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index e83136e3a80..2abe632add9 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -283,7 +283,6 @@ def test_builtin_to_custom_crossing_is_rejected_before_mutation(self): pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") for builtin in ( - recipe.MXFP8BlockScaling(fp8_dpa=True), # stateless recipe.Float8CurrentScaling(fp8_dpa=True), # delayed S/dP half recipe.DelayedScaling(fp8_dpa=True), ): @@ -321,6 +320,29 @@ def counting_factory(role): assert dpa._quantization_runtime is None assert dpa.fp8_meta["scaling_fwd"] is old_state + def test_stateless_builtin_dpa_accepts_a_custom_recipe(self): + """D5/H2: an inert or stateless built-in DPA must still accept a recipe.""" + if not mxfp8_available: + pytest.skip(f"MXFP8: {reason_for_no_mxfp8}") + + for builtin in ( + recipe.Float8CurrentScaling(), # quantization-inert DPA (fp8_dpa=False) + recipe.MXFP8BlockScaling(fp8_dpa=True), # stateless attention quantizers + ): + dpa = self._make_dpa() + with autocast(enabled=True, recipe=builtin): + dpa.init_fp8_metadata(num_gemms=3) + assert dpa._quantization_runtime is None + + custom = recipe.CustomRecipe( + qfactory=mxfp8_factory, + qfactory_key=("dpa-stateless-migration", 1), + fp8_dpa=True, + ) + with autocast(enabled=True, recipe=custom): + dpa.get_qkv_quantization_capabilities() + assert dpa._quantization_runtime is not None + def test_builtin_dpa_recipe_change_is_rejected(self): """Any built-in recipe change is rejected instead of silently ignored.""" # CurrentScaling: the recipe object compared on re-entry is the fake @@ -598,8 +620,9 @@ def test_all_current_scaling_update_is_rejected_before_commit(self): fp8_dpa=True, ) + # apply_recipe annotates the owner's exception; it does not rewrite its type. with pytest.raises( - RuntimeError, + TypeError, match="Float8CurrentScaling DPA requires delayed scaling for S and dP", ): te.apply_recipe(dpa, invalid_recipe) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 343cb425ffe..767b140865d 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -15,6 +15,7 @@ Float8CurrentScaling, MXFP8BlockScaling, NVFP4BlockScaling, + QParams, ) from transformer_engine.pytorch import ( DotProductAttention, @@ -32,7 +33,10 @@ is_nvfp4_available, quantized_model_init, ) -from transformer_engine.pytorch.custom_recipes.quantizer_factories import current_scaling_factory +from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( + current_scaling_factory, + delayed_scaling_factory, +) from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.quantization import ( DelayedScalingRequest, @@ -1793,7 +1797,8 @@ def test_apply_recipe_rejects_e2e_mha_reproducer_during_planning(): old_global_state = _global_recipe_state() try: - with pytest.raises(RuntimeError, match="O quantizer is NVFP4Quantizer"): + # The owner raises TypeError; apply_recipe annotates but does not rewrite it. + with pytest.raises(TypeError, match="O quantizer is NVFP4Quantizer"): apply_recipe(module, recipe) assert _global_recipe_state() == old_global_state @@ -2451,11 +2456,14 @@ def failing_factory(role): qfactory=failing_factory, qfactory_key=("apply-failure", 2, failure_index), ) - with pytest.raises( - RuntimeError, - match=rf"planning module '{failure_index}': model-wide factory failure", - ): + with pytest.raises(RuntimeError, match="model-wide factory failure") as failure: apply_recipe(model, replacement_recipe) + # The owner's own exception type survives; the module is named in a note. + assert type(failure.value) is RuntimeError + assert any( + f"while planning module '{failure_index}'" in note + for note in getattr(failure.value, "__notes__", ()) + ) assert _global_recipe_state() == old_global_state for module, expected_views, workspace in zip(modules, old_views, workspaces): @@ -2500,15 +2508,23 @@ def test_apply_recipe_deduplicates_shared_runtime_owner(): FP8GlobalStateManager.reset() -def test_apply_recipe_rejects_fusible_owner_before_factory(): - """Excluded fusible owners are discovered before any participant planning.""" +def test_apply_recipe_rejects_fusible_owner_without_committing_anything(): + """A fusible owner rejects the model, and nothing is committed on the way there. + + Owners are planned as they are discovered, so a participant ahead of the + fusible owner may have its factory invoked. Planning publishes no state, so + the documented guarantee -- no module and no global recipe changes -- holds. + """ FP8GlobalStateManager.reset() - def unexpected_factory(_role): - raise AssertionError("excluded model invoked qfactory") + calls = [] + + def counting_factory(role): + calls.append(role) + return IdentityQuantizer() recipe = CustomRecipe( - qfactory=unexpected_factory, + qfactory=counting_factory, qfactory_key=("apply-fusible", 1), ) model = torch.nn.ModuleList( @@ -2528,6 +2544,41 @@ def unexpected_factory(_role): FP8GlobalStateManager.reset() +def test_apply_recipe_skips_zero_quantizer_fusible_owner(): + """An operation that builds no quantizers is inert, not an excluded owner.""" + FP8GlobalStateManager.reset() + + calls = [] + + def counting_factory(role): + calls.append(role) + return IdentityQuantizer() + + recipe = CustomRecipe( + qfactory=counting_factory, + qfactory_key=("apply-zero-quantizer", 1), + ) + model = torch.nn.ModuleList( + [ + Linear(16, 16, bias=False, device="cuda", name="linear"), + te_ops.LayerNorm(16, device="cuda"), + ] + ) + try: + apply_recipe(model, recipe) + assert model[0]._quantization_runtime is not None + assert FP8GlobalStateManager.quantization_state.fp8_recipe is recipe + + # A fusible owner that does build quantizers still rejects. + with pytest.raises(RuntimeError, match="does not support fusible operations"): + apply_recipe( + torch.nn.ModuleList([te_ops.Quantize()]), + CustomRecipe(qfactory=counting_factory, qfactory_key=("apply-zero-q", 2)), + ) + finally: + FP8GlobalStateManager.reset() + + def test_apply_recipe_rejects_legacy_dpa_path(): """Model-wide application does not alter the built-in NVTE_DPA_* mechanism.""" available, reason = is_fp8_available(return_reason=True) @@ -2541,16 +2592,134 @@ def test_apply_recipe_rejects_legacy_dpa_path(): attention_dropout=0.0, name="dpa", ).cuda() + linear = Linear(16, 16, bias=False, device="cuda", name="linear") + model = torch.nn.ModuleList([linear, dpa]) old_global_state = _global_recipe_state() try: - with pytest.raises(RuntimeError, match="only through CustomRecipe"): - apply_recipe(dpa, Float8CurrentScaling()) + # Quantization-inert for attention: the DPA is skipped, not rejected, so a + # stock TransformerLayer-shaped model can still be updated model-wide. + inert = Float8CurrentScaling() + apply_recipe(model, inert) + assert linear._quantization_runtime is not None assert dpa._quantization_runtime is None + assert FP8GlobalStateManager.quantization_state.fp8_recipe is inert + + # Asking the built-in path to quantize attention is still rejected. + for attention_recipe in ( + Float8CurrentScaling(fp8_dpa=True), + Float8CurrentScaling(fp8_dpa=True, fp8_mha=True), + ): + with pytest.raises(RuntimeError, match="only through CustomRecipe"): + apply_recipe(model, attention_recipe) + assert dpa._quantization_runtime is None + + # A bare inert DPA has nothing to update at all. + with pytest.raises(ValueError, match="found no Transformer Engine runtime owners"): + apply_recipe(dpa, Float8CurrentScaling()) + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_dispatches_on_the_owner_protocol_not_concrete_classes(): + """Any owner implementing the private plan/apply pair participates.""" + FP8GlobalStateManager.reset() + + class _StubOwner(torch.nn.Module): + """A non-TE owner that plugs into te.apply_recipe().""" + + def __init__(self): + super().__init__() + self.planned = None + self.applied = None + + def _plan_recipe_update(self, recipe, *, diagnostic_name): + self.planned = (recipe, diagnostic_name) + return "stub-update" + + def _apply_recipe_update(self, update): + # Record the call itself, so a skipped owner is distinguishable from + # one committed with a None update. + self.applied = ("applied", update) + + class _InertOwner(_StubOwner): + def _plan_recipe_update(self, recipe, *, diagnostic_name): + super()._plan_recipe_update(recipe, diagnostic_name=diagnostic_name) + return None + + stub, inert = _StubOwner(), _InertOwner() + model = torch.nn.ModuleList([stub, inert]) + recipe = _make_counting_recipe(("apply-duck-typing", 1), []) + try: + apply_recipe(model, recipe) + assert stub.planned == (recipe, "0") + assert stub.applied == ("applied", "stub-update") + assert inert.planned == (recipe, "1") + assert inert.applied is None + assert FP8GlobalStateManager.quantization_state.fp8_recipe is recipe + finally: + FP8GlobalStateManager.reset() + + +def test_apply_recipe_rejects_custom_delayed_dpa_crossing_before_commit(): + """A delayed custom DPA runtime cannot be handed to the built-in path.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + dpa = DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + name="dpa", + ).cuda() + linear = Linear(16, 16, bias=False, device="cuda", name="linear") + model = torch.nn.ModuleList([linear, dpa]) + try: + with autocast( + enabled=True, + recipe=CustomRecipe( + qfactory=delayed_scaling_factory, + qfactory_key=("apply-crossing-warm", 1), + fp8_dpa=True, + ), + ): + dpa.get_qkv_quantization_capabilities() + runtime = dpa._quantization_runtime + assert runtime is not None + + if not TransformerEngineBaseModule._runtime_has_delayed_scaling(runtime): + pytest.skip("this factory produced no delayed DPA state") + + old_global_state = _global_recipe_state() + with pytest.raises(RuntimeError, match="frozen after initialization"): + apply_recipe(model, Float8CurrentScaling()) + assert dpa._quantization_runtime is runtime assert _global_recipe_state() == old_global_state finally: FP8GlobalStateManager.reset() +def test_zero_quantizer_norm_survives_same_class_change_on_lazy_path(): + """te.ops norms build no quantizers, so a same-class change cannot invalidate them.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + model = te_ops.Sequential(te_ops.RMSNorm(16, device="cuda")) + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + try: + first = Float8CurrentScaling() + second = Float8CurrentScaling() + second.fp8_quant_fwd_inp = QParams(power_2_scale=True) + for active in (first, second, first): + with torch.no_grad(), autocast(enabled=True, recipe=active): + model(inp) + finally: + FP8GlobalStateManager.reset() + + def test_apply_recipe_rejects_active_autocast_and_graph_capture(monkeypatch): """Explicit model-wide application is limited to the documented safe boundary.""" FP8GlobalStateManager.reset() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 647e6da4222..bfe43e05792 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -33,6 +33,7 @@ Float8BlockScalingRecipeState, _QuantizationRuntime, _QuantizationRuntimeKey, + _is_delayed_scaling_state, ) from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -1375,6 +1376,29 @@ def _rehash_autocast_arguments(fp8_recipe, fp8_recipe_dpa, fp8_group) -> None: fp8_group, ) + def _plan_recipe_update( + self, + recipe: Recipe, + *, + diagnostic_name: str, + ) -> Optional[object]: + """Leave the built-in path on its lazy route; plan only CustomRecipe here.""" + if not recipe.custom(): + if recipe.fp8_dpa or recipe.fp8_mha: + raise RuntimeError( + "te.apply_recipe() supports DotProductAttention only through CustomRecipe; " + "the built-in NVTE_DPA_* path remains lazy and unchanged for " + f"{diagnostic_name!r}." + ) + runtime = getattr(self, "_quantization_runtime", None) + if runtime is not None and self._runtime_has_delayed_scaling(runtime): + # Committing would hand this module to the built-in path while it + # owns delayed state registered for reduction. + self._reject_builtin_dpa_update() + # Quantization-inert for attention: nothing here to update. + return None + return super()._plan_recipe_update(recipe, diagnostic_name=diagnostic_name) + def _check_quantization_update_supported( self, *, @@ -1382,12 +1406,18 @@ def _check_quantization_update_supported( requested_key: _QuantizationRuntimeKey, num_gemms: int, ) -> None: - """Also reject a CustomRecipe update onto an initialized built-in DPA.""" - if getattr(self, "_quantization_runtime", None) is None and self.fp8_initialized: - # The built-in path owns this module's state, and its delayed half is - # registered in the global reduction buckets. Running before super() - # means both the lazy forward and apply_recipe planning reject here, - # before any candidate is built. + """Also reject a CustomRecipe update onto built-in delayed state.""" + if getattr(self, "_quantization_runtime", None) is None and any( + _is_delayed_scaling_state(self.fp8_meta.get(key)) + for key in ("scaling_fwd", "scaling_bwd") + ): + # The built-in path registered this delayed state in the global + # reduction buckets, so a CustomRecipe runtime must not retarget it. + # Stateless built-in state is rebuilt instead, which is what the lazy + # path does for every other module and what D5/H2 needs: an inert + # built-in DPA inside a TransformerLayer must still accept a recipe. + # Running before super() means both the lazy forward and apply_recipe + # planning reject here, before any candidate is built. self._reject_builtin_dpa_update() super()._check_quantization_update_supported( recipe=recipe, diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 234384afe1f..08d464387c5 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1463,6 +1463,29 @@ def _plan_quantization_update( role_revision=role_revision, ) + def _plan_recipe_update( + self, + recipe: Recipe, + *, + diagnostic_name: str, + ) -> Optional[object]: + """Plan a model-wide recipe update for this owner, or ``None`` to skip it. + + One half of the private owner protocol :func:`te.apply_recipe` dispatches + on, so it never needs to import concrete owner classes. An owner outside + this hierarchy participates by implementing this and + :meth:`_apply_recipe_update`. + """ + del diagnostic_name + return self._plan_quantization_update( + recipe=recipe, + num_gemms=self._get_quantization_runtime_num_gemms(), + ) + + def _apply_recipe_update(self, update: object) -> None: + """Commit a plan produced by :meth:`_plan_recipe_update`.""" + self._apply_quantization_update(update) + def _apply_quantization_update( self, update: _QuantizationUpdate, diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index a22501b8886..1e3b626b4ea 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -648,6 +648,10 @@ def maybe_fuse_ops( if ( recipe_type is self.recipe_type and self.recipe_config != recipe_config + # A pipeline whose ops build no quantizers holds nothing a recipe change + # could invalidate. Evaluated only on the rare mismatch path. Interim + # form; WP8 folds it into the transition check with the op-level guard. + and any(op._builds_quantizers() for op in self._basic_ops) and not _is_preserved_fusible_recipe_transition( recipe, self.recipe_config, diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index b7218f696c8..f0ff13f2622 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -113,6 +113,19 @@ class FusibleOperation(torch.nn.Module, metaclass=abc.ABCMeta): def is_fused_op(self) -> bool: """Whether this op is the fusion of one or more basic ops""" + def _plan_recipe_update( + self, + recipe: Recipe, + *, + diagnostic_name: str, + ) -> Optional[object]: + """Reject a model-wide recipe update: fusible ops are recreated, not updated.""" + del recipe + raise RuntimeError( + "te.apply_recipe() does not support fusible operations yet; " + f"recreate or update this owner separately: {diagnostic_name!r}." + ) + def pre_first_fuser_forward(self) -> None: """Preprocessing before first fuser forward pass""" @@ -235,6 +248,25 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta): # Number of extra tensor outputs num_extra_outputs: int = 0 + def _builds_quantizers(self) -> bool: + """Whether this operation constructs any quantizers. + + An operation that builds none (e.g. ``te.ops.LayerNorm`` / ``RMSNorm``, + Megatron's ``TENorm``) holds no state a recipe change could invalidate. + """ + return bool(self.num_quantizers("forward") or self.num_quantizers("backward")) + + def _plan_recipe_update( + self, + recipe: Recipe, + *, + diagnostic_name: str, + ) -> Optional[object]: + """Skip an operation that builds no quantizers; it holds nothing to update.""" + if not self._builds_quantizers(): + return None + return super()._plan_recipe_update(recipe, diagnostic_name=diagnostic_name) + def __init__(self) -> None: super().__init__() @@ -387,6 +419,8 @@ def reset_recipe_state( self._recipe_type is recipe_type and self._recipe_config is not None and self._recipe_config != recipe_config + # Interim form; WP8 moves this into the transition check. + and self._builds_quantizers() and not _is_preserved_fusible_recipe_transition( recipe, self._recipe_config, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 1855a20f01f..12fe7417292 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1005,66 +1005,35 @@ def apply_recipe(model: torch.nn.Module, recipe: Recipe) -> None: check_recipe_support(recipe) recipe.quantizer_config() - # Import locally to keep quantization.py independent from module/base.py - # during package initialization. - from .attention.dot_product_attention.dot_product_attention import DotProductAttention - from .module.base import TransformerEngineBaseModule - from .ops.op import FusibleOperation - - participants = [] - fusible_owners = [] - legacy_dpa_owners = [] - seen = set() + # Dispatch on a private owner protocol rather than concrete classes, so this + # function needs no imports from the module/attention/ops packages and any + # owner implementing the pair can participate. + plans = [] for fqn, module in model.named_modules(): - module_id = id(module) - if module_id in seen: + planner = getattr(module, "_plan_recipe_update", None) + if planner is None: continue - seen.add(module_id) diagnostic_name = fqn or "" - if isinstance(module, FusibleOperation): - fusible_owners.append(diagnostic_name) - if isinstance(module, DotProductAttention) and not recipe.custom(): - legacy_dpa_owners.append(diagnostic_name) - if isinstance(module, TransformerEngineBaseModule): - participants.append((diagnostic_name, module)) - - # Reject excluded owners after one complete discovery pass and before any - # qfactory invocation or candidate construction. - if fusible_owners: - owners = ", ".join(repr(name) for name in fusible_owners) - raise RuntimeError( - "te.apply_recipe() does not support fusible operations yet; " - f"recreate or update these owners separately: {owners}." - ) - if legacy_dpa_owners: - owners = ", ".join(repr(name) for name in legacy_dpa_owners) - raise RuntimeError( - "te.apply_recipe() supports DotProductAttention only through CustomRecipe; " - f"the built-in NVTE_DPA_* path remains lazy and unchanged for: {owners}." - ) - if not participants: - raise ValueError("te.apply_recipe() found no Transformer Engine runtime owners in model.") - - updates = [] - for fqn, module in participants: try: - # pylint: disable-next=protected-access - num_gemms = module._get_quantization_runtime_num_gemms() - update = module._plan_quantization_update( # pylint: disable=protected-access - recipe=recipe, - num_gemms=num_gemms, - ) + update = planner(recipe, diagnostic_name=diagnostic_name) except Exception as exc: - raise RuntimeError( - f"te.apply_recipe() failed while planning module {fqn!r}: {exc}" - ) from exc - updates.append((module, update)) + # Preserve the owner's exception type; only annotate it. + if hasattr(exc, "add_note"): # Python 3.11+ + exc.add_note(f"te.apply_recipe(): while planning module {diagnostic_name!r}") + raise + if update is None: + # The owner has nothing this recipe would change. + continue + plans.append((module, update)) + + if not plans: + raise ValueError("te.apply_recipe() found no Transformer Engine runtime owners in model.") # Applying an update only publishes state that was fully constructed and # validated above. Publish the manager state last so planning failures # cannot expose a requested recipe globally. - for module, update in updates: - module._apply_quantization_update(update) # pylint: disable=protected-access + for module, update in plans: + module._apply_recipe_update(update) # pylint: disable=protected-access FP8GlobalStateManager.quantization_state.fp8_recipe = recipe From ba9b2b75cdd5df0023fca020553a162407a18070 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 11 Sep 2026 13:22:25 +0000 Subject: [PATCH 92/97] Tie GroupedLinear's derived state to its runtime Signed-off-by: Evgeny --- tests/pytorch/test_hybrid_quantization.py | 241 ++++++++++----- tests/pytorch/test_recipe.py | 4 +- tests/pytorch/test_recipe_update.py | 29 +- transformer_engine/pytorch/module/base.py | 3 +- .../pytorch/module/grouped_linear.py | 278 ++++++------------ transformer_engine/pytorch/quantization.py | 3 + .../tensor/storage/grouped_tensor_storage.py | 29 +- 7 files changed, 293 insertions(+), 294 deletions(-) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index 2abe632add9..ecb58dd408c 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -5500,30 +5500,161 @@ def test_validation_rejects_same_family_config_mismatch(self): operand_name="input", ) - def test_builtin_recipe_skips_grouped_quantizer_validation(self, monkeypatch): - def unexpected_validation(*_args, **_kwargs): - pytest.fail("built-in recipes must not run grouped-quantizer validation") + def test_bf16_forward_after_quantized_forward_has_no_stale_classification(self): + """Grouped traits die with the runtime, so the high-precision path is clean.""" + torch.manual_seed(1234) + common = dict(bias=False, params_dtype=torch.bfloat16, save_original_input=True) + model = GroupedLinear(2, 128, 128, **common).cuda() + reference = GroupedLinear(2, 128, 128, **common).cuda() + with torch.no_grad(): + for i in range(2): + getattr(reference, f"weight{i}").copy_(getattr(model, f"weight{i}")) + + m_splits = torch.tensor([64, 64], dtype=torch.int64) + warmup = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + with torch.no_grad(), autocast(enabled=True, recipe=recipe.DelayedScaling()): + model(warmup, m_splits) + assert model._quantization_runtime.owner_traits is not None + + # A high-precision training step must not read the quantized run's + # classification, and must match a module that never saw a recipe. + def step(module): + inp = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda", requires_grad=True) + torch.manual_seed(99) + out = module(inp, m_splits) + out.backward(torch.ones_like(out)) + return inp.grad, [getattr(module, f"weight{i}").grad for i in range(2)] + + torch.manual_seed(7) + actual_input_grad, actual_weight_grads = step(model) + torch.manual_seed(7) + expected_input_grad, expected_weight_grads = step(reference) + + torch.testing.assert_close(actual_input_grad, expected_input_grad, rtol=0.0, atol=0.0) + for actual, expected in zip(actual_weight_grads, expected_weight_grads): + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + def test_grouped_tensor_recipe_is_rejected_before_any_state_changes(self): + """The device/recipe half of the grouped-tensor check runs at validation.""" + if not fp8_block_scaling_available: + pytest.skip(f"FP8 block scaling: {reason_for_no_fp8_block_scaling}") + from transformer_engine.pytorch.utils import get_device_compute_capability + + if not (10, 0) <= get_device_compute_capability() <= (11, 0): + pytest.skip("the grouped-tensor block-scaling restriction is Blackwell-only") - monkeypatch.setattr( - split_quantization, - "validate_grouped_quantizer_list", - unexpected_validation, - ) model = GroupedLinear( - 2, - 128, - 128, - bias=False, - params_dtype=torch.bfloat16, + 2, 128, 128, bias=False, params_dtype=torch.bfloat16, use_grouped_tensor=True ).cuda() tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + m_splits = torch.tensor([64, 64], dtype=torch.int64, device="cuda") + with torch.no_grad(), autocast(enabled=True, recipe=recipe.MXFP8BlockScaling()): + model(tensor, m_splits) + active_runtime = model._quantization_runtime + + with pytest.raises(RuntimeError, match="does not support the FP8 block-scaling recipe"): + te.apply_recipe(model, recipe.Float8BlockScaling()) + assert model._quantization_runtime is active_runtime + + def test_packed_grouped_storage_rejects_unsupported_quantizers(self): + """One owner of the packed-storage contract, stated as an allowlist.""" + from transformer_engine.pytorch.tensor.storage.grouped_tensor_storage import ( + GroupedTensorStorage, + ) + + for unsupported in (IdentityQuantizer(), _make_hybrid_quantizer_fp8_row_fp4_col()): + with pytest.raises(NotImplementedError, match="GroupedTensorStorage does not support"): + GroupedTensorStorage.make_grouped_tensor( + 2, + None, + None, + 64, + 64, + quantizer=unsupported, + device="cuda", + dtype=torch.bfloat16, + ) + # A supported quantizer is allowed through the same gate. + GroupedTensorStorage.make_grouped_tensor( + 2, + None, + None, + 64, + 64, + quantizer=_make_fp8_quantizer(), + device="cuda", + dtype=torch.bfloat16, + ) + + @staticmethod + def _track_validation(monkeypatch): + """Record every grouped-operand validation call.""" + real_validate = split_quantization.validate_grouped_quantizer_list + calls = [] + + def tracked(quantizers, *, operand_name="operand"): + calls.append(operand_name) + return real_validate(quantizers, operand_name=operand_name) + + monkeypatch.setattr(split_quantization, "validate_grouped_quantizer_list", tracked) + return calls + + @pytest.mark.parametrize( + "make_recipe,expected", + [ + pytest.param(recipe.DelayedScaling, [], id="builtin-delayed"), + pytest.param(recipe.Float8CurrentScaling, [], id="builtin-current"), + pytest.param( + lambda: recipe.CustomRecipe( + qfactory=delayed_scaling_factory, + qfactory_key=("grouped-validation-count", 1), + ), + ["grad_output", "input", "weight"], + id="custom", + ), + ], + ) + def test_only_custom_recipes_pay_for_grouped_operand_validation( + self, monkeypatch, make_recipe, expected + ): + """Built-in recipes are homogeneous by construction, so they skip the scan. + + A CustomRecipe checks input, weight and grad_output exactly once per + commit; steady state revalidates nothing either way. + """ + calls = self._track_validation(monkeypatch) + model = GroupedLinear(2, 128, 128, bias=False, params_dtype=torch.bfloat16).cuda() + tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") m_splits = torch.tensor([64, 64], dtype=torch.int64) + active = make_recipe() - with torch.no_grad(), autocast(enabled=True, recipe=recipe.DelayedScaling()): + with torch.no_grad(), autocast(enabled=True, recipe=active): model(tensor, m_splits) + assert sorted(calls) == expected + + for grad_mode in (torch.no_grad, torch.enable_grad): + with grad_mode(), autocast(enabled=True, recipe=active): + model(tensor, m_splits) + assert sorted(calls) == expected + + def test_steady_state_forward_does_not_query_grad_mode_for_validation(self, monkeypatch): + """The deleted generation guard was the only steady-state grad-mode query.""" + model = GroupedLinear(2, 128, 128, bias=False, params_dtype=torch.bfloat16).cuda() + tensor = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + m_splits = torch.tensor([64, 64], dtype=torch.int64) + active = recipe.DelayedScaling() + + with torch.no_grad(), autocast(enabled=True, recipe=active): model(tensor, m_splits) - def test_validation_runs_only_with_quantizer_generation(self, monkeypatch): + calls = self._track_validation(monkeypatch) + with torch.no_grad(), autocast(enabled=True, recipe=active): + model._get_quantizers() + assert not calls + + def test_runtime_traits_are_replaced_atomically_with_the_runtime(self, monkeypatch): + """Traits live on the runtime, so a rejected candidate leaves them alone.""" + def make_qfactory(columnwise_source): def qfactory(_role): return HybridQuantizer( @@ -5541,59 +5672,22 @@ def qfactory(_role): m_splits = torch.tensor([64, 64], dtype=torch.int64) original_recipe = recipe.CustomRecipe(qfactory=make_qfactory("original")) - real_validate = split_quantization.validate_grouped_quantizer_list - validation_calls = [] - - def tracked_validate(quantizers, *, operand_name="operand"): - validation_calls.append((operand_name, id(quantizers[0]))) - return real_validate(quantizers, operand_name=operand_name) - - monkeypatch.setattr( - split_quantization, - "validate_grouped_quantizer_list", - tracked_validate, - ) - + calls = self._track_validation(monkeypatch) with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): model(tensor, m_splits) - first_call_count = len(validation_calls) - first_generation = model._validated_quantizer_generations["scaling_fwd"] - assert first_call_count > 0 - # Candidate validation covers both directions before the commit, so the - # first forward validates the backward operand too, whatever the grad mode. - assert {name for name, _ in validation_calls} == {"input", "weight", "grad_output"} - - with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): - model(tensor, m_splits) - assert len(validation_calls) == first_call_count - assert model._validated_quantizer_generations["scaling_fwd"] is first_generation - - # Selecting quantizers with gradients enabled hits the O(1) generation - # guard: the committed runtime already recorded both generations, so - # nothing is revalidated. Do not run a full forward here: this validation - # test intentionally uses a columnwise-only configuration that is not - # supported by the split-quantization kernel. - with torch.enable_grad(), autocast(enabled=True, recipe=original_recipe): - model._get_quantizers() - assert len(validation_calls) == first_call_count - assert ( - model._validated_quantizer_generations["scaling_bwd"] is model.quantizers["scaling_bwd"] - ) - - def unexpected_grad_mode_query(): - pytest.fail("cached validation must not query grad mode") - - with monkeypatch.context() as context: - context.setattr(torch, "is_grad_enabled", unexpected_grad_mode_query) - model._validate_quantizer_generation(False) + assert sorted(calls) == ["grad_output", "input", "weight"] + first_runtime = model._quantization_runtime + first_traits = first_runtime.owner_traits + assert first_traits is not None + # A supported rebuild replaces runtime and traits together. rebuilt_recipe = recipe.CustomRecipe(qfactory=make_qfactory("rowwise_dequantized")) with torch.no_grad(), autocast(enabled=True, recipe=rebuilt_recipe): model(tensor, m_splits) - rebuilt_generation = model._validated_quantizer_generations["scaling_fwd"] - assert len(validation_calls) > first_call_count - assert rebuilt_generation is not first_generation - assert rebuilt_generation[0].columnwise_source == "rowwise_dequantized" + rebuilt_runtime = model._quantization_runtime + assert rebuilt_runtime is not first_runtime + assert rebuilt_runtime.owner_traits is not first_traits + assert model.quantizers["scaling_fwd"][0].columnwise_source == "rowwise_dequantized" input_count = 0 @@ -5613,27 +5707,18 @@ def mixed_source_qfactory(role): qfactory=mixed_source_qfactory, qfactory_key=("test_mixed_columnwise_sources", 1), ) - # A failed candidate is neither committed nor marked validated. The - # active runtime stays behind the requested revision, so each retry - # prepares and validates a fresh candidate. - rebuilt_runtime = model._quantization_runtime - rebuilt_backward_generation = model._validated_quantizer_generations["scaling_bwd"] - rebuilt_delayed_quantizer = model._delayed_scaling_input_quantizer - rebuilt_unsafe_quantizer = model._unsafe_requantization_input_quantizer - failed_validation_call_count = len(validation_calls) + # A failed candidate is neither committed nor marked validated, so each + # retry prepares and validates a fresh candidate. + rebuilt_traits = rebuilt_runtime.owner_traits + call_count = len(calls) for _ in range(2): with pytest.raises(ValueError, match="mixed columnwise source policies"): with torch.no_grad(), autocast(enabled=True, recipe=mixed_recipe): model(tensor, m_splits) assert model._quantization_runtime is rebuilt_runtime - assert model._validated_quantizer_generations["scaling_fwd"] is rebuilt_generation - assert ( - model._validated_quantizer_generations["scaling_bwd"] is rebuilt_backward_generation - ) - assert model._delayed_scaling_input_quantizer is rebuilt_delayed_quantizer - assert model._unsafe_requantization_input_quantizer is rebuilt_unsafe_quantizer - assert len(validation_calls) > failed_validation_call_count - failed_validation_call_count = len(validation_calls) + assert rebuilt_runtime.owner_traits is rebuilt_traits + assert len(calls) > call_count + call_count = len(calls) # Stale invalid recipe metadata must not affect the non-quantized path. with torch.no_grad(): diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 6b839e52945..ad9dd3ee02b 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -712,9 +712,7 @@ def make_module_owner(module_type, num_gemms): "grad_input": 1, }, _num_fp8_tensors_per_gemm={"fwd": 3, "bwd": 2}, - _validated_quantizer_generations={}, - _delayed_scaling_input_quantizer=None, - _unsafe_requantization_input_quantizer=None, + use_grouped_tensor=False, ) GroupedLinear.set_meta_tensor(grouped, True, recipe) GroupedLinear.set_meta_tensor(grouped, False, recipe) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 767b140865d..7bfbb259946 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -2817,9 +2817,7 @@ def qfactory(role): old_backward_state = module.fp8_meta["scaling_bwd"] old_forward_quantizers = module.quantizers["scaling_fwd"] old_backward_quantizers = module.quantizers["scaling_bwd"] - old_validated_generations = module._validated_quantizer_generations - old_delayed_quantizer = module._delayed_scaling_input_quantizer - old_unsafe_quantizer = module._unsafe_requantization_input_quantizer + old_traits = old_runtime.owner_traits invalid_recipe = make_recipe( ("grouped-atomic", "invalid", mismatched_tensor_type), @@ -2833,9 +2831,7 @@ def qfactory(role): assert module.fp8_meta["scaling_bwd"] is old_backward_state assert module.quantizers["scaling_fwd"] is old_forward_quantizers assert module.quantizers["scaling_bwd"] is old_backward_quantizers - assert module._validated_quantizer_generations is old_validated_generations - assert module._delayed_scaling_input_quantizer is old_delayed_quantizer - assert module._unsafe_requantization_input_quantizer is old_unsafe_quantizer + assert old_runtime.owner_traits is old_traits replacement_recipe = make_recipe( ("grouped-atomic", "replacement", mismatched_tensor_type), @@ -2843,21 +2839,12 @@ def qfactory(role): unsafe_inputs=True, ) update = _prepare_runtime_update(module, replacement_recipe, num_gemms=2) - assert module._validated_quantizer_generations is old_validated_generations - assert module._delayed_scaling_input_quantizer is old_delayed_quantizer - assert module._unsafe_requantization_input_quantizer is old_unsafe_quantizer + # Planning derives traits for the candidate without touching the active runtime. + assert old_runtime.owner_traits is old_traits assert module._apply_quantization_update(update) # pylint: disable=protected-access replacement_runtime = module._quantization_runtime # pylint: disable=protected-access assert replacement_runtime is not old_runtime - assert ( - module._validated_quantizer_generations["scaling_fwd"] - is replacement_runtime.forward_quantizers - ) - assert ( - module._validated_quantizer_generations["scaling_bwd"] - is replacement_runtime.backward_quantizers - ) - assert module._delayed_scaling_input_quantizer is None - assert ( - module._unsafe_requantization_input_quantizer is replacement_runtime.forward_quantizers[0] - ) + traits = replacement_runtime.owner_traits + assert traits is not old_traits + assert traits.delayed_scaling_input_quantizer is None + assert traits.unsafe_requantization_input_quantizer is replacement_runtime.forward_quantizers[0] diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 08d464387c5..3ce318169cc 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1317,7 +1317,8 @@ def _activate_quantization_runtime( validation_result: Any = None, ) -> None: """Publish a validated replacement runtime and its compatibility views.""" - del validation_result + # Owner-derived state lives on the runtime, so it cannot outlive it. + candidate.owner_traits = validation_result recipe = candidate.recipe forward_state = candidate.forward_state backward_state = candidate.backward_state diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index d7bd1ef562d..117f396c1dc 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -5,6 +5,7 @@ """GroupedLinear API""" from typing import Union, Optional, Callable, Tuple, List +import dataclasses from itertools import chain import os import warnings @@ -75,7 +76,6 @@ Float8CurrentScalingQuantizer, Float8Quantizer, HybridQuantizer, - IdentityQuantizer, MXFP8Quantizer, ) from ..quantized_tensor import ( @@ -90,6 +90,40 @@ __all__ = ["GroupedLinear", "is_module_grouped_tensor_path_supported"] +@dataclasses.dataclass(frozen=True) +class _GroupedRuntimeTraits: + """Grouped-kernel classification derived when a runtime is validated. + + Published on the runtime and discarded with it, so a high-precision forward + after a quantized one cannot read a classification that no longer applies. + """ + + delayed_scaling_input_quantizer: Optional[Quantizer] = None + unsafe_requantization_input_quantizer: Optional[Quantizer] = None + + +# Read by the high-precision path, which has no runtime to ask. +_NO_GROUPED_RUNTIME_TRAITS = _GroupedRuntimeTraits() + + +def _check_grouped_tensor_recipe(recipe: Optional[Recipe], use_grouped_tensor: bool) -> None: + """Reject a recipe the native grouped-tensor path cannot run on this device. + + Depends only on the recipe and the device, so it is checked when a runtime is + validated as well as on the forward path that would dispatch the kernel. + """ + if not use_grouped_tensor or recipe is None or not recipe.float8_block_scaling(): + return + if (10, 0) <= get_device_compute_capability() <= (11, 0): + raise RuntimeError( + "use_grouped_tensor=True does not support the FP8 block-scaling recipe on " + "Blackwell GPUs: the native grouped FP8 block-scaling path is Hopper-only. " + "Set use_grouped_tensor=False, or unset " + "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM if it enabled this path, to use " + "the MXFP8-emulated path on Blackwell." + ) + + def is_module_grouped_tensor_path_supported( recipe: Optional[Recipe], dtype: torch.dtype, @@ -756,18 +790,8 @@ def forward( or cpu_offloading or any(q is not None for q in output_quantizers) ): - if ( - fp8 - and recipe.float8_block_scaling() - and (10, 0) <= get_device_compute_capability() <= (11, 0) - ): - raise RuntimeError( - "use_grouped_tensor=True does not support the FP8 block-scaling recipe on " - "Blackwell GPUs: the native grouped FP8 block-scaling path is Hopper-only. " - "Set use_grouped_tensor=False, or unset " - "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM if it enabled this path, to use " - "the MXFP8-emulated path on Blackwell." - ) + if fp8: + _check_grouped_tensor_recipe(recipe, use_grouped_tensor=True) grouped_tensor_supported = is_module_grouped_tensor_path_supported( recipe, activation_dtype, @@ -1714,10 +1738,6 @@ def __init__( "fwd": 3, "bwd": 2, } - self._uses_custom_recipe = False - self._validated_quantizer_generations = {} - self._delayed_scaling_input_quantizer = None - self._unsafe_requantization_input_quantizer = None if tp_group is None: self.tp_size = tp_size @@ -1797,33 +1817,42 @@ def __init__( if name in (f"weight{i}", f"bias{i}"): param.skip_backward_post_hook = True - def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: - """Init scales and amaxes for fwd | bwd.""" - if recipe.float8_current_scaling() and self.tp_size > 1: + def _validate_quantization_runtime( + self, + candidate: _QuantizationRuntime, + ) -> "_GroupedRuntimeTraits": + """Validate both grouped directions and derive the runtime's traits. + + Only a ``CustomRecipe`` is checked: the built-in recipes build their own + expert quantizers and are homogeneous by construction, so scanning the + operands would spend CPU confirming something it cannot disprove. + """ + super()._validate_quantization_runtime(candidate) + if ( + isinstance(candidate.forward_state, Float8CurrentScalingRecipeState) + and self.tp_size > 1 + ): raise ValueError( "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " "Because the TP communication is handled outside of this module." ) - super().set_meta_tensor(fwd, recipe) - if self._uses_custom_recipe: - self._validate_quantizer_generation(fwd) + _check_grouped_tensor_recipe(candidate.recipe, self.use_grouped_tensor) - def _validate_forward_quantizer_generation( - self, - generation: List[Quantizer], - *, - num_gemms: int, - validate_grouped_quantizers: bool, - ) -> Tuple[Optional[Quantizer], Optional[Quantizer]]: - """Validate candidate forward quantizers and derive grouped runtime state.""" - stride = self._num_fp8_tensors_per_gemm["fwd"] - first_input_quantizer = generation[self._offsets["input"]] - if validate_grouped_quantizers: + num_gemms = candidate.num_gemms + forward = candidate.forward_quantizers + forward_stride = self._num_fp8_tensors_per_gemm["fwd"] + + if candidate.recipe.custom(): + backward_stride = self._num_fp8_tensors_per_gemm["bwd"] input_quantizers = tuple( - generation[self._offsets["input"] + i * stride] for i in range(num_gemms) + forward[self._offsets["input"] + i * forward_stride] for i in range(num_gemms) ) weight_quantizers = tuple( - generation[self._offsets["weight"] + i * stride] for i in range(num_gemms) + forward[self._offsets["weight"] + i * forward_stride] for i in range(num_gemms) + ) + grad_output_quantizers = tuple( + candidate.backward_quantizers[self._offsets["grad_output"] + i * backward_stride] + for i in range(num_gemms) ) _split_quantization.validate_grouped_quantizer_list( input_quantizers, operand_name="input" @@ -1831,133 +1860,27 @@ def _validate_forward_quantizer_generation( _split_quantization.validate_grouped_quantizer_list( weight_quantizers, operand_name="weight" ) + _split_quantization.validate_grouped_quantizer_list( + grad_output_quantizers, operand_name="grad_output" + ) else: # Native recipe construction guarantees homogeneous expert quantizers. - input_quantizers = (first_input_quantizer,) - delayed_scaling_input_quantizer = next( - (q for q in input_quantizers if isinstance(q, Float8Quantizer)), - None, - ) - unsafe_requantization_input_quantizer = next( - ( - q - for q in input_quantizers - if q is not None and not can_reconstruct_wgrad_input_from_original(q) - ), - None, - ) - return ( - delayed_scaling_input_quantizer, - unsafe_requantization_input_quantizer, - ) - - def _validate_backward_quantizer_generation( - self, - generation: List[Quantizer], - *, - num_gemms: int, - validate_grouped_quantizers: bool, - ) -> None: - """Validate candidate backward quantizers for grouped-kernel compatibility.""" - if not validate_grouped_quantizers: - return - stride = self._num_fp8_tensors_per_gemm["bwd"] - grad_output_quantizers = tuple( - generation[self._offsets["grad_output"] + i * stride] for i in range(num_gemms) - ) - _split_quantization.validate_grouped_quantizer_list( - grad_output_quantizers, - operand_name="grad_output", - ) - - def _validate_quantization_runtime( - self, - candidate: _QuantizationRuntime, - ) -> Tuple[Optional[Quantizer], Optional[Quantizer]]: - """Validate both grouped directions without changing live module state.""" - super()._validate_quantization_runtime(candidate) - if ( - isinstance(candidate.forward_state, Float8CurrentScalingRecipeState) - and self.tp_size > 1 - ): - raise ValueError( - "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " - "Because the TP communication is handled outside of this module." - ) - - validate_grouped_quantizers = candidate.recipe.custom() - validation_result = self._validate_forward_quantizer_generation( - candidate.forward_quantizers, - num_gemms=candidate.num_gemms, - validate_grouped_quantizers=validate_grouped_quantizers, - ) - self._validate_backward_quantizer_generation( - candidate.backward_quantizers, - num_gemms=candidate.num_gemms, - validate_grouped_quantizers=validate_grouped_quantizers, - ) - return validation_result - - def _activate_quantization_runtime( - self, - candidate: _QuantizationRuntime, - *, - validation_result: Tuple[Optional[Quantizer], Optional[Quantizer]], - ) -> None: - """Publish validated grouped-derived state with the candidate runtime.""" - delayed_scaling_input_quantizer, unsafe_requantization_input_quantizer = validation_result - validated_generations = { - "scaling_fwd": candidate.forward_quantizers, - "scaling_bwd": candidate.backward_quantizers, - } + input_quantizers = (forward[self._offsets["input"]],) - super()._activate_quantization_runtime( - candidate, - validation_result=validation_result, - ) - self.fast_setattr( - "_delayed_scaling_input_quantizer", - delayed_scaling_input_quantizer, - ) - self.fast_setattr( - "_unsafe_requantization_input_quantizer", - unsafe_requantization_input_quantizer, + return _GroupedRuntimeTraits( + delayed_scaling_input_quantizer=next( + (q for q in input_quantizers if isinstance(q, Float8Quantizer)), + None, + ), + unsafe_requantization_input_quantizer=next( + ( + q + for q in input_quantizers + if q is not None and not can_reconstruct_wgrad_input_from_original(q) + ), + None, + ), ) - self.fast_setattr("_uses_custom_recipe", candidate.recipe.custom()) - self.fast_setattr("_validated_quantizer_generations", validated_generations) - - def _validate_quantizer_generation(self, fwd: bool) -> None: - """Validate grouped-kernel invariants once per quantizer generation.""" - # Recipe state replaces this list object only when it constructs a new - # quantizer generation. The O(1) identity guard keeps validation off the - # steady-state forward path. Record a generation only after all of its - # operand roles pass, so a failed recipe transition is retried. - meta_key = "scaling_fwd" if fwd else "scaling_bwd" - generation = self.quantizers.get(meta_key) - if generation is None: - return - if self._validated_quantizer_generations.get(meta_key) is generation: - return - - if fwd: - ( - delayed_scaling_input_quantizer, - unsafe_requantization_input_quantizer, - ) = self._validate_forward_quantizer_generation( - generation, - num_gemms=self.num_gemms, - validate_grouped_quantizers=True, - ) - self._delayed_scaling_input_quantizer = delayed_scaling_input_quantizer - self._unsafe_requantization_input_quantizer = unsafe_requantization_input_quantizer - else: - self._validate_backward_quantizer_generation( - generation, - num_gemms=self.num_gemms, - validate_grouped_quantizers=True, - ) - - self._validated_quantizer_generations[meta_key] = generation def get_quantizer_roles( self, @@ -1995,21 +1918,8 @@ def make_grouped_weights(self, defer_init=False) -> None: return weight_quantizers = self._get_weight_quantizers() - # TODO(#3158): Support Identity/Hybrid single grouped weights. - unsupported_quantizers = tuple( - type(quantizer).__name__ - for quantizer in weight_quantizers - if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer)) - ) - if unsupported_quantizers: - quantizer_names = ", ".join(dict.fromkeys(unsupported_quantizers)) - raise NotImplementedError( - "GroupedLinear(single_grouped_weight=True) does not support " - f"{quantizer_names} weight quantizers yet. Set " - "single_grouped_weight=False or unset " - "NVTE_GROUPED_LINEAR_SINGLE_PARAM. See #3158." - ) - + # Unsupported weight quantizers are rejected by GroupedTensorStorage, + # which owns the packed-storage contract. recipe = ( weight_quantizers[0]._get_compatible_recipe() if weight_quantizers and weight_quantizers[0] is not None @@ -2377,6 +2287,15 @@ def forward( linear_fn = _GroupedLinear.forward autograd_ctx = [None] + # Traits belong to the active runtime, so the high-precision path + # cannot read a classification derived for a quantized one. + runtime = self._quantization_runtime if self.fp8 else None + runtime_traits = ( + runtime.owner_traits + if runtime is not None and runtime.owner_traits is not None + else _NO_GROUPED_RUNTIME_TRAITS + ) + cache_weight = is_first_microbatch is not None if self.single_grouped_weight: weight_workspaces = [self._fp8_workspaces.get("weight")] if cache_weight else [None] @@ -2408,8 +2327,8 @@ def forward( cache_weight, skip_fp8_weight_update, self.save_original_input, - self._delayed_scaling_input_quantizer, - self._unsafe_requantization_input_quantizer, + runtime_traits.delayed_scaling_input_quantizer, + runtime_traits.unsafe_requantization_input_quantizer, debug, self.single_grouped_weight, use_grouped_bias, @@ -2538,13 +2457,6 @@ def _get_weight_quantizers(self) -> List[Quantizer]: return weight_quantizers def _get_quantizers(self): - if self.fp8 and self._uses_custom_recipe: - # Normally validated before candidate commit. Keep this O(1) - # generation guard for legacy/direct quantizer-list replacement. - self._validate_quantizer_generation(True) - if torch.is_grad_enabled(): - self._validate_quantizer_generation(False) - weight_quantizers = self._get_weight_quantizers() input_quantizers, output_quantizers = ( [None] * self.num_gemms, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 12fe7417292..9e8b035411b 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -143,6 +143,9 @@ class _QuantizationRuntime: backward_state: "RecipeState" forward_quantizers: List["Quantizer"] backward_quantizers: List["Quantizer"] + # Owner-specific state derived during validation, published with the runtime + # and discarded with it. Owners that need none leave it ``None``. + owner_traits: Any = None @dataclasses.dataclass(frozen=True) diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 3473024c03f..e8a0566f70e 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -655,15 +655,28 @@ def make_grouped_tensor( """ if quantizer is not None: - # TODO(#3158): Support Identity/Hybrid packed grouped storage. - from ..hybrid_tensor import HybridQuantizer - from ..identity_tensor import IdentityQuantizer - - if isinstance(quantizer, (IdentityQuantizer, HybridQuantizer)): + # TODO(#3158): widen this allowlist as packed grouped storage gains + # support, rather than enumerating what is unsupported. + from ...tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, + ) + from ...tensor.mxfp8_tensor import MXFP8Quantizer + from ...tensor.nvfp4_tensor import NVFP4Quantizer + from ...tensor.float8_blockwise_tensor import Float8BlockQuantizer + + supported = ( + Float8Quantizer, + Float8CurrentScalingQuantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, + ) + if not isinstance(quantizer, supported): raise NotImplementedError( - "GroupedTensorStorage does not support IdentityQuantizer or " - "HybridQuantizer yet. Use separate tensors, or set " - "GroupedLinear(single_grouped_weight=False). See #3158." + f"GroupedTensorStorage does not support {type(quantizer).__name__} yet. " + "Use separate tensors, set GroupedLinear(single_grouped_weight=False), " + "or unset NVTE_GROUPED_LINEAR_SINGLE_PARAM. See #3158." ) # Set device From 5fb9af1bdfbd9825c585aa809fb62adba00a90fb Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 11 Sep 2026 13:41:39 +0000 Subject: [PATCH 93/97] Treat restores and resizes as replacements, not updates Signed-off-by: Evgeny --- tests/pytorch/test_recipe_update.py | 77 +++++++++++++++++++++++ transformer_engine/pytorch/module/base.py | 22 +++++++ 2 files changed, 99 insertions(+) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 7bfbb259946..0302af79297 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -33,6 +33,7 @@ is_nvfp4_available, quantized_model_init, ) +from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV from transformer_engine.pytorch.custom_recipes.quantizer_factories import ( current_scaling_factory, delayed_scaling_factory, @@ -2620,6 +2621,82 @@ def test_apply_recipe_rejects_legacy_dpa_path(): FP8GlobalStateManager.reset() +def test_graph_capture_resizes_amax_history_and_leaves_a_consistent_runtime(): + """A direct resize bypasses the planner, so it must re-key the runtime itself.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + FP8GlobalStateManager.reset() + try: + with quantized_model_init(enabled=True, recipe=DelayedScaling(amax_history_len=1024)): + module = Linear( + 16, 16, bias=False, params_dtype=torch.bfloat16, device="cuda", name="linear" + ) + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + + import transformer_engine.pytorch as te + + te.make_graphed_callables( + module, + (inp,), + fp8_enabled=True, + fp8_recipe=DelayedScaling(amax_history_len=16), + ) + + assert module.fp8_meta["scaling_fwd"].amax_history.shape[0] == 16 + runtime = module._quantization_runtime # pylint: disable=protected-access + assert runtime.recipe.amax_history_len == 16 + assert runtime.key.recipe_config == runtime.recipe.quantizer_config() + assert module.fp8_meta["recipe"] is runtime.recipe + finally: + FP8GlobalStateManager.reset() + + +@pytest.mark.parametrize( + "live_recipe", + [ + pytest.param(DelayedScaling(amax_history_len=8), id="different-history"), + pytest.param(DelayedScaling(amax_history_len=4, margin=1), id="different-margin"), + pytest.param(MXFP8BlockScaling(), id="different-recipe"), + ], +) +def test_checkpoint_restore_adopts_the_checkpoints_recipe(monkeypatch, live_recipe): + """Loading a checkpoint replaces state; it is not an update from the live recipe.""" + # Delayed-scaling extra state is pickled, which is gated off by default. + monkeypatch.setenv(UNSAFE_PICKLE_EXTRA_STATE_ENV, "1") + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + if live_recipe.mxfp8() and not is_mxfp8_available(): + pytest.skip("MXFP8 is required for this variant") + + FP8GlobalStateManager.reset() + try: + + def warmed(active): + module = Linear( + 64, 64, bias=False, params_dtype=torch.bfloat16, device="cuda", name="linear" + ) + with autocast(enabled=True, recipe=active): + module(torch.randn(64, 64, device="cuda", dtype=torch.bfloat16)) + return module + + saved = warmed(DelayedScaling(amax_history_len=4)) + state_dict = saved.state_dict() + expected_history = saved.fp8_meta["scaling_fwd"].amax_history.clone() + + destination = warmed(live_recipe) + destination.load_state_dict(state_dict) + + state = destination.fp8_meta["scaling_fwd"] + assert state.amax_history.shape[0] == 4 + assert torch.equal(state.amax_history, expected_history) + assert destination.fp8_meta["recipe"].delayed() + finally: + FP8GlobalStateManager.reset() + + def test_apply_recipe_dispatches_on_the_owner_protocol_not_concrete_classes(): """Any owner implementing the private plan/apply pair participates.""" FP8GlobalStateManager.reset() diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 3ce318169cc..3238384908f 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -5,6 +5,7 @@ """Base modules and utilities for TransformerEngine PyTorch API""" import copy +import dataclasses import io import math import os @@ -1183,6 +1184,20 @@ def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> meta_key ].amax_history[0] + # Callers that resize directly (CUDA graph capture, the built-in DPA + # path) bypass the planner, so re-key the runtime to the length that is + # now installed. Otherwise the next forward compares against a stale + # configuration and rejects a module it has already resized. + runtime = getattr(self, "_quantization_runtime", None) + if runtime is not None and runtime.recipe.delayed() and not runtime.recipe.custom(): + recipe = copy.copy(runtime.recipe) + recipe.amax_history_len = length + runtime.recipe = recipe + runtime.key = dataclasses.replace(runtime.key, recipe_config=recipe.quantizer_config()) + runtime.forward_quantizers = self.quantizers["scaling_fwd"] + runtime.backward_quantizers = self.quantizers["scaling_bwd"] + self.fp8_meta["recipe"] = recipe + def _resolve_boundary_quantizer_role( self, *, @@ -1830,6 +1845,13 @@ def set_extra_state(self, state: torch.Tensor) -> None: if "global_fp8_buffer_pos_fwd_recompute" in self.fp8_meta: del self.fp8_meta["global_fp8_buffer_pos_fwd_recompute"] + # A checkpoint replaces this module's quantization state wholesale; it is + # not a mid-training update from whatever recipe happens to be active. + # Dropping the runtime first makes the load a first activation, so the + # checkpoint's recipe is adopted instead of being rejected as a + # transition away from the live one. + self.fast_setattr("_quantization_runtime", None) + # Initialize before loading self.init_fp8_meta_tensors(self.fp8_meta["recipe"]) From f8d86f1179c877489c2083c23d6cbd51d895c2b3 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Fri, 11 Sep 2026 15:07:06 +0000 Subject: [PATCH 94/97] Reject recipe changes a CUDA graph would ignore Signed-off-by: Evgeny --- tests/pytorch/test_recipe_update.py | 46 +++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 9 +++++ transformer_engine/pytorch/module/base.py | 10 +++++ 3 files changed, 65 insertions(+) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 0302af79297..d665fc04fd4 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -2653,6 +2653,52 @@ def test_graph_capture_resizes_amax_history_and_leaves_a_consistent_runtime(): FP8GlobalStateManager.reset() +def test_live_cuda_graph_rejects_a_recipe_update_instead_of_ignoring_it(): + """A captured module replays captured kernels, so an update must not be accepted.""" + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + import transformer_engine.pytorch as te + + FP8GlobalStateManager.reset() + try: + module = Linear( + 64, 64, bias=False, params_dtype=torch.bfloat16, device="cuda", name="linear" + ) + inp = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + captured = Float8CurrentScaling() + graphed = te.make_graphed_callables(module, (inp,), fp8_enabled=True, fp8_recipe=captured) + assert module._graph_lease_count == 1 # pylint: disable=protected-access + runtime = module._quantization_runtime # pylint: disable=protected-access + + changed = Float8CurrentScaling() + changed.fp8_quant_fwd_inp = QParams(power_2_scale=True) + + # apply_recipe used to be accepted and then silently ignored by replay. + with pytest.raises(RuntimeError, match="captured by a live CUDA graph"): + apply_recipe(module, changed) + # Replaying under a different recipe was already refused by the graph's own + # capture check; that path is unchanged. + with pytest.raises(RuntimeError, match="differs from the CUDA graph capture recipe"): + with autocast(enabled=True, recipe=changed): + module(inp) + assert module._quantization_runtime is runtime # pylint: disable=protected-access + + # Re-entering the captured recipe is not a change and stays allowed. + with autocast(enabled=True, recipe=captured): + graphed(inp) + + # Releasing the graph releases the lease. + graphed.reset() + assert module._graph_lease_count == 0 # pylint: disable=protected-access + apply_recipe(module, changed) + # pylint: disable-next=protected-access + assert module._quantization_runtime is not runtime + finally: + FP8GlobalStateManager.reset() + + @pytest.mark.parametrize( "live_recipe", [ diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 952d4d433e8..c62c7bb15c1 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -1238,6 +1238,12 @@ def make_graphed_attribute_functions(graph_idx, helpers): need_bwd_dw = need_bwd_dw_graph.get(graph_idx, False) te_modules = visited_te_modules.get(graph_idx, set()) + # A captured module executes replayed kernels, so a later recipe update + # would be accepted and then silently ignored. Hold a lease for as long + # as this graph is live; the planner refuses to update a leased module. + for module in te_modules: + module._graph_lease_count = getattr(module, "_graph_lease_count", 0) + 1 + # Attach backward_dw as an attribute to the graphed callable. def backward_dw(): helpers.ensure_not_reset() @@ -1256,6 +1262,9 @@ def backward_dw(): def reset(): nonlocal fwd_graph, bwd_graph, bwd_dw_graph, te_modules + for module in te_modules: + module._graph_lease_count = max(0, getattr(module, "_graph_lease_count", 0) - 1) + for graph in (fwd_graph, bwd_graph, bwd_dw_graph): if graph is not None: graph.reset() diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 3238384908f..937c4c59ead 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1453,6 +1453,16 @@ def _plan_quantization_update( role_revision=role_revision, ) + # An unchanged recipe is fine above; anything past this point would + # replace state a live CUDA graph has already captured, which replay + # would then ignore. Reject instead of accepting a silent no-op. + if getattr(self, "_graph_lease_count", 0): + raise RuntimeError( + "This module's state is captured by a live CUDA graph, so a recipe, " + "quantizer role or GEMM layout change cannot take effect. Call reset() " + "on the graphed callable and recapture before changing it." + ) + self._check_quantization_update_supported( recipe=recipe, requested_key=requested_key, From 001597e9062a1ef9e33431cf1a844bc9c441740d Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 14 Sep 2026 10:24:30 +0000 Subject: [PATCH 95/97] Minor fixes Signed-off-by: Evgeny --- tests/pytorch/test_custom_recipe.py | 150 +++++++++++++++++++++ tests/pytorch/test_recipe_update.py | 2 + transformer_engine/pytorch/__init__.py | 4 +- transformer_engine/pytorch/module/base.py | 33 ++--- transformer_engine/pytorch/quantization.py | 75 ++++------- 5 files changed, 197 insertions(+), 67 deletions(-) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 258e39ef471..c8f23763b1f 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -2,6 +2,8 @@ # # See LICENSE for license information. +from dataclasses import replace + import pytest import torch @@ -1644,6 +1646,154 @@ def test_fp8block_recipe_state_unknown_or_none_role_falls_back_positionally(): assert quantizers[2].block_scaling_dim == fp8_recipe.x_block_scaling_dim +def test_fp8block_recipe_state_per_slot_qparams_are_honoured(): + """Each slot reads its own qparams, not just its own block-scaling dim. + + The recipe defaults give every slot the same ``amax_epsilon`` / + ``power_2_scale``, so the dispatch tests above cannot tell a correct + per-slot mapping from a uniform one. Distinct values per slot can. + """ + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + from transformer_engine.pytorch.quantization import Float8BlockScalingRecipeState + + fp8_recipe = recipe.Float8BlockScaling() + fp8_recipe.fp8_quant_fwd_inp = replace( + fp8_recipe.fp8_quant_fwd_inp, amax_epsilon=0.125, power_2_scale=True + ) + fp8_recipe.fp8_quant_fwd_weight = replace( + fp8_recipe.fp8_quant_fwd_weight, amax_epsilon=0.25, power_2_scale=False + ) + fp8_recipe.fp8_quant_bwd_grad = replace( + fp8_recipe.fp8_quant_bwd_grad, amax_epsilon=0.5, power_2_scale=True + ) + + forward = Float8BlockScalingRecipeState( + fp8_recipe, + mode="forward", + num_quantizers=3, + roles=[ + _fp8block_role("input"), + _fp8block_role("weight"), + _fp8block_role("output"), + ], + ).make_quantizers() + # input, then weight, then the output boundary slot mirroring input. + assert [q.amax_epsilon for q in forward] == [0.125, 0.25, 0.125] + assert [q.force_pow_2_scales for q in forward] == [True, False, True] + + backward = Float8BlockScalingRecipeState( + fp8_recipe, + mode="backward", + num_quantizers=2, + roles=[ + _fp8block_role("grad_output"), + _fp8block_role("grad_input"), + ], + ).make_quantizers() + assert [q.amax_epsilon for q in backward] == [0.5, 0.5] + assert [q.force_pow_2_scales for q in backward] == [True, True] + + +# ---------------------------------------------------------------------- +# Role-aware dispatch in the current-scaling recipe state +# ---------------------------------------------------------------------- + + +def _current_scaling_role(tensor_type): + """QuantizerRole helper for current-scaling tests.""" + return QuantizerRole(module_type="linear", tensor_type=tensor_type, name="t") + + +def test_current_scaling_custom_qparams_are_honoured(): + """Input, weight and grad-output slots each carry their own customized qparams.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + + from transformer_engine.pytorch.quantization import Float8CurrentScalingRecipeState + + cs_recipe = recipe.Float8CurrentScaling() + # Assigned after any shorthand flag so the per-slot values are what survives. + cs_recipe.fp8_quant_fwd_inp = replace( + cs_recipe.fp8_quant_fwd_inp, amax_epsilon=0.125, power_2_scale=True + ) + cs_recipe.fp8_quant_fwd_weight = replace( + cs_recipe.fp8_quant_fwd_weight, amax_epsilon=0.25, power_2_scale=False + ) + cs_recipe.fp8_quant_bwd_grad = replace( + cs_recipe.fp8_quant_bwd_grad, amax_epsilon=0.5, power_2_scale=True + ) + + forward = Float8CurrentScalingRecipeState( + cs_recipe, + mode="forward", + num_quantizers=2, + roles=[ + _current_scaling_role("input"), + _current_scaling_role("weight"), + ], + ).make_quantizers() + assert [q.amax_epsilon for q in forward] == [0.125, 0.25] + assert [q.force_pow_2_scales for q in forward] == [True, False] + + backward = Float8CurrentScalingRecipeState( + cs_recipe, + mode="backward", + num_quantizers=1, + roles=[_current_scaling_role("grad_output")], + ).make_quantizers() + assert backward[0].amax_epsilon == 0.5 + assert backward[0].force_pow_2_scales is True + + +def test_current_scaling_boundary_slots_keep_constructor_defaults(): + """``output`` / ``grad_input`` do not borrow a neighboring slot's qparams. + + Current scaling differs from the block-scaling states here: a boundary slot + falls back to ``use_power_2_scales`` and ``amax_epsilon=0.0`` rather than + reusing the input or gradient configuration. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + + from transformer_engine.pytorch.quantization import Float8CurrentScalingRecipeState + + cs_recipe = recipe.Float8CurrentScaling(use_power_2_scales=False) + cs_recipe.fp8_quant_fwd_inp = replace( + cs_recipe.fp8_quant_fwd_inp, amax_epsilon=0.125, power_2_scale=True + ) + cs_recipe.fp8_quant_bwd_grad = replace( + cs_recipe.fp8_quant_bwd_grad, amax_epsilon=0.5, power_2_scale=True + ) + + output_slot = Float8CurrentScalingRecipeState( + cs_recipe, + mode="forward", + num_quantizers=3, + roles=[ + _current_scaling_role("input"), + _current_scaling_role("weight"), + _current_scaling_role("output"), + ], + ).make_quantizers()[2] + assert output_slot.amax_epsilon == 0.0 + assert output_slot.force_pow_2_scales is False + + grad_input_slot = Float8CurrentScalingRecipeState( + cs_recipe, + mode="backward", + num_quantizers=2, + roles=[ + _current_scaling_role("grad_output"), + _current_scaling_role("grad_input"), + ], + ).make_quantizers()[1] + assert grad_input_slot.amax_epsilon == 0.0 + assert grad_input_slot.force_pow_2_scales is False + + def _nvfp4_role(tensor_type): return QuantizerRole(module_type="linear", tensor_type=tensor_type, name="t") diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index d665fc04fd4..f5a03f44e1f 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -416,6 +416,8 @@ def test_delayed_runtime_rejects_role_and_slot_layout_updates(): _ensure_runtime(module, recipe, num_gemms=2) assert all(current is old for current, old in zip(_runtime_views(module), old_views)) + +def test_stateless_runtime_cannot_acquire_delayed_state(): """An active stateless runtime cannot acquire delayed state mid-training.""" module = Linear(16, 16, bias=False, device="cuda", name="linear") assert _ensure_runtime(module, Float8CurrentScaling()) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 4eb5d136c53..8ca94ef0eaa 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -43,10 +43,10 @@ moe_sort_chunks_by_index, moe_sort_chunks_by_index_with_probs, ) -from transformer_engine.pytorch.quantization import apply_recipe -from transformer_engine.pytorch.quantization import autocast from transformer_engine.pytorch.quantization import fp8_autocast from transformer_engine.pytorch.quantization import fp8_model_init +from transformer_engine.pytorch.quantization import autocast +from transformer_engine.pytorch.quantization import apply_recipe from transformer_engine.pytorch.quantization import quantized_model_init from transformer_engine.pytorch.quantization import is_fp8_available from transformer_engine.pytorch.quantization import is_mxfp8_available diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 937c4c59ead..a90b514efc7 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -904,24 +904,25 @@ class TransformerEngineBaseModule(torch.nn.Module, ABC): Private quantization runtime lifecycle -------------------------------------- ``init_fp8_metadata`` calls ``_ensure_quantization_runtime``. An unchanged - forward returns from its constant-time revision fast path. A cold-path - mismatch follows this staged workflow:: - - _ensure_quantization_runtime - -> _plan_quantization_update - -> resolve roles and requested runtime key - -> recognize an equal semantic runtime - -> reject unsupported transitions - -> _build_quantization_runtime - -> _validate_quantization_runtime - -> _apply_quantization_update - -> synchronize revisions for an equal runtime, or - -> _activate_quantization_runtime for a replacement + forward returns from a constant-time check: the module's role revision, its + GEMM count and the recipe's own configuration object (compared by identity, + then by equality) all still match the committed runtime. Anything else goes + through the cold path:: + + _plan_quantization_update # builds and validates, publishes nothing + -> resolve roles and the requested runtime key + -> return early for a semantically equal runtime + -> reject a change a live CUDA graph has already captured + -> reject unsupported transitions + -> _build_quantization_runtime + -> _validate_quantization_runtime + _apply_quantization_update # cannot fail + -> _activate_quantization_runtime, or re-mark an equal runtime Planning constructs and validates both forward and backward state without - publishing it. Applying a validated update is limited to revision updates - or no-fail publication of precomputed module state. Subclasses customize - candidate validation and activation through their corresponding hooks. + publishing any of it, so a rejected update leaves the module untouched. + Subclasses customize candidate validation and activation through their + corresponding hooks. """ def __init__( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 9e8b035411b..4ce8632dbc4 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -20,7 +20,6 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import ( Recipe, - QParams, DelayedScaling, Format, MXFP8BlockScaling, @@ -106,11 +105,11 @@ def __str__(self) -> str: class _QuantizationRuntimeKey: """Immutable semantic request for one module or operation's quantizers. - ``recipe_config`` is the recipe-owned semantic configuration cached by - :class:`FP8GlobalStateManager`. The role tuples preserve the ordered - forward and backward quantizer-slot layouts for this particular runtime - owner. ``None`` is a meaningful boundary slot and is therefore retained - rather than filtered out. + ``recipe_config`` is the semantic configuration the recipe itself owns and + caches, so it stays identity-stable until the recipe is mutated. The role + tuples preserve the ordered forward and backward quantizer-slot layouts for + this particular runtime owner. ``None`` is a meaningful boundary slot and + is therefore retained rather than filtered out. This is intentionally private: runtime owners use it to decide whether a candidate runtime would construct the same quantizers as the active one. @@ -1510,27 +1509,6 @@ def _slot_tensor_type(self, idx: int) -> str: ) return default_tensor_types[idx % len(default_tensor_types)] - @staticmethod - def _qparams_for_tensor_type( - tensor_type: str, - *, - input_qparams: QParams, - weight_qparams: QParams, - grad_qparams: QParams, - output_uses_input_qparams: bool = True, - grad_input_uses_grad_qparams: bool = True, - ) -> Optional[QParams]: - """Select a recipe qparam bundle for a canonical tensor type.""" - if tensor_type == "weight": - return weight_qparams - if tensor_type == "input" or (tensor_type == "output" and output_uses_input_qparams): - return input_qparams - if tensor_type == "grad_output" or ( - tensor_type == "grad_input" and grad_input_uses_grad_qparams - ): - return grad_qparams - return None - @staticmethod def create( recipe: Recipe, @@ -1737,14 +1715,16 @@ def make_quantizers(self) -> list: from .tensor.float8_tensor import Float8CurrentScalingQuantizer def _make(tensor_type: str) -> Float8CurrentScalingQuantizer: - qparams = self._qparams_for_tensor_type( - tensor_type, - input_qparams=self.recipe.fp8_quant_fwd_inp, - weight_qparams=self.recipe.fp8_quant_fwd_weight, - grad_qparams=self.recipe.fp8_quant_bwd_grad, - output_uses_input_qparams=False, - grad_input_uses_grad_qparams=False, - ) + if tensor_type == "input": + qparams = self.recipe.fp8_quant_fwd_inp + elif tensor_type == "weight": + qparams = self.recipe.fp8_quant_fwd_weight + elif tensor_type == "grad_output": + qparams = self.recipe.fp8_quant_bwd_grad + else: + # "output" and "grad_input" are boundary slots: they keep the + # constructor defaults instead of borrowing a neighbor's qparams. + qparams = None force_pow_2_scales = self.recipe.use_power_2_scales amax_epsilon = 0.0 @@ -1877,22 +1857,18 @@ def make_quantizers(self) -> list: from .tensor.float8_blockwise_tensor import Float8BlockQuantizer def _make(tensor_type: str) -> Float8BlockQuantizer: - qparams = self._qparams_for_tensor_type( - tensor_type, - input_qparams=self.recipe.fp8_quant_fwd_inp, - weight_qparams=self.recipe.fp8_quant_fwd_weight, - grad_qparams=self.recipe.fp8_quant_bwd_grad, - ) - assert qparams is not None if tensor_type == "weight": + qparams = self.recipe.fp8_quant_fwd_weight fp8_dtype = self.qw_dtype block_scaling_dim = self.recipe.w_block_scaling_dim elif tensor_type in ("grad_output", "grad_input"): + qparams = self.recipe.fp8_quant_bwd_grad fp8_dtype = self.qgrad_dtype block_scaling_dim = self.recipe.grad_block_scaling_dim else: # "input", "output", or any unknown forward type fall back to # the input config, matching the legacy positional behavior. + qparams = self.recipe.fp8_quant_fwd_inp fp8_dtype = self.qx_dtype block_scaling_dim = self.recipe.x_block_scaling_dim return Float8BlockQuantizer( @@ -1963,13 +1939,14 @@ def make_quantizers(self) -> list: from .tensor.nvfp4_tensor import NVFP4Quantizer def _make(tensor_type: str) -> NVFP4Quantizer: - qparams = self._qparams_for_tensor_type( - tensor_type, - input_qparams=self.recipe.fp4_quant_fwd_inp, - weight_qparams=self.recipe.fp4_quant_fwd_weight, - grad_qparams=self.recipe.fp4_quant_bwd_grad, - ) - assert qparams is not None + if tensor_type == "weight": + qparams = self.recipe.fp4_quant_fwd_weight + elif tensor_type in ("grad_output", "grad_input"): + qparams = self.recipe.fp4_quant_bwd_grad + else: + # "input", "output", or any unknown forward type fall back to + # the input config, matching the legacy positional behavior. + qparams = self.recipe.fp4_quant_fwd_inp nvfp4_use_4over6 = False if tensor_type not in ("grad_output", "grad_input"): if self.recipe.nvfp4_4over6 == "all": From c2946a1899f821281c94d1d790ce844a9c9b1a65 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Mon, 14 Sep 2026 16:12:03 +0000 Subject: [PATCH 96/97] Misc fixes Signed-off-by: Evgeny --- tests/pytorch/test_custom_recipe.py | 6 +- tests/pytorch/test_recipe.py | 95 +++- tests/pytorch/test_recipe_update.py | 405 +++++++----------- transformer_engine/common/recipe/__init__.py | 236 ++++------ .../dot_product_attention.py | 76 ++-- .../pytorch/attention/multi_head_attention.py | 22 +- transformer_engine/pytorch/module/base.py | 42 +- .../pytorch/module/grouped_linear.py | 4 +- transformer_engine/pytorch/ops/fuser.py | 9 +- transformer_engine/pytorch/ops/op.py | 6 + transformer_engine/pytorch/quantization.py | 43 +- 11 files changed, 416 insertions(+), 528 deletions(-) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index c8f23763b1f..701e9b37eeb 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -66,7 +66,9 @@ def test_first_party_qfactories_have_canonical_qfactory_keys(qfactory): assert len(qfactory_key) == 2 assert isinstance(qfactory_key[0], str) and qfactory_key[0] assert isinstance(qfactory_key[1], int) and qfactory_key[1] > 0 - assert recipe.CustomRecipe(qfactory=qfactory).qfactory_key == qfactory_key + custom_recipe = recipe.CustomRecipe(qfactory=qfactory) + assert custom_recipe.qfactory_key is None + assert dict(custom_recipe.quantizer_config())["qfactory_key"] == qfactory_key @pytest.mark.parametrize("module_type", ["Linear", "LayerNormLinear", "OpsLinear"]) @@ -1674,6 +1676,7 @@ def test_fp8block_recipe_state_per_slot_qparams_are_honoured(): fp8_recipe, mode="forward", num_quantizers=3, + device=torch.device("cpu"), roles=[ _fp8block_role("input"), _fp8block_role("weight"), @@ -1688,6 +1691,7 @@ def test_fp8block_recipe_state_per_slot_qparams_are_honoured(): fp8_recipe, mode="backward", num_quantizers=2, + device=torch.device("cpu"), roles=[ _fp8block_role("grad_output"), _fp8block_role("grad_input"), diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index ad9dd3ee02b..8dbdcac3db9 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -382,8 +382,8 @@ def test_invalid_field_assignment_rejected_before_config( assert recipe_instance.quantizer_config() == valid_config -def test_high_level_recipe_flags_configure_concrete_quantizers_at_construction(): - """Constructor flags must reach derived QParams and concrete quantizers.""" +def test_legacy_scaling_flags_keep_import_time_qparams_defaults(): + """Keep constructor behavior stable until the shorthand semantics change separately.""" current = Float8CurrentScaling(use_power_2_scales=True) current_quantizers = Float8CurrentScalingRecipeState( current, @@ -392,7 +392,10 @@ def test_high_level_recipe_flags_configure_concrete_quantizers_at_construction() device=torch.device("cpu"), roles=[QuantizerRole(tensor_type="input"), QuantizerRole(tensor_type="weight")], ).make_quantizers() - assert all(quantizer.force_pow_2_scales for quantizer in current_quantizers) + assert all( + quantizer.force_pow_2_scales == current.fp8_quant_fwd_inp.power_2_scale + for quantizer in current_quantizers + ) block = Float8BlockScaling(use_f32_scales=True) block_quantizers = Float8BlockScalingRecipeState( @@ -402,7 +405,7 @@ def test_high_level_recipe_flags_configure_concrete_quantizers_at_construction() device=torch.device("cpu"), roles=[QuantizerRole(tensor_type="input"), QuantizerRole(tensor_type="weight")], ).make_quantizers() - assert all(not quantizer.force_pow_2_scales for quantizer in block_quantizers) + assert all(quantizer.force_pow_2_scales for quantizer in block_quantizers) nvfp4 = NVFP4BlockScaling( disable_rht=True, @@ -429,23 +432,23 @@ def test_high_level_recipe_flags_configure_concrete_quantizers_at_construction() assert not nvfp4_backward[0].stochastic_rounding -def test_high_level_recipe_mutation_preserves_unowned_qparams_fields(): - """Convenience fields update their traits without discarding nested customization.""" +def test_high_level_recipe_mutation_preserves_legacy_qparams_semantics(): + """Legacy inert shorthands do not overwrite QParams; NVFP4 shorthands still do.""" current = Float8CurrentScaling() current.fp8_quant_fwd_inp = replace(current.fp8_quant_fwd_inp, amax_epsilon=0.125) current.use_power_2_scales = True - assert current.fp8_quant_fwd_inp.power_2_scale + assert not current.fp8_quant_fwd_inp.power_2_scale assert current.fp8_quant_fwd_inp.amax_epsilon == 0.125 - assert current.fp8_quant_fwd_weight.power_2_scale - assert current.fp8_quant_bwd_grad.power_2_scale + assert not current.fp8_quant_fwd_weight.power_2_scale + assert not current.fp8_quant_bwd_grad.power_2_scale block = Float8BlockScaling() block.fp8_quant_fwd_weight = replace(block.fp8_quant_fwd_weight, amax_epsilon=0.25) block.use_f32_scales = True - assert not block.fp8_quant_fwd_inp.power_2_scale - assert not block.fp8_quant_fwd_weight.power_2_scale + assert block.fp8_quant_fwd_inp.power_2_scale + assert block.fp8_quant_fwd_weight.power_2_scale assert block.fp8_quant_fwd_weight.amax_epsilon == 0.25 - assert not block.fp8_quant_bwd_grad.power_2_scale + assert block.fp8_quant_bwd_grad.power_2_scale block_quantizer = Float8BlockScalingRecipeState( block, mode="forward", @@ -453,7 +456,7 @@ def test_high_level_recipe_mutation_preserves_unowned_qparams_fields(): device=torch.device("cpu"), roles=[QuantizerRole(tensor_type="weight")], ).make_quantizers()[0] - assert not block_quantizer.force_pow_2_scales + assert block_quantizer.force_pow_2_scales nvfp4 = NVFP4BlockScaling() nvfp4.fp4_quant_fwd_inp = replace(nvfp4.fp4_quant_fwd_inp, amax_epsilon=0.375) @@ -523,6 +526,17 @@ def test_current_scaling_recipe_state_configures_roles_and_preserves_boundary_de (recipe.use_power_2_scales, 0.0), ] + legacy_dpa_quantizers = Float8CurrentScalingRecipeState( + recipe, + mode="forward", + num_quantizers=3, + device=torch.device("cpu"), + roles=None, + ).make_quantizers() + assert [(q.force_pow_2_scales, q.amax_epsilon) for q in legacy_dpa_quantizers] == [ + (recipe.use_power_2_scales, 0.0), + ] * 3 + @pytest.mark.parametrize( "recipe", @@ -741,8 +755,11 @@ def make_module_owner(module_type, num_gemms): basic_linear = te_ops.BasicLinear(4, 4, device="meta") basic_linear.reset_recipe_state(recipe=recipe) - assert settings(basic_linear._quantizers["forward"]) == [inp, weight] - assert settings(basic_linear._quantizers["backward"]) == [grad] + assert settings(basic_linear._quantizers["forward"]) == [ + (inp[0], 0.0), + (weight[0], 0.0), + ] + assert settings(basic_linear._quantizers["backward"]) == [(grad[0], 0.0)] assert basic_linear.get_quantizer("forward", 0).internal assert basic_linear.get_quantizer("forward", 0).optimize_for_gemm assert basic_linear.get_quantizer("forward", 1).internal @@ -751,8 +768,16 @@ def make_module_owner(module_type, num_gemms): basic_grouped = te_ops.GroupedLinear(2, 4, 4, device="meta") basic_grouped.reset_recipe_state(recipe=recipe) - assert settings(basic_grouped._quantizers["forward"]) == [inp, weight, inp, weight] - assert settings(basic_grouped._quantizers["backward"]) == [grad, grad] + assert settings(basic_grouped._quantizers["forward"]) == [ + (inp[0], 0.0), + (weight[0], 0.0), + (inp[0], 0.0), + (weight[0], 0.0), + ] + assert settings(basic_grouped._quantizers["backward"]) == [ + (grad[0], 0.0), + (grad[0], 0.0), + ] for quantizer in basic_grouped._quantizers["forward"]: assert quantizer.internal for quantizer in basic_grouped._quantizers["backward"]: @@ -843,7 +868,7 @@ def factory(role): # pylint: disable=unused-argument calls.append(("attached", role)) attached = CustomRecipe(qfactory=factory) - assert attached.qfactory_key == ("attached_policy", 1) + assert attached.qfactory_key is None assert dict(attached.quantizer_config())["qfactory_key"] == ("attached_policy", 1) explicit = CustomRecipe(qfactory=factory, qfactory_key=("explicit_policy", 2)) @@ -902,7 +927,7 @@ def test_legacy_pickled_recipe_gains_defaults_and_canonical_attention_flags(): assert restored.backward_override is None assert restored.fp8_dpa is True - assert restored.fp8_quant_fwd_inp == QParams(power_2_scale=True, amax_epsilon=0.0) + assert restored.fp8_quant_fwd_inp == Float8CurrentScaling().fp8_quant_fwd_inp assert ( restored.quantizer_config() == Float8CurrentScaling(use_power_2_scales=True, fp8_mha=True).quantizer_config() @@ -957,13 +982,13 @@ def second_factory(role): def unkeyed_factory(role): raise AssertionError("factory must not be called during key resolution") - # An attached key is picked up at construction and re-derived on replacement. + # An attached key is resolved when semantic configuration is constructed. recipe_instance = CustomRecipe(qfactory=first_factory) - assert recipe_instance.qfactory_key == ("first_factory", 1) + assert recipe_instance.qfactory_key is None first_config = recipe_instance.quantizer_config() recipe_instance.qfactory = second_factory - assert recipe_instance.qfactory_key == ("second_factory", 1) + assert recipe_instance.qfactory_key is None assert recipe_instance.quantizer_config() != first_config # An explicit key describes the factory it was passed with, so it does not @@ -976,6 +1001,31 @@ def unkeyed_factory(role): explicit.quantizer_config() +def test_reassigning_qfactory_refreshes_changed_attached_key_without_losing_explicit_key(): + """Factory metadata and an explicit recipe key cannot make the cache disagree.""" + + @quantizer_factory(key=("decorated", 1)) + def factory(role): + raise AssertionError("factory must not be called during key resolution") + + decorated = CustomRecipe(qfactory=factory) + first_config = decorated.quantizer_config() + factory.qfactory_key = ("decorated", 2) + decorated.qfactory = factory + + assert decorated.qfactory_key is None + assert decorated.quantizer_config() != first_config + assert dict(decorated.quantizer_config())["qfactory_key"] == ("decorated", 2) + + explicit = CustomRecipe(qfactory=factory, qfactory_key=("explicit", 1)) + explicit_config = explicit.quantizer_config() + factory.qfactory_key = ("decorated", 3) + explicit.qfactory = factory + + assert explicit.qfactory_key == ("explicit", 1) + assert explicit.quantizer_config() is explicit_config + + def test_unkeyed_qfactory_error_names_both_ways_to_supply_a_key(): """The rejection is the whole migration path, so it must be actionable.""" @@ -1068,7 +1118,6 @@ def _autocast_activation_state(): qstate.is_first_fp8_module, qstate.fp8_graph_capturing, qstate.autocast_depth, - qstate.abort_amax_reduction, tuple( (key, id(recipe), id(group)) for key, (recipe, group) in qstate.autocast_arguments.items() diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index f5a03f44e1f..581aa23cb93 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -502,149 +502,24 @@ def test_same_delayed_recipe_object_mutation_keeps_committed_snapshot(): assert committed_recipe.margin == 0 -def test_rejected_delayed_update_aborts_autocast_reduction(): - """A caught activation failure cannot update the old delayed tensors on exit.""" - available, reason = is_fp8_available(return_reason=True) - if not available: - pytest.skip(reason) - - module = Linear( - 16, - 16, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - name="linear", - ) - inp = torch.randn(8, 16, device="cuda", dtype=torch.bfloat16) - recipe = DelayedScaling(amax_history_len=4, margin=0) - with autocast(enabled=True, recipe=recipe): - module(inp) - - state = module.fp8_meta["scaling_fwd"] - state.scale.fill_(3) - state.amax_history.fill_(7) - expected_scale = state.scale.clone() - expected_history = state.amax_history.clone() - - recipe.margin = 1 - with autocast(enabled=True, recipe=recipe): - with pytest.raises(RuntimeError, match="do not support delayed scaling"): - module(inp) - - assert torch.equal(state.scale, expected_scale) - assert torch.equal(state.amax_history, expected_history) - - -def test_failed_stateless_activation_aborts_a_registered_delayed_reduction(): - """A caught failure in one module must not update another module's delayed bucket.""" - available, reason = is_fp8_available(return_reason=True) - if not available: - pytest.skip(reason) - - FP8GlobalStateManager.reset() - try: - inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) - delayed_module = Linear( - 16, - 16, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - name="delayed", - ) - with autocast(enabled=True, recipe=DelayedScaling(amax_history_len=4, margin=0)): - delayed_module(inp) - - state = delayed_module.fp8_meta["scaling_fwd"] - state.scale.fill_(3) - state.amax_history.fill_(7) - expected_scale = state.scale.clone() - expected_history = state.amax_history.clone() - - def failing_qfactory(role): - del role - raise RuntimeError("deliberate qfactory failure") - - stateless_module = Linear( - 16, - 16, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - name="stateless", - ) - recipe = CustomRecipe(qfactory=failing_qfactory, qfactory_key=("abort-stateless", 1)) - with autocast(enabled=True, recipe=recipe): - with pytest.raises(RuntimeError, match="deliberate qfactory failure"): - stateless_module(inp) - - assert torch.equal(state.scale, expected_scale) - assert torch.equal(state.amax_history, expected_history) - finally: - FP8GlobalStateManager.reset() - - -def test_uncaught_region_failure_aborts_the_autocast_reduction(): - """An exception that escapes the region leaves its delayed tensors untouched.""" - available, reason = is_fp8_available(return_reason=True) - if not available: - pytest.skip(reason) - - FP8GlobalStateManager.reset() - try: - inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) - module = Linear( - 16, - 16, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - name="linear", - ) - recipe = DelayedScaling(amax_history_len=4, margin=0) - with autocast(enabled=True, recipe=recipe): - module(inp) - - state = module.fp8_meta["scaling_fwd"] - state.scale.fill_(3) - state.amax_history.fill_(7) - expected_scale = state.scale.clone() - expected_history = state.amax_history.clone() - - with pytest.raises(RuntimeError, match="deliberate region failure"): - with autocast(enabled=True, recipe=recipe): - raise RuntimeError("deliberate region failure") - - assert torch.equal(state.scale, expected_scale) - assert torch.equal(state.amax_history, expected_history) - finally: - FP8GlobalStateManager.reset() - - -def test_rejected_fusible_update_aborts_autocast_reduction(monkeypatch): - """The fusible rejection fires before the op reset loop, so it must still abort.""" +def test_region_failure_does_not_skip_autocast_reduction(monkeypatch): + """All ranks retain the DelayedScaling collective schedule after an exception.""" available, reason = is_fp8_available(return_reason=True) if not available: pytest.skip(reason) FP8GlobalStateManager.reset() try: - model = te_ops.Sequential(te_ops.Linear(16, 16, bias=False, device="cuda")) - inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) - with autocast(enabled=True, recipe=DelayedScaling(amax_history_len=4, margin=0)): - model(inp) - reductions = [] monkeypatch.setattr( FP8GlobalStateManager, "reduce_and_update_fp8_tensors", classmethod(lambda _cls, forward=True: reductions.append(forward)), ) - with autocast(enabled=True, recipe=DelayedScaling(amax_history_len=4, margin=1)): - with pytest.raises(RuntimeError, match="not supported for fusible operations"): - model(inp) - assert not reductions + with pytest.raises(RuntimeError, match="deliberate region failure"): + with autocast(enabled=True, recipe=DelayedScaling()): + raise RuntimeError("deliberate region failure") + assert reductions == [True] finally: FP8GlobalStateManager.reset() @@ -910,77 +785,47 @@ def test_same_recipe_object_semantic_mutation_rebuilds_runtime(): assert len(calls) == 10 -def test_current_scaling_flag_mutation_rebuilds_concrete_quantizers(): - """A same-object power-of-two update must change the active quantizers.""" - available, reason = is_fp8_available(return_reason=True) - if not available: - pytest.skip(reason) - - FP8GlobalStateManager.reset() - recipe = Float8CurrentScaling(use_power_2_scales=False) - module = Linear( - 128, - 128, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - name="current", - ) - try: - apply_recipe(module, recipe) - old_runtime = module._quantization_runtime # pylint: disable=protected-access - assert not old_runtime.forward_quantizers[0].force_pow_2_scales - assert not old_runtime.forward_quantizers[1].force_pow_2_scales - assert not old_runtime.backward_quantizers[0].force_pow_2_scales - - recipe.use_power_2_scales = True - apply_recipe(module, recipe) - runtime = module._quantization_runtime # pylint: disable=protected-access - assert runtime is not old_runtime - assert runtime.forward_quantizers[0].force_pow_2_scales - assert runtime.forward_quantizers[1].force_pow_2_scales - assert runtime.backward_quantizers[0].force_pow_2_scales - - inp = torch.randn(128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True) - _run_update_step(module, recipe, inp) - assert module._quantization_runtime is runtime # pylint: disable=protected-access - finally: - FP8GlobalStateManager.reset() - - -@pytest.mark.parametrize(("initial", "target"), ((False, True), (True, False))) -def test_current_scaling_same_object_mutation_matches_fresh_recipe(initial, target): - """Same-object flag updates match fresh target recipes in both directions.""" - available, reason = is_fp8_available(return_reason=True) +def test_float8_block_same_object_layout_mutation_matches_fresh_recipe(): + """A warmed Float8BlockScaling layout update matches a fresh target.""" + available, reason = is_fp8_block_scaling_available(return_reason=True) if not available: pytest.skip(reason) FP8GlobalStateManager.reset() try: _check_same_object_mutation_against_fresh_recipe( - Float8CurrentScaling(use_power_2_scales=initial), - Float8CurrentScaling(use_power_2_scales=target), - {"use_power_2_scales": target}, - seed=6200 + int(target), + Float8BlockScaling(x_block_scaling_dim=1, w_block_scaling_dim=2), + Float8BlockScaling(x_block_scaling_dim=2, w_block_scaling_dim=1), + {"x_block_scaling_dim": 2, "w_block_scaling_dim": 1}, + seed=6250, ) finally: FP8GlobalStateManager.reset() -def test_float8_block_same_object_layout_mutation_matches_fresh_recipe(): - """A warmed Float8BlockScaling layout update matches a fresh target.""" +def test_apply_recipe_rejects_unsupported_block_scales_before_commit(): + """Blackwell scale restrictions are checked before an existing runtime is replaced.""" available, reason = is_fp8_block_scaling_available(return_reason=True) if not available: pytest.skip(reason) + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("FP32 block scales are supported before Blackwell") FP8GlobalStateManager.reset() + module = Linear(128, 128, bias=False, device="cuda", name="block") + initial_recipe = Float8BlockScaling() try: - _check_same_object_mutation_against_fresh_recipe( - Float8BlockScaling(x_block_scaling_dim=1, w_block_scaling_dim=2), - Float8BlockScaling(x_block_scaling_dim=2, w_block_scaling_dim=1), - {"x_block_scaling_dim": 2, "w_block_scaling_dim": 1}, - seed=6250, - ) + apply_recipe(module, initial_recipe) + old_runtime = module._quantization_runtime # pylint: disable=protected-access + old_global_state = _global_recipe_state() + + unsupported_recipe = Float8BlockScaling() + unsupported_recipe.fp8_quant_fwd_inp = QParams(power_2_scale=False) + for _ in range(2): + with pytest.raises(RuntimeError, match="requires power-of-two scales"): + apply_recipe(module, unsupported_recipe) + assert module._quantization_runtime is old_runtime # pylint: disable=protected-access + assert _global_recipe_state() == old_global_state finally: FP8GlobalStateManager.reset() @@ -1060,42 +905,6 @@ def test_native_recipe_transition_matches_fresh_target(): FP8GlobalStateManager.reset() -def test_float8_block_scaling_flag_mutation_rebuilds_concrete_quantizers(): - """A same-object FP32-scale update must change the active quantizers.""" - available, reason = is_fp8_block_scaling_available(return_reason=True) - if not available: - pytest.skip(reason) - if torch.cuda.get_device_capability() >= (10, 0): - pytest.skip("Blackwell FP8 block-scaling emulation requires power-of-two scales") - - FP8GlobalStateManager.reset() - recipe = Float8BlockScaling(use_f32_scales=False) - module = Linear( - 128, - 128, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - name="block", - ) - try: - apply_recipe(module, recipe) - old_runtime = module._quantization_runtime # pylint: disable=protected-access - assert old_runtime.forward_quantizers[0].force_pow_2_scales - assert old_runtime.forward_quantizers[1].force_pow_2_scales - assert old_runtime.backward_quantizers[0].force_pow_2_scales - - recipe.use_f32_scales = True - apply_recipe(module, recipe) - runtime = module._quantization_runtime # pylint: disable=protected-access - assert runtime is not old_runtime - assert not runtime.forward_quantizers[0].force_pow_2_scales - assert not runtime.forward_quantizers[1].force_pow_2_scales - assert not runtime.backward_quantizers[0].force_pow_2_scales - finally: - FP8GlobalStateManager.reset() - - def test_nvfp4_flag_mutations_rebuild_concrete_quantizers(): """Every mutable NVFP4 convenience flag must change its concrete trait.""" if not is_nvfp4_available(): @@ -1549,6 +1358,70 @@ def test_composed_module_executes_after_recipe_update(module_factory): _assert_numerical_traces_match(second_trace, first_trace) +@pytest.mark.parametrize("eager_apply", (False, True), ids=("lazy", "apply-recipe")) +def test_inert_builtin_dpa_recipe_update_matches_fresh_target(eager_apply): + """A warmed BF16-attention MHA accepts an outer built-in recipe update.""" + if not is_nvfp4_available(): + pytest.skip("NVFP4 is not available") + available, reason = is_mxfp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + def make_mha(): + return MultiheadAttention( + hidden_size=128, + num_attention_heads=2, + attention_dropout=0.0, + attn_mask_type="no_mask", + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + name="mha", + ) + + FP8GlobalStateManager.reset() + torch.manual_seed(2026) + switching = make_mha() + reference = make_mha() + reference.load_state_dict(switching.state_dict()) + old_recipe = NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + ) + target_recipe = MXFP8BlockScaling() + + try: + warmup_inp = torch.randn( + 128, + 2, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + _run_numerical_step(switching, old_recipe, warmup_inp) + assert switching.core_attention.fp8_initialized + + if eager_apply: + apply_recipe(switching, target_recipe) + + base_inp = torch.randn(128, 2, 128, device="cuda", dtype=torch.bfloat16) + actual, output_grad = _run_numerical_step( + switching, + target_recipe, + base_inp.detach().clone().requires_grad_(True), + ) + expected, _ = _run_numerical_step( + reference, + target_recipe, + base_inp.detach().clone().requires_grad_(True), + output_grad=output_grad, + ) + _assert_numerical_traces_match(actual, expected) + finally: + FP8GlobalStateManager.reset() + + @pytest.mark.parametrize( "module_factory", ( @@ -1822,8 +1695,8 @@ def test_apply_recipe_rejects_e2e_mha_reproducer_during_planning(): "name", (pytest.param(None, id="unnamed"), pytest.param("mha", id="named")), ) -def test_mha_constructor_declares_literal_dpa_qkv_role(attention_type, input_layernorm, name): - """Every MHA variant constructs and declares the QKV boundary role by literal name. +def test_mha_constructor_uses_consistent_dpa_qkv_role(attention_type, input_layernorm, name): + """Every MHA variant uses one effective name on both sides of its boundaries. Construction alone is the regression guard: ``dpa_name`` used to be referenced before assignment, so every one of these variants raised @@ -1841,12 +1714,10 @@ def test_mha_constructor_declares_literal_dpa_qkv_role(attention_type, input_lay name=name, ) - # Literal, not derived from the DPA child: an unnamed MHA declares "" like - # qkv_name/proj_name do, a named one declares the child's composed name. expected_qkv = QuantizerRole( module_type="dpa", tensor_type="qkv", - name="mha.core_attention" if name is not None else "", + name=module.core_attention.name, ) if attention_type == "self": qkv_producers = [module.layernorm_qkv if input_layernorm else module.qkv] @@ -1866,16 +1737,9 @@ def test_mha_constructor_declares_literal_dpa_qkv_role(attention_type, input_lay num_quantizers=9, boundary_role=None, )[0] - if name is not None: - # Producer-declared and DPA own-slot roles agree on both sides. - assert dpa_own_qkv_role == expected_qkv - else: - # WP9(b) will give unnamed modules "" on both sides. Today the DPA child - # still auto-names itself in __init__, so the two sides disagree; pin that - # so the equality above can be extended deliberately. - assert module.core_attention.name.startswith("Layer_") - assert dpa_own_qkv_role.name == module.core_attention.name - assert dpa_own_qkv_role != expected_qkv + assert dpa_own_qkv_role == expected_qkv + assert module.core_attention.name == f"{module.name}.core_attention" + assert module.proj.name == f"{module.name}.proj" @pytest.mark.parametrize( @@ -2040,9 +1904,9 @@ def runtime_object_ids(owner): initialized_owners = _active_runtime_owners(module) assert initialized_owners == [qkv, module.proj] assert all(owner.primary_weights_in_fp8 for owner in initialized_owners) - semantic_dpa_name = "" if name is None else module.core_attention.name - semantic_qkv_name = "" if name is None else qkv.name - semantic_proj_name = "" if name is None else module.proj.name + semantic_dpa_name = module.core_attention.name + semantic_qkv_name = qkv.name + semantic_proj_name = module.proj.name assert all( owner._role_revision == owner._quantization_runtime.role_revision == 0 for owner in initialized_owners @@ -2623,6 +2487,67 @@ def test_apply_recipe_rejects_legacy_dpa_path(): FP8GlobalStateManager.reset() +def test_rejected_builtin_dpa_update_keeps_planning_and_forward_state_atomic(monkeypatch): + """Unsupported built-in FP8-attention changes never publish partial state.""" + import importlib + + available, reason = is_fp8_available(return_reason=True) + if not available: + pytest.skip(reason) + + dpa_module = importlib.import_module( + "transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention" + ) + monkeypatch.setattr(dpa_module, "_dpa_fp8_recipe", "Float8CurrentScaling") + + FP8GlobalStateManager.reset() + dpa = DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + name="dpa", + ).cuda() + linear = Linear(16, 16, bias=False, device="cuda", name="linear") + active_recipe = Float8CurrentScaling(fp8_dpa=True) + + try: + with autocast(enabled=True, recipe=active_recipe): + dpa.init_fp8_metadata() + assert dpa.fp8_initialized + + metadata_keys = ( + "global_recipe", + "local_recipes", + "recipe", + "scaling_fwd", + "scaling_bwd", + ) + metadata_before = {key: dpa.fp8_meta.get(key) for key in metadata_keys} + + def assert_metadata_unchanged(): + assert all(dpa.fp8_meta.get(key) is value for key, value in metadata_before.items()) + + global_before = _global_recipe_state() + + # Planning an active-attention -> BF16-attention transition rejects before + # the sibling Linear candidate or global recipe can be committed. + with pytest.raises(RuntimeError, match="only through CustomRecipe"): + apply_recipe(torch.nn.ModuleList([linear, dpa]), Float8CurrentScaling()) + assert linear._quantization_runtime is None # pylint: disable=protected-access + assert _global_recipe_state() == global_before + assert_metadata_unchanged() + + changed_recipe = Float8CurrentScaling(fp8_dpa=True) + changed_recipe.fp8_quant_fwd_inp = QParams(amax_epsilon=0.25) + for _ in range(2): + with pytest.raises(RuntimeError, match="built-in DotProductAttention recipe path"): + with autocast(enabled=True, recipe=changed_recipe): + dpa.init_fp8_metadata() + assert_metadata_unchanged() + finally: + FP8GlobalStateManager.reset() + + def test_graph_capture_resizes_amax_history_and_leaves_a_consistent_runtime(): """A direct resize bypasses the planner, so it must re-key the runtime itself.""" available, reason = is_fp8_available(return_reason=True) @@ -2649,6 +2574,8 @@ def test_graph_capture_resizes_amax_history_and_leaves_a_consistent_runtime(): assert module.fp8_meta["scaling_fwd"].amax_history.shape[0] == 16 runtime = module._quantization_runtime # pylint: disable=protected-access assert runtime.recipe.amax_history_len == 16 + assert runtime.forward_state.recipe is runtime.recipe + assert runtime.backward_state.recipe is runtime.recipe assert runtime.key.recipe_config == runtime.recipe.quantizer_config() assert module.fp8_meta["recipe"] is runtime.recipe finally: diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index e4f4385955d..c19921da901 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -149,11 +149,6 @@ def _algorithm_config( _DerivedQParams = Dict[str, Tuple[Tuple[str, str, Callable[[Any], Any]], ...]] -def _unchanged(value: Any) -> Any: - """Derived ``QParams`` field takes the recipe option's value unchanged.""" - return value - - def _negated(value: Any) -> bool: """Derived ``QParams`` field takes the negation of the recipe option's value.""" return not value @@ -169,7 +164,8 @@ def _attached_qfactory_key(qfactory: Any) -> Optional[Hashable]: def _config_factory_key(recipe: Any, value: Any) -> Hashable: """Configuration entry for the quantizer factory: its semantic key alone.""" - if value is None: + key = value if value is not None else _attached_qfactory_key(recipe.qfactory) + if key is None: qfactory = getattr(recipe, "qfactory", None) name = getattr(qfactory, "__qualname__", None) or repr(qfactory) raise ValueError( @@ -183,83 +179,7 @@ def _config_factory_key(recipe: Any, value: Any) -> Hashable: "Bump the revision whenever the factory's behavior changes, so quantizers built from" " the previous behavior are rebuilt instead of silently reused." ) - return value - - -def _config_algorithm(field_name: str, *, half: int) -> Callable[[Any, Any], Any]: - """Converter for one half of a literal-or-callable algorithm field. - - ``half=0`` is the algorithm entry, ``half=1`` its semantic key; each keeps its - own declaration position in the configuration. - """ - key_name = f"{field_name}_key" - - def convert(recipe: Any, value: Any) -> Any: - algorithm = value if half == 0 else getattr(recipe, field_name) - return _algorithm_config(algorithm, getattr(recipe, key_name), field_name=field_name)[half] - - return convert - - -# Semantic-configuration fields per recipe class, populated on first use. Kept in a -# module dict rather than on the class: writing to a class ``__dict__`` is not -# traceable by torch.compile, and a first activation can happen inside a compiled -# region. -_QUANTIZER_CONFIG_FIELDS: Dict[type, tuple] = {} - - -# Converters for field values that are not their own configuration, by value type. -_VALUE_CONFIG_CONVERTERS = ( - (Enum, lambda recipe, value: value.name), - (QParams, lambda recipe, value: _qparams_config(value)), - (MMParams, lambda recipe, value: _mmparams_config(value)), -) - - -def _build_quantizer_config_fields(recipe: Any) -> tuple: - """Resolve how each dataclass field contributes to the semantic configuration. - - The result depends only on the recipe class, so it is computed once per class: - per instance the work is then one attribute read and one ``is None`` test per - field. Entries are ``(entry_name, source_field, converter)`` in declaration - order; a ``None`` converter means the value is already its own configuration. - """ - # ``fields`` is given the instance, not the class: torch.compile can trace the - # instance form, and a recipe's first activation may happen inside a compiled - # region. - recipe_fields = fields(recipe) - field_names = {f.name for f in recipe_fields} - config_fields = [] - for recipe_field in recipe_fields: - name = recipe_field.name - - # The quantizer factory is represented by its semantic key alone: the - # callable itself is deliberately not part of the configuration. - if name == "qfactory": - config_fields.append(("qfactory_key", "qfactory_key", _config_factory_key)) - continue - if name == "qfactory_key": - continue - - # A field with a sibling ``_key`` is a literal-or-callable algorithm. - if f"{name}_key" in field_names: - config_fields.append((name, name, _config_algorithm(name, half=0))) - continue - if name.endswith("_key") and name[: -len("_key")] in field_names: - config_fields.append((name, name, _config_algorithm(name[: -len("_key")], half=1))) - continue - - value = getattr(recipe, name) - converter = next( - ( - convert - for value_type, convert in _VALUE_CONFIG_CONVERTERS - if isinstance(value, value_type) - ), - None, - ) - config_fields.append((name, name, converter)) - return tuple(config_fields) + return _validate_qfactory_key(key, source="CustomRecipe qfactory_key") def _recipe_type_entry(cls: type) -> tuple: @@ -488,19 +408,35 @@ def _make_quantizer_config(self) -> Hashable: get this for free; adding a field to a recipe therefore cannot leave it out of the semantic configuration. """ - cls = type(self) - cached_fields = _QUANTIZER_CONFIG_FIELDS.get(cls) - if cached_fields is None: - cached_fields = (_recipe_type_entry(cls), _build_quantizer_config_fields(self)) - _QUANTIZER_CONFIG_FIELDS[cls] = cached_fields - recipe_type_entry, config_fields = cached_fields - - values = self.__dict__ - items = [recipe_type_entry] + recipe_fields = fields(self) + field_names = {recipe_field.name for recipe_field in recipe_fields} + items = [_recipe_type_entry(type(self))] append = items.append - for entry_name, source, convert in config_fields: - value = values[source] - append((entry_name, value if convert is None else convert(self, value))) + for recipe_field in recipe_fields: + name = recipe_field.name + + if name == "qfactory": + append(("qfactory_key", _config_factory_key(self, self.qfactory_key))) + continue + if name == "qfactory_key": + continue + + value = getattr(self, name) + key_name = f"{name}_key" + if key_name in field_names: + value = _algorithm_config(value, getattr(self, key_name), field_name=name)[0] + elif name.endswith("_key") and name[: -len("_key")] in field_names: + algorithm_name = name[: -len("_key")] + value = _algorithm_config( + getattr(self, algorithm_name), value, field_name=algorithm_name + )[1] + elif isinstance(value, Enum): + value = value.name + elif isinstance(value, QParams): + value = _qparams_config(value) + elif isinstance(value, MMParams): + value = _mmparams_config(value) + append((name, value)) return tuple(items) @abc.abstractmethod @@ -695,9 +631,15 @@ class Float8CurrentScaling(Recipe): use_power_2_scales: bool = os.getenv("NVTE_FP8_CURRENT_SCALING_POWER_2_SCALES", "0") == "1" fp8_format: Format = Format.HYBRID - fp8_quant_fwd_inp: QParams = field(init=False) - fp8_quant_fwd_weight: QParams = field(init=False) - fp8_quant_bwd_grad: QParams = field(init=False) + fp8_quant_fwd_inp: QParams = field( + default=QParams(power_2_scale=use_power_2_scales, amax_epsilon=0.0), init=False + ) + fp8_quant_fwd_weight: QParams = field( + default=QParams(power_2_scale=use_power_2_scales, amax_epsilon=0.0), init=False + ) + fp8_quant_bwd_grad: QParams = field( + default=QParams(power_2_scale=use_power_2_scales, amax_epsilon=0.0), init=False + ) fp8_gemm_fprop: MMParams = MMParams(use_split_accumulator=False) fp8_gemm_dgrad: MMParams = MMParams(use_split_accumulator=True) fp8_gemm_wgrad: MMParams = MMParams(use_split_accumulator=True) @@ -705,14 +647,6 @@ class Float8CurrentScaling(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) - _DERIVED_QPARAMS: ClassVar[_DerivedQParams] = { - "use_power_2_scales": ( - ("fp8_quant_fwd_inp", "power_2_scale", _unchanged), - ("fp8_quant_fwd_weight", "power_2_scale", _unchanged), - ("fp8_quant_bwd_grad", "power_2_scale", _unchanged), - ), - } - def validate(self) -> None: super().validate() assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." @@ -720,21 +654,6 @@ def validate(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." - def __post_init__(self) -> None: - self.fp8_quant_fwd_inp = QParams( - power_2_scale=self.use_power_2_scales, - amax_epsilon=0.0, - ) - self.fp8_quant_fwd_weight = QParams( - power_2_scale=self.use_power_2_scales, - amax_epsilon=0.0, - ) - self.fp8_quant_bwd_grad = QParams( - power_2_scale=self.use_power_2_scales, - amax_epsilon=0.0, - ) - super().__post_init__() - def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -848,9 +767,15 @@ class Float8BlockScaling(Recipe): use_f32_scales: bool = os.getenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") == "1" fp8_format: Format = Format.E4M3 - fp8_quant_fwd_inp: QParams = field(init=False) - fp8_quant_fwd_weight: QParams = field(init=False) - fp8_quant_bwd_grad: QParams = field(init=False) + fp8_quant_fwd_inp: QParams = field( + default=QParams(power_2_scale=not use_f32_scales, amax_epsilon=0.0), init=False + ) + fp8_quant_fwd_weight: QParams = field( + default=QParams(power_2_scale=not use_f32_scales, amax_epsilon=0.0), init=False + ) + fp8_quant_bwd_grad: QParams = field( + default=QParams(power_2_scale=not use_f32_scales, amax_epsilon=0.0), init=False + ) x_block_scaling_dim: int = 1 w_block_scaling_dim: int = 2 grad_block_scaling_dim: int = 1 @@ -861,14 +786,6 @@ class Float8BlockScaling(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) - _DERIVED_QPARAMS: ClassVar[_DerivedQParams] = { - "use_f32_scales": ( - ("fp8_quant_fwd_inp", "power_2_scale", _negated), - ("fp8_quant_fwd_weight", "power_2_scale", _negated), - ("fp8_quant_bwd_grad", "power_2_scale", _negated), - ), - } - def validate(self) -> None: super().validate() assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x" @@ -894,22 +811,6 @@ def validate(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." - def __post_init__(self) -> None: - power_2_scale = not self.use_f32_scales - self.fp8_quant_fwd_inp = QParams( - power_2_scale=power_2_scale, - amax_epsilon=0.0, - ) - self.fp8_quant_fwd_weight = QParams( - power_2_scale=power_2_scale, - amax_epsilon=0.0, - ) - self.fp8_quant_bwd_grad = QParams( - power_2_scale=power_2_scale, - amax_epsilon=0.0, - ) - super().__post_init__() - def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " @@ -1178,17 +1079,37 @@ def validate(self) -> None: raise ValueError("CustomRecipe quantization_alignment must be positive.") def __setattr__(self, name: str, value: Any) -> None: + if name != "qfactory" or "qfactory" not in self.__dict__: + super().__setattr__(name, value) + return + + current = self.__dict__["qfactory"] + if current is value or current == value: + # Decorator metadata may have changed in place. In that case assigning + # the same factory is the explicit signal to rebuild its semantic config. + if self.qfactory_key is None: + object.__setattr__(self, "_cached_repr", None) + try: + attached_key = _attached_qfactory_key(value) + except (TypeError, ValueError): + object.__setattr__(self, "_cached_quantizer_config", None) + object.__setattr__(self, "_has_cached_quantizer_config", False) + raise + if self._has_cached_quantizer_config: + cached_key = dict(self._cached_quantizer_config)["qfactory_key"] + if cached_key != attached_key: + object.__setattr__(self, "_cached_quantizer_config", None) + object.__setattr__(self, "_has_cached_quantizer_config", False) + return + + # An explicit key describes the factory it was supplied with. Replacing + # that factory requires a new explicit key or decorator metadata. super().__setattr__(name, value) - # A key stands in for one factory's behavior, so it must not outlive the - # factory it came from: replacing the factory re-derives it, and a - # replacement without a key is reported at the next activation. During - # construction the key is resolved by ``__post_init__`` instead. - if name == "qfactory" and "qfactory_key" in self.__dict__: - object.__setattr__(self, "qfactory_key", _attached_qfactory_key(value)) + super().__setattr__("qfactory_key", None) def __post_init__(self) -> None: if self.qfactory_key is None: - self.qfactory_key = _attached_qfactory_key(self.qfactory) + _attached_qfactory_key(self.qfactory) else: self.qfactory_key = _validate_qfactory_key( self.qfactory_key, @@ -1197,10 +1118,13 @@ def __post_init__(self) -> None: super().__post_init__() def _make_repr(self) -> str: + qfactory_key = self.qfactory_key + if qfactory_key is None: + qfactory_key = _attached_qfactory_key(self.qfactory) return ( f"recipe_type={self.__class__.__name__}, " f"qfactory={self.qfactory}, " - f"qfactory_key={self.qfactory_key}, " + f"qfactory_key={qfactory_key}, " f"backward_override={self.backward_override}, " f"quantization_alignment={self.quantization_alignment}" ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 36f02073646..5d989ee3e22 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1143,14 +1143,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: super().init_fp8_metadata(num_gemms=num_gemms) return - try: - self._init_builtin_fp8_metadata(num_gemms, fp8_recipe, _original_recipe) - except BaseException: - # The built-in path bypasses the base init_fp8_metadata abort wrapper, - # so its rejections must still protect a reduction this region may - # complete with delayed tensors this module can no longer update. - FP8GlobalStateManager.abort_current_amax_reduction() - raise + self._init_builtin_fp8_metadata(num_gemms, fp8_recipe, _original_recipe) def _init_builtin_fp8_metadata( self, @@ -1295,18 +1288,14 @@ def _init_builtin_fp8_metadata( # assume attention uses the same fp8_group as GEMMs fp8_group = FP8GlobalStateManager.get_fp8_group() - self.fast_setattr("fp8_parameters", FP8GlobalStateManager.with_fp8_parameters()) - self.fast_setattr("fp8", FP8GlobalStateManager.is_fp8_enabled()) - self.fast_setattr("fp8_calibration", FP8GlobalStateManager.is_fp8_calibration()) - fp8_enabled = self.fp8 or self.fp8_calibration - self.fp8_meta["fp8_checkpoint"] = self.fp8 or self.fp8_calibration + fp8_parameters = FP8GlobalStateManager.with_fp8_parameters() + fp8 = FP8GlobalStateManager.is_fp8_enabled() + fp8_calibration = FP8GlobalStateManager.is_fp8_calibration() + fp8_enabled = fp8 or fp8_calibration local_recipes = fp8_recipes if isinstance(fp8_recipes, List) else [fp8_recipes] previous_local_recipes = self.fp8_meta.get("local_recipes") - if self.fp8_parameters or fp8_enabled: - self.fp8_meta["global_recipe"] = fp8_recipe - self.fp8_meta["local_recipes"] = local_recipes - if self.fp8_parameters or fp8_enabled: + if fp8_parameters or fp8_enabled: # Taking ownership here would strand a CustomRecipe runtime still # pointing at this module's state. The two paths do not hand state # back and forth; with quantization off nothing is published, so this @@ -1319,8 +1308,23 @@ def _init_builtin_fp8_metadata( self.fp8_initialized and previous_local_recipes and previous_local_recipes != local_recipes + and ( + self._builtin_recipes_quantize_attention(previous_local_recipes) + or self._builtin_recipes_quantize_attention(local_recipes) + ) ): self._reject_builtin_dpa_update() + + # Publish only after transition validation. A rejected update must not + # change the baseline used by the next attempt. + self.fast_setattr("fp8_parameters", fp8_parameters) + self.fast_setattr("fp8", fp8) + self.fast_setattr("fp8_calibration", fp8_calibration) + self.fp8_meta["fp8_checkpoint"] = fp8_enabled + + if self.fp8_parameters or fp8_enabled: + self.fp8_meta["global_recipe"] = fp8_recipe + self.fp8_meta["local_recipes"] = local_recipes if self.fp8_initialized and fp8_recipe_dpa == self.fp8_meta["recipe"]: # FP8 init has already been run and recipe is the same, don't do anything. return @@ -1373,6 +1377,11 @@ def _reject_builtin_dpa_update() -> None: "Recreate the module, or use a CustomRecipe throughout." ) + @staticmethod + def _builtin_recipes_quantize_attention(recipes: Optional[List[Recipe]]) -> bool: + """Return whether resolved built-in recipes enable attention quantization.""" + return bool(recipes) and any(recipe.fp8_dpa or recipe.fp8_mha for recipe in recipes) + @staticmethod def _rehash_autocast_arguments(fp8_recipe, fp8_recipe_dpa, fp8_group) -> None: """Re-register reduction arguments when DPA resolved its own recipe.""" @@ -1399,10 +1408,19 @@ def _plan_recipe_update( f"{diagnostic_name!r}." ) runtime = getattr(self, "_quantization_runtime", None) - if runtime is not None and self._runtime_has_delayed_scaling(runtime): - # Committing would hand this module to the built-in path while it - # owns delayed state registered for reduction. + if runtime is not None: + # The forward path cannot hand CustomRecipe-owned state to the + # built-in path. Reject during planning before sibling owners commit. self._reject_builtin_dpa_update() + previous_local_recipes = self.fp8_meta.get("local_recipes") + if self.fp8_initialized and self._builtin_recipes_quantize_attention( + previous_local_recipes + ): + raise RuntimeError( + "te.apply_recipe() supports DotProductAttention only through CustomRecipe; " + "the built-in NVTE_DPA_* path remains lazy and unchanged for " + f"{diagnostic_name!r}." + ) # Quantization-inert for attention: nothing here to update. return None return super()._plan_recipe_update(recipe, diagnostic_name=diagnostic_name) @@ -1438,20 +1456,12 @@ def _invalidate_dpa_runtime_caches() -> None: """Force backend reselection after a committed quantizer replacement.""" _attention_backends["backend_selection_requires_update"] = True - def _activate_quantization_runtime( - self, - candidate: _QuantizationRuntime, - *, - validation_result: Optional[List[Recipe]] = None, - ) -> None: + def _activate_quantization_runtime(self, candidate: _QuantizationRuntime) -> None: """Publish a CustomRecipe DPA runtime and its derived cache state.""" - super()._activate_quantization_runtime( - candidate, - validation_result=validation_result, - ) + super()._activate_quantization_runtime(candidate) self._invalidate_dpa_runtime_caches() - local_recipes = validation_result + local_recipes = candidate.owner_traits if local_recipes is None: # Do not leave labels from an earlier supported quantizer family # attached after a rebuild to an unsupported family. @@ -1462,7 +1472,7 @@ def _activate_quantization_runtime( def _validate_quantization_runtime( self, candidate: _QuantizationRuntime, - ) -> Optional[List[Recipe]]: + ) -> None: """Validate an FP8-attention candidate before publishing it.""" super()._validate_quantization_runtime(candidate) if candidate.recipe.fp8_dpa: @@ -1504,7 +1514,7 @@ def _validate_quantization_runtime( "for DPA tensor types 's' and 'dp'." ) - return _infer_custom_dpa_local_recipes( + candidate.owner_traits = _infer_custom_dpa_local_recipes( candidate.recipe, {"scaling_fwd": candidate.forward_state}, { diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index c928c1868ce..d818c2d5f29 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -372,7 +372,10 @@ def __init__( # quantized_model_init eagerly creates runtimes for weight-owning children. # Pass the final composed topology into their constructors so each initial # runtime is created with the same roles it will use during execution. - dpa_name = name + ".core_attention" if name is not None else "" + # ``_validate_name`` gives unnamed MHA modules a stable effective name. + # Compose every child and boundary role from it so factories see the same + # consumer identity on both sides of each boundary. + dpa_name = f"{self.name}.core_attention" qkv_output_role = QuantizerRole( module_type="dpa", tensor_type="qkv", @@ -407,7 +410,7 @@ def __init__( ub_overlap_ag=ub_overlap_ag, normalization=normalization, ub_name="qkv", - name=name + ".layernorm_linear_qkv" if name is not None else None, + name=f"{self.name}.layernorm_linear_qkv", _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) @@ -421,7 +424,7 @@ def __init__( return_bias=False, parallel_mode=qkv_parallel_mode, parameters_split=parameters_split, - name=name + ".linear_qkv" if name is not None else None, + name=f"{self.name}.linear_qkv", _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) @@ -445,7 +448,7 @@ def __init__( ub_overlap_ag=ub_overlap_ag, normalization=normalization, ub_name="qkv", - name=name + ".layernorm_linear_q" if name is not None else None, + name=f"{self.name}.layernorm_linear_q", _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) @@ -457,6 +460,7 @@ def __init__( bias=bias, return_bias=False, parallel_mode=qkv_parallel_mode, + name=f"{self.name}.linear_q", _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) @@ -468,7 +472,7 @@ def __init__( return_bias=False, parallel_mode=qkv_parallel_mode, parameters_split=("key", "value") if not fuse_qkv_params else None, - name=name + ".linear_kv" if name is not None else None, + name=f"{self.name}.linear_kv", _declared_output_quantizer_role=qkv_output_role, **common_gemm_kwargs, ) @@ -477,8 +481,8 @@ def __init__( self.key_value, ] - qkv_name = qkv_producers[0].name if name is not None else "" - proj_name = name + ".proj" if name is not None else "" + qkv_name = qkv_producers[0].name + proj_name = f"{self.name}.proj" # Attention. self.core_attention = DotProductAttention( @@ -494,7 +498,7 @@ def __init__( layer_number=self.layer_number, attention_type=self.attention_type, softmax_type=self.softmax_type, - name=name + ".core_attention" if name is not None else None, + name=dpa_name, _declared_output_quantizer_role=QuantizerRole( module_type="linear", tensor_type="input", @@ -518,7 +522,7 @@ def __init__( ub_overlap_rs=ub_overlap_rs, ub_overlap_ag=ub_overlap_ag, ub_name="proj", - name=name + ".proj" if name is not None else None, + name=proj_name, _declared_grad_input_quantizer_role=QuantizerRole( module_type="dpa", tensor_type="do", diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index a90b514efc7..55f85bcedcd 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -87,7 +87,6 @@ class _QuantizationUpdate: """A validated runtime update that is ready to commit.""" candidate: Optional[_QuantizationRuntime] - validation_result: Any role_revision: int @@ -1194,6 +1193,8 @@ def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> recipe = copy.copy(runtime.recipe) recipe.amax_history_len = length runtime.recipe = recipe + runtime.forward_state.recipe = recipe + runtime.backward_state.recipe = recipe runtime.key = dataclasses.replace(runtime.key, recipe_config=recipe.quantizer_config()) runtime.forward_quantizers = self.quantizers["scaling_fwd"] runtime.backward_quantizers = self.quantizers["scaling_bwd"] @@ -1322,19 +1323,12 @@ def _build_quantization_runtime( backward_quantizers=backward_quantizers, ) - def _validate_quantization_runtime(self, candidate: _QuantizationRuntime) -> Any: - """Validate a built candidate and return optional precomputed activation data.""" + def _validate_quantization_runtime(self, candidate: _QuantizationRuntime) -> None: + """Validate a built candidate and attach any owner-specific traits to it.""" del candidate - def _activate_quantization_runtime( - self, - candidate: _QuantizationRuntime, - *, - validation_result: Any = None, - ) -> None: + def _activate_quantization_runtime(self, candidate: _QuantizationRuntime) -> None: """Publish a validated replacement runtime and its compatibility views.""" - # Owner-derived state lives on the runtime, so it cannot outlive it. - candidate.owner_traits = validation_result recipe = candidate.recipe forward_state = candidate.forward_state backward_state = candidate.backward_state @@ -1450,7 +1444,6 @@ def _plan_quantization_update( if active is not None and active.key == requested_key and active.num_gemms == num_gemms: return _QuantizationUpdate( candidate=None, - validation_result=None, role_revision=role_revision, ) @@ -1480,13 +1473,9 @@ def _plan_quantization_update( ) if active is not None and self._runtime_has_delayed_scaling(candidate): self._reject_delayed_scaling_update() - # Subclasses may return transient commit data even though the base implementation - # returns None, which pylint cannot infer through dynamic dispatch. - # pylint: disable-next=assignment-from-no-return - validation_result = self._validate_quantization_runtime(candidate) + self._validate_quantization_runtime(candidate) return _QuantizationUpdate( candidate=candidate, - validation_result=validation_result, role_revision=role_revision, ) @@ -1525,10 +1514,7 @@ def _apply_quantization_update( active.role_revision = update.role_revision return False - self._activate_quantization_runtime( - candidate, - validation_result=update.validation_result, - ) + self._activate_quantization_runtime(candidate) return True def _ensure_quantization_runtime( @@ -1950,16 +1936,10 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: original_recipe = meta.get("recipe") recipe = FP8GlobalStateManager.get_fp8_recipe() - try: - runtime_changed = self._ensure_quantization_runtime( - recipe=recipe, - num_gemms=num_gemms, - ) - except BaseException: - # Any failure here leaves this module's delayed state unusable for a - # reduction the region may still complete. - FP8GlobalStateManager.abort_current_amax_reduction() - raise + runtime_changed = self._ensure_quantization_runtime( + recipe=recipe, + num_gemms=num_gemms, + ) if fp8_enabled: meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() if fp8_enabled: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 2326f0ce47b..84c56bba0b0 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1825,7 +1825,7 @@ def __init__( def _validate_quantization_runtime( self, candidate: _QuantizationRuntime, - ) -> "_GroupedRuntimeTraits": + ) -> None: """Validate both grouped directions and derive the runtime's traits. Only a ``CustomRecipe`` is checked: the built-in recipes build their own @@ -1872,7 +1872,7 @@ def _validate_quantization_runtime( # Native recipe construction guarantees homogeneous expert quantizers. input_quantizers = (forward[self._offsets["input"]],) - return _GroupedRuntimeTraits( + candidate.owner_traits = _GroupedRuntimeTraits( delayed_scaling_input_quantizer=next( (q for q in input_quantizers if isinstance(q, Float8Quantizer)), None, diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 1e3b626b4ea..bb87520d873 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -759,13 +759,8 @@ def __call__( recipe = FP8GlobalStateManager.get_fp8_recipe() is_grad_enabled = torch.is_grad_enabled() - # Attempt to fuse operations if neccesary. The recipe transition inside - # can reject before the op reset loop runs, so guard the whole call. - try: - self.maybe_fuse_ops(is_grad_enabled, recipe, input, basic_op_extra_inputs) - except BaseException: - FP8GlobalStateManager.abort_current_amax_reduction() - raise + # Attempt to fuse operations if neccesary + self.maybe_fuse_ops(is_grad_enabled, recipe, input, basic_op_extra_inputs) # Initialization before forward for idx, op in enumerate(self._basic_ops): diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index f0ff13f2622..42978a6716a 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -486,6 +486,12 @@ def reset_recipe_state( # Construct builder class for quantized tensors self._quantizers[mode] = recipe_state.make_quantizers() + if recipe.float8_current_scaling(): + # Preserve the legacy te.ops behavior. Its current-scaling + # customization wrote the wrong attribute, so amax epsilon + # remained at the quantizer default. + for quantizer in self._quantizers[mode]: + quantizer.amax_epsilon = 0.0 else: # Update quantization recipe states for mode in ("forward", "backward"): diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 4fceacf6c5c..d69d407f2ab 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -455,7 +455,6 @@ class FP8GlobalState: is_first_fp8_module: bool = False fp8_graph_capturing: bool = False autocast_depth: int = 0 - abort_amax_reduction: bool = False global_amax_buffer: Dict[str, list] = field(default_factory=dict) global_amax_history_buffer: Dict[str, list] = field(default_factory=dict) global_scale_buffer: Dict[str, list] = field(default_factory=dict) @@ -755,8 +754,6 @@ def reduce_and_update_fp8_tensors( """Delayed scaling only. Concatenate, reduce, and split amaxes in the global buffer.""" # global_amax_buffer should only be non-empty for fp8 delayed scaling qstate = cls.quantization_state - if qstate.abort_amax_reduction: - return for ( buffer_key, amax_buffer, @@ -862,36 +859,18 @@ def autocast_enter( qstate.fp8_graph_capturing = _graph if qstate.autocast_depth == 0: - qstate.abort_amax_reduction = False qstate.is_first_fp8_module = True qstate.autocast_depth += 1 - @classmethod - def abort_current_amax_reduction(cls) -> None: - """Prevent delayed-state updates when the active autocast has failed.""" - qstate = cls.quantization_state - if qstate.autocast_depth > 0: - qstate.abort_amax_reduction = True - @classmethod def autocast_exit(cls, enabled: bool, _graph: bool) -> None: """Set state and tracking variables for exit from FP8 region.""" qstate = cls.quantization_state qstate.autocast_depth -= 1 - outermost = qstate.autocast_depth == 0 # Reduce only the non-FP8 weight modules here. # FP8 weight modules are reduced at the end of the optimizer # step after the weight amax is populated. - should_reduce = ( - enabled - and outermost - and not qstate.abort_amax_reduction - and not _graph - and torch.is_grad_enabled() - ) - if outermost: - qstate.abort_amax_reduction = False - if should_reduce: + if enabled and qstate.autocast_depth == 0 and not _graph and torch.is_grad_enabled(): # Delayed scaling only function. For other recipes this is a # no-op because the global amax buffer is empty. cls.reduce_and_update_fp8_tensors(forward=True) @@ -1264,8 +1243,6 @@ def __enter__(self) -> "autocast": return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: - if exc_type is not None: - FP8GlobalStateManager.abort_current_amax_reduction() try: FP8GlobalStateManager.set_autocast_state(self._fp8_state) FP8GlobalStateManager.autocast_exit(self._enabled, _graph=self._graph) @@ -1709,13 +1686,15 @@ def make_quantizers(self) -> list: Input, weight, and grad-output slots use their corresponding recipe qparams. Output and grad-input boundary slots retain the legacy - constructor defaults. Missing or non-canonical roles use the common - positional fallback provided by :meth:`RecipeState._slot_tensor_type`. + constructor defaults. A missing role list also retains constructor + defaults, matching the built-in DPA path. """ from .tensor.float8_tensor import Float8CurrentScalingQuantizer def _make(tensor_type: str) -> Float8CurrentScalingQuantizer: - if tensor_type == "input": + if self.roles is None: + qparams = None + elif tensor_type == "input": qparams = self.recipe.fp8_quant_fwd_inp elif tensor_type == "weight": qparams = self.recipe.fp8_quant_fwd_weight @@ -1832,6 +1811,16 @@ def __init__( if device is None: device = torch.device("cuda") self.device = device + if self.device.type == "cuda" and torch.cuda.get_device_capability(self.device) >= (10, 0): + qparams = ( + recipe.fp8_quant_fwd_inp, + recipe.fp8_quant_fwd_weight, + recipe.fp8_quant_bwd_grad, + ) + if any(not params.power_2_scale for params in qparams): + raise RuntimeError( + "Float8BlockScaling requires power-of-two scales on Blackwell and newer GPUs." + ) def make_quantizers(self) -> list: """Build one ``Float8BlockQuantizer`` per slot, dispatched by tensor type. From e6a841d689079d214d2bbedf25fcf9da45a2191c Mon Sep 17 00:00:00 2001 From: Evgeny Date: Tue, 15 Sep 2026 12:50:04 +0000 Subject: [PATCH 97/97] Fix ci Signed-off-by: Evgeny --- tests/pytorch/test_fusible_ops.py | 4 ++++ tests/pytorch/test_hybrid_quantization.py | 6 +++-- tests/pytorch/test_onnx_export.py | 2 +- tests/pytorch/test_recipe_update.py | 26 +++++++++++++++++++++- transformer_engine/pytorch/quantization.py | 21 ++++++++++++----- 5 files changed, 49 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 5b1aef1c911..0ca79616328 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -4395,6 +4395,10 @@ def test_linear( quantized_compute = quantization is not None maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) maybe_skip_quantization(quantization, dims=out_shape) + if quantized_weight: + fp8_available, reason = te.is_fp8_available(return_reason=True) + if not fp8_available: + pytest.skip(reason) # Construct model recipe = make_recipe(quantization) diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index ecb58dd408c..ccdb6570c68 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -506,8 +506,10 @@ def test_direct_runtime_commit_updates_dpa_derived_state(self): assert dpa._apply_quantization_update(update) assert dpa._quantization_runtime is update.candidate - assert dpa.fp8_meta["local_recipes"] is update.validation_result - assert [type(item).__name__ for item in update.validation_result] == ["MXFP8BlockScaling"] + assert dpa.fp8_meta["local_recipes"] is update.candidate.owner_traits + assert [type(item).__name__ for item in update.candidate.owner_traits] == [ + "MXFP8BlockScaling" + ] assert _canonical_qkv_quantizer(dpa) is not old_qkv_quantizer def test_candidate_failure_preserves_active_runtime_and_caches(self): diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py index 6ddf9cbbbc7..abed418c9fc 100644 --- a/tests/pytorch/test_onnx_export.py +++ b/tests/pytorch/test_onnx_export.py @@ -584,7 +584,7 @@ def _test_export_layernorm_linear( normalization=normalization, ).to(device="cuda") if fp8_recipe is not None: - set_layer_scale(model, scale_factor, num_gemms=2) + set_layer_scale(model, scale_factor, num_gemms=1) do_export(model, inp, fname, fp8_recipe) te_outputs = te_infer(model, inp, is_fp8=fp8_recipe is not None, fp8_recipe=fp8_recipe) diff --git a/tests/pytorch/test_recipe_update.py b/tests/pytorch/test_recipe_update.py index 581aa23cb93..c34011c25ed 100644 --- a/tests/pytorch/test_recipe_update.py +++ b/tests/pytorch/test_recipe_update.py @@ -47,6 +47,8 @@ from transformer_engine.pytorch.tensor.identity_tensor import IdentityQuantizer pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) +requires_fp8 = pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) def _make_counting_recipe(key, calls, *, fail_on_grad_output=False): @@ -736,6 +738,7 @@ def test_module_owns_runtime_slot_layout(module_factory, expected_num_gemms): ) +@requires_fp8 def test_unchanged_forward_uses_revision_hot_path(monkeypatch): """Repeated forwards do not enter any runtime-construction cold path.""" calls = [] @@ -1130,6 +1133,7 @@ def test_runtime_update_workspace_lifecycle( assert not module._fp8_workspaces # pylint: disable=protected-access +@requires_fp8 def test_forward_recipe_mutation_breaks_committed_key_invariant(): """A forward path must not change its committed recipe snapshot.""" @@ -1203,6 +1207,7 @@ def forward(self, inp, *args, **kwargs): ), ), ) +@requires_fp8 def test_module_family_executes_after_recipe_update( module_factory, input_factory, @@ -1456,6 +1461,7 @@ def make_mha(): ) def test_apply_recipe_before_first_mha_forward_has_stable_boundary_topology( module_factory, + monkeypatch, ): """E2E TE-1 regression: first forward must not migrate unwarmed MHA roles.""" from transformer_engine.pytorch.utils import get_device_compute_capability @@ -1466,6 +1472,9 @@ def test_apply_recipe_before_first_mha_forward_has_stable_boundary_topology( compute_capability = get_device_compute_capability() if compute_capability < (9, 0) or compute_capability >= (12, 0): pytest.skip("FP8 attention is not supported on this compute capability") + # This test exercises runtime topology, not backend availability. Keep an + # executable fallback on CI configurations without FP8 fused attention. + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") FP8GlobalStateManager.reset() module = module_factory() @@ -1535,7 +1544,7 @@ def test_apply_recipe_before_first_mha_forward_has_stable_boundary_topology( FP8GlobalStateManager.reset() -def test_mha_warmed_and_unwarmed_topology_matches_and_rope_does_not_migrate_runtime(): +def test_mha_warmed_and_unwarmed_topology_matches_and_rope_does_not_migrate_runtime(monkeypatch): """Cover the TE-1 invariants not asserted by existing attention numerics tests.""" from transformer_engine.pytorch import RotaryPositionEmbedding from transformer_engine.pytorch.utils import get_device_compute_capability @@ -1546,6 +1555,9 @@ def test_mha_warmed_and_unwarmed_topology_matches_and_rope_does_not_migrate_runt compute_capability = get_device_compute_capability() if compute_capability < (9, 0) or compute_capability >= (12, 0): pytest.skip("FP8 attention is not supported on this compute capability") + # This test exercises runtime topology, not backend availability. Keep an + # executable fallback on CI configurations without FP8 fused attention. + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") FP8GlobalStateManager.reset() recipe = _make_compatible_fp8_mha_recipe(("mha-warm-rope-topology", 1)) @@ -1952,6 +1964,7 @@ def runtime_object_ids(owner): FP8GlobalStateManager.reset() +@requires_fp8 def test_explicit_mha_boundary_override_wins_and_can_restore_declared_role(): """Public role overrides retain revision semantics above declared topology.""" FP8GlobalStateManager.reset() @@ -2064,6 +2077,7 @@ def test_shipped_dpa_factory_unwarmed_mha_keeps_bf16_boundaries_stable(): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_rejects_quantized_primary_update_before_candidate_construction(): """The model-wide API preserves the quantized-model-init support boundary.""" FP8GlobalStateManager.reset() @@ -2104,6 +2118,7 @@ def test_apply_recipe_rejects_quantized_primary_update_before_candidate_construc FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_success_noop_mutation_and_forward_fast_path(): """Explicit application publishes once and matching forwards remain lazy-path no-ops.""" FP8GlobalStateManager.reset() @@ -2152,6 +2167,7 @@ def test_apply_recipe_success_noop_mutation_and_forward_fast_path(): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_computes_config_once(monkeypatch): """All participant plans receive one shared semantic configuration.""" FP8GlobalStateManager.reset() @@ -2177,6 +2193,7 @@ def counted_make_config(recipe): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_reconstructs_stateless_runtime_after_checkpoint_restore(): """A restored model can rebuild its committed recipe before its first forward.""" FP8GlobalStateManager.reset() @@ -2219,6 +2236,7 @@ def test_apply_recipe_reconstructs_stateless_runtime_after_checkpoint_restore(): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_includes_unexecuted_conditional_branch(): """Traversal updates owners independently of the branch executed by forward.""" @@ -2260,6 +2278,7 @@ def forward(self, inp, *, use_unexecuted=False): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_supports_custom_recipe_dpa(): """Runtime-managed CustomRecipe DPA participates in model-wide application.""" FP8GlobalStateManager.reset() @@ -2282,6 +2301,7 @@ def test_apply_recipe_supports_custom_recipe_dpa(): @pytest.mark.parametrize("failure_index", (0, 1, 2)) +@requires_fp8 def test_apply_recipe_planning_failure_is_model_wide_atomic(failure_index): """A factory failure in any participant leaves modules and manager unchanged.""" FP8GlobalStateManager.reset() @@ -2357,6 +2377,7 @@ def failing_factory(role): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_deduplicates_shared_runtime_owner(): """A shared module is planned and applied exactly once.""" FP8GlobalStateManager.reset() @@ -2375,6 +2396,7 @@ def test_apply_recipe_deduplicates_shared_runtime_owner(): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_rejects_fusible_owner_without_committing_anything(): """A fusible owner rejects the model, and nothing is committed on the way there. @@ -2411,6 +2433,7 @@ def counting_factory(role): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_skips_zero_quantizer_fusible_owner(): """An operation that builds no quantizers is inert, not an excluded owner.""" FP8GlobalStateManager.reset() @@ -2672,6 +2695,7 @@ def warmed(active): FP8GlobalStateManager.reset() +@requires_fp8 def test_apply_recipe_dispatches_on_the_owner_protocol_not_concrete_classes(): """Any owner implementing the private plan/apply pair participates.""" FP8GlobalStateManager.reset() diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index d69d407f2ab..65b8f105c68 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -292,6 +292,17 @@ def check_recipe_support(recipe: Recipe) -> None: raise RuntimeError(unsupported_reason) +def _validate_recipe_activation(recipe: Recipe) -> None: + """Validate recipe semantics before checking platform support. + + Semantic validation is platform-independent and must not be hidden by a + hardware error. In particular, an invalid ``CustomRecipe`` should report + its missing factory key consistently on every GPU architecture. + """ + recipe.quantizer_config() + check_recipe_support(recipe) + + def get_default_fp8_recipe() -> Recipe: """FP8 recipe with default args.""" assert not torch.compiler.is_compiling(), ( @@ -828,12 +839,12 @@ def _prepare_autocast_enter( ) -> Recipe: """Resolve and validate an autocast activation without publishing state.""" fp8_recipe = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe - if enabled or calibrating: - check_recipe_support(fp8_recipe) if enabled or calibrating or cls.quantization_state.fp8_parameters: # Building the configuration validates the recipe. A region that # quantizes nothing must not reject a recipe it will never use. fp8_recipe.quantizer_config() + if enabled or calibrating: + check_recipe_support(fp8_recipe) return fp8_recipe @classmethod @@ -983,8 +994,7 @@ def apply_recipe(model: torch.nn.Module, recipe: Recipe) -> None: if FP8GlobalStateManager.fp8_graph_capturing(): raise RuntimeError("te.apply_recipe() must be called outside CUDA graph capture.") - check_recipe_support(recipe) - recipe.quantizer_config() + _validate_recipe_activation(recipe) # Dispatch on a private owner protocol rather than concrete classes, so this # function needs no imports from the module/attention/ops packages and any @@ -1104,8 +1114,7 @@ def quantized_model_init( # cannot use must leave the process as it found it. recipe = get_default_fp8_recipe() if recipe is None else recipe if enabled: - check_recipe_support(recipe) - recipe.quantizer_config() + _validate_recipe_activation(recipe) _fp8_parameters = qstate.fp8_parameters _fp8_recipe = qstate.fp8_recipe