Enable FSDP2 cpu offloading #624
Conversation
facebook-github-bot
added
the
CLA Signed
train.py
Outdated
| model.to_empty(device=init_device) | ||
| model.init_weights() | ||
| if job_config.training.enable_cpu_offload: | ||
| with torch.device("cuda"): |
There was a problem hiding this comment.
I somehow think the model.init_weights(buffer_device="cuda") change sounds better. It is straightforward on what we want to achieve, while making minimum change to code.
tianyu-l
left a comment
tianyu-l
left a comment
There was a problem hiding this comment.
Looks good in general. Had some inline comments. In particular, let's figure out if CPU offloading and PP should coexist; if so we should add support for that as well.
train.py
Outdated
| init_device = ( | ||
| "cpu" | ||
| if job_config.checkpoint.create_seed_checkpoint | ||
| or job_config.training.enable_cpu_offload |
There was a problem hiding this comment.
Don't we need to do the same for the PP case (several lines above)? Or are we assuming if PP is used, CPU offloading is not an option?
There was a problem hiding this comment.
I test it with pp.
Based on llama2_7b.toml, cpu_offload also works but has a much obvious latency in training.
| "Enable CPU Offload with PP", | ||
| "enable_cpu_offload+PP", | ||
| ngpu=4, | ||
| ), |
There was a problem hiding this comment.
test with pp, could remove pp later if not necessary in the CI test
resolve pytorch#620 Add config: `--training.enable_cpu_offload` Command: `CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.sh` For non-pp case: <img width="611" alt="Screenshot 2024-10-23 at 1 45 56 PM" src="https://github.com/user-attachments/assets/8692f8a6-c0f3-460e-8eb6-7f7195bed370"> For pp case: <img width="587" alt="cpu offload+pp" src="https://github.com/user-attachments/assets/73e40861-47e2-4845-a41c-4bfea2860109">
3 participants
resolve #620
Add config:
--training.enable_cpu_offloadCommand:
CONFIG_FILE="./train_configs/llama3_8b.toml" ./run_llama_train.shFor non-pp case:
For pp case:
