[BUG] Step3 RLHF Training failed with GPTJ 6B on 8x32GB V100

Describe the bug
This may not be bug but to ask people for their suggestion if anyone has been successful in fixing it.
I am running step 3 RLHF training with both actor and critic GPTJ6B. I used the following command:

deepspeed --master_port 12346 main.py --data_path Dahoas/rm-static --data_split 2,4,4 --actor_model_name_or_path EleutherAI/gpt-j-6b --critic_model_name_or_path /mnt/data/ds-chat-step2GPTJoutput --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 50 --max_prompt_seq_len 256 --actor_learning_rate 5e-4 --critic_learning_rate 5e-6 --actor_weight_decay 0.1 --critic_weight_decay 0.1 --num_train_epochs 1 --lr_scheduler_type cosine --gradient_accumulation_steps 1 --num_warmup_steps 100 --deepspeed --seed 1234 --enable_hybrid_engine --inference_tp_size 1 --actor_zero_stage 3 --critic_zero_stage 3 --actor_gradient_checkpointing --critic_gradient_checkpointing --actor_lora_dim 64 --actor_lora_module_name decoder.layers. --critic_lora_dim 64 --critic_lora_module_name decoder.layers. --output_dir /mnt/data/ds-chat-step3output

Log output
RuntimeError: CUDA out of memory. Tried to allocate 2.72 GiB (GPU 0; 31.75 GiB total capacity; 23.68 GiB already allocated; 2.36 GiB free; 27.80 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

To Reproduce
Steps to reproduce the behavior:

1. Install the requirements in deepspeed-chat repo
2. Run the command above on Azure ML

Expected behavior
Run successfully

Screenshots

System info (please complete the following information):

• Python 3.8, pytorch 1.12.1 + cuda 11.6
• GPU count and types: 8x32GV100
• (if applicable) what DeepSpeed-MII version are you using: deepspeed 0.9.0

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.

Start with main.py and reproduce the supplied Step 3 RLHF command using the listed Python, PyTorch, CUDA, DeepSpeed, GPT-J, and 8x32GB V100 setup. Investigate the reported CUDA out-of-memory failure during training; the issue does not identify a target file or test, and completion would require defining a successful run or a confirmed fix.