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 optimizations #194

Merged

epwalsh merged 3 commits into v2 from epwalsh/v2-fsdp-optimizations

Mar 7, 2025

Member

epwalsh commented Mar 7, 2025 •

edited

Loading

Improve FSDP settings for gradient accumulation. For example, with HSDP we delay the gradient all-reduce until the final micro-batch.
Always wrap the embeddings separately but don't unshard them in the backwards pass since embedding params aren't needed for the gradient computation.

epwalsh added 3 commits

March 7, 2025 10:42


          improve FSDP settings for gradient accumulation

70ddda5


          wrap embeddings seperate

e42e807


          clean up

eebc709

epwalsh merged commit f52d11f into v2

14 checks passed

epwalsh deleted the epwalsh/v2-fsdp-optimizations branch

March 7, 2025 19:32

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet