Repository navigation
Conversation
Backward unshards MXFP8 weights with columnwise planes only and relies on a recomputed FsdpModule's pre_forward to gather rowwise planes. A checkpoint can recompute a module whose parameters another FsdpModule owns without re-entering that owner's pre_forward, e.g. MTP eh_proj under the root, or TransformerLayer MLP linears under selective MLP recompute, so TE fails with missing row-scaled FP8 data. Register a forward pre-hook on each module directly holding an FsdpModule's parameters that unshards the owner for forward compute during recomputation. Groups serving both phases are already ready; sharded MXFP8 groups gather their rowwise planes. Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>
Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>
Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>
Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>
…e-enter the module Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>
…cks in tests Signed-off-by: Xiaoda Zhang <xiaodaz@nvidia.com>
Signed-off-by: Jack Chang <jianbinc@nvidia.com>
copy-pr-bot
Bot
force-pushed
the
pull-request/8432
branch
from
November 1, 2026 01:28
7133fc6 to
efe6b4a
Compare
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
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.
Enable #8432's backward-plane policy in the MCore adapter when MXFP8 parameter gathering and activation recomputation are configured. Selective MLA and root-owned MTP recomputation can invoke child modules without re-entering their owning FSDP unit, so backward must also make the forward weight representation available.
Depends on #8432; this PR targets
pull-request/8432and contains only the adapter integration and its regression test. The policy conservatively gathers both MXFP8 layouts for recomputing models. Models without recomputation retain the existing per-phase behavior. Finer per-unit selection can be optimized separately.Validation:
git diff --checkpassed. The formatter's advisory mypy step reports missing Torch/TE imports in the workstation environment.