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

Add FP8 TransformerEngine activation checkpointing #3156

Merged
merged 16 commits into from
Apr 9, 2024
Merged

Add FP8 TransformerEngine activation checkpointing #3156

merged 16 commits into from
Apr 9, 2024

Conversation

cli99
Copy link
Contributor

@cli99 cli99 commented Mar 29, 2024 •
edited
Loading

What does this PR do?

This PR adds FSDP activation checkpointing support to TransformerEngine FP8
example run: 3-mpt-30b-ba1280-actfull-fp8-J6LQuV

What issue(s) does this change relate to?

FSDP activation checkpointing currently does not work with FP8 TE

Before submitting

  • Have you read the contributor guidelines?
  • Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
  • Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
  • Did you update any related docs and document your change?
  • Did you update any related tests and add any new tests related to your change? (see testing)
  • Did you run the tests locally to make sure they pass?
  • Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

@cli99 cli99 requested a review from j316chuck April 1, 2024 18:28
@cli99 cli99 marked this pull request as ready for review April 1, 2024 18:28
j316chuck
j316chuck approved these changes Apr 1, 2024
View reviewed changes
Copy link
Contributor

@j316chuck j316chuck left a comment •
edited
Loading

Choose a reason for hiding this comment

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

Thanks for adding this Cheng! Mind pasting the llm-foundry run name so we can reproduce this result easily?

composer/trainer/dist_strategy.py Show resolved Hide resolved
@cli99 cli99 requested a review from j316chuck April 1, 2024 21:01
@cli99 cli99 merged commit a471278 into mosaicml:dev Apr 9, 2024
14 checks passed
staghado pushed a commit to lightonai/composer that referenced this pull request Apr 13, 2024
@cli99 @j316chuck
@dakinggg @staghado
Add FP8 TransformerEngine activation checkpointing (mosaicml#3156)
3331293 
* add te checkpoint wrapper

* remove extra wrapper

* add option to shard fp8 weights

* rename

* rename

* add log info

* fix te checkpoint

* update format

* add comment

---------

Co-authored-by: Charles Tang <j316chuck@users.noreply.github.com>
Co-authored-by: Daniel King <43149077+dakinggg@users.noreply.github.com>
staghado pushed a commit to lightonai/composer that referenced this pull request Apr 13, 2024
@cli99 @j316chuck
@dakinggg @staghado
Add FP8 TransformerEngine activation checkpointing (mosaicml#3156)
f4827e0 
* add te checkpoint wrapper

* remove extra wrapper

* add option to shard fp8 weights

* rename

* rename

* add log info

* fix te checkpoint

* update format

* add comment

---------

Co-authored-by: Charles Tang <j316chuck@users.noreply.github.com>
Co-authored-by: Daniel King <43149077+dakinggg@users.noreply.github.com>
j316chuck added a commit that referenced this pull request May 16, 2024
@cli99 @j316chuck @dakinggg
Add FP8 TransformerEngine activation checkpointing (#3156)
af1b353 
* add te checkpoint wrapper

* remove extra wrapper

* add option to shard fp8 weights

* rename

* rename

* add log info

* fix te checkpoint

* update format

* add comment

---------

Co-authored-by: Charles Tang <j316chuck@users.noreply.github.com>
Co-authored-by: Daniel King <43149077+dakinggg@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@j316chuck j316chuck j316chuck approved these changes

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
@cli99 @j316chuck @dakinggg