[PyTorch] Support FP8 weight caching in compiled Linear (#3505)

* [PyTorch] Support FP8 weight caching in compiled Linear

Create a fresh FP8 workspace on the first microbatch so the functional Linear op does not mutate a cached input. Reuse the workspace on the remaining microbatches and preserve the fused weight-gradient accumulation fallback.

Add repeated optimizer-step comparisons, deferred backward coverage, and CUDA Graph replay assertions. GPU validation is pending in this draft.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [PyTorch] Remove redundant microbatch cache identity test

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [PyTorch] Measure microbatch cache memory and mark CUDA Graph minibatches

Add an isolated-process memory example covering eager, functional cache refresh, compile and CUDA Graph replay. Mark each complete minibatch as one CUDA Graph iteration so cached weights survive the intervening backward calls.

Linear tests: 18 passed, 12 skipped, 1 existing XPASS on RTX Ada. Four 4096-wide layers peak at 452 MiB in eager and 539 MiB in compile after warmup; creating fresh caches in eager does not raise that global peak.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* [PyTorch] Keep memory diagnostics outside the microbatch PR

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Author: pggPL

tests/pytorch/test_torch_compile.py

@@ -2238,81 +2238,117 @@ def fn(inp):
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
-@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-def test_te_linear_compile_is_first_microbatch():
-    """te.Linear with ``is_first_microbatch`` under torch.compile: FP8 weight
-    caching updates the cached workspace in place, which the functional custom
-    op can't express, so the schedule must fall back to eager -- warning +
-    numerics identical to eager, cache reused in place across steps. The eager
-    reference runs on a separate module so it cannot mask a corrupted or
-    rebuilt cache."""
-    dtype = torch.bfloat16
-    device = "cuda"
-    fp8_recipe = recipe.Float8CurrentScaling()
-    model = te.Linear(64, 32, params_dtype=dtype, device=device)
-    ref_model = te.Linear(64, 32, params_dtype=dtype, device=device)
-    with torch.no_grad():
-        ref_model.weight.copy_(model.weight)
-        ref_model.bias.copy_(model.bias)
-
-    schedule = [True, False, False]
-    is_first = schedule[0]  # rebound each step; closed over by the fns.
+@pytest.mark.parametrize("compile_mode", _compile_modes)
+@pytest.mark.parametrize("deferred_backward", [False, True])
+@pytest.mark.parametrize(
+    "fp8_recipe",
+    [
+        pytest.param(
+            recipe.Float8CurrentScaling(),
+            marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8),
+        ),
+        pytest.param(
+            recipe.Float8BlockScaling(),
+            marks=pytest.mark.skipif(
+                not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling
+            ),
+        ),
+        pytest.param(
+            recipe.MXFP8BlockScaling(),
+            marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8),
+        ),
+        pytest.param(
+            recipe.NVFP4BlockScaling(disable_stochastic_rounding=True),
+            marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4),
+        ),
+    ],
+    ids=recipe_id,
+)
+def test_te_linear_compile_is_first_microbatch(
+    fp8_recipe, compile_mode, deferred_backward, monkeypatch
+):
+    """Reuse cached weights across microbatches and refresh after optimizer updates."""
+    dtype, device = torch.bfloat16, "cuda"
+    model = te.Linear(128, 128, params_dtype=dtype, device=device)
+    ref_model = te.Linear(128, 128, params_dtype=dtype, device=device)
+    ref_model.load_state_dict(model.state_dict())
 
-    def fn(inp):
+    def fn(inp, is_first):
         with te.autocast(recipe=fp8_recipe):
             return model(inp, is_first_microbatch=is_first)
 
-    def ref_fn(inp):
+    def ref_fn(inp, is_first):
         with te.autocast(recipe=fp8_recipe):
             return ref_model(inp, is_first_microbatch=is_first)
 
