Skip to content
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

Tensor Parallelism Integration #3269

Merged
merged 89 commits into from
May 24, 2024
Merged

Tensor Parallelism Integration #3269

merged 89 commits into from
May 24, 2024

Conversation

mvpatel2000
Copy link
Contributor

@mvpatel2000 mvpatel2000 commented May 8, 2024 •
edited
Loading

What does this PR do?

This PR adds Tensor Parallelism (TP) integration. As part of this, we simplify the Trainer interface for parallelism.

Current limitations:

  • TP requires FSDP
  • TP means no load_monolith_rank0_only
  • TP does not support monolith checkpointing
  • TP must be all layers with same TP world size

Task List

  • Add TensorParallelism integration
  • Switch to parallelism_config interface
  • Code cleanup with respect to where device_mesh is tracked
  • Make device_mesh/DTensor on by default with FSDP
  • Docs update
  • Move all distributed work to a new folder

Test List

  • Add TP train test
  • Add TP sharded resumption test
  • Add TP full state dict resumption test

Follow-on PR:

  • Switch to Dataclass
  • Load monolith rank0 only support
  • Train TP only
  • Add TP vs. FSDP determinism test (blocked by no monolith checkpointing)

Note:

  • simulate_tp.py will be deleted before merge. It is currently used for testing.

@dakinggg dakinggg marked this pull request as draft May 8, 2024 21:01
mvpatel2000 and others added 21 commits May 8, 2024 14:05
@mvpatel2000
lint and pytorch pinning
1dd4245
@mvpatel2000
results
d38c254
tweak warnings
2c4cfc3
fix lint
bb74919
@mvpatel2000
lint
8ebc28a
@mvpatel2000
Merge branch 'mvpatel2000/nd-parallelism' of github.com-mvpatel2000:m…
e8253f5 
…vpatel2000/composer into mvpatel2000/nd-parallelism
filter
3cdf544
Merge branch 'mvpatel2000/nd-parallelism' of github.com-mvpatel2000:m…
e491ebc 
…vpatel2000/composer into mvpatel2000/nd-parallelism
add ckpt
1ec42c6
checkdown
d74082b
@mvpatel2000
Merge branch 'dev' into mvpatel2000/nd-parallelism
a47d3df
fix tests
da05e5d
fix
5d628d0
fix test
998da6b
fix tests
ee185e1
@mvpatel2000
Merge branch 'dev' into mvpatel2000/nd-parallelism
8b6664c
fix tests
50ff5bb
Merge branch 'mvpatel2000/nd-parallelism' of github.com-mvpatel2000:m…
3964e87 
…vpatel2000/composer into mvpatel2000/nd-parallelism
paralleli config
6f6d0d9
@mvpatel2000
Merge branch 'dev' into mvpatel2000/nd-parallelism
b150a33
fix some arg parsing
69111b4
b-chu
b-chu reviewed May 23, 2024
View reviewed changes
composer/core/state.py Show resolved Hide resolved
composer/core/state.py Show resolved Hide resolved
composer/core/state.py Show resolved Hide resolved
composer/core/state.py Show resolved Hide resolved
composer/core/state.py Outdated Show resolved Hide resolved
composer/trainer/trainer.py Show resolved Hide resolved
composer/trainer/trainer.py Show resolved Hide resolved
composer/trainer/trainer.py Show resolved Hide resolved
composer/trainer/trainer.py Show resolved Hide resolved
composer/trainer/trainer.py Outdated Show resolved Hide resolved
b-chu
b-chu approved these changes May 24, 2024
View reviewed changes
Copy link
Contributor

@b-chu b-chu left a comment

Choose a reason for hiding this comment

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

Minor comments, but looks good! Thanks for the PR!

composer/core/state.py Show resolved Hide resolved
composer/core/state.py Show resolved Hide resolved
composer/distributed/mosaic_fsdp.py Outdated Show resolved Hide resolved
@mvpatel2000 mvpatel2000 merged commit 09f14f9 into mosaicml:dev May 24, 2024
15 checks passed
dakinggg added a commit to dakinggg/composer that referenced this pull request May 25, 2024
@dakinggg
Revert "Tensor Parallelism Integration (mosaicml#3269)"
6dad3c8 
This reverts commit 09f14f9.
dakinggg added a commit that referenced this pull request May 25, 2024
@dakinggg
Revert TP integration (#3328)
f154ea6 
* Revert "Bugfixes to FSDP + TP (#3323)"

This reverts commit 79e79eb.

* Revert "Tensor Parallelism Integration (#3269)"

This reverts commit 09f14f9.
@mvpatel2000 mvpatel2000 mentioned this pull request May 28, 2024
Merged
8 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@mihir-db mihir-db mihir-db left review comments

@b-chu b-chu b-chu approved these changes

@j316chuck j316chuck Awaiting requested review from j316chuck

@eracah eracah Awaiting requested review from eracah

@XiaohanZhangCMU XiaohanZhangCMU Awaiting requested review from XiaohanZhangCMU

@snarayan21 snarayan21 Awaiting requested review from snarayan21

@dakinggg dakinggg Awaiting requested review from dakinggg

@bigning bigning Awaiting requested review from bigning

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
6 participants
@mvpatel2000 @bigning @b-chu @snarayan21 @dakinggg @mihir-db