[PyTorch] Support dynamic recipe updates during training

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)
+.. autoapifunction:: transformer_engine.common.recipe.quantizer_factory
+
+.. autoapiclass:: transformer_engine.common.recipe.CustomRecipe(qfactory, fp8_dpa=False, fp8_mha=False, qfactory_key=None)

docs/api/pytorch.rst

@@ -43,6 +43,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

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)

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)

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"

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

tests/pytorch/hybrid_quantization_utils.py

@@ -37,6 +37,7 @@ def _make_mxfp8(*, fp8_dtype=te.DType.kFloat8E4M3):
     return te.MXFP8Quantizer(fp8_dtype=fp8_dtype)
 
 
+@recipe.quantizer_factory(key=("test_nvfp4_linear_mxfp8_dpa", 1))
 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"
@@ -211,7 +212,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, *, qfactory_key):
     """Build a CustomRecipe with hybrid forward and configurable grad quantizers."""
     if grad_factory is None:
         grad_factory = col_factory
@@ -227,9 +228,10 @@ 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, 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:
@@ -240,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:
@@ -252,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:
@@ -262,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:
@@ -272,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)
@@ -294,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:
@@ -304,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:
@@ -316,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:
@@ -328,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:
@@ -343,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:
@@ -353,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:
@@ -365,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:

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)

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)

tests/pytorch/test_cuda_graphs.py

@@ -753,6 +753,56 @@ 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() -> None:
+    """Accept replays under an equal configuration and reject changed ones."""
+    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,
+    )
+
+    def replay():
+        output = model(generate_data(model_config, dtype))
+        output.backward(generate_data(model_config, dtype, requires_grad=False))
+
+    try:
+        # 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):
+            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)
+
+
 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()