[PyTorch] Support FP8 weight caching in compiled Linear

tests/pytorch/test_torch_compile.py

@@ -2142,81 +2142,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)
-
-        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()
+    def replay(graph):
+        nonlocal replay_count
+        replay_count += 1
+        return original_replay(graph)
 
-        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

@@ -210,11 +210,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 (
@@ -2363,6 +2360,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()
 
@@ -2389,7 +2393,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.
@@ -2434,9 +2438,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
@@ -2633,7 +2638,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,
@@ -2672,8 +2677,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