Repository navigation
[MXFP8] Fix compile test configuration - #5809
Conversation
|
|
|
@claude review |
There was a problem hiding this comment.
Verdict: Approve -- please address the items below before merging. This test-only fix is at the right layer: the library raises on purpose when in-place WGRAD accumulation runs under tracing, so the test has to opt out.
What this PR does. Linear.Config defaults to inplace_wgrad_accum=True. uses_inplace_wgrad_accum raises RuntimeError when it sees torch.compiler.is_compiling(). Before this PR, both compile cases of test_mxfp8_linear_runs_outside_plain_eager failed at that guard and never reached the gradient assertions. The PR turns the option off only for the compile cases. The activation-checkpoint cases still exercise the in-place path.
Example.
# on main, any GPU at SM100 or newer
linear = _make_mxfp8_linear(bias=False) # inplace_wgrad_accum=True
torch.compile(linear, fullgraph=True)(x)
# -> RuntimeError: MXFP8Linear has inplace_wgrad_accum=True but is being traced ...
# with this PR: linear.inplace_wgrad_accum = False first -> forward and backward run
Must fix
tests/unit_tests/gpu/quantization/test_mxfp8_linear.py:349-- the comment says "compiled autograd", but the guard fires becausetorch.compiletraces the forward. See the inline suggestion.
Should consider
- PR description -- the failure is not specific to SM103 or any other hardware. On any GPU at SM100 or newer, the test fails on main because of the default config, so the "on SM103" framing and the "does not add hardware support" disclaimer point the reader the wrong way. A clearer summary: "the test compiles a layer built with the default
inplace_wgrad_accum=True, which the library rejects under tracing."
Not verified: the GPU run (needs an SM100+ device). The failure on main follows directly from inplace_wgrad_accum.py:51-59 and linear.py:56.
| ) | ||
|
|
||
| if execution_mode == "compile": | ||
| # Compiled autograd requires the non-in-place WGRAD path. |
There was a problem hiding this comment.
This comment names the wrong mechanism. "Compiled autograd" is a separate feature, and this test doesn't use it. The error comes from uses_inplace_wgrad_accum (torchtitan/models/common/linear/inplace_wgrad_accum.py:54), which raises when torch.compiler.is_compiling() is true. That happens because torch.compile traces the forward, and a traced backward can't read and clear weight.grad. A future reader might go looking for compiled autograd and wonder why it isn't enabled.
| # Compiled autograd requires the non-in-place WGRAD path. | |
| # A traced forward cannot add WGRAD into weight.grad in place. |
Summary (author-provided intent, AI-edited)
Why
What changed
Testing