[spmd_types] SPMD support for ChunkedCELoss - #3472
Open
pianpwk wants to merge 71 commits into
Open
Conversation
This was referenced May 31, 2026
meta-cla
Bot
added
the
CLA Signed
This was referenced Jun 1, 2026
pianpwk
changed the title
pianpwk
marked this pull request as ready for review
June 4, 2026 23:34
submartingales
pushed a commit
to botcanlearn/torchtitan-upstream
that referenced
this pull request
Jun 18, 2026
Pull Request resolved: pytorch/torchtitan#3472 ghstack-source-id: 34283ce
tianyu-l
reviewed
Jun 18, 2026
Comment on lines
+721
to
+722
| # The hidden activation type depends on the model's sharding config; | ||
| # we'll defer hidden_states typechecking to previous module boundaries. |
Contributor
There was a problem hiding this comment.
what are the possibilities?
Contributor
Author
There was a problem hiding this comment.
hidden_states is directly from decoder, so we'd have to condition this on SP (I or S(1))
Comment on lines
+645
to
+650
| if get_spmd_backend() == "spmd_types": | ||
| spmd.assert_type( | ||
| accumulated_grad, | ||
| {"dp": spmd.V, "cp": spmd.V, "tp": spmd.R}, | ||
| spmd.PartitionSpec("dp", "cp", None), | ||
| ) |
Contributor
There was a problem hiding this comment.
This is gradient, why would it need a type?
Also, why "tp" dim is spmd.R?
Contributor
Author
There was a problem hiding this comment.
fair point, removed typechecking for this
| # PP: only the last stage has lm_head; non-last stages skip this. | ||
| if isinstance(self.loss_fn, ChunkedCELoss): | ||
| assert isinstance(self.model_config, Decoder.Config) | ||
| self.loss_fn.global_vocab_size = self.model_config.vocab_size |
Contributor
There was a problem hiding this comment.
global_vocab_size should be a config of ChunkedCELoss, set in config_registry.py's not rather than patched in trainer
Contributor
Author
There was a problem hiding this comment.
hardcoded values in model registries
submartingales
pushed a commit
to botcanlearn/torchtitan-upstream
that referenced
this pull request
Jun 18, 2026
Pull Request resolved: pytorch/torchtitan#3472 ghstack-source-id: 98e6e13
tianyu-l
reviewed
Jun 18, 2026
| def qwen3_debugmodel() -> Trainer.Config: | ||
| return Trainer.Config( | ||
| loss=ChunkedCELoss.Config(), | ||
| loss=ChunkedCELoss.Config(global_vocab_size=2048), |
Contributor
There was a problem hiding this comment.
should come from model_spec rather than hardcoding
pianpwk
added a commit
that referenced
this pull request
Jun 18, 2026
loss migration: loss-parallel custom autograd handling, DTensor local_map / unwrapping, now lives under cross_entropy_loss fn loss-parallel flag deletion: - for trainers, TP on = LP on, uses custom autograd impl instead of pytorch's implicit `loss_parallel()` context. Sets model sharding configs to return vocab-sharded logits: `tp_gather_logits=False` - RL usecases require full logits from lm_head, `tp_gather_logits=True` - GraphTrainer TP on = LP on Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0) (oldest at bottom): * #3690 * #3673 * #3655 * #3664 * #3654 * #3278 * #3472 * __->__ #3694
tianyu-l
approved these changes
Jun 18, 2026
tianyu-l
left a comment
tianyu-l
left a comment
Contributor
There was a problem hiding this comment.
see #3472 (comment)
let's not hardcode vocab_size
submartingales
pushed a commit
to botcanlearn/torchtitan-upstream
that referenced
this pull request
Jun 18, 2026
Pull Request resolved: pytorch/torchtitan#3472 ghstack-source-id: bd941ec
submartingales
pushed a commit
to botcanlearn/torchtitan-upstream
that referenced
this pull request
Jun 19, 2026
Pull Request resolved: pytorch/torchtitan#3472 ghstack-source-id: fa65d5a
tianyu-l
reviewed
Jun 20, 2026
Comment on lines
+146
to
+148
| loss_config = getattr(config, "loss", None) | ||
| if isinstance(loss_config, ChunkedCELoss.Config): | ||
| loss_config.global_vocab_size = self.vocab_size |
Contributor
There was a problem hiding this comment.
This function is to give CLI config a chance to override python file config. The vocab size shouldn't be adjusted by CLI, so we could've set it entirely in config_registry files.
If you think the radius is too large, I'm fine with this for now but please add a NOTE / TODO.
submartingales
pushed a commit
to botcanlearn/torchtitan-upstream
that referenced
this pull request
Jun 20, 2026
Pull Request resolved: pytorch/torchtitan#3472 ghstack-source-id: 6ade7ab
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
3 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.
Moves the seq-dim allgather onto the pre-lm-head norm, making it ShardingConfig-based. Annotates types in ChunkedCELoss impl.
Stack from ghstack (oldest at bottom):