feat(fsdp): add memory_efficient_load option for large model initialization#862
Conversation
Summary of ChangesHello @rchardx, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses Out-Of-Memory (OOM) issues that can occur during the initialization of very large models in a Fully Sharded Data Parallel (FSDP) setup. It introduces a new configuration option, Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
The pull request introduces a memory_efficient_load option for FSDP, which allows initializing large models on the CPU and then broadcasting weights after FSDP sharding. This is a valuable performance improvement for reducing peak GPU memory during initialization. The changes are well-implemented across the configuration, engine logic, and documentation, ensuring consistency and clarity. The new feature is clearly explained in the best practices and CLI reference documentation.
There was a problem hiding this comment.
Pull request overview
This PR adds a memory_efficient_load configuration option for FSDP training to address out-of-memory (OOM) issues during large model initialization, specifically fixing issue #845 regarding Qwen3 32B OOM with FSDP2.
Changes:
- Added
fsdp.memory_efficient_loadboolean config option that enables memory-efficient model loading - Modified model initialization to create model structure on CPU and load pretrained weights only on rank 0, then broadcast to all ranks after FSDP sharding
- Updated documentation to explain the new feature and its benefits for handling OOM during initialization
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| areal/api/cli_args.py | Added memory_efficient_load field to FSDPEngineConfig with appropriate metadata |
| areal/engine/fsdp_engine.py | Implemented memory-efficient loading logic: model creation on CPU, conditional weight loading on rank 0, and broadcasting after FSDP sharding |
| docs/cli_reference.md | Added documentation for the new memory_efficient_load parameter in the FSDP configuration table |
| docs/best_practices/handling_oom.md | Added section explaining how to use memory-efficient loading to resolve initialization OOM errors |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
rchardx
force-pushed
the
rchardx/fsdp_oom
branch
from
ca08863 to
0bfada4
Compare
rchardx
temporarily deployed
to
AReaL-unittests
GitHub Actions
Inactive
…zation When training large models, loading weights directly onto each GPU can cause OOM during initialization. This adds `fsdp.memory_efficient_load` config option that initializes model structure on CPU, then rank 0 loads pretrained weights and broadcasts to all ranks after FSDP sharding. This reduces peak GPU memory during model loading at the cost of slightly longer initialization time.
rchardx
force-pushed
the
rchardx/fsdp_oom
branch
from
0bfada4 to
4047986
Compare
rchardx
added
safe-to-test
rchardx
temporarily deployed
to
AReaL-unittests
GitHub Actions
Inactive
…zation (areal-project#862) When training large models, loading weights directly onto each GPU can cause OOM during initialization. This adds `fsdp.memory_efficient_load` config option that initializes model structure on CPU, then rank 0 loads pretrained weights and broadcasts to all ranks after FSDP sharding. This reduces peak GPU memory during model loading at the cost of slightly longer initialization time.
…zation (areal-project#862) When training large models, loading weights directly onto each GPU can cause OOM during initialization. This adds `fsdp.memory_efficient_load` config option that initializes model structure on CPU, then rank 0 loads pretrained weights and broadcasts to all ranks after FSDP sharding. This reduces peak GPU memory during model loading at the cost of slightly longer initialization time.
Description
When training large models, loading weights directly onto each GPU can cause OOM during initialization. This adds
fsdp.memory_efficient_loadconfig option that initializes model structure on CPU, then rank 0 loads pretrained weights and broadcasts to all ranks after FSDP sharding.This reduces peak GPU memory during model loading at the cost of slightly longer initialization time.
Related Issue
Fixes #845
Type of Change
work as expected)
Checklist
jb build docs/gemini review)