-    # Eager priming: FP8 state must exist before tracing (creating quantizers
-    # in-graph breaks later recompiles; upstream Dynamo bug).
-    is_first = None
-    fn(torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True))
-    is_first = schedule[0]
+    for forward in (fn, ref_fn):
+        inp = torch.randn(128, 128, dtype=dtype, device=device, requires_grad=True)
+        forward(inp, None).sum().backward()
 
-    torch._dynamo.reset()
-    compiled = torch.compile(fn)
+    replay_count = 0
+    original_replay = torch.cuda.CUDAGraph.replay
 
-    cached_workspace = None
-    for step, is_first in enumerate(schedule):
-        base = torch.randn(32, 64, dtype=dtype, device=device)
+    def replay(graph):
+        nonlocal replay_count
+        replay_count += 1
+        return original_replay(graph)
 
-        inp_ref = base.detach().clone().requires_grad_(True)
-        ref_model.zero_grad(set_to_none=True)
-        out_ref = ref_fn(inp_ref)
-        out_ref.sum().backward()
-
-        inp = base.detach().clone().requires_grad_(True)
-        model.zero_grad(set_to_none=True)
-        if step == 0:
-            with pytest.warns(
-                UserWarning, match="Falling back to eager execution under torch.compile"
-            ):
-                out = compiled(inp).clone()
-        else:
-            out = compiled(inp).clone()
-        out.sum().backward()
+    if compile_mode == "reduce-overhead":
+        monkeypatch.setattr(torch.cuda.CUDAGraph, "replay", replay)
+    torch._dynamo.reset()
+    compiled = torch.compile(fn, fullgraph=True, mode=compile_mode)
 
