Automatically omit unused columnwise primary weights for backward overrides

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(

transformer_engine/pytorch/module/base.py

 with_columnwise_usage = torch.is_grad_enabled() 
 quantizer.set_usage( 
     rowwise=True, 
     columnwise=torch.is_grad_enabled(), 
 ) 
 with_columnwise_usage = torch.is_grad_enabled() 
@@ -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):

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

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

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