[muse_glimmer] Enable CP for the multimodal model - #4387
Open
jinsooihm wants to merge 1 commit into
Open
Conversation
Build packed vision source indices before context-parallel token sharding and carry them through the same load-balancing permutation. Keep vision compute replicated across CP, then use a stateless sharding boundary to gather each local token shard while preserving TP and FSDP gradient semantics. Generalize the shared vision attention layouts, keep multimodal embedding preparation on pipeline stage 0, and support the combined TP+CP+PP+SP topology without changing existing non-CP behavior or checkpoint keys. Test Plan: 71 focused CPU tests; 2 distributed GPU tests; 10-step CP2xPP2 and TP2xCP2xPP2+SP real-process-group runs; exact 10-step deterministic loss and grad-norm parity; repository pre-commit hooks excluding the separately documented optional-dependency Pyrefly blocker.
meta-cla
Bot
added
the
CLA Signed
jinsooihm
marked this pull request as ready for review
September 1, 2026 00:59
jinsooihm
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
September 1, 2026 00:59
tianyu-l
reviewed
Sep 1, 2026
| return vision_bank_indices_T.masked_fill(~vision_mask_T, -1) | ||
|
|
||
|
|
||
| class VisionScatter(Module): |
Contributor
There was a problem hiding this comment.
Question 1: are you putting it here because other multimodal models could share? Is there a vlm model in torchtitan that cannot use this to implement CP?
Question 2: do you create this stateless module only because you want _vision_scatter_config to perform spmd collectives at the module boundary? Fwiw we are moving away from such pattern because they couldn't express fused comm + computation. Could you put collectives INSIDE model code instead? I think we at least we could have a function that does vision scattering, instead of a module.
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
2 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.
Summary
This PR enables CP for the multimodal muse_glimmer.
MuseGlimmerModel.preprocess_inputsbuildsvision_bank_indices_T: [T]from unsharded text tokens, which maps each token position to the index of vision encoder output that should be placed there.Then,
vision_bank_indices_Tgets permuted and sharded byprepare_context_parallel_inputjointly with tokens, so it can be used in the model forward to gather the vision encoder output (which is replicated on CP) after token is embedded.Test Plan
TP2+SP vs TP2+CP2+SP: max loss difference 1.38e-5; max grad_norm difference 5.04e-4 from running
loss_compare.py