[PyTorch] Scope the quantized-param caching flag to its own graph capture

tests/pytorch/test_cuda_graphs.py

@@ -1209,3 +1209,69 @@ def test_make_graphed_callables_with_interleaved_pipeline_parallelism_reused_buf
     )
     assert_all_equal(outputs, graph_outputs)
     assert_all_equal(weights, graph_weights)
+
+
+@pytest.mark.skipif(not fp8_available, reason="FP8 is not supported")
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id)
+def test_make_graphed_callables_weight_caching_does_not_leak(
+    *,
+    fp8_recipe: recipe.Recipe,
+    dtype: torch.dtype = torch.bfloat16,
+) -> None:
+    """A capture with quantized-param caching must not affect later captures.
+
+    The flag tensor gating quantized weight updates is process-global, so a capture
+    without caching that picks it up bakes in a flag nothing ever writes, and then
+    silently reuses a stale quantized weight.
+    """
+    reset_rng_states()
+    FP8GlobalStateManager.reset()
+    hidden, seqlen = 32, 32  # MXFP8 needs both dims divisible by 32
+
+    def build() -> torch.nn.Module:
+        torch.manual_seed(0)
+        model = Linear(hidden, hidden, bias=False, params_dtype=dtype, device="cuda")
+        for param in model.parameters():
+            param.grad = torch.empty_like(param)
+        return model
+
+    def data() -> torch.Tensor:
+        return torch.randn((seqlen, hidden), dtype=dtype, device="cuda")
+
+    def capture(model: torch.nn.Module, caching: bool) -> torch.nn.Module:
+        return make_graphed_callables(
+            model,
+            (data(),),
+            num_warmup_iters=3,
+            enabled=True,
+            recipe=fp8_recipe,
+            cache_quantized_params=caching,
+        )
+
+    def train(model: torch.nn.Module) -> List[torch.Tensor]:
+        outputs = []
+        optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
+        for step in range(2):
+            optimizer.zero_grad(set_to_none=False)
+            for microbatch in range(2):
+                torch.manual_seed(step * 10 + microbatch)
+                with autocast(enabled=True, recipe=fp8_recipe):
+                    out = model(data())
+                out.backward(torch.ones_like(out))
+                outputs.append(out.detach().clone())
+            optimizer.step()
+        return outputs
+
+    # Leave the global flag asking for the weight update to be skipped. Each
+    # microbatch needs its own autocast: the flag is only written for the first
+    # module of a context.
+    cached_model = capture(build(), caching=True)
+    for is_first_microbatch in (True, False):
+        with autocast(enabled=True, recipe=fp8_recipe):
+            out = cached_model(data(), is_first_microbatch=is_first_microbatch)
+        out.backward(torch.ones_like(out))
+
+    # A later capture without caching must still refresh its quantized weight.
+    graphed = train(capture(build(), caching=False))
+    ungraphed = train(build())
+    assert_all_equal(graphed, ungraphed)

transformer_engine/pytorch/graph.py

