From 068bf5d18702c65992aa138df5e4cd6aeed216c7 Mon Sep 17 00:00:00 2001 From: xiuhu17 Date: Wed, 2 Sep 2026 00:07:33 -0700 Subject: [PATCH 1/3] Avoid unused columnwise primary weights Backward overrides run dgrad with high-precision or dequantized weights, so primary quantized parameters do not need a columnwise representation. Initialize only rowwise storage for these modes while preserving bidirectional storage for quantized backward. Add coverage for MXFP8 and 1D NVFP4 before and after backward. Signed-off-by: xiuhu17 --- tests/pytorch/test_backward_override.py | 53 +++++++++++++++++++++++ transformer_engine/pytorch/module/base.py | 10 ++++- 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index c0acf2e6b3..9c64fb10ef 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -90,6 +90,19 @@ ), ] +_primary_weight_recipe_list = [ + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), + id="MXFP8BlockScaling", + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), + id="NVFP4BlockScaling1D", + ), +] + @pytest.fixture(autouse=True) def _reset_global_fp8_state(): @@ -858,6 +871,46 @@ def test_backward_override_recipe_matches_requested_mode( assert quant_recipe.backward_override is None +@pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) +@pytest.mark.parametrize("backward_override", (None, *_BACKWARD_OVERRIDES)) +def test_primary_weight_layout_respects_backward_override( + recipe_name: str, + backward_override: Optional[str], +) -> None: + """Primary weights only keep representations consumed by the configured backward.""" + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override("linear", mode_recipe, backward_override) + + with te.quantized_model_init(enabled=True, recipe=mode_recipe): + module = te.Linear( + 64, + 64, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + ) + + weight = module.weight + expect_columnwise = backward_override is None + + def _check_weight_layout() -> None: + assert weight._rowwise_data is not None + assert weight._rowwise_scale_inv is not None + assert (weight._columnwise_data is not None) == expect_columnwise + assert (weight._columnwise_scale_inv is not None) == expect_columnwise + if hasattr(weight, "_amax_columnwise"): + assert (weight._amax_columnwise is not None) == expect_columnwise + + _check_weight_layout() + + x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + with te.autocast(enabled=True, recipe=mode_recipe): + y = module(x) + y.sum().backward() + + _check_weight_layout() + + @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) @pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias")) def test_linear_backward_override_dequantized_ignores_save_original_input( diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 59a4d7e08a..d41c540e33 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1845,7 +1845,15 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: quantizer = self.quantizers["scaling_fwd"][fp8_meta_index] if quantizer is None: raise RuntimeError("Weight quantizer has not been initialized") - quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) + # Backward overrides consume high-precision or dequantized weights, + # so they do not need a quantized columnwise representation. + quantizer.set_usage( + rowwise=True, + columnwise=( + torch.is_grad_enabled() + and self.fp8_meta["recipe"].backward_override is None + ), + ) quantizer.internal = False # HybridQuantizer is included so its current-scaling / NVFP4 # sub-quantizers get the same cross-shard amax reduction as the From bae76d173fa569d05719114e3fd4eef7f46efc12 Mon Sep 17 00:00:00 2001 From: xiuhu17 Date: Wed, 2 Sep 2026 00:26:12 -0700 Subject: [PATCH 2/3] Make row-only primary weights opt-in Signed-off-by: xiuhu17 --- tests/pytorch/test_backward_override.py | 82 ++++++++++++++++++++-- transformer_engine/pytorch/module/base.py | 14 ++-- transformer_engine/pytorch/quantization.py | 30 +++++++- 3 files changed, 115 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index 9c64fb10ef..1632eb112a 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -872,16 +872,26 @@ def test_backward_override_recipe_matches_requested_mode( @pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) -@pytest.mark.parametrize("backward_override", (None, *_BACKWARD_OVERRIDES)) -def test_primary_weight_layout_respects_backward_override( +@pytest.mark.parametrize("backward_override", _BACKWARD_OVERRIDES) +@pytest.mark.parametrize( + "omit_columnwise_primary_weight_storage", + (False, True), + ids=("default_storage", "rowwise_only"), +) +def test_primary_weight_layout_with_backward_override( recipe_name: str, - backward_override: Optional[str], + backward_override: str, + omit_columnwise_primary_weight_storage: bool, ) -> None: - """Primary weights only keep representations consumed by the configured backward.""" + """Columnwise primary-weight storage is omitted only when explicitly requested.""" mode_recipe = make_recipe(recipe_name, backward_override=backward_override) skip_unsupported_backward_override("linear", mode_recipe, backward_override) - with te.quantized_model_init(enabled=True, recipe=mode_recipe): + with te.quantized_model_init( + enabled=True, + recipe=mode_recipe, + omit_columnwise_primary_weight_storage=omit_columnwise_primary_weight_storage, + ): module = te.Linear( 64, 64, @@ -891,7 +901,7 @@ def test_primary_weight_layout_respects_backward_override( ) weight = module.weight - expect_columnwise = backward_override is None + expect_columnwise = not omit_columnwise_primary_weight_storage def _check_weight_layout() -> None: assert weight._rowwise_data is not None @@ -911,6 +921,66 @@ def _check_weight_layout() -> None: _check_weight_layout() +@pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) +def test_default_primary_weight_storage_allows_quantized_backward_switch( + recipe_name: str, +) -> None: + """Default storage preserves runtime switches from an override to quantized backward.""" + mode_recipe = make_recipe(recipe_name, backward_override="dequantized") + default_recipe = make_recipe(recipe_name) + + with te.quantized_model_init(enabled=True, recipe=mode_recipe): + module = te.Linear( + 64, + 64, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + ) + + x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + with te.autocast(enabled=True, recipe=default_recipe): + y = module(x) + y.sum().backward() + + +@pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) +def test_rowwise_only_primary_weight_rejects_quantized_backward(recipe_name: str) -> None: + """A rowwise-only primary weight fails before quantized backward requests columnwise data.""" + mode_recipe = make_recipe(recipe_name, backward_override="dequantized") + default_recipe = make_recipe(recipe_name) + + with te.quantized_model_init( + enabled=True, + recipe=mode_recipe, + omit_columnwise_primary_weight_storage=True, + ): + module = te.Linear( + 64, + 64, + bias=False, + params_dtype=torch.bfloat16, + device="cuda", + ) + + x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + with pytest.raises(RuntimeError, match="without columnwise storage"): + with te.autocast(enabled=True, recipe=default_recipe): + module(x) + + +@pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) +def test_rowwise_only_primary_weight_requires_backward_override(recipe_name: str) -> None: + """The rowwise-only opt-in rejects recipes that need quantized backward.""" + with pytest.raises(ValueError, match="requires a recipe with backward_override"): + with te.quantized_model_init( + enabled=True, + recipe=make_recipe(recipe_name), + omit_columnwise_primary_weight_storage=True, + ): + pass + + @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) @pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias")) def test_linear_backward_override_dequantized_ignores_save_original_input( diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index d41c540e33..ebb7fda310 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -910,6 +910,9 @@ def __init__(self, name: Optional[str] = None) -> None: self.param_init_meta = {} self.primary_weights_in_fp8 = FP8GlobalStateManager.with_fp8_parameters() self.preserve_high_precision_init_val = FP8GlobalStateManager.with_high_precision_init_val() + self.omit_columnwise_primary_weight_storage = ( + FP8GlobalStateManager.should_omit_columnwise_primary_weight_storage() + ) self.fsdp_wrapped = False self.fsdp_group = None self._fp8_workspaces: Dict[str, QuantizedTensor] = {} @@ -1845,13 +1848,10 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: quantizer = self.quantizers["scaling_fwd"][fp8_meta_index] if quantizer is None: raise RuntimeError("Weight quantizer has not been initialized") - # Backward overrides consume high-precision or dequantized weights, - # so they do not need a quantized columnwise representation. quantizer.set_usage( rowwise=True, columnwise=( - torch.is_grad_enabled() - and self.fp8_meta["recipe"].backward_override is None + torch.is_grad_enabled() and not self.omit_columnwise_primary_weight_storage ), ) quantizer.internal = False @@ -2069,6 +2069,12 @@ def _check_weight_tensor_recipe_correspondence(self) -> None: return recipe = self.fp8_meta["recipe"] + if self.omit_columnwise_primary_weight_storage and recipe.backward_override is None: + raise RuntimeError( + "Primary weights were initialized without columnwise storage, but the current " + "recipe uses quantized backward. Recreate the model with columnwise primary-weight " + "storage or keep backward_override set to 'high_precision' or 'dequantized'." + ) weight_tensors = [getattr(self, name) for name in self.weight_names] for i, tensor in enumerate(weight_tensors): if isinstance(tensor, QuantizedTensorStorage): diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 98c67be922..278ebf060f 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -399,6 +399,7 @@ class FP8GlobalState: fp8_distributed_group: Optional[dist_group_type] = None fp8_parameters: bool = False high_precision_init_val: bool = False + omit_columnwise_primary_weight_storage: bool = False is_first_fp8_module: bool = False fp8_graph_capturing: bool = False autocast_depth: int = 0 @@ -584,6 +585,11 @@ def with_high_precision_init_val(cls) -> bool: """Should the high precision initial values be stored with FP8 parameters""" return cls.quantization_state.high_precision_init_val + @classmethod + def should_omit_columnwise_primary_weight_storage(cls) -> bool: + """Should quantized primary weights omit columnwise storage""" + return cls.quantization_state.omit_columnwise_primary_weight_storage + @classmethod def fp8_graph_capturing(cls) -> bool: """Is CUDA graph capture under way?""" @@ -879,6 +885,7 @@ def quantized_model_init( enabled: bool = True, recipe: Optional[Recipe] = None, preserve_high_precision_init_val: bool = False, + omit_columnwise_primary_weight_storage: bool = False, ) -> None: """ Context manager for initialization of quantized parameters. @@ -920,6 +927,14 @@ def quantized_model_init( using quantized parameters directly. Note that after the master weights are initialized, users should call `clear_high_precision_init_val()` to release this CPU memory. + This functionality is *EXPERIMENTAL*. + omit_columnwise_primary_weight_storage : bool, default = False + when enabled, initialize quantized primary weights with rowwise storage only. + This requires a recipe with ``backward_override`` set to ``"high_precision"`` + or ``"dequantized"`` and is intended for fixed-mode workloads such as frozen + LoRA base weights. A model initialized this way cannot later use quantized + backward without reconstructing its primary weights with columnwise storage. + This functionality is *EXPERIMENTAL*. """ @@ -927,15 +942,28 @@ def quantized_model_init( _fp8_parameters = qstate.fp8_parameters _fp8_recipe = qstate.fp8_recipe _high_precision_init_val = qstate.high_precision_init_val + _omit_columnwise_primary_weight_storage = qstate.omit_columnwise_primary_weight_storage + resolved_recipe = get_default_fp8_recipe() if recipe is None else recipe + if ( + enabled + and omit_columnwise_primary_weight_storage + and resolved_recipe.backward_override is None + ): + raise ValueError( + "omit_columnwise_primary_weight_storage requires a recipe with " + "backward_override='high_precision' or 'dequantized'" + ) qstate.fp8_parameters = enabled - qstate.fp8_recipe = get_default_fp8_recipe() if recipe is None else recipe + qstate.fp8_recipe = resolved_recipe qstate.high_precision_init_val = preserve_high_precision_init_val + qstate.omit_columnwise_primary_weight_storage = omit_columnwise_primary_weight_storage try: yield finally: qstate.fp8_parameters = _fp8_parameters qstate.fp8_recipe = _fp8_recipe qstate.high_precision_init_val = _high_precision_init_val + qstate.omit_columnwise_primary_weight_storage = _omit_columnwise_primary_weight_storage def fp8_autocast( From 7ac8a08eaf78fcfdd4ee6c7e1338412c1c3aef31 Mon Sep 17 00:00:00 2001 From: xiuhu17 Date: Sat, 5 Sep 2026 22:19:46 -0500 Subject: [PATCH 3/3] Infer primary weight storage from backward recipe --- tests/pytorch/test_backward_override.py | 100 +++++++++--------- transformer_engine/pytorch/module/base.py | 14 +-- .../pytorch/ops/basic/basic_linear.py | 15 ++- .../pytorch/ops/basic/grouped_linear.py | 17 ++- transformer_engine/pytorch/quantization.py | 30 +----- 5 files changed, 90 insertions(+), 86 deletions(-) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index 1632eb112a..297821f69e 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -872,36 +872,26 @@ def test_backward_override_recipe_matches_requested_mode( @pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) -@pytest.mark.parametrize("backward_override", _BACKWARD_OVERRIDES) -@pytest.mark.parametrize( - "omit_columnwise_primary_weight_storage", - (False, True), - ids=("default_storage", "rowwise_only"), -) +@pytest.mark.parametrize("backward_override", (None, *_BACKWARD_OVERRIDES)) +@pytest.mark.parametrize("module_kind", ("linear", "basic_linear")) def test_primary_weight_layout_with_backward_override( recipe_name: str, - backward_override: str, - omit_columnwise_primary_weight_storage: bool, + backward_override: Optional[str], + module_kind: str, ) -> None: - """Columnwise primary-weight storage is omitted only when explicitly requested.""" + """The recipe determines primary storage, which survives forward/backward.""" mode_recipe = make_recipe(recipe_name, backward_override=backward_override) - skip_unsupported_backward_override("linear", mode_recipe, backward_override) + if backward_override is not None: + skip_unsupported_backward_override("linear", mode_recipe, backward_override) - with te.quantized_model_init( - enabled=True, - recipe=mode_recipe, - omit_columnwise_primary_weight_storage=omit_columnwise_primary_weight_storage, - ): - module = te.Linear( - 64, - 64, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - ) + with te.quantized_model_init(enabled=True, recipe=mode_recipe): + if module_kind == "linear": + module = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16, device="cuda") + else: + module = te_ops.BasicLinear(64, 64, dtype=torch.bfloat16, device="cuda") weight = module.weight - expect_columnwise = not omit_columnwise_primary_weight_storage + expect_columnwise = backward_override is None def _check_weight_layout() -> None: assert weight._rowwise_data is not None @@ -925,11 +915,11 @@ def _check_weight_layout() -> None: def test_default_primary_weight_storage_allows_quantized_backward_switch( recipe_name: str, ) -> None: - """Default storage preserves runtime switches from an override to quantized backward.""" + """Weights initialized for quantized backward can enter and leave override mode.""" mode_recipe = make_recipe(recipe_name, backward_override="dequantized") default_recipe = make_recipe(recipe_name) - with te.quantized_model_init(enabled=True, recipe=mode_recipe): + with te.quantized_model_init(enabled=True, recipe=default_recipe): module = te.Linear( 64, 64, @@ -939,29 +929,26 @@ def test_default_primary_weight_storage_allows_quantized_backward_switch( ) x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) - with te.autocast(enabled=True, recipe=default_recipe): - y = module(x) - y.sum().backward() + for runtime_recipe in (mode_recipe, default_recipe): + with te.autocast(enabled=True, recipe=runtime_recipe): + y = module(x) + y.sum().backward() @pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) -def test_rowwise_only_primary_weight_rejects_quantized_backward(recipe_name: str) -> None: +@pytest.mark.parametrize("module_kind", ("linear", "basic_linear")) +def test_rowwise_only_primary_weight_rejects_quantized_backward( + recipe_name: str, module_kind: str +) -> None: """A rowwise-only primary weight fails before quantized backward requests columnwise data.""" mode_recipe = make_recipe(recipe_name, backward_override="dequantized") default_recipe = make_recipe(recipe_name) - with te.quantized_model_init( - enabled=True, - recipe=mode_recipe, - omit_columnwise_primary_weight_storage=True, - ): - module = te.Linear( - 64, - 64, - bias=False, - params_dtype=torch.bfloat16, - device="cuda", - ) + with te.quantized_model_init(enabled=True, recipe=mode_recipe): + if module_kind == "linear": + module = te.Linear(64, 64, bias=False, params_dtype=torch.bfloat16, device="cuda") + else: + module = te_ops.BasicLinear(64, 64, dtype=torch.bfloat16, device="cuda") x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) with pytest.raises(RuntimeError, match="without columnwise storage"): @@ -969,16 +956,27 @@ def test_rowwise_only_primary_weight_rejects_quantized_backward(recipe_name: str module(x) -@pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) -def test_rowwise_only_primary_weight_requires_backward_override(recipe_name: str) -> None: - """The rowwise-only opt-in rejects recipes that need quantized backward.""" - with pytest.raises(ValueError, match="requires a recipe with backward_override"): - with te.quantized_model_init( - enabled=True, - recipe=make_recipe(recipe_name), - omit_columnwise_primary_weight_storage=True, - ): - pass +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("backward_override", (None, *_BACKWARD_OVERRIDES)) +def test_grouped_op_primary_weight_layout(backward_override: Optional[str]) -> None: + """Packed op weights use the same recipe-driven allocation policy as modules. + + This checks allocation, not support for grouped-op override backward. + """ + mode_recipe = make_recipe("mxfp8", backward_override=backward_override) + with te.quantized_model_init(recipe=mode_recipe): + module = te_ops.GroupedLinear(2, 64, 64, bias=False, dtype=torch.bfloat16, device="cuda") + for idx in range(2): + weight = getattr(module, f"weight{idx}") + assert weight._rowwise_data is not None + assert weight._rowwise_scale_inv is not None + assert (weight._columnwise_data is not None) == (backward_override is None) + assert (weight._columnwise_scale_inv is not None) == (backward_override is None) + + if backward_override is not None: + with te.autocast(recipe=make_recipe("mxfp8")): + with pytest.raises(RuntimeError, match="without columnwise storage"): + module.pre_fuser_forward(requires_grad=True) @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index ebb7fda310..ec67795d4f 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -910,9 +910,7 @@ def __init__(self, name: Optional[str] = None) -> None: self.param_init_meta = {} self.primary_weights_in_fp8 = FP8GlobalStateManager.with_fp8_parameters() self.preserve_high_precision_init_val = FP8GlobalStateManager.with_high_precision_init_val() - self.omit_columnwise_primary_weight_storage = ( - FP8GlobalStateManager.should_omit_columnwise_primary_weight_storage() - ) + self._primary_weights_rowwise_only = False self.fsdp_wrapped = False self.fsdp_group = None self._fp8_workspaces: Dict[str, QuantizedTensor] = {} @@ -1848,11 +1846,13 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: quantizer = self.quantizers["scaling_fwd"][fp8_meta_index] if quantizer is None: raise RuntimeError("Weight quantizer has not been initialized") + self._primary_weights_rowwise_only = ( + FP8GlobalStateManager.get_fp8_recipe().backward_override + in ("high_precision", "dequantized") + ) quantizer.set_usage( rowwise=True, - columnwise=( - torch.is_grad_enabled() and not self.omit_columnwise_primary_weight_storage - ), + columnwise=torch.is_grad_enabled() and not self._primary_weights_rowwise_only, ) quantizer.internal = False # HybridQuantizer is included so its current-scaling / NVFP4 @@ -2069,7 +2069,7 @@ def _check_weight_tensor_recipe_correspondence(self) -> None: return recipe = self.fp8_meta["recipe"] - if self.omit_columnwise_primary_weight_storage and recipe.backward_override is None: + if self._primary_weights_rowwise_only and recipe.backward_override is None: raise RuntimeError( "Primary weights were initialized without columnwise storage, but the current " "recipe uses quantized backward. Recreate the model with columnwise primary-weight " diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index cb429055a4..7077cbb76f 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -328,9 +328,13 @@ def reset_parameters(self) -> None: "within quantized_model_init, but the forward pass was not " "performed within autocast." ) + self._primary_weight_rowwise_only = ( + FP8GlobalStateManager.get_fp8_recipe().backward_override + in ("high_precision", "dequantized") + ) quantizer.set_usage( rowwise=True, - columnwise=torch.is_grad_enabled(), + columnwise=torch.is_grad_enabled() and not self._primary_weight_rowwise_only, ) quantizer.internal = False with torch.no_grad(): @@ -347,6 +351,15 @@ def pre_first_fuser_forward(self) -> None: self.reset_parameters() def pre_fuser_forward(self, *, requires_grad: bool) -> None: + if ( + FP8GlobalStateManager.is_fp8_enabled() + and getattr(self, "_primary_weight_rowwise_only", False) + and FP8GlobalStateManager.get_fp8_recipe().backward_override is None + ): + raise RuntimeError( + "Primary weights were initialized without columnwise storage; " + "keep backward_override set to 'high_precision' or 'dequantized'." + ) super().pre_fuser_forward(requires_grad=requires_grad) if FP8GlobalStateManager.is_fp8_enabled(): # Configure quantizer usages diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 9551650045..e22071d95e 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -463,7 +463,13 @@ def reset_parameters(self) -> None: self.get_quantizer("forward", 2 * idx + 1) for idx in range(self.num_groups) ] with_rowwise_usage = True - with_columnwise_usage = torch.is_grad_enabled() + self._primary_weights_rowwise_only = ( + FP8GlobalStateManager.get_fp8_recipe().backward_override + in ("high_precision", "dequantized") + ) + with_columnwise_usage = ( + torch.is_grad_enabled() and not self._primary_weights_rowwise_only + ) for quantizer in quantizers: if quantizer is None: raise RuntimeError( @@ -756,6 +762,15 @@ def pre_first_fuser_forward(self) -> None: ) def pre_fuser_forward(self, *, requires_grad: bool) -> None: + if ( + FP8GlobalStateManager.is_fp8_enabled() + and getattr(self, "_primary_weights_rowwise_only", False) + and FP8GlobalStateManager.get_fp8_recipe().backward_override is None + ): + raise RuntimeError( + "Primary weights were initialized without columnwise storage; " + "keep backward_override set to 'high_precision' or 'dequantized'." + ) super().pre_fuser_forward(requires_grad=requires_grad) if FP8GlobalStateManager.is_fp8_enabled(): # Assume weights have consistent grad requirement diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 278ebf060f..ad0ee99fc3 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -399,7 +399,6 @@ class FP8GlobalState: fp8_distributed_group: Optional[dist_group_type] = None fp8_parameters: bool = False high_precision_init_val: bool = False - omit_columnwise_primary_weight_storage: bool = False is_first_fp8_module: bool = False fp8_graph_capturing: bool = False autocast_depth: int = 0 @@ -585,11 +584,6 @@ def with_high_precision_init_val(cls) -> bool: """Should the high precision initial values be stored with FP8 parameters""" return cls.quantization_state.high_precision_init_val - @classmethod - def should_omit_columnwise_primary_weight_storage(cls) -> bool: - """Should quantized primary weights omit columnwise storage""" - return cls.quantization_state.omit_columnwise_primary_weight_storage - @classmethod def fp8_graph_capturing(cls) -> bool: """Is CUDA graph capture under way?""" @@ -885,7 +879,6 @@ def quantized_model_init( enabled: bool = True, recipe: Optional[Recipe] = None, preserve_high_precision_init_val: bool = False, - omit_columnwise_primary_weight_storage: bool = False, ) -> None: """ Context manager for initialization of quantized parameters. @@ -928,42 +921,27 @@ def quantized_model_init( users should call `clear_high_precision_init_val()` to release this CPU memory. This functionality is *EXPERIMENTAL*. - omit_columnwise_primary_weight_storage : bool, default = False - when enabled, initialize quantized primary weights with rowwise storage only. - This requires a recipe with ``backward_override`` set to ``"high_precision"`` - or ``"dequantized"`` and is intended for fixed-mode workloads such as frozen - LoRA base weights. A model initialized this way cannot later use quantized - backward without reconstructing its primary weights with columnwise storage. - This functionality is *EXPERIMENTAL*. + Recipes with ``backward_override="high_precision"`` or ``"dequantized"`` + automatically omit columnwise primary-weight storage. Such weights must be + reconstructed with columnwise storage before switching to quantized backward + or using an external optimizer that requires both storage directions. """ qstate = FP8GlobalStateManager.quantization_state _fp8_parameters = qstate.fp8_parameters _fp8_recipe = qstate.fp8_recipe _high_precision_init_val = qstate.high_precision_init_val - _omit_columnwise_primary_weight_storage = qstate.omit_columnwise_primary_weight_storage resolved_recipe = get_default_fp8_recipe() if recipe is None else recipe - if ( - enabled - and omit_columnwise_primary_weight_storage - and resolved_recipe.backward_override is None - ): - raise ValueError( - "omit_columnwise_primary_weight_storage requires a recipe with " - "backward_override='high_precision' or 'dequantized'" - ) qstate.fp8_parameters = enabled qstate.fp8_recipe = resolved_recipe qstate.high_precision_init_val = preserve_high_precision_init_val - qstate.omit_columnwise_primary_weight_storage = omit_columnwise_primary_weight_storage try: yield finally: qstate.fp8_parameters = _fp8_parameters qstate.fp8_recipe = _fp8_recipe qstate.high_precision_init_val = _high_precision_init_val - qstate.omit_columnwise_primary_weight_storage = _omit_columnwise_primary_weight_storage def fp8_autocast(