Skip to content

[WIP] Sharing work on FSDP2 x QLoRA profiling #1

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 55 commits into
base: main
Choose a base branch
from
Open

[WIP] Sharing work on FSDP2 x QLoRA profiling #1

wants to merge 55 commits into from

Conversation

@janeyx99
Copy link
Owner

@janeyx99 janeyx99 commented May 16, 2024

Context

What is the purpose of this PR? Is it to

  • add a new feature
  • fix a bug
  • update tests and/or documentation
  • other (please add here)

Please link to any issues this PR addresses.

Changelog

What are the changes made in this PR?

Test plan

Please make sure to do each of the following if applicable to your PR. (If you're not sure about any one of these just ask and we will happily help.)

  • run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
  • add unit tests for any new functionality
  • update docstrings for any new or updated methods or classes
  • run unit tests via pytest tests
  • run recipe tests via pytest tests -m integration_test
  • manually run any new or modified recipes with sufficient proof of correctness
    • include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

weifengpy and others added 30 commits April 23, 2024 17:45
@weifengpy
enable LoRA + FSDP2
e5826a1 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
reset params for lora weights and rope
64fc870 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
support lora weights checkpoint and checkpoint utils
0cd21c6 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
fix lora meta device bug
589191e 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
save optim state dict
c801f26 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
mark TODO
19a2d70 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
optimizer foreach=True for DTensor
441da10 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
clip grad norm
750b9e5 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
switch to ptd state dict api
3d632d5 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
add profiler
cb3abb3 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
qlora 7b config
dfcdde3 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
use torchao copy_
e68804a 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
Merge pull request #1 from weifengpy/fsdp2
b6fad93 
use torchao copy_
@weifengpy
enable saving checkpoint
d6af9a2 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
Merge pull request meta-pytorch#2 from weifengpy/fsdp2
7bbe522 
enable saving checkpoint
@weifengpy
optimizer state dict: load on rank0 and broadcast
b616394 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
import Optimizer
a400497 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
resume training
e9de63c 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
prepare for full test
05d3895 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
prepare for full test
7a5bb80 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
remove profiler
64bf49c 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
passed integration test
cb1bba4 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
remove uncesssary change
ac516e9 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
Merge branch 'main' into fsdp2
bfde704
@weifengpy
bring back state dict validation
102db31 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
align indent on comment
0b66651 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
remove unused import
672aabb 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
switch to ptd state dict and keep self implemented in record
6af2723 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
clean unused code
42ad99c 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
remove cuda value error
74f6175 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
weifengpy and others added 20 commits May 8, 2024 00:19
@weifengpy
comment on to_empty
f1b8a5e 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
fix memory issues by switching model state dict api
36e6829 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
clean for review
08cd1fd 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
Merge branch 'main' into fsdp2
559bc4d
@weifengpy
fix linter
2333134 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
fix checkpoint loading
49a0364 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
expecttest CI depedency
dc2ce02 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
ci depdencecy
0a604aa 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
fix CI issue
fa83140 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
Merge branch 'main' into qlora
6203a1f
@weifengpy
Merge branch 'pytorch:main' into fsdp2
4b5a895
@weifengpy
Merge branch 'fsdp2' into qlora
1080e2c
@weifengpy
rebase qlora
1a70498 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
rebase qlora
cb862e9 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
sync lora changes
21f5458 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
push qlora for perf measurement
33773bd 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
fix meta init + cpu offloading
483028b 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
init RotaryPositionalEmbeddings in both fresh training and resume
cf42618 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
import cpu offloading when needed
b519d50 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
FSDP(CheckpointWrapper(Model))
8600ced 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
weifengpy and others added 5 commits May 22, 2024 14:37
@weifengpy
bring back cpu offloading
b2fd531 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
remove model.to
bb8a8bc 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
apply nf4 when loading model state dict
db71c5c 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@weifengpy
move lora to cpu when cpu offloading
16bf2de 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:
@janeyx99
added compile
4908773
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

@janeyx99 @weifengpy