KEMBAR78
Updating memory optimization overview by SalmanMohammadi · Pull Request #1916 · meta-pytorch/torchtune · GitHub
Skip to content

Updating memory optimization overview #1916

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
Merged
merged 7 commits into from
Oct 29, 2024
Merged

Updating memory optimization overview #1916

merged 7 commits into from
Oct 29, 2024

Conversation

SalmanMohammadi
Copy link
Contributor

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)

Updating the memory optimization tutorial with CPUOffloadingOptimizer,

                                      .,'
                                   .'`.'
                                  .' .'
                      _.ood0Pp._ ,'  `.~ .q?00doo._
                  .od00Pd0000Pdb._. . _:db?000b?000bo.
                .?000Pd0000PP?000PdbMb?000P??000b?0000b.
              .d0000Pd0000P'  `?0Pd000b?0'  `?000b?0000b.
             .d0000Pd0000?'     `?d000b?'     `?00b?0000b.
             d00000Pd0000Pd0000Pd00000b?00000b?0000b?0000b
             ?00000b?0000b?0000b?b    dd00000Pd0000Pd0000P
             `?0000b?0000b?0000b?0b  dPd00000Pd0000Pd000P'
              `?0000b?0000b?0000b?0bd0Pd0000Pd0000Pd000P'
                `?000b?00bo.   `?P'  `?P'   .od0Pd000P'
                  `~?00b?000bo._  .db.  _.od000Pd0P~'
                      `~?0b?0b?000b?0Pd0Pd000PdP~'

@pytorch-bot
Copy link

pytorch-bot bot commented Oct 29, 2024 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1916

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8cafa6f with merge base 1bbd749 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 29, 2024
@SalmanMohammadi SalmanMohammadi changed the title Upadting Updating memory optimization overview Oct 29, 2024
joecummings
joecummings approved these changes Oct 29, 2024
View reviewed changes
Copy link
Member

@joecummings joecummings left a comment

Choose a reason for hiding this comment

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

I am the King of All Nits, bow before me

docs/source/tutorials/memory_optimizations.rst

* The CPU optimizer step is often the bottleneck when optimizer CPU offload is used. To minimize the slowdown, it is recommended to (1) use full ``bf16`` training so that parameters, gradients, and optimizer states are in ``bf16``; and (2) give GPU more work per optimizer step (e.g. larger batch size with activation checkpointing, gradient accumulation).
* Gradient accumulation should always be set to 1 when ``offload_gradients=True``, as gradients are cleared on GPU every backward pass.
* This optimizer works by keeping a copy of parameters and pre-allocating gradient memory on CPU. Therefore, expect your RAM usage to increase by 4x model size.
Copy link
Member

Choose a reason for hiding this comment

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

Is it exactly 4x model size? Or roughly?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

According to the ao docs yeah

To minimize the amount of CPU<->GPU data transfer, we keep a copy of parameters and pre-allocate gradients memory on CPU. Therefore, expect your RAM usage to increase by 2x model size + optimizer state (which is 2x model size for Adam).

and since we always use adam it is 4x

@SalmanMohammadi SalmanMohammadi merged commit 0cc1b1f into meta-pytorch:main Oct 29, 2024
16 checks passed
@SalmanMohammadi SalmanMohammadi deleted the cpu_offload_docs branch October 29, 2024 14:29
@SalmanMohammadi SalmanMohammadi mentioned this pull request Oct 30, 2024
Closed
34 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@joecummings joecummings joecummings approved these changes

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@SalmanMohammadi @joecummings @facebook-github-bot