Skip to content

[MXFP8] Fix compile test configuration - #5809

Merged
tianyu-l merged 1 commit into
pytorch:mainfrom
0z5a:fix/mxfp8-compile-test
Nov 6, 2026
Merged

tianyu-l merged 1 commit into
pytorch:mainfrom
0z5a:fix/mxfp8-compile-test

Conversation

@0z5a

@0z5a 0z5a commented Nov 4, 2026

Copy link
Copy Markdown
Contributor

Summary (author-provided intent, AI-edited)

Improve MXFP8 compile-test execution on SM103 by selecting the supported
weight-gradient accumulation mode. This test-only change also applies to
other SM10x devices and does not add hardware support.

Why

AI-assisted technical notes: test_mxfp8_linear_runs_outside_plain_eager
compiles an MXFP8Linear with its default inplace_wgrad_accum=True.
Traced backward explicitly rejects that configuration, so the two compiled
cases fail before checking the weight gradient. The compile cases need the
supported non-in-place path; the checkpoint cases can retain their current path.

What changed

AI-assisted technical notes: Set linear.inplace_wgrad_accum = False only
when this test chooses torch.compile. This changes the test setup and does
not change the production configuration or the checkpoint cases.

Testing

AI-assisted test notes: On B300 (SM 10.3), with PyTorch
2.16.0.dev20261102+cu130 and Triton 3.9.0, the two compile cases on base
485c065cb248863bf2e2c65588becd6d7a8bcf83 fail at the in-place WGRAD guard.
With this change, the four compile/checkpoint combinations pass, and the
complete test file reports 25 passed. The tested file was a byte-identical
copy of the changed file.

python -m pytest -p no:cacheprovider -q candidate_test_mxfp8_linear.py::test_mxfp8_linear_runs_outside_plain_eager
# 4 passed
python -m pytest -p no:cacheprovider -q candidate_test_mxfp8_linear.py
# 25 passed

The local pre-commit run --all-files run passed the formatting, lint and
other available hooks; the link checker skipped because lychee was absent.
Pyrefly could not pass locally because that environment lacked the project
dependencies. The same Pyrefly 0.45.1 check with project dependencies on an
isolated source copy completed with 0 errors and 2 warnings.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Nov 4, 2026
@0z5a
0z5a marked this pull request as ready for review November 5, 2026 22:01
@pytorch-bot pytorch-bot Bot added the ciflow/fake-pg Run 1-GPU Fake PG integration tests label Nov 5, 2026
@pytorch-bot

pytorch-bot Bot commented Nov 5, 2026 •

Copy link
Copy Markdown

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

@tianyu-l

tianyu-l commented Nov 6, 2026

Copy link
Copy Markdown
Contributor

@claude review

@claude

claude Bot commented Nov 6, 2026 •

Copy link
Copy Markdown
Contributor

Claude finished @tianyu-l's task in 1m 17s —— View job


Approve (please fix the code comment first) -- review

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

  1. tests/unit_tests/gpu/quantization/test_mxfp8_linear.py:349 -- the comment says "compiled autograd", but the guard fires because torch.compile traces 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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
# Compiled autograd requires the non-in-place WGRAD path.
# A traced forward cannot add WGRAD into weight.grad in place.

@tianyu-l
tianyu-l merged commit 40d524a into pytorch:main Nov 6, 2026
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/fake-pg Run 1-GPU Fake PG integration tests CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants