Skip to content

[spmd_types] helion rope & fused_swiglu rules - #3741

Merged
pianpwk merged 47 commits into
mainfrom
gh/pianpwk/50/head
Jun 26, 2026
Merged

[spmd_types] helion rope & fused_swiglu rules#3741
pianpwk merged 47 commits into
mainfrom
gh/pianpwk/50/head

Conversation

@pianpwk

@pianpwk pianpwk commented Jun 22, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

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

pianpwk added 2 commits June 22, 2026 16:06
@pianpwk
Update
f9617f0 
[ghstack-poisoned]
@pianpwk
Update (base update)
f0e688a 
[ghstack-poisoned]
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jun 22, 2026
This was referenced Jun 22, 2026
Open
Merged
Merged
Merged
Merged
Merged
Merged
@pianpwk pianpwk changed the title [spmd_types] helion rope & fused_swiglu override rules [spmd_types] helion rope & fused_swiglu rules Jun 22, 2026
@pianpwk pianpwk mentioned this pull request Jun 23, 2026
Draft
pianpwk added 6 commits June 22, 2026 18:01
@pianpwk
Update
35a342e 
[ghstack-poisoned]
@pianpwk
Update (base update)
8232e03 
[ghstack-poisoned]
@pianpwk
Update
f0094ab 
[ghstack-poisoned]
@pianpwk
Update (base update)
fb87a72 
[ghstack-poisoned]
@pianpwk
Update
5410e44 
[ghstack-poisoned]
@pianpwk
Update (base update)
b22a3c5 
[ghstack-poisoned]
tianyu-l
tianyu-l reviewed Jun 23, 2026
View reviewed changes
Comment thread torchtitan/overrides/fused_swiglu.py Outdated
Comment on lines +418 to +422
if get_spmd_backend() == "spmd_types" and spmd.is_type_checking():
spmd.assert_type(
out,
{"dp": spmd.S(0), "cp": spmd.S(1), "tp": spmd.S(2)},
)

@tianyu-l tianyu-l Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since with spmd.no_typecheck() and with spmd.local() don't provide an "official" closure to annotate exiting tensors, I wonder if we could in torchtitan always put these back-to-global assert_type inside the with context instead of outside.

Main reason is for readability.

Comment thread torchtitan/overrides/helion_rope.py
xq_out, xk_out = _helion_rope_fwd(xq, xk, cache, pos)
# TODO(pianpwk): Migrate this local_map workaround to a custom op SPMD
# propagation rule registration system.
xq_out, xk_out = spmd.local_map(

@tianyu-l tianyu-l Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

a bit confused -- why the swiglu ops use no_typechecking but here we use spmd.local_map?

The only difference may be whether we are checking against in_types? Or are you saying we are actually checking something local in _helion_rope_fwd?

@pianpwk pianpwk Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will switch , I thought we needed no_typecheck but maybe not . with spmd.local we just check the custom op has valid local type interaction (no I mixing non-I, etc.)

@aditvenk aditvenk Jun 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@pianpwk -- I think PR description / commit message needs to be updated

pianpwk added 2 commits June 23, 2026 15:20
@pianpwk
Update
a37e4be 
[ghstack-poisoned]
@pianpwk
Update (base update)
94ba354 
[ghstack-poisoned]
tianyu-l
tianyu-l reviewed Jun 23, 2026
View reviewed changes
Comment thread torchtitan/overrides/fused_swiglu.py Outdated

# TODO(pianpwk): Migrate this local workaround to a custom op SPMD
# propagation rule registration system.
with spmd.local():

@tianyu-l tianyu-l Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

change to local_map

pianpwk added 2 commits June 23, 2026 17:20
@pianpwk
Update
165723f 
[ghstack-poisoned]
@pianpwk
Update (base update)
2969a0f 
[ghstack-poisoned]
@pianpwk
pianpwk requested a review from tianyu-l June 24, 2026 00:21
pianpwk added 22 commits June 25, 2026 10:16
@pianpwk
Update (base update)
fbcba4c 
[ghstack-poisoned]
@pianpwk
Update
0b89fe2 
[ghstack-poisoned]
@pianpwk
Update (base update)
cf8b38f 
[ghstack-poisoned]
@pianpwk
Update
fc4784b 
[ghstack-poisoned]
@pianpwk
Update (base update)
dc6e8d4 
[ghstack-poisoned]
@pianpwk
Update
4d4c1d4 
[ghstack-poisoned]
@pianpwk
Update (base update)
7995bf7 
[ghstack-poisoned]
@pianpwk
Update
a217da5 
[ghstack-poisoned]
@pianpwk
Update (base update)
4e2771c 
[ghstack-poisoned]
@pianpwk
Update
e7e5b95 
[ghstack-poisoned]
@pianpwk
Update (base update)
a9e8c64 
[ghstack-poisoned]
@pianpwk
Update
4fe094e 
[ghstack-poisoned]
@pianpwk
Update (base update)
c4c1ac3 
[ghstack-poisoned]
@pianpwk
Update
2fdad36 
[ghstack-poisoned]
@pianpwk
Update (base update)
e061f25 
[ghstack-poisoned]
@pianpwk
Update
91d554f 
[ghstack-poisoned]
@pianpwk
Update (base update)
b9ed654 
[ghstack-poisoned]
@pianpwk
Update
bc42b6f 
[ghstack-poisoned]
@pianpwk
Update (base update)
d816af2 
[ghstack-poisoned]
@pianpwk
Update
ba79383 
[ghstack-poisoned]
@pianpwk
Update (base update)
baf184f 
[ghstack-poisoned]
@pianpwk
Update
d077023 
[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)
6445cd5 
[ghstack-poisoned]
@pianpwk
Update
e033727 
[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
pianpwk changed the base branch from gh/pianpwk/50/base to main June 26, 2026 02:15
@pianpwk
pianpwk merged commit fbf79bc into main Jun 26, 2026
8 of 11 checks passed
submartingales pushed a commit to botcanlearn/torchtitan-upstream that referenced this pull request Jun 26, 2026
@pianpwk
[spmd_types] helion rope & fused_swiglu override rules
eed1187 
Wrap Helion RoPE and dense fused SwiGLU calls with explicit spmd.local_map contracts so the custom ops can run under spmd_types typechecking. Grouped fused SwiGLU runs directly in the existing local SPMD expert region.

These workarounds should be migrated to a custom op SPMD propagation rule registration system once that API is available.

Test Plan:

- python -m py_compile torchtitan/overrides/helion_rope.py torchtitan/overrides/fused_swiglu.py

- git diff --check -- torchtitan/overrides/helion_rope.py torchtitan/overrides/fused_swiglu.py

- 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

ghstack-source-id: a1ae35a
Pull Request resolved: pytorch/torchtitan#3741
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@aditvenk aditvenk aditvenk left review comments
@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/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.

3 participants

@pianpwk @aditvenk @tianyu-l