Skip to content

Expert Parallelism #361

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 95 commits into
base: dev
Choose a base branch
from
Open

Expert Parallelism #361

wants to merge 95 commits into from

Conversation

xrsrke
Copy link
Member

@xrsrke xrsrke commented Apr 29, 2025

What does this PR do?

Fixes # (issue)

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guidelines?
  • Did you write any new necessary tests?
  • Did you log the throughput and loss you get to ensure the PR works as expected in actual training?
  • Did you log the memory usage? you can use this tool to understand the memory usage breakdown in nanotron.
  • If you modified anything related to checkpoints, did you verify that saving and reloading checkpoints still works correctly?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

NouamaneTazi and others added 30 commits April 14, 2025 16:17
@NouamaneTazi
can only merge to main from dev (#348)
d22eb80
@xrsrke
move moe from qwen modeling to src/nn
f1160f1
@xrsrke
add groupedmlp
bb8ac96
@xrsrke
add token permute and unpermute
ebdd55d
@xrsrke
fix num_tokens_per_expert counting < num_experts
3178944
@NouamaneTazi @c8ef @eliebak
fix init and init scaling factor and run evals in background (#353)
37cc8b0 
* can only merge to main from dev

* Fix UnBoundLocalError in `clm_collator.py` (#339)

* Update clm_collator.py

* can only merge to main from dev (#348)

---------

Co-authored-by: Nouamane Tazi <nouamane@huggingface.co>

* fix init and init scaling factor and run evals in background (#349)

* InitScalingMethod

* InitScalingMethod

* run evals in background (#352)

* eval

* try adding lightevalrunner to trainer

* amend

* amend

* amend

* amend

* amend

* amend

* .

* amend

* amend

* .

* qos to low

* add nanotron_path

* some fix: logs, and config

* cp instead of sync

* eval_interval

* serialize sanity checks

* add output dir and s3_save path in the config

* fix s3 only if define

* fixes

---------

Co-authored-by: elie <97572401+eliebak@users.noreply.github.com>
Co-authored-by: “eliebak” <elie.bakouch@huggingface.co>

---------

Co-authored-by: elie <97572401+eliebak@users.noreply.github.com>
Co-authored-by: “eliebak” <elie.bakouch@huggingface.co>

---------

Co-authored-by: Connector Switch <c8ef@outlook.com>
Co-authored-by: elie <97572401+eliebak@users.noreply.github.com>
Co-authored-by: “eliebak” <elie.bakouch@huggingface.co>
@zzhhjjj
inference qwen moe seems to work
4c39f62 
inference seems good rn
@zzhhjjj
update readme
20f0196
@xrsrke
fix router's weight initialization and wrong hidden size for non-moe …
e472f77 
…mlp in qwen
@xrsrke
add source for router weight and router logits in float32
76ff5c7
@xrsrke
Merge branch 'main' into dev_qwen_moe
8851256
@NouamaneTazi
Merge branch 'dev' of https://github.com/huggingface/nanotron into de…
863f45f 
…v_qwen_moe
@NouamaneTazi
fixes
11f7997
@NouamaneTazi
.
81c31ba
@NouamaneTazi
.
2e79d24
@xrsrke
add parametrize grouped mlp in column and row linear
bcf1658
@xrsrke
add logging per-param grad norm
2b3a59d
@xrsrke
fix conversation fail due to buffer on cpu
abbb4a1
@xrsrke
add ep, ep+tp, ep+dp process group initialization
e04c5f1
@xrsrke
add all-to-all, but not finished tests
009b8f7
@xrsrke
remove ep dimension from non-moe
309470f
@xrsrke
add the first all-to-all of ep
4cdba21
@xrsrke
add 2nd all-to-all
a48aca2
@xrsrke
fix ep inference (not correct outputs), shard expert weights in load_…
fd740e0 
…checkpoints
@xrsrke
fix runtime generation with ep=1
2539352
@xrsrke
refactor token dispatching out of moe
adbe7db
@xrsrke
partially working dispatching (incorrect order in the resulting dispa…
fb20017 
…tched input)
@xrsrke
add and fix all-to-all tests for custom input/output splits
97c56a4
@xrsrke
add moe token dispatcher test
628b53e
@xrsrke
fix all-to-all token dispatching with permute before all-to-all
5d513ac
xrsrke and others added 30 commits May 9, 2025 19:27
@xrsrke
remove moe_logging, model grads logging, clean up
c8239fb
@xrsrke
Merge branch 'dev' into dev_qwen_moe_with_ep
246604c
@xrsrke
rename use_haojun_permute to use_torch_permute
eadf97f
@NouamaneTazi
add te implem
84b3f55
@NouamaneTazi
Merge branch 'dev_qwen_moe_with_ep' of https://github.com/huggingface…
4a64da4 
…/nanotron into dev_qwen_moe_with_ep
@NouamaneTazi
init works
543072e
@NouamaneTazi
training works
ab649f0
@xrsrke
add moe ep
d99aa4b
@xrsrke
add saving moe checkpoints
aa4e754
@xrsrke
add resume moe checkpoints
f0d410f
@xrsrke
adapt edp to te moe
e04f6be
@xrsrke
fix edp+ep saving/resume checkpoint
d36905b
@xrsrke
use all_gather_into_tensor in token dispatcher
d24e603
@xrsrke
use fused permute, megablock's grouped gemm, and fuse multiplying exp…
559d532 
…ert output in unpermute. tested correctness with ep + edp
@NouamaneTazi
check init, fwd in case of ep=1
5d36aef
@NouamaneTazi
add all2all
df8cd20
@NouamaneTazi
loss goes down after fixing data labels
135c715
@NouamaneTazi
fix bias_activation_fusion and set to True by default
98df8cc
@NouamaneTazi
disable collect_env
88c2ee0
@NouamaneTazi
preprocess_data
acad09b
@NouamaneTazi
fix bug
1a8fb35
@NouamaneTazi
assert in case of top_k=1
d678ae7
@NouamaneTazi
add SimpleTokenDataset
e233a7d
@xrsrke
fix all-to-all
55fb95f
@xrsrke
scripts/scaling_moe_benchmark.py
aea7910
@xrsrke
fix no gradients, and expert device has no tokens
d30ae0f
@xrsrke
add compute moe params
88d4140
@xrsrke
fix sync_tied_weights_gradients without gradient accumulator
9e687c4
@xrsrke
fix sync_tied_weights_gradients
06faa91
@xrsrke
fix cannot import name 'Qwen2Config'
ce5f61c
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@xrsrke @NouamaneTazi @zzhhjjj