[BUG] Fix ds_chat regression #7014
Description
The module injection refactoring in PR #7004 broke existing code leading to ds-chat CI failures as below:
[rank2]: File "/scratch/azureml/cr/j/1df407cd106142c2afb0cda94a311736/exe/wd/actions-runner/_work/DeepSpeed/DeepSpeed/unit-test-venv/lib/python3.10/site-packages/deepspeed/runtime/hybrid_engine.py", line 293, in create_inference_containers
[rank2]: self._other_layers.append(self.inference_policies[child.__class__][0](
[rank2]: TypeError: LinearLayer.__init__() missing 2 required positional arguments: 'module' and 'mp_group'To Reproduce
Run ds-chat workflow. Or after environment setup, you the following command will reproduce the issue.
deepspeed --num_gpus=1 main.py --data_path Dahoas/rm-static --data_split 2,4,4 --actor_model_name_or_path facebook/opt-125m --critic_model_name_or_path /tmp/ds_chat --num_padding_at_beginning 1 --per_device_generation_batch_size 4 --per_device_training_batch_size 4 --generation_batches 1 --ppo_epochs 1 --max_answer_seq_len 256 --max_prompt_seq_len 256 --actor_learning_rate 9.65e-6 --critic_learning_rate 5e-6 --num_train_epochs 1 --lr_scheduler_type cosine --gradient_accumulation_steps 1 --actor_weight_decay 0 --critic_weight_decay 0 --num_warmup_steps 100 --deepspeed --seed 1234 --actor_dropout 0.0 --print_answers --actor_zero_stage 3 --critic_zero_stage 3 --output_dir z3_he_true_offload_true_lora_true --enable_hybrid_engine --offload --actor_lora_dim 128 --actor_lora_module_name decoder.layers. --enable_test_mode --test_stop_step 5
Expected behavior
ds-chat CI should work.