search

princeton-nlp / LLM-Shearing

[ICLR 2024] Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning
https://arxiv.org/abs/2310.06694
MIT License
533 stars 39 forks source link

Avoid OOM using deepspeed zero-stage #47

Open gywlssww opened 7 months ago

gywlssww commented 7 months ago

When pruning or continue training, I got OOM error even with batch size 1 & 8 GPUs. I'm trying to use deepspeed zero stage or FSDP to avoid OOM. However, it seems your code isn't compatible well with deepspeed. Could you help me with these OOM errors?

xiamengzhou commented 7 months ago

Hi!! The current code supports FSDP as it is, and should be able to run on 8 GPUs. Are you using a separate implementation?

gywlssww commented 7 months ago

I'm stuck with OOM when using only FSDP too. Do I really need to run the code based on slurm?? I think it is okay to run the code based on vanilla composer without slurm. And also when I use 16 GPUs, gpuutil doesn't rise and only takes up the memory.

xiamengzhou commented 7 months ago

@gywlssww The code can be absolutely run without slurm! Do you mean you still get OOM issues even if you use 16 GPUs for pruning?