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

[GasolSun36]

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.

[imraviagrawal]

Are you able to solve this error??

[Ricardokevins]

similar issue https://github.com/microsoft/DeepSpeedExamples/issues/448