-        torch.testing.assert_close(out, out_ref.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
-        torch.testing.assert_close(inp.grad, inp_ref.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
-        torch.testing.assert_close(
-            model.weight.grad, ref_model.weight.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL
-        )
+    with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
+        for step in range(5):
+            if compile_mode == "reduce-overhead":
+                torch.compiler.cudagraph_mark_step_begin()
+            model.zero_grad(set_to_none=True)
+            ref_model.zero_grad(set_to_none=True)
+            pending = []
+            replays_before = replay_count
+            for microbatch in range(3):
+                base = torch.randn(128, 128, dtype=dtype, device=device)
+                inp_ref = base.detach().clone().requires_grad_(True)
+                inp = base.detach().clone().requires_grad_(True)
+                is_first = microbatch == 0
+                out_ref = ref_fn(inp_ref, is_first)
+                out = compiled(inp, is_first).clone()
+                torch.testing.assert_close(out, out_ref, atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
+
+                workspace = model._fp8_workspaces["weight"]
+                if is_first:
+                    cached_workspace = workspace
+                else:
+                    assert workspace is cached_workspace
+                torch.testing.assert_close(
+                    workspace.dequantize(),
+                    ref_model._fp8_workspaces["weight"].dequantize(),
+                    atol=_EAGER_ATOL,
+                    rtol=_EAGER_RTOL,
+                )
 
-        workspace = model._fp8_workspaces.get("weight")
-        assert workspace is not None, f"no cached FP8 weight after step {step}"
-        if step == 0:
-            cached_workspace = workspace
-        else:
-            assert workspace is cached_workspace, f"cache rebuilt at step {step}"
+                pending.append((out, out_ref, inp, inp_ref))
+                if not deferred_backward or microbatch == 2:
+                    for result, ref_result, input_tensor, ref_input in pending:
+                        ref_result.sum().backward()
+                        result.sum().backward()
+                        torch.testing.assert_close(
+                            input_tensor.grad, ref_input.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL
+                        )
+                        for param, ref_param in zip(model.parameters(), ref_model.parameters()):
+                            torch.testing.assert_close(
+                                param.grad, ref_param.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL
+                            )
+                    pending.clear()
+            if step == 4 and compile_mode == "reduce-overhead":
+                assert replay_count > replays_before, "CUDA graphs were recorded but never replayed"
+            with torch.no_grad():
+                for param, ref_param in zip(model.parameters(), ref_model.parameters()):
+                    param.add_(param.grad, alpha=-0.001)
+                    ref_param.add_(ref_param.grad, alpha=-0.001)
+            del cached_workspace, workspace, out, out_ref, result, ref_result
 
     torch._dynamo.reset()
-    compiled_fg = torch.compile(fn, fullgraph=True)
-    is_first = True
-    with pytest.raises(Exception, match=re.escape("FP8 weight caching")):
-        compiled_fg(torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True))
 
 
 @pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")

transformer_engine/pytorch/module/linear.py

@@ -211,11 +211,8 @@ def compile_unsupported_reason(self) -> Optional[str]:
             # A quantized dgrad can't cross the op boundary: grads are packed
             # one plain Tensor[] slot each (_pack_bwd_result).
             return "a quantized input grad (fp8_grad=True)"
-        if self.cache_weight and self.fp8:
-            # The cached workspace is updated in place on the first microbatch,
-            # which the functional op (mutates_args=()) can't express. Without
-            # FP8 no workspace exists, so is_first_microbatch is inert.
-            return "FP8 weight caching (is_first_microbatch)"
+        if self.fp8 and self.skip_fp8_weight_update is not None:
+            return "GPU-controlled FP8 weight cache updates"
         if self.fuse_wgrad_accumulation:
             return "fuse_wgrad_accumulation (main_grad)"
         for quantizer in (
@@ -2389,6 +2386,13 @@ def forward(
                              * it also allows skipping gradient accumulation during the
                                first microbatch (since it is the first gradient being
                                produced)
+
+                             Under ``torch.compile``, the FP8 weight cache is replaced on
+                             the first microbatch and reused on subsequent microbatches.
+                             ``fuse_wgrad_accumulation=True`` remains unsupported.
+                             With ``mode="reduce-overhead"``, call
+                             ``torch.compiler.cudagraph_mark_step_begin()`` once before
+                             each minibatch to keep the cache live across its microbatches.
         """
         is_grad_enabled = torch.is_grad_enabled()
 
@@ -2415,7 +2419,7 @@ def forward(
 
         if torch.compiler.is_compiling() and _linear_op is not None:
             reason = self._compile_eager_fallback_reason(
-                inp, is_first_microbatch, fp8_output, fp8_grad, is_grad_enabled, debug
+                inp, skip_fp8_weight_update, fp8_output, fp8_grad, is_grad_enabled, debug
             )
             if reason is not None:
                 # A break inside the try/finally below would skip the whole frame.
@@ -2460,9 +2464,10 @@ def forward(
                         set_quantizer_amax_reduction_group(quantizer, None)
 
             cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight"
-            weight_workspace = (
-                self._fp8_workspaces.get(cache_name) if cache_name is not None else None
-            )
+            # A compiled refresh must not pass the old cache into the functional op.
+            weight_workspace = None
+            if cache_name is not None and not (use_compiled_op and is_first_microbatch):
+                weight_workspace = self._fp8_workspaces.get(cache_name)
 
             dgrad_use_split_accumulator = _2X_ACC_DGRAD
             wgrad_use_split_accumulator = _2X_ACC_WGRAD
@@ -2659,7 +2664,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage
     def _compile_eager_fallback_reason(
         self,
         inp: torch.Tensor,
-        is_first_microbatch: Optional[bool],
+        skip_fp8_weight_update: Optional[torch.Tensor],
         fp8_output: bool,
         fp8_grad: bool,
         is_grad_enabled: bool,
@@ -2698,8 +2703,8 @@ def _compile_eager_fallback_reason(
             and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad)
         ):
             return "a quantized input grad (fp8_grad=True)"
-        if fp8 and is_first_microbatch is not None and not self.is_fsdp2:
-            return "FP8 weight caching (is_first_microbatch)"
+        if fp8 and skip_fp8_weight_update is not None:
+            return "GPU-controlled FP8 weight cache updates"
         return None
 
     @torch._dynamo.disable