[mxfp8 moe training][docs] add tutorial for training with MXFP8 expert parallel#3752
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3752
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 49737db with merge base 30fcb15 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
806eb90 to
a54d592
Compare
302d0b4 to
4270ba0
Compare
This was referenced Jan 29, 2026
a54d592 to
8282223
Compare
8282223 to
3f8c54a
Compare
3f8c54a to
761c09d
Compare
761c09d to
cbe10e7
Compare
021de8c to
9d79704
Compare
9d79704 to
67b3b96
Compare
67b3b96 to
0c0e066
Compare
Contributor
|
will need a rebase on #3769 |
c449c6d to
cc3891f
Compare
0c0e066 to
4c607e4
Compare
svekars
reviewed
Jan 30, 2026
| ^^^^^^^^^^^^^ | ||
|
|
||
| 1. (Recommended) Create a new virtual environment with conda or venv. | ||
| 2. `Install torchao <https://github.com/pytorch/ao/tree/main?tab=readme-ov-file#installation>`__ nightly build (required for CUDA 12.8+ support). |
Contributor
There was a problem hiding this comment.
Do you need torch too? Which versions of all these dependencies? For example, 2.10 or later?
svekars
reviewed
Jan 30, 2026
| | 131072 | 2048 | 1408 | 8 | 3.278 | 2.913 | 1.13x | 4.934 | 3.881 | 1.27x | 1.21x | | ||
| +----------+-------+--------------+---------------+---------------+----------------+---------------+---------------+----------------+---------------+-----------------+ | ||
|
|
||
| As shown, using MXFP8 for all-to-all communications achieves **1.14-1.25x total speedup** versus only quantizing directly before the grouped GEMMs. |
Contributor
There was a problem hiding this comment.
Can you add a conclusion? Like what did the user learn in the tutorial and where else to find more info?
Contributor
Author
There was a problem hiding this comment.
good point, done!
ezyang
reviewed
Jan 30, 2026
| # 🔥 happens in BF16 as well. | ||
| # 🔥 In the backward pass, the incoming upstream gradients will be the MXTensor outputs of the | ||
| # 🔥 MXFP8 all-to-all combine backward pass, so this unpermute autograd func accepts MXTensor | ||
| # 🔥 inputs and performs the reordering in MXFP8. |
Contributor
Author
There was a problem hiding this comment.
ok haha i was trying to draw the reader's eye to the most important parts... will de-emojify
|
How is the code tested? |
Contributor
Author
@ezyang we have unit tests for this pattern of chaining the autograd functions (with eager and compile), and we also integrated into Torchtitan and did large scale convergence + performance testing on an external cloud partner's B200 cluster (joint blog post on this coming soon!) |
…t parallel stack-info: PR: #3752, branch: danielvegamyhre/stack/127
9 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked PRs:
[mxfp8 moe training][docs] add tutorial for training with MXFP8 expert parallel