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

Fsdp 2.0 integration #206

Draft
wants to merge 44 commits into
base: main
Choose a base branch
from
Draft

Fsdp 2.0 integration #206

wants to merge 44 commits into from

Conversation

le1nux
Copy link
Member

@le1nux le1nux commented Aug 4, 2024 •
edited
Loading

What does this PR do?

This PR ..

General Changes

  • ..

Breaking Changes

  • ..

TODOs

  • Need to check if we want to use tied weights for the embedding and LM head and how this would be parallelized with FSDP and tensor parallelism

Checklist before submitting final PR

  • My PR is minimal and addresses one issue in isolation
  • I have merged the latest version of the target branch into this feature branch
  • I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
  • I have run a sample config for model training
  • I have checked that all tests run through (python tests/tests.py)

@le1nux
feat: added FSDP2MixedPrecisionSettings for FSDP2
e8f73bc
@le1nux
feat: drafted FSDP2
e290385
@le1nux
feat: added FSDP 2 Wrapper Config
79cd87f
@le1nux
chore: Added ModelStateError class
4ed5013
@le1nux
feat: wired up FSDP 2 draft
9190468
@le1nux
feat: added assertion that each parameter belongs to exactly one para…
ae8b337 
…meter group
@le1nux
refactor: layer norm is now instantiated within the GPT2 model
1fe8b6b
@le1nux
refactor: added fsdp 1 and fsdp 2 support side by side
f13ccb4
@le1nux
refactor: upated lorem ipsum for meta device support and layer norm i…
3e6e830 
…nit in gp2 model
@le1nux
feat: added config lorem ipsum fsdp 2
d3fb538
@le1nux
feat: added device mesh component
939beff
@le1nux
feat: integrated device mesh into fsdp 2 call
29977e4
@le1nux
refactor: changed gradient clipping to match FSDP 2 implementation
2c07800
@le1nux
refactor: model forward pass must now be initiated via __call__ and n…
102c99a 
…ot via forward method
@le1nux
refactor: removed T_co generics from dataloader
7b3c9f9
@le1nux
refactor: removed device mesh from collator
64d401d
@le1nux
chore: cleaned up fsdp2 integration
737cecf
@le1nux
fix: fixed failing test_evaluate_cpu test
60316b3
@le1nux
fix: fixed more failing tests
d7ac011
@le1nux
chore: Merge branch 'main' into fsdp_2.0_integration
b41f01e
@le1nux
refactor: replaced norm dependency with norm config in the test confi…
f718201 
…g lorem ipsum
@le1nux
refactor: adapted the test configs w.r.t. meta device and layer norm …
080e200 
…configs
@le1nux
feat: added weight init to model after fsdp2 sharding
05eae90
@le1nux
feat: added torch compile
d6deb22
@le1nux
feat: added data_parallel_degree, tensor_parallel_degree and pipeline…
ceb701d 
…_parallel_degree to DeviceMesh
@le1nux
feat: updated fsdp 2 config yaml
d2d3e72
@le1nux
feat: added logging to device mesh init
02996c8
@le1nux
feat: added validation check for enable_loss_parallel in device mesh
80c9fa9
@le1nux
feat: added initial tensor parallelization plan
a46c9e3
@le1nux
refactor: added lm_head to tp plan
eba6cbc
@le1nux
feat: added model parallelization
3b77516
@le1nux
feat: wired up tensor parallelism component
4e5b5b9
@le1nux
feat: now using foreach flag in clip_grad_norm_
bb32750
@le1nux
refactor: moved meta device initialization to ModelFactory.get_weight…
a906483 
…_initalized_model
@le1nux
refactor: updated TP plans for sequence parallelism
507ccfb
@le1nux
refactor: moved lm_head into transformer module dict
30e2ff7
@le1nux
refactor: implemented tensor parallelism for GPT2 model
c1df139
@le1nux
refactor: added local variable for simplified tensor parallelism debu…
9dd33fe 
…gging
@le1nux
refactor: refined sequence and tensor parallelism plan
bac4b9f
@le1nux
feat: added function that recursively instantiates all norms by calli…
8842a4f 
…ng reset_parameters
@le1nux
refactor: separated ModelFactory into GenericModelFactory and GPT2Mod…
1bf6808 
…elFactory
@le1nux
feat: added initial fsdp2 unit test
87a0b6c
@le1nux
feat: added weight equivalence test
f96ef5e
@le1nux
chore: moved reshard_after_forward calculation outside
0375df5
@le1nux le1nux added the enhancement New feature or request label Aug 23, 2024
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
enhancement New feature or request
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@le1nux