Skip to content

[spmd_types] qwen3 sharding - #3653

Merged
pianpwk merged 22 commits into
mainfrom
gh/pianpwk/36/head
Jun 17, 2026
Merged

[spmd_types] qwen3 sharding#3653
pianpwk merged 22 commits into
mainfrom
gh/pianpwk/36/head

Conversation

@pianpwk

@pianpwk pianpwk commented Jun 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

pianpwk added 2 commits June 13, 2026 01:26
@pianpwk
Update
51b99b7 
[ghstack-poisoned]
@pianpwk
Update (base update)
357378d 
[ghstack-poisoned]
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jun 13, 2026
This was referenced Jun 13, 2026
Merged
Open
Open
Merged
Merged
Merged
pianpwk added 2 commits June 15, 2026 00:01
@pianpwk
Update
8cc668a 
[ghstack-poisoned]
@pianpwk
Update (base update)
c286d54 
[ghstack-poisoned]
This was referenced Jun 15, 2026
Merged
Merged
Merged
@pianpwk pianpwk changed the title Add Qwen3 SPMD sharding configs [spmd_types] qwen3 sharding Jun 15, 2026
pianpwk added 2 commits June 15, 2026 10:18
@pianpwk
Update
8b2e948 
[ghstack-poisoned]
@pianpwk
Update (base update)
7c59d97 
[ghstack-poisoned]
pianpwk
pianpwk commented Jun 15, 2026
View reviewed changes
Comment thread torchtitan/models/common/moe_sharding.py Outdated
dense_sequence_parallel_placement()
if enable_ep and enable_sp
else dense_activation_placement(tp=spmd.R)
)

@pianpwk pianpwk Jun 15, 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.

I think this is outdated, on main both shared & shared.colwise have this src->dst redistribution

@tianyu-l tianyu-l Jun 16, 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.

on main both shared & shared.colwise have this src->dst redistribution

wouldn't DTensor error out in this case? as the in_src is different from actual sharding

@pianpwk pianpwk Jun 16, 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.

it's ignored if input is already DTensor it seems, we should be stricter. I can put up a separate PR:

torchtitan/torchtitan/protocols/module.py

Lines 594 to 601 in d226342

if not isinstance(value, DTensor) and src_spmd_layout is not None:
layout = resolve_placements(src_spmd_layout, mesh)
value = DTensor.from_local(
value,
mesh,
layout,
run_check=False,
)

dst_placements is fine, we already match from outer-level redistribute so no-op there

@pianpwk pianpwk Jun 16, 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.

put up at #3682

@pianpwk
Update
beb7775 
[ghstack-poisoned]
@pianpwk
pianpwk marked this pull request as ready for review June 15, 2026 20:13
tianyu-l
tianyu-l reviewed Jun 16, 2026
View reviewed changes
Comment thread torchtitan/models/common/decoder_sharding.py Outdated
import spmd_types as spmd

from torchtitan.distributed.parallel_dims import MeshAxisName
from torchtitan.distributed.utils import get_spmd_backend

@tianyu-l tianyu-l Jun 15, 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.

accidental?

@pianpwk pianpwk Jun 16, 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.

ah yes stale

Comment thread torchtitan/models/common/moe_sharding.py Outdated
dense_sequence_parallel_placement()
if enable_ep and enable_sp
else dense_activation_placement(tp=spmd.R)
)

@tianyu-l tianyu-l Jun 16, 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.

on main both shared & shared.colwise have this src->dst redistribution

wouldn't DTensor error out in this case? as the in_src is different from actual sharding

Comment thread torchtitan/models/common/moe_sharding.py

@tianyu-l tianyu-l Jun 16, 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 quick scanning looks correct to me. But

  1. since MoE / EP is not supported yet, spmd_types enablement is not there yet?
  2. Was full dtensor working? Several annotations were wrong?

@pianpwk pianpwk Jun 16, 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.

is it not? #3447 landed, I was able to fully typecheck

@tianyu-l tianyu-l Jun 16, 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.

are you typechecking EP as well? Not sure because I saw #3654

@pianpwk pianpwk Jun 16, 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.

maybe am confused, this PR has no effect until we enable qwen3 + spmd backend, 2 PRs from now

pianpwk added 4 commits June 15, 2026 23:54
@pianpwk
Update (base update)
7d683b0 
[ghstack-poisoned]
@pianpwk
Update
4cf333e 
[ghstack-poisoned]
@pianpwk
Update (base update)
206f23c 
[ghstack-poisoned]
@pianpwk
Update
ff49eb4 
[ghstack-poisoned]
pianpwk added 2 commits June 16, 2026 10:47
@pianpwk
Update (base update)
5480972 
[ghstack-poisoned]
@pianpwk
Update
0bc0eed 
[ghstack-poisoned]
submartingales pushed a commit to botcanlearn/torchtitan-upstream that referenced this pull request Jun 16, 2026
@pianpwk
Add Qwen3 SPMD sharding configs
fb0b729 
ghstack-source-id: 4ec04ec
Pull Request resolved: pytorch/torchtitan#3653
pianpwk added a commit that referenced this pull request Jun 16, 2026
@pianpwk
[spmd_types] switch full_dtensor to LP autograd fn (#3671)
05d3c36 
Moves full_dtensor backend off pytorch's loss_parallel -> our custom
autograd fn

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.12.0)
(oldest at bottom):
* #3673
* #3655
* #3654
* #3653
* #3278
* #3672
* #3472
* #3632
* __->__ #3671
@pianpwk pianpwk mentioned this pull request Jun 16, 2026
Merged
pianpwk added 2 commits June 16, 2026 16:22
@pianpwk
Update (base update)
c9113cf 
[ghstack-poisoned]
@pianpwk
Update
ae3d5db 
[ghstack-poisoned]
@pianpwk pianpwk mentioned this pull request Jun 16, 2026
Merged
pianpwk added 4 commits June 16, 2026 16:31
@pianpwk
Update (base update)
c2577d1 
[ghstack-poisoned]
@pianpwk
Update
62ada6f 
[ghstack-poisoned]
@pianpwk
Update (base update)
68b97bc 
[ghstack-poisoned]
@pianpwk
Update
2ae6577 
[ghstack-poisoned]
pianpwk added 2 commits June 16, 2026 20:37
@pianpwk
Update (base update)
bf820e5 
[ghstack-poisoned]
@pianpwk
Update
4b5c70a 
[ghstack-poisoned]
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
Update
0cad716 
[ghstack-poisoned]
@pianpwk
pianpwk changed the base branch from gh/pianpwk/36/base to main June 17, 2026 04:07
@pianpwk
pianpwk merged commit e5cc365 into main Jun 17, 2026
7 of 11 checks passed
pianpwk added a commit that referenced this pull request Jun 17, 2026
@pianpwk
[spmd_types] spmd.local region in ChunkedCELoss (#3632)
feb8cdd 
No immediate effect, just prepping for next commit.

decided not to use `spmd.local_map` to avoid larger refactor

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
submartingales pushed a commit to botcanlearn/torchtitan-upstream that referenced this pull request Jun 17, 2026
@pianpwk
Add Qwen3 SPMD sharding configs
f26110a 
ghstack-source-id: 82640f9
Pull Request resolved: pytorch/torchtitan#3653
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/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