diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index c0acf2e6b3..297821f69e 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,114 @@ 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)) +@pytest.mark.parametrize("module_kind", ("linear", "basic_linear")) +def test_primary_weight_layout_with_backward_override( + recipe_name: str, + backward_override: Optional[str], + module_kind: str, +) -> None: + """The recipe determines primary storage, which survives forward/backward.""" + mode_recipe = make_recipe(recipe_name, backward_override=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): + 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 = 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", _primary_weight_recipe_list) +def test_default_primary_weight_storage_allows_quantized_backward_switch( + recipe_name: str, +) -> None: + """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=default_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) + 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) +@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): + 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"): + with te.autocast(enabled=True, recipe=default_recipe): + module(x) + + +@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) @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..ec67795d4f 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -910,6 +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._primary_weights_rowwise_only = False self.fsdp_wrapped = False self.fsdp_group = None self._fp8_workspaces: Dict[str, QuantizedTensor] = {} @@ -1845,7 +1846,14 @@ 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()) + 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._primary_weights_rowwise_only, + ) quantizer.internal = False # HybridQuantizer is included so its current-scaling / NVFP4 # sub-quantizers get the same cross-shard amax reduction as the @@ -2061,6 +2069,12 @@ def _check_weight_tensor_recipe_correspondence(self) -> None: return recipe = self.fp8_meta["recipe"] + 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 " + "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/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 98c67be922..ad0ee99fc3 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -921,14 +921,20 @@ def quantized_model_init( users should call `clear_high_precision_init_val()` to release this CPU memory. 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 + resolved_recipe = get_default_fp8_recipe() if recipe is None else recipe 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 try: yield