diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index c0acf2e6b3b..726c2e18064 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,158 @@ 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)) +@pytest.mark.parametrize("freeze_weight", (False, True)) +def test_grouped_op_primary_weight_layout( + backward_override: Optional[str], freeze_weight: bool +) -> None: + """Grouped op dgrad still needs both directions, including for a frozen base.""" + 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 + assert weight._columnwise_scale_inv is not None + weight.requires_grad_(not freeze_weight) + + x = torch.randn(64, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + splits = torch.tensor([32, 32], dtype=torch.int32, device="cuda") + for runtime_recipe in (mode_recipe, make_recipe("mxfp8")): + x.grad = None + module.zero_grad(set_to_none=True) + with te.autocast(recipe=runtime_recipe): + y = module(x, splits) + y.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + for idx in range(2): + weight = getattr(module, f"weight{idx}") + assert weight._columnwise_data is not None + assert weight._columnwise_scale_inv is not None + assert (weight.grad is None) == freeze_weight + + +@pytest.mark.parametrize("recipe_name", _primary_weight_recipe_list) +@pytest.mark.parametrize("module_kind", ("linear", "basic_linear")) +@pytest.mark.parametrize("inference_context", (torch.no_grad, torch.inference_mode)) +def test_rowwise_only_primary_weight_allows_inference_recipe_switch( + recipe_name: str, module_kind: str, inference_context, backward_override: str +) -> None: + """An inference-only recipe switch must not require columnwise storage.""" + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override("linear", mode_recipe, backward_override) + with te.quantized_model_init(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 inference_context(), te.autocast(recipe=make_recipe(recipe_name)): + y = module(x) + assert not y.requires_grad + assert torch.isfinite(y).all() + assert module.weight._columnwise_data is None + assert module.weight._columnwise_scale_inv is None + + # Inference must not silently change the primary layout or allow quantized training. + with pytest.raises(RuntimeError, match="without columnwise storage"): + with te.autocast(recipe=make_recipe(recipe_name)): + module(x) + with te.autocast(recipe=mode_recipe): + y = module(x) + y.sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + @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 a3131f7436d..c39a04a0d0d 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,16 @@ def _check_weight_tensor_recipe_correspondence(self) -> None: return recipe = self.fp8_meta["recipe"] + if ( + torch.is_grad_enabled() + and 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 cb429055a4f..468e43a5745 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,16 @@ def pre_first_fuser_forward(self) -> None: self.reset_parameters() def pre_fuser_forward(self, *, requires_grad: bool) -> None: + if ( + requires_grad + and 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 95516500450..74506964193 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -463,6 +463,7 @@ def reset_parameters(self) -> None: self.get_quantizer("forward", 2 * idx + 1) for idx in range(self.num_groups) ] with_rowwise_usage = True + # This op still uses quantized dgrad even under a backward override. with_columnwise_usage = torch.is_grad_enabled() for quantizer in quantizers: if quantizer is None: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 34ae20b4984..11e01c35e16 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -921,14 +921,22 @@ 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 in modules and ops that + support override backward. Op-fuser GroupedLinear retains columnwise storage + because its backward still uses quantized GEMMs. Rowwise-only 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