@@ -371,6 +371,9 @@ def _make_graphed_callables(
                     consumed_sample_q[sample_keys].append(per_callable_fwd_idx)
                 fwd_sample_qs[m_chunk] = fwd_sample_qs[m_chunk][num_consumed_samples:]
 
+    # Scope caching to this capture: the flag tensor below outlives it (replay fills
+    # it in place), so a later capture must not bake in this one's leftover flag.
+    FP8GlobalStateManager.set_caching_quantized_params(cache_quantized_params)
     if cache_quantized_params:
         # Initialize flag that controls FP8 weight updates
         FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(False)
@@ -1609,23 +1612,26 @@ def call_func(self, *args, **kwargs):
     else:
         original_rng_states = torch.cuda.get_rng_state()
 
-    graphed_callables = _make_graphed_callables(
-        forward_funcs,
-        sample_args,
-        num_warmup_iters=num_warmup_iters,
-        allow_unused_input=allow_unused_input,
-        cache_quantized_params=cache_quantized_params,
-        sample_kwargs=sample_kwargs,
-        _order=_order,
-        _num_layers_per_chunk=_num_layers_per_chunk,
-        pool=pool,
-        retain_graph_in_backward=retain_graph_in_backward,
-        _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers,
-        clone_param_grads_on_return=clone_param_grads_on_return,
-        pre_warmup_hook=pre_warmup_hook,
-        post_warmup_hook=post_warmup_hook,
-        capture_time_hooks=capture_time_hooks,
-    )
+    try:
+        graphed_callables = _make_graphed_callables(
+            forward_funcs,
+            sample_args,
+            num_warmup_iters=num_warmup_iters,
+            allow_unused_input=allow_unused_input,
+            cache_quantized_params=cache_quantized_params,
+            sample_kwargs=sample_kwargs,
+            _order=_order,
+            _num_layers_per_chunk=_num_layers_per_chunk,
+            pool=pool,
+            retain_graph_in_backward=retain_graph_in_backward,
+            _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers,
+            clone_param_grads_on_return=clone_param_grads_on_return,
+            pre_warmup_hook=pre_warmup_hook,
+            post_warmup_hook=post_warmup_hook,
+            capture_time_hooks=capture_time_hooks,
+        )
+    finally:
+        FP8GlobalStateManager.set_caching_quantized_params(False)
 
     # Ensures warmup does not affect numerics for ops such as dropout.
     if graph_safe_rng_available():

transformer_engine/pytorch/module/grouped_linear.py

@@ -1872,9 +1872,7 @@ def forward(
         num_gemms = self.num_gemms
 
         if FP8GlobalStateManager.fp8_graph_capturing():
-            skip_fp8_weight_update = (
-                FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
-            )
+            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
         else:
             skip_fp8_weight_update = None
         if skip_fp8_weight_update is not None:
@@ -1893,9 +1891,7 @@ def forward(
             )
 
         if FP8GlobalStateManager.fp8_graph_capturing():
-            skip_fp8_weight_update = (
-                FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
-            )
+            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
         else:
             skip_fp8_weight_update = None
         if skip_fp8_weight_update is not None:

transformer_engine/pytorch/module/layernorm_linear.py

@@ -1706,9 +1706,7 @@ def forward(
         debug = self.is_debug_iter()
 
         if FP8GlobalStateManager.fp8_graph_capturing():
-            skip_fp8_weight_update = (
-                FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
-            )
+            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
         else:
             skip_fp8_weight_update = None
         if skip_fp8_weight_update is not None:

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -2291,9 +2291,7 @@ def forward(
         debug = self.is_debug_iter()
 
         if FP8GlobalStateManager.fp8_graph_capturing():
-            skip_fp8_weight_update = (
-                FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
-            )
+            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
         else:
             skip_fp8_weight_update = None
         if skip_fp8_weight_update is not None:

transformer_engine/pytorch/module/linear.py

@@ -1881,9 +1881,7 @@ def forward(
         debug = self.is_debug_iter()
 
         if FP8GlobalStateManager.fp8_graph_capturing():
-            skip_fp8_weight_update = (
-                FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
-            )
+            skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor()
         else:
             skip_fp8_weight_update = None
         if skip_fp8_weight_update is not None:

transformer_engine/pytorch/quantization.py

@@ -400,6 +400,9 @@ class FP8GlobalState:
         default_factory=dict
     )
     skip_fp8_weight_update_tensor: Optional[torch.Tensor] = None
+    # Whether the graph capture in progress caches quantized params. Scopes the flag
+    # tensor above, which outlives capture and would otherwise leak into later ones.
+    caching_quantized_params: bool = False
 
 
 class FP8GlobalStateManager:
@@ -418,9 +421,19 @@ def set_skip_fp8_weight_update_tensor(cls, skip: bool) -> None:
             )
         cls.quantization_state.skip_fp8_weight_update_tensor.fill_(skip)
 
+    @classmethod
+    def set_caching_quantized_params(cls, caching: bool) -> None:
+        """Mark whether the graph capture in progress caches quantized params"""
+        cls.quantization_state.caching_quantized_params = caching
+
     @classmethod
     def get_skip_fp8_weight_update_tensor(cls) -> Optional[torch.Tensor]:
-        """Get the skip fp8 weight update tensor"""
+        """Get the skip fp8 weight update tensor
+
+        ``None`` unless the capture in progress caches quantized params.
+        """
+        if not cls.quantization_state.caching_quantized_params:
+            return None
         return cls.quantization_state.skip_fp8_weight_update_tensor
 
     @classmethod