Repository: huggingface/trl

Issue:
PPO: DeepSpeed ZeRO 3 fails if reference model is created via `trl.create_reference_model()`
Recent changes have added DeepSpeed ZeRO 3 support for `PPOTrainer`. This can be seen in the `sentiment_tuning.py` example. If we modify the script to create the reference model after the regular model via `trl.create_reference_model(model, num_shared_layers)`, we encounter the following error:

```
Traceback (most recent call last):                                                                                                                                                                  
  File "/home/xxx/repos/llm-training/trl-latest/examples/scripts/sentiment_tuning.py", line 210, in <module>                                                                         
    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)                                                                                                                              
  File "/home/xxx/.conda/envs/trl/lib/python3.10/contextlib.py", line 79, in inner                                                                                                   
    return func(*args, **kwds)                                                                                                                                                                      
  File "/home/xxx/repos/llm-training/trl-latest/trl/trainer/ppo_trainer.py", line 696, in step                                                                                       
    ref_logprobs, ref_logits_or_none, _, _ = self.batched_forward_pass(                                                                                                                             
  File "/home/xxx/.conda/envs/trl/lib/python3.10/contextlib.py", line 79, in inner                                                                                                   
    return func(*args, **kwds)                                                                                                                                                                      
  File "/home/xxx/repos/llm-training/trl-latest/trl/trainer/ppo_trainer.py", line 955, in batched_forward_pass                                                                       
    logits, _, values = model(**input_kwargs)                                                                                                                                                       
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                 
    return forward_call(*args, **kwargs)                                                                                                                                                            
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn                                                                      
    ret_val = func(*args, **kwargs)                                                                                                                                                                 
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1801, in forward                                                                   
    loss = self.module(*inputs, **kwargs)                                                                                                                                                           
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl                                                                 
    result = forward_call(*args, **kwargs)                                                                                                                                                          
  File "/home/xxx/repos/llm-training/trl-latest/trl/models/modeling_value_head.py", line 165, in forward                                                                             
    base_model_output = self.pretrained_model(                                                                                                                                                      
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl                                                                 
    result = forward_call(*args, **kwargs)                                                                                                                                                          
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 1076, in forward                                                     
    transformer_outputs = self.transformer(                                                                                                                                                         
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl                                                                 
    result = forward_call(*args, **kwargs)                                                                                                                                                          
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 843, in forward                                                      
    inputs_embeds = self.wte(input_ids)                                                                                                                                                             
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl                                                                 
    result = hook(self, args)                                                                                                                                                                       
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn                                                                      
    ret_val = func(*args, **kwargs)                                                                                                                                                                 
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 392, in _pre_forward_module_hook
    self.pre_sub_module_forward_function(module)
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/deepspeed/runtime/zero/parameter_offload.py", line 505, in pre_sub_module_forward_function
    param_coordinator.fetch_sub_module(sub_module, forward=prev_grad_state)
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/xxx/.conda/envs/trl/lib/python3.10/site-packages/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 310, in fetch_sub_module
    assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
AssertionError: {'id': 151, 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 0, 'shape': (0,), 'ds_shape': (0,), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules':
 {171}, 'ds_tensor.shape': torch.Size([0])}
```

This occurs with the following changes to `sentiment_tuning.py`. Note that it occurs when setting `num_shared_layers=0` as well.

```
diff --git a/trl-latest/examples/scripts/sentiment_tuning.py b/trl-latest/examples/scripts/sentiment_tuning.py
index df0576d..f13d3f4 100644
--- a/trl-latest/examples/scripts/sentiment_tuning.py
+++ b/trl-latest/examples/scripts/sentiment_tuning.py
@@ -23,7 +23,8 @@ from peft import LoraConfig
 from tqdm import tqdm
 from transformers import AutoTokenizer, pipeline
 
-from trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead, PPOConfig, PPOTrainer, set_seed
+from trl import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead, PPOConfig, PPOTrainer, set_seed, \
+    create_reference_model
 from trl.core import LengthSampler
 
 
@@ -121,24 +122,33 @@ def collator(data):
 # set seed before initializing value head for deterministic eval
 set_seed(args.ppo_config.seed)
 print(f"Args: {args}")
-# Now let's build the model, the reference model, and the tokenizer.
-if not args.use_peft:
-    ref_model = trl_model_class.from_pretrained(args.ppo_config.model_name, trust_remote_code=True)
-    device_map = None
-    peft_config = None
-else:
-    peft_config = args.peft_config
-    ref_model = None
-    # Copy the model to each device
-    device_map = {"": Accelerator().local_process_index}
 
 model = trl_model_class.from_pretrained(
     args.ppo_config.model_name,
     trust_remote_code=True,
-    device_map=device_map,
-    peft_config=peft_config,
+    device_map=None,
+    peft_config=None,
 )
 
+ref_model = create_reference_model(model, num_shared_layers=5)
+# # Now let's build the model, the reference model, and the tokenizer.
+# if not args.use_peft:
+#     ref_model = trl_model_class.from_pretrained(args.ppo_config.model_name, trust_remote_code=True)
+#     device_map = None
+#     peft_config = None
+# else:
+#     peft_config = args.peft_config
+#     ref_model = None
+#     # Copy the model to each device
+#     device_map = {"": Accelerator().local_process_index}
+
+# model = trl_model_class.from_pretrained(
+#     args.ppo_config.model_name,
+#     trust_remote_code=True,
+#     device_map=device_map,
+#     peft_config=peft_config,
+# )
+
```


I am launching the script via ` CUDA_VISIBLE_DEVICES=0,1,2 && accelerate launch --config_file=examples/accelerate_configs/deepspeed_zero3.yaml --main_process_port 25678 examples/scripts/sentiment_tuning.py`

**Expected behavior:**
The script should either function as it does when instantiating the reference model before the regular model, or, if creating the reference model this way is somehow incompatible with ZeRO 3, it should fail early with an appropriate error message.

@lewtun 


Discussion:
Indeed, I think sharing the same LLM backbone for the reference and active model could be a challenge. @lewtun have you tried this configuration before?
No I haven't tried this and I am not sure it would be easy to handle in DeepSpeed without some low-level changes. For now I'll open a PR to raise an error is this combination of shared layer + ZeRO-3 is used

Resolving patch:
diff --git a/trl/models/modeling_base.py b/trl/models/modeling_base.py
--- a/trl/models/modeling_base.py
+++ b/trl/models/modeling_base.py
@@ -22,6 +22,7 @@
 from huggingface_hub import hf_hub_download
 from huggingface_hub.utils import EntryNotFoundError, HFValidationError, LocalEntryNotFoundError
 from transformers import PreTrainedModel
+from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
 
 from ..import_utils import is_peft_available
 
@@ -519,6 +520,10 @@ def create_reference_model(
     Returns
         `PreTrainedModelWrapper`
     """
+    if is_deepspeed_zero3_enabled():
+        raise ValueError(
+            "DeepSpeed ZeRO-3 is enabled and is not compatible with `create_reference_model()`. Please instantiate your reference model directly with `AutoCausalLM.from_pretrained()`."
+        )
 
     parameter_names = [n for n, _ in model.named_parameters()]
     ref_model = deepcopy(model)


Test patch: