Step3: 8 * A100-40G training LLAMA2-7B and OPT-350M out of memory.

Hi,
In step 3, run the following command and getting "OOM" when Initializing Ref Model (Actor Model initialized perfectly):

Actor_Lr=9.65e-6
Critic_Lr=5e-6
deepspeed --master_port 12346 main.py
--data_path Dahoas/rm-static
--data_split 2,4,4
--actor_model_name_or_path LLAMA-2-7B-path
--critic_model_name_or_path OPT-350M-path
--num_padding_at_beginning 1
--per_device_generation_batch_size 1
--per_device_training_batch_size 1
--generation_batches 1
--ppo_epochs 1
--max_answer_seq_len 256
--max_prompt_seq_len 256
--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 1
--actor_gradient_checkpointing
--critic_gradient_checkpointing
--disable_actor_dropout
--num_warmup_steps 100
--deepspeed
--seed 1234
--actor_zero_stage 2
--critic_zero_stage 2
--enable_hybrid_engine
--offload_reference_model
--output_dir training_log_output

When I use zero_stage 3 like:

--actor_zero_stage 3
--critic_zero_stage3
--offload \

I got :

Traceback (most recent call last):
File "DeepSpeedExamples-master/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py", line 631, in
main()
File "DeepSpeedExamples-master/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py", line 447, in main
rlhf_engine = DeepSpeedRLHFEngine(
File "DeepSpeedExamples-master/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py", line 48, in init
self.actor = self._init_actor(
File "DeepSpeedExamples-master/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py", line 125, in init_actor
actor_engine, * = deepspeed.initialize(model=actor_model,
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/init.py", line 153, in initialize
engine = DeepSpeedHybridEngine(args=args,
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py", line 52, in init
self.create_inference_module()
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py", line 359, in create_inference_module
self.create_inference_containers(self.module)
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py", line 308, in create_inference_containers
self.create_inference_containers(child, layer_id=layer_id)
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py", line 308, in create_inference_containers
self.create_inference_containers(child, layer_id=layer_id)
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py", line 288, in create_inference_containers
self._inference_containers.append(self.inference_policies[child.class][0](
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py", line 109, in new_inference_container
_container.create_ds_model_config()
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/module_inject/containers/base.py", line 84, in create_ds_model_config
self.set_hidden_heads(*self.policy.get_hidden_heads())
File "anaconda3/envs/llm/lib/python3.10/site-packages/deepspeed/module_inject/containers/llama.py", line 119, in get_hidden_heads
return self.client_module.self_attn.q_proj.weight.shape[1],
IndexError: tuple index out of range

even before Initializing the Actor Model.

Start with training/step3_rlhf_finetuning/main.py and rlhf_engine.py, then reproduce the Step 3 command with the reported zero-stage and offload settings. Trace initialization through the DeepSpeed hybrid engine and compare the OOM path with the IndexError path; done means the configuration initializes and runs without either failure.