Skip to content

[spmd_types] GPT-OSS enablement - #3690

Merged
pianpwk merged 95 commits into
mainfrom
gh/pianpwk/43/head
Jun 26, 2026
Merged

[spmd_types] GPT-OSS enablement#3690
pianpwk merged 95 commits into
mainfrom
gh/pianpwk/43/head

Conversation

@pianpwk

@pianpwk pianpwk commented Jun 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

pianpwk added 2 commits June 16, 2026 14:48
@pianpwk
Update (base update)
ac555e4 
[ghstack-poisoned]
@pianpwk
Update
c3b4749 
[ghstack-poisoned]
This was referenced Jun 16, 2026
Merged
Open
Merged
Open
Merged
This was referenced Jun 16, 2026
Merged
Merged
Merged
Merged
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jun 16, 2026
@pianpwk pianpwk changed the title Enable GPT-OSS SPMD backend integration [spmd_types] GPT-OSS enablement Jun 16, 2026
pianpwk added 2 commits June 16, 2026 16:22
@pianpwk
Update (base update)
3c77d07 
[ghstack-poisoned]
@pianpwk
Update
45007ed 
[ghstack-poisoned]
pianpwk added a commit that referenced this pull request Jun 16, 2026
@pianpwk
Enable GPT-OSS SPMD backend integration
6d1aad4 
ghstack-source-id: cc74dcc
Pull Request resolved: #3690
@pianpwk pianpwk mentioned this pull request Jun 16, 2026
Merged
pianpwk added 2 commits June 16, 2026 16:31
@pianpwk
Update (base update)
16440a4 
[ghstack-poisoned]
@pianpwk
Update
6f28c18 
[ghstack-poisoned]
pianpwk added a commit that referenced this pull request Jun 16, 2026
@pianpwk
Enable GPT-OSS SPMD backend integration
e705cbb 
ghstack-source-id: 3c1f7ab
Pull Request resolved: #3690
pianpwk added 2 commits June 16, 2026 17:08
@pianpwk
Update (base update)
1716d4d 
[ghstack-poisoned]
@pianpwk
Update
2dc66c9 
[ghstack-poisoned]
pianpwk added a commit that referenced this pull request Jun 17, 2026
@pianpwk
Enable GPT-OSS SPMD backend integration
69cdd0f 
ghstack-source-id: 30d8436
Pull Request resolved: #3690
pianpwk added 2 commits June 16, 2026 20:37
@pianpwk
Update (base update)
16c5509 
[ghstack-poisoned]
@pianpwk
Update
1b741b8 
[ghstack-poisoned]
pianpwk added a commit that referenced this pull request Jun 17, 2026
@pianpwk
Enable GPT-OSS SPMD backend integration
3fd8ff7 
ghstack-source-id: 0f6f491
Pull Request resolved: #3690
pianpwk added a commit that referenced this pull request Jun 17, 2026
@pianpwk
[spmd_types] raise on SAC + Flex + typechecking (#3672)
eb27b0f 
SAC + Flex, or SAC + typecheck + non-flex, or typecheck + Flex + full/no AC is fine, but this is not (likely SAC + regional inductor issue). Raise error asking user to use one of the other cominbations

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3690
* #3673
* #3655
* #3664
* #3654
* #3278
* #3694
* #3472
* #3632
* #3653
* __->__ #3672
pianpwk added 2 commits June 16, 2026 20:45
@pianpwk
Update (base update)
889a440 
[ghstack-poisoned]
@pianpwk
Update
3969d0b 
[ghstack-poisoned]
pianpwk added 10 commits June 25, 2026 14:55
@pianpwk
Update (base update)
6ca3ccc 
[ghstack-poisoned]
@pianpwk
Update
663ed67 
[ghstack-poisoned]
@pianpwk
Update (base update)
dad84dc 
[ghstack-poisoned]
@pianpwk
Update
7fe1431 
[ghstack-poisoned]
@pianpwk
Update (base update)
34e1886 
[ghstack-poisoned]
@pianpwk
Update
895c3aa 
[ghstack-poisoned]
@pianpwk
Update (base update)
509040a 
[ghstack-poisoned]
@pianpwk
Update
5648632 
[ghstack-poisoned]
@pianpwk
Update (base update)
c7b7a93 
[ghstack-poisoned]
@pianpwk
Update
967e1e8 
[ghstack-poisoned]
submartingales pushed a commit to botcanlearn/torchtitan-upstream that referenced this pull request Jun 25, 2026
@pianpwk
Enable GPT-OSS SPMD backend integration
7a33660 
ghstack-source-id: f91a8b9
Pull Request resolved: pytorch/torchtitan#3690
pianpwk added 2 commits June 25, 2026 16:43
@pianpwk
Update (base update)
797e2dd 
[ghstack-poisoned]
@pianpwk
Update
48d00e3 
[ghstack-poisoned]
pianpwk added a commit that referenced this pull request Jun 26, 2026
@pianpwk
[spmd_types] llama3 enablement (#3763)
0a4aa45 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3740
* #3690
* #3673
* #3655
* #3741
* #3664
* #3654
* __->__ #3763
pianpwk added 2 commits June 25, 2026 18:38
@pianpwk
Update (base update)
3c7643c 
[ghstack-poisoned]
@pianpwk
Update
099edd7 
[ghstack-poisoned]
pianpwk added a commit that referenced this pull request Jun 26, 2026
@pianpwk
[spmd_types] MoE sparse mesh transitions (#3654)
b052f36 
`set_current_spmd_mesh(sparse_mesh)` transitions for MoE region, spmd.all_to_all switch, plus
mesh reinterprets, assert_type, mutate_type

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3740
* #3690
* #3673
* #3655
* #3741
* #3664
* __->__ #3654
pianpwk added a commit that referenced this pull request Jun 26, 2026
@pianpwk
[spmd_types] fix MoE sequence-sharded combine (#3664)
565cb81 
Fix spmd_types SP-padding support, by requiring evenly sharded seq-dim
for MoE input activations, calculating post-combine output size based on
local tensor shapes.

Raises in trainer for MoE + uneven seq-dim sharding.

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3740
* #3690
* #3673
* #3655
* #3741
* __->__ #3664
pianpwk added a commit that referenced this pull request Jun 26, 2026
@pianpwk
[spmd_types] helion rope & fused_swiglu rules (#3741)
fbf79bc 
Wrap Helion RoPE and fused SwiGLU with explicit spmd.local_map calls

Test Plan:

- NGPU=8 LOG_RANK=0 MODULE=qwen3 CONFIG=qwen3_debugmodel_fused_qkv
./run_train.sh --parallelism.spmd_backend spmd_types
--debug.spmd_typechecking --parallelism.data_parallel_shard_degree 2
--parallelism.context_parallel_degree 2
--parallelism.tensor_parallel_degree 2 --training.steps 1
--override.imports torchtitan.overrides.helion_rope
activation-checkpoint:none

- NGPU=8 LOG_RANK=0 MODULE=deepseek_v3 CONFIG=deepseek_v3_debugmodel
./run_train.sh --parallelism.spmd_backend spmd_types
--debug.spmd_typechecking --parallelism.data_parallel_shard_degree 2
--parallelism.context_parallel_degree 2
--parallelism.tensor_parallel_degree 2
--parallelism.expert_parallel_degree 2 --training.steps 1
--override.imports torchtitan.overrides.fused_swiglu
activation-checkpoint:none

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3740
* #3690
* #3673
* #3655
* __->__ #3741
pianpwk added 2 commits June 25, 2026 20:17
@pianpwk
Update (base update)
388024f 
[ghstack-poisoned]
@pianpwk
Update
dcfec29 
[ghstack-poisoned]
pianpwk added a commit that referenced this pull request Jun 26, 2026
@pianpwk
[spmd_types] qwen3 enablement (#3655)
9b91827 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3740
* #3690
* #3673
* __->__ #3655
pianpwk added 2 commits June 25, 2026 22:17
@pianpwk
Update (base update)
fa18dc4 
[ghstack-poisoned]
@pianpwk
Update
d68d3a7 
[ghstack-poisoned]
pianpwk added a commit that referenced this pull request Jun 26, 2026
@pianpwk
[spmd_types] dsv3 enablement (#3673)
622c1ac 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3740
* #3690
* __->__ #3673
pianpwk added 2 commits June 26, 2026 00:22
@pianpwk
Update (base update)
4648ab7 
[ghstack-poisoned]
@pianpwk
Update
7af611e 
[ghstack-poisoned]
@pianpwk
pianpwk changed the base branch from gh/pianpwk/43/base to main June 26, 2026 07:23
@pianpwk
pianpwk merged commit 7144431 into main Jun 26, 2026
7 of 11 checks passed
submartingales pushed a commit to botcanlearn/torchtitan-upstream that referenced this pull request Jun 26, 2026
@pianpwk
Enable GPT-OSS SPMD backend integration
697be4a 
ghstack-source-id: 4a12303
Pull Request resolved: pytorch/torchtitan#3690
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tianyu-l tianyu-l tianyu-l approved these changes
@fegin fegin Awaiting requested review from fegin fegin is a code owner
@wwwjn wwwjn Awaiting requested review from wwwjn wwwjn is a code owner
@wconstab wconstab Awaiting requested review from wconstab wconstab is a code owner

Assignees

No one assigned

Labels

ciflow/rl ciflow/8gpu CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@pianpwk @tianyu-l