transformers
[Umbrella] FSDP2 + expert-parallel token dispatch, with sharded-loading prefetch
#48204
Open
Open
[Umbrella] FSDP2 + expert-parallel token dispatch, with sharded-loading prefetch #48204
qgallouedec
wants to merge
316 commits into
main
from
ep-fsdp-2d-mesh
qgallouedec
changed the title
qgallouedec
changed the title
Merge remote-tracking branch 'origin/fix-ep-training-gradients' into …
2d8108d5
Merge remote-tracking branch 'origin/ep-fsdp-2d-mesh' into ep-fsdp-2d…
b13f35b6
Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh
e3a43b26
Gather the sharded state dict through the model's own save helpers
71fcd09a
Document combining expert parallelism with FSDP2
bba3125e
Address review: drop the dead expert-gradient sync, torch-native mixe…
ad5df191
Give each device mesh its own optimizer param group instead of steppi…
c18e5d99
Test the 2-D mesh with a batch split across the fsdp ranks and compar…
508996b6
Refresh the throughput table
ef412569
Cosmetic: docstring layout, comment placement, doc link
393af534
Merge branch 'fsdp2-ep-2d-mesh' into ep-fsdp-2d-mesh
724c6f91
Restore the shard prefetch helper dropped in the merge
b4045930
Gate the per-matmul sentinel masks on expert parallelism (#48201 x #4…
51177c78
Load: convert transposed and RoPE-permuted weights in full before tak…
bb3a775d
Let conversion ops declare the dims they move; defer the shard only w…
9e2bca95
Keep the decision in one helper
d83a1f72
Merge branch 'main' into moe-skip-ep-sentinels-v2
9987b909
Merge branch 'main' into load-shard-prefetch
1bd524ae
Merge branch 'main' into fix-empty-fsdp-shards
8aa54334
Merge branch 'main' into fix-empty-fsdp-shards
84c5e320
Merge branch 'main' into load-shard-prefetch
14f0a81c
Merge branch 'main' into moe-skip-ep-sentinels-v2
2e4efe73
Merge branch 'main' into ep-fsdp-2d-mesh
75b74f8e
Merge branch 'main' into load-shard-prefetch
07eb49e4
Merge branch 'main' into fix-empty-fsdp-shards
a6b88157
Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh
2414cedf
Gather the sharded state dict through the model's own save helpers
4914e5ae
Document combining expert parallelism with FSDP2
489714fe
Fix NaN gradients in expert-parallel training: mask uninitialized gro…
91ad4af2
Address review: drop the dead expert-gradient sync, torch-native mixe…
f9e24a39
Fix wrong gradients for all non-expert parameters in expert-parallel …
50d95d25
Give each device mesh its own optimizer param group instead of steppi…
683176d4
Trim comments
86dbd5a4
Test the 2-D mesh with a batch split across the fsdp ranks and compar…
a8301666
Gate the router-score backward all-reduce on grad mode; drop the post…
036caf30
Refresh the throughput table
ef86f595
Fix expert parallelism through Trainer (#48208)
0e5d1ebc
Cosmetic: docstring layout, comment placement, doc link
d9394c9e
Carry is_expert_parallel on the test experts stand-ins
eb9dafc6
Read _device_mesh directly and hoist the itertools import
c64443d6
Merge remote-tracking branch 'origin/main' into ep-fsdp-2d-mesh
28d58f03
Merge remote-tracking branch 'pr/48201' into ep-fsdp-2d-mesh
4f51f0d5
Merge remote-tracking branch 'pr/48227' into ep-fsdp-2d-mesh
e2bf5ee6
Merge remote-tracking branch 'pr/48237' into ep-fsdp-2d-mesh
d0e12fff
Merge remote-tracking branch 'pr/48205' into ep-fsdp-2d-mesh
09db09fd
Merge remote-tracking branch 'pr/48516' into ep-fsdp-2d-mesh
36662856
Merge branch 'main' into moe-skip-ep-sentinels-v2
2ec1d59c
Merge branch 'main' into load-shard-prefetch
9059c09e
Merge branch 'main' into fix-empty-fsdp-shards
9c00001a
Merge branch 'main' into load-convert-before-shard
324b3116
Prefetch only the byte spans each rank will read
b01da96b
Cover per-expert checkpoints, and compare against a fair share of the…
69e08c23
Warm short ranges with kernel readahead, and bridge cheap gaps
862086dc
Take the meta state dict again, not named_parameters
cca1d864
Decide readahead on the total a rank warms, not the size of each range
f2c47436
Set is_expert_parallel at the call site instead of threading a parameter
a66940a0
Merge branch 'main' into moe-skip-ep-sentinels-v2
7fe66827
Set the flag where the experts are sharded, and keep the router path …
e9f3a20c
Drop the dead default on _dispatch_routed_input
b22c3025
Fix NaN gradients in expert-parallel training: mask uninitialized gro…
8af1a002
Fix wrong gradients for all non-expert parameters in expert-parallel …
1f0b7bf6
Trim comments
4636d425
Gate the router-score backward all-reduce on grad mode; drop the post…
89fa55ca
Fix expert parallelism through Trainer (#48208)
96769d99
Enable FSDP2 + expert parallelism via a 2-D (fsdp, tp) device mesh
c74c1228
Gather the sharded state dict through the model's own save helpers
1c66acbc
Document combining expert parallelism with FSDP2
38aefbf2
Address review: drop the dead expert-gradient sync, torch-native mixe…
e72b117d
Give each device mesh its own optimizer param group instead of steppi…
449acd21
Test the 2-D mesh with a batch split across the fsdp ranks and compar…
9d17bd7a
Refresh the throughput table
6f2d9c4b
Cosmetic: docstring layout, comment placement, doc link
3a8da119
Expert-parallel token dispatch: each rank trains on its own part of t…
1b071993
Make the dispatch handling of the experts explicit in apply_fully_sha…
7063c456
Combine the dispatched expert outputs in the routing weights' precisi…
dc1fdeee
Split the batches across every rank with a DistributedSampler under t…
f2a3c26d
Simplify the dispatch: torch's differentiable all-to-all, the grouped…
3ab94ff1
Dispatch as plan-level styles; one notion of the ranks sharing a batc…
d1881444
Select the dispatch styles in the mixin, guard the torch version and …
b9bf58bd
Rewrite the expert parallel plan the model actually reads for the dis…
39cdce50
Seed the dispatch sampler with data_seed; require the grouped_mm expe…
9f862308
Run the local experts directly: the grouped forward's masking and com…
930c146a
Run the local experts through the module's own forward
15588e00
Dispatch: one host sync per layer, comments say what the code does
51f281d9
Dispatch: drop the router style, reuse the named mesh and is_dtensor,…
101bfea9
Dispatch: skip the mixin tests for models without an FSDP plan, drop …
2d9046bd
Fix the _has_fsdp_plan docstring
42606afb
Drop the FSDP plan skip from the mixin tests: every config carries a …
d8eb1fc1
Format
23920e3a
Keep the base install_forward signature on the dispatch experts style
e1498696
Docs: shorter token dispatch overview
e139e8a5
Replicate the rest of the plan under expert parallel token dispatch
cf74929d
Merge remote-tracking branch 'upstream/main' into fsdp2-ep-2d-mesh
0e21b598
Narrow to the RoPE permutation, leave transposes to shard-on-read
f4562cf1
Merge branch 'main' into load-convert-before-shard
d8c8613a
Merge branch 'main' into load-convert-before-shard
7d960108
Merge branch 'fsdp2-ep-2d-mesh' into ep-token-dispatch
18144abb
Update docs/source/en/expert_parallelism.md
c9a42167
Test the byte spans a rank warms
e21db3dc
Make local ranks pick the same prefetch plan
20896a7b
Reduce over the accelerator the file already binds to
b191be0e
Merge branch 'main' into load-shard-prefetch
19f7ae1d
Update docs/source/en/expert_parallelism.md
8f2f9cb1
Update docs/source/en/expert_parallelism.md
3c8f600e
Update docs/source/en/expert_parallelism.md
e7be30a9
Update docs/source/en/expert_parallelism.md
8dbccad4
Merge branch 'main' into load-shard-prefetch
9f432f54
Bind the prefetch barrier to the accelerator
9e8a65d2
Merge branch 'load-shard-prefetch' into prefetch-rank-spans
82a4f5ba
Merge branch 'main' into load-convert-before-shard
717a3a95
Merge branch 'main' into fix-empty-fsdp-shards
f436d9b4
Merge branch 'main' into fsdp2-ep-2d-mesh
35ed430b
Merge remote-tracking branch 'upstream/main' into moe-skip-ep-sentine…
1103b0ec
Merge branch 'fsdp2-ep-2d-mesh' into ep-token-dispatch
79fe11d4
Merge branch 'main' into qwen3_vl_moe_tp_plan
298ffa4d
Merge remote-tracking branch 'upstream/main' into umb48204
3e6e8755
Merge branch 'pr48516' into umb48204
82e8fca8
Merge branch 'pr48201' into umb48204
a8916bef
Merge branch 'pr48227' into umb48204
55b35321
Merge branch 'pr48237' into umb48204
4460294e
Merge branch 'pr48519' into umb48204
43a10a22
Merge branch 'pr48518' into umb48204
b3119027
Drop the HF_EP_DISPATCH prototype call sites superseded by #48518
b89b0e67
Merge branch 'pr48373' into umb48204
fba560cb
Keep replicated_with_grad_allreduce under expert-parallel dispatch
c8d3107f
Keep replicated_with_grad_allreduce under expert-parallel dispatch
0817429f
Compare gradients, not only the loss, in the EP backward test
34e63c9d
Merge branch 'pr48518' into umb48204
caac0ad5
Merge branch 'main' into ep-fsdp-2d-mesh
fccd0e7f
Merge branch 'main' into load-convert-before-shard
5892c331
Merge branch 'main' into fix-empty-fsdp-shards
7b09389e
Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFE…
153a3994
style
cf674aa3
Move the shard prefetch into distributed/utils.py
77ad02d0
Apply ruff format to the prefetch helper
70c6ac4b
Bind the prefetch barrier to the accelerator
f70364f3
Merge branch 'main' into moe-skip-ep-sentinels-v2
e3e2d58d
Update docs/source/en/expert_parallelism.md
75fff716
Update docs/source/en/expert_parallelism.md
acba6bc9
Update docs/source/en/expert_parallelism.md
69668c3c
Update docs/source/en/expert_parallelism.md
9cfb2d3b
Merge branch 'main' into fsdp2-ep-2d-mesh
2884e2f1
Merge branch 'main' into qwen3_vl_moe_tp_plan
86fc3945
Merge branch 'fsdp2-ep-2d-mesh' into ep-token-dispatch
4ce47004
Merge remote-tracking branch 'upstream/load-shard-prefetch' into pref…
a687e4da
Flag the empty shard in set_param_for_module
fee68450
Only build the sharding op when it is used
cb61b1aa
Declare dispatch_experts_forward next to its only caller
28cb4827
Merge branch 'main' into fix-empty-fsdp-shards
f22f18a7
Let each op say whether it breaks shard-on-read
071db775
Merge branch 'main' into load-convert-before-shard
5fa9a03d
Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFE…
11bca8c0
style
04bedbe5
Move the shard prefetch into distributed/utils.py
8c873dc2
Apply ruff format to the prefetch helper
2a42ab85
Bind the prefetch barrier to the accelerator
31a256b5
Name the dispatch strategy instead of switching it with a bool
a6e36c09
Merge branch 'main' into moe-skip-ep-sentinels-v2
70bd69ef
Merge remote-tracking branch 'upstream/load-shard-prefetch' into umb4…
ad3b02ff
Merge remote-tracking branch 'upstream/prefetch-rank-spans' into umb4…
9505dee1
Merge remote-tracking branch 'upstream/moe-skip-ep-sentinels-v2' into…
bbf73973
Merge remote-tracking branch 'upstream/fix-empty-fsdp-shards' into um…
18ef7f71
Merge remote-tracking branch 'upstream/fsdp2-ep-2d-mesh' into umb48204
6a207528
Merge remote-tracking branch 'upstream/ep-token-dispatch' into umb48204
7683dddb
Merge remote-tracking branch 'upstream/load-convert-before-shard' int…
b669fc20
Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFE…
6e40b0b1
style
9a61e7b7
Move the shard prefetch into distributed/utils.py
e8b13c27
Apply ruff format to the prefetch helper
eed24d1e
Bind the prefetch barrier to the accelerator
a98cde20
Merge branch 'main' into fix-empty-fsdp-shards
758ed92f
Merge branch 'main' into fsdp2-ep-2d-mesh
1be3a7f5
Merge branch 'main' into load-convert-before-shard
10f60888
Merge remote-tracking branch 'upstream/load-shard-prefetch' into pref…
62e85eab
Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFE…
72548d08
style
dd8804bd
Move the shard prefetch into distributed/utils.py
acae96cf
Apply ruff format to the prefetch helper
7db2467b
Bind the prefetch barrier to the accelerator
68bc14bf
Merge remote-tracking branch 'upstream/load-shard-prefetch' into pref…
e514c7fd
Merge branch 'main' into load-convert-before-shard
14fe0333
Merge remote-tracking branch 'upstream/fsdp2-ep-2d-mesh' into ep-toke…
db788602
Merge branch 'main' into fsdp2-ep-2d-mesh
a6f5824a
Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFE…
5489cde1
style
0b5d4402
Move the shard prefetch into distributed/utils.py
170161c5
Apply ruff format to the prefetch helper
2f54f17c
Bind the prefetch barrier to the accelerator
a1846bc2
Merge branch 'main' into fix-empty-fsdp-shards
d9183573
Merge branch 'fsdp2-ep-2d-mesh' into ep-token-dispatch
b56e1d68
Merge branch 'main' into load-convert-before-shard
65b6e406
Merge branch 'main' into qwen3_vl_moe_tp_plan
6fbcb300
Merge remote-tracking branch 'upstream/main' into umb48204
dc7143f5
Merge remote-tracking branch 'upstream/load-shard-prefetch' into umb4…
a58566db
Merge remote-tracking branch 'upstream/prefetch-rank-spans' into umb4…
56699c82
Merge remote-tracking branch 'upstream/fsdp2-ep-2d-mesh' into umb48204
fef0435d
Merge remote-tracking branch 'upstream/ep-token-dispatch' into umb48204
21e12768
Merge remote-tracking branch 'upstream/fix-empty-fsdp-shards' into um…
cfc7f175
Merge branch 'pr48373-head' into umb48204
d12ff1fb
Merge remote-tracking branch 'upstream/load-convert-before-shard' int…
2a7a88d6
Merge branch 'main' into fsdp2-ep-2d-mesh
cb02b27f
Pass the ignored params through the policy helper
6a86e181
Merge remote-tracking branch 'upstream/load-shard-prefetch' into pref…
37c8ee6f
Merge remote-tracking branch 'upstream/main' into fsdp2-ep-2d-mesh
79988c35
Drop the comment on the parallel sizes
51207488
Trim the parallelism config comment
0147840f
Warm the page cache before sharded loading (opt-in via HF_SHARD_PREFE…
b1744994
style
48ce8e26
Move the shard prefetch into distributed/utils.py
3e51ffbc
Apply ruff format to the prefetch helper
ef758c4e
Bind the prefetch barrier to the accelerator
cfdc8bab
Merge remote-tracking branch 'upstream/fsdp2-ep-2d-mesh' into ep-toke…
d604a2a2
Merge remote-tracking branch 'upstream/main' into ep-token-dispatch
b169a75d
Address review: guard dispatch + PP, derive the mesh name, move the g…
3439765f
Keep reading the tensor parallel size from accelerate in `get_tp_size`
4f2a6b91
Give the token dispatch weight comparison its own tolerance
86561b69
Separate expert dispatch mesh from tensor parallel size
df8fa771
Support trunk tensor parallelism with folded expert dispatch
60655f89
rm config test
d7753400
Simplify expert dispatch and rename edp to efsdp
a49aa6d1
Merge branch 'ep-token-dispatch' into ep-mesh
b1fdbdf3
Merge origin/ep-mesh and restore expert-parallel mesh support
6d8f54de
Unify dense and expert mesh construction
ea65cfb3
remove
b33e45e0
Encapsulate distributed mesh views behind named accessors
0eb9ffa8
Simplify MeshManager access and keep mesh views private
4d63a32c
Simplify expert parallel configuration and explicit plan selection
d4da6e0b
Resolve TP and EP plans before sharding and derive expert dispatch
6a1059ee
Merge branch 'main' into ep-token-dispatch
56047f4c
Merge branch 'ep-token-dispatch' into ep-mesh
00717373
Inline legacy expert parallel compatibility handling
c8a05bc6
Simplify parallel plan resolution and clarify TP and EP meshes
cc2237c5
remove
fc37645e
Remove auto EP plans and clarify parallel plan tests
251b2b94
Clarify expert parallel docs and explicit sharding mesh selection
8c2588da
Always wrap dispatched experts on the expert FSDP mesh
19118585
Clarify TP and EP layouts in FSDP comparison test
49f432cd
Clarify all-reduce and dispatch layouts in EP tests
6d231256
Take TP and EP meshes in apply_tensor_parallelism
e1d33d25
add default all-to-all plan for qwen3_moe only + derive rule from pl…
79fe4a43
raise in `resolve_parallel_plans` when an override key matches no mod…
2dc99524
better distinguish masked/all-to-all path for nicer API
b2b0d324
Merge remote-tracking branch 'upstream/ep-mesh' into ep-umbrella
45ad8d1e
Merge remote-tracking branch 'upstream/load-shard-prefetch' into ep-u…
6e4964ff
Merge remote-tracking branch 'upstream/prefetch-rank-spans' into ep-u…
092a8e31
Trainer: average the gradients of parameters FSDP2 does not shard
f4582493
Trainer: broadcast the parameters FSDP2 does not shard before training
9d0eb661
Style
65781c6e
Merge branch 'ep-umbrella' into ep-fsdp-2d-mesh
36ec32db
Drop #48237, #48519 and #48373 from the umbrella
ee8ecdf0
qgallouedec
changed the title
Login to write a write a comment.
Login via GitHub