Fix fp8_linear compilability - #47623
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| def test_linear_rejects_float32_scales_on_sm100(self): | ||
| # Regression for #47030: rounding a checkpoint's float32 scales up to UE8M0 without requantizing | ||
| # silently corrupts the output on SM100. The forward must reject them before the kernel runs (so | ||
| # `fp8_linear` falls back to Triton) — no kernel op should execute. | ||
| input = torch.randn(4, 128, dtype=torch.bfloat16, device=torch_device) | ||
| weight = torch.randn(256, 128, device=torch_device).to(torch.float8_e4m3fn) | ||
| weight_scale = torch.ones(2, 1, dtype=torch.float32, device=torch_device) | ||
| with self._bundle(is_sm100=True) as captured: | ||
| with self.assertRaisesRegex(NotImplementedError, "float32 scale-factor path"): | ||
| deepgemm_fp8_fp4_linear(input, weight, weight_scale, block_size=(128, 128)) |
There was a problem hiding this comment.
also proposing a regression test for #47030 which was already fixed
There was a problem hiding this comment.
Gotcha, yea fair enough, let's close after merging this PR please
|
@askserge review this pr |
![sergereview[bot]](https://avatars.githubusercontent.com/in/3673521?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤗 Serge says:
This PR cleanly fixes the fp8_linear compilability issue by replacing the fragile allow_in_graph + module-global double-hop for is_sm100() with a direct @torch._dynamo.assume_constant_result decorator, and applies the same decorator to is_deepgemm_loadable() and is_sonicmoe_loadable(). The new SM100 float32-scale guard in fp8_linear correctly routes around the silent-corruption path from #47030, and the added regression tests cover both compile-safety and the backstop rejection. The changes are focused, well-documented, and consistent with the existing assume_constant_result usage in moe.py.
serge v0.1.0 · model: moonshotai/Kimi-K2.7-Code · 32 LLM turns · 36 tool calls · 41.6s · 878536 in / 3872 out tokens
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Smaller comments but overall good, pls double check whether older torch really have that as well 😬
| _IS_SM100 = torch.cuda.get_device_capability()[0] >= 10 | ||
|
|
||
|
|
||
| @torch._dynamo.assume_constant_result |
There was a problem hiding this comment.
Wow I didnt know that existed, is it safe for the older torch versions as well? (does it exist)
There was a problem hiding this comment.
yes it's safe, we also use it in moe.py
| def test_linear_rejects_float32_scales_on_sm100(self): | ||
| # Regression for #47030: rounding a checkpoint's float32 scales up to UE8M0 without requantizing | ||
| # silently corrupts the output on SM100. The forward must reject them before the kernel runs (so | ||
| # `fp8_linear` falls back to Triton) — no kernel op should execute. | ||
| input = torch.randn(4, 128, dtype=torch.bfloat16, device=torch_device) | ||
| weight = torch.randn(256, 128, device=torch_device).to(torch.float8_e4m3fn) | ||
| weight_scale = torch.ones(2, 1, dtype=torch.float32, device=torch_device) | ||
| with self._bundle(is_sm100=True) as captured: | ||
| with self.assertRaisesRegex(NotImplementedError, "float32 scale-factor path"): | ||
| deepgemm_fp8_fp4_linear(input, weight, weight_scale, block_size=(128, 128)) |
There was a problem hiding this comment.
Gotcha, yea fair enough, let's close after merging this PR please
* assumle the loadability util a constant and tests * fix * test for 47030 * fix
* assumle the loadability util a constant and tests * fix * test for 47030 * fix
3 participants
What does this PR do?
is_deepgemm_loadable does not use the double hop allow_in_graph nor assume_constant_result because i thought i only used it in the already protected loader utils, but it's also used in the fp8_linear path which makes it break it during compilation. In this PR i tag every util that returns a trace time python constant with assume_constant_result and add regression tests.
Code Agent Policy
The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. These often are low-quality, or fix extremely minor issues that occur rarely or never in practice.
As a result, we're instituting a rule that first-time contributors should not use code agents to submit PRs or issues.
We'd also ask autonomous "OpenClaw"-like agents not to open any PRs or issues.
Issues/PRs from first-time contributors that violate this rule will probably just be closed without review, and we
might block you, especially if you open more than one or appear to be deliberately ignoring this. We especially do not
want new contributors to jump in on random issues to contribute an agent-written fix. This creates lots of noise
for reviewers and other users and will almost certainly get you blocked.
For more information, please read
CONTRIBUTING.md.Before submitting
Pull Request checks?
to it if that's the case.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.