From 4fdb04c8fcbe8d7a3c8c914920827301ded457b1 Mon Sep 17 00:00:00 2001 From: zhihaow6 Date: Sat, 1 Aug 2026 03:41:37 -0700 Subject: [PATCH] [PyTorch] Scope the quantized-param caching flag to its own graph capture `make_graphed_callables(cache_quantized_params=True)` creates a process-global flag tensor that gates quantized weight updates, and never scopes it. The modules only check `fp8_graph_capturing()` before reading it, so any *later* capture in the same process picks up the leftover tensor and bakes it in as its own quantize noop flag. Nothing ever writes that flag for a callable graphed without caching, so it keeps whatever the earlier capture left there: if that was "skip", the second module silently reuses a stale quantized weight for the rest of training while its master weight keeps being updated. The flag tensor cannot be cleared or reallocated -- already-captured graphs bake in its address and replay fills it in place -- so gate the read instead. Track whether the capture in progress requested caching and hand out the tensor only then, clearing the scope once capture finishes. Affects every recipe, not a specific one. TE's own graph tests run each case in a fresh process, which is why this never showed up in CI. Signed-off-by: zhihaow6 --- tests/pytorch/test_cuda_graphs.py | 66 +++++++++++++++++++ transformer_engine/pytorch/graph.py | 40 ++++++----- .../pytorch/module/grouped_linear.py | 8 +-- .../pytorch/module/layernorm_linear.py | 4 +- .../pytorch/module/layernorm_mlp.py | 4 +- transformer_engine/pytorch/module/linear.py | 4 +- transformer_engine/pytorch/quantization.py | 15 ++++- 7 files changed, 108 insertions(+), 33 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 5a848dc0e8..78872e3ac2 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/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) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index b298b3d8ff..95070b3a4a 100644 --- a/transformer_engine/pytorch/graph.py +++ b/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(): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 0b51ca6a1e..48725f5776 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/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: diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a588e21a0c..13810b201f 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/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: diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 19e20d775c..e5da435d4e 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/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: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 6b0f941bc4..9b6f11a466 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/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: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index f1bffb122f..f4fe8dce96 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/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