A100 40 GB: OOM on step-3 for opt-6.7B

[akashsaravanan-georgian]

Hi, I managed to train step 1 and step 2 for a 6.7B actor model and 350m reward model but I keep running into an out of memory issue for step 3. I was wondering what config was used in your tests with this setup?

[puyuanOT]

I guess you could use LoRA or gradient checkpoint?

[akashsaravanan-georgian]

I guess you could use LoRA or gradient checkpoint?

I've tried both based on the message that pops up when the OOM error occurs. I also tried reducing the batch size and that didn't help either. So I was wondering how the folks here got it running.

[tiandongtao]

same OOM question my step3 single_gpu config:

deepspeed main.py
--data_path wangrui6/Zhihu-KOL
--data_split 2,4,4
--actor_model_name_or_path $ACTOR_MODEL_PATH
--critic_model_name_or_path $CRITIC_MODEL_PATH
--num_padding_at_beginning 1
--per_device_train_batch_size 1
--per_device_mini_train_batch_size 1
--generation_batch_numbers 1
--ppo_epochs 1
--max_answer_seq_len 256
--max_prompt_seq_len 256
--ppo_epochs 1
--actor_learning_rate ${Actor_Lr}
--critic_learning_rate ${Critic_Lr}
--actor_weight_decay 0.1
--critic_weight_decay 0.1
--num_train_epochs 1
--lr_scheduler_type cosine
--gradient_accumulation_steps 16
--num_warmup_steps 100
--deepspeed --seed 1234
${ACTOR_ZERO_STAGE}
${CRITIC_ZERO_STAGE} ${OFFLOAD}
--actor_lora_dim 128
--actor_gradient_checkpointing
--critic_gradient_checkpointing
--enable_hybrid_engine
--offload_reference_model
--output_dir $OUTPUT
&> $OUTPUT/training.log

[akashsaravanan-georgian]

As an update, I used the default configuration for 6.7B (run_6.7b_lora.sh). I made the following changes:

• reduce per_device_train_batch_size and per_device_mini_train_batch_size to 1.
• offload reference model (--offload_reference_model)
• use gradient checkpointing (--actor_gradient_checkpointing --critic_gradient_checkpointing)
• use lora instead of the above (--only_optimize_lora)

All of these are on a standard A100 40GB instance with 12 CPUs and 85 GB of RAM. In all cases, I ran into a CUDA OOM issue. I assumed that this was a supported configuration since Table 3 here says so but it seems like it is wrong.

Tagging @tjruwase and @yaozhewei