GPU memory requirements of train_ppo

[skepsun]

I tried to perform sft and ppo on a 3b llama model on 4xA100 (40g). In ppo stage (using bloom-1b7-rm as reward model), I always got OOM errors even after using very small batch sizes (1 for actor, 4 for reward model). I'm curious about how much GPU memory should I prepare for training ppo on a 3b model?

[NicholasCao]

Without trying the 3B model, I followed my hyperparameters (batch_size, seq_len) to train the 1.7B model with A100 (40G)*4.
One suggestion is to reduce the seq_len.

llama 3b? A100 16G? 🤣

[skepsun]

Without trying the 3B model, I followed my hyperparameters (batch_size, seq_len) to train the 1.7B model with A100 (40G)*4. One suggestion is to reduce the seq_len.

llama 3b? A100 16G? 🤣

Sorry for typo (not 16g, 40g)🤣. I tried open-llama-3b. Maybe I should reduce seq_len?

[NicholasCao]

The default length is 1024, which can be reduced appropriately to get the program running so that the required GPU memory can be estimated