Avoid OOM using deepspeed zero-stage

[gywlssww]

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]

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]

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]

@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?