Add flop formula for _scaled_mm #144973
Add flop formula for _scaled_mm #144973
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/144973
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit b8ea47e with merge base 62ce3e6 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
lw
added
ciflow/trunk
| def test_scaled_mm(self): | ||
| dtype = torch.float8_e4m3fnuz if torch.version.hip else torch.float8_e4m3fn | ||
| with FlopCounterMode() as mode: | ||
| torch._scaled_mm( | ||
| torch.randn((3 * 16, 5 * 16), device="cuda").to(dtype), | ||
| torch.randn((7 * 16, 5 * 16), device="cuda").to(dtype).t(), | ||
| scale_a=torch.ones((), device="cuda"), | ||
| scale_b=torch.ones((), device="cuda"), | ||
| out_dtype=torch.bfloat16, | ||
| ) |
There was a problem hiding this comment.
Similar change suggested in previous PR.
docs/source/conf.py
Outdated
| "mm_flop", | ||
| "normalize_tuple", | ||
| "register_flop_formula", | ||
| "scaled_mm_flop", |
There was a problem hiding this comment.
Adding new entry to this file and the allowlist below are NEVER ok.
If you want this function to be public, you need to properly expose and document it.
My guess is that in this case you don't want it to be public so you should prepend an _* before it!
There was a problem hiding this comment.
I'm following the convention of all other functions in that file. Did I miss some aspects of it or are you saying that all those functions are also "bad" and I should rather break the convention?
|
@pytorchbot merge -f "Still failing because of torchtune, unrelated" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
4 participants
Stack from ghstack (oldest at bottom):
This will make it work correctly with the partitioner's AutoAC