internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: huggingface/trl

Issue:
dpo_llama2.py: peft on top of peft?
https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py#L127
https://github.com/huggingface/trl/blob/main/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py#L186

It looks like the model for DPO training is a PeftModel built upon another PeftModel without weight merging. Is this expected? 

```
(Pdb) PeftModelForCausalLM(
  (base_model): LoraModel(
    (model): PeftModelForCausalLM(
      (base_model): LoraModel(
        (model): LlamaForCausalLM(
          (model): LlamaModel(
            (embed_tokens): Embedding(32000, 4096)
            (layers): ModuleList(
            ......
```


Discussion:
cc @kashif @younesbelkada 
Hi @ZHZisZZ 
Indeed this is not expected, all instances of `AutoPeftModelForCausalLM` should be in fact `AutoModelForCausalLM` here. I think what happened is that we first fine-tuned with SFT + PEFT and directly loaded that un-merged peft model for DPO and attached new adapters on top of it. 

Resolving patch:
diff --git a/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py b/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py
--- a/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py
+++ b/examples/research_projects/stack_llama_2/scripts/dpo_llama2.py
@@ -5,8 +5,8 @@
 
 import torch
 from datasets import Dataset, load_dataset
-from peft import AutoPeftModelForCausalLM, LoraConfig
-from transformers import AutoTokenizer, HfArgumentParser, TrainingArguments
+from peft import LoraConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, TrainingArguments
 
 from trl import DPOTrainer
 
@@ -124,7 +124,7 @@ def return_prompt_and_responses(samples) -> Dict[str, str]:
     script_args = parser.parse_args_into_dataclasses()[0]
 
     # 1. load a pretrained model
-    model = AutoPeftModelForCausalLM.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         script_args.model_name_or_path,
         low_cpu_mem_usage=True,
         torch_dtype=torch.float16,
@@ -138,7 +138,7 @@ def return_prompt_and_responses(samples) -> Dict[str, str]:
             name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool
         ]
 
-    model_ref = AutoPeftModelForCausalLM.from_pretrained(
+    model_ref = AutoModelForCausalLM.from_pretrained(
         script_args.model_name_or_path,
         low_cpu_mem_usage=True,
         torch_dtype=torch.float16,


Test patch:
Repository: huggingface/trl

Issue:
Wrong trigger for warning "UserWarning: The passed formatting_func has more than one argument."
In `trl/trainer/utils.py` the code `func.__code__.co_varnames` is being used to check if the user passed a **formatting_func** with more than 1 parameter. This code actually counts the function variables rather than function parameters.

For instance
```
def add_v1(a):
  return a + 1

def add_v2(a):
  b = a + 1
  return b

print("Number of parameters for add_v1 is", len(add_v1.__code__.co_varnames))
print("Number of parameters for add_v2 is", len(add_v2.__code__.co_varnames))
```
outputs
```
Number of parameters for add_v1 is 1
Number of parameters for add_v2 is 2
```

The proposed fix is to change the following:

```python
formatting_func_signature = formatting_func.__code__.co_varnames
if len(formatting_func_signature) > 1:
    warnings.warn(
        "The passed formatting_func has more than one argument. Usually that function should have a single argument `example`"
        " which corresponds to the dictionary returned by each element of the dataset. Make sure you know what you are doing."
    )
```
to:
```python
if formatting_func.__code__.co_argcount > 1:
    warnings.warn(
        "The passed formatting_func has more than one argument. Usually that function should have a single argument `example`"
        " which corresponds to the dictionary returned by each element of the dataset. Make sure you know what you are doing."
    )
```

Tested on python Python 2.7.5 and Python 3.6.8


Discussion:

Resolving patch:
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -543,8 +543,7 @@ def __init__(
             self.formatting_func = formatting_func
 
         if formatting_func is not None:
-            formatting_func_signature = formatting_func.__code__.co_varnames
-            if len(formatting_func_signature) > 1:
+            if formatting_func.__code__.co_argcount > 1:
                 warnings.warn(
                     "The passed formatting_func has more than one argument. Usually that function should have a single argument `example`"
                     " which corresponds to the dictionary returned by each element of the dataset. Make sure you know what you are doing."


Test patch:
Repository: huggingface/trl

Issue:
Deprecated option `optimize_cuda_cache` warning on import of trl
When you just import trl, we get the following warning:
```
$ python3
Python 3.10.8 (main, Nov 24 2022, 14:13:03) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import trl
/opt/conda/lib/python3.10/site-packages/trl/trainer/ppo_config.py:141: UserWarning: The `optimize_cuda_cache` arguement will be deprecated soon, please use `optimize_device_cache` instead.
  warnings.warn(
```

Unless this is left on purpose, I believe that this line https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_config.py#L107, should be changed from:
```
optimize_cuda_cache: bool = False
```
to:
```
optimize_cuda_cache: Optional[bool] = None
```
to get rid of the warning.

I can open a PR if this warning is not needed on default config creation.


Discussion:
Indeed, that's correct! A PR would be great, thank you!

Resolving patch:
diff --git a/trl/trainer/ppo_config.py b/trl/trainer/ppo_config.py
--- a/trl/trainer/ppo_config.py
+++ b/trl/trainer/ppo_config.py
@@ -104,7 +104,7 @@ class PPOConfig:
     """Number of optimisation epochs per batch of samples"""
     max_grad_norm: Optional[float] = None
     """Maximum gradient norm for gradient clipping"""
-    optimize_cuda_cache: bool = False
+    optimize_cuda_cache: Optional[bool] = None
     """DEPRECATED: use `optimize_device_cache` instead, which does the same thing."""
     optimize_device_cache: Optional[bool] = False
     """Optimize device cache for slightly more memory-efficient training"""


Test patch:
Repository: huggingface/trl

Issue:
`_prepare_packed_dataloader` and `_prepare_non_packed_dataloader` have inconsistent keys
`_prepare_packed_dataloader`  outputs a `ConstantLengthDataset` with [keys](https://github.com/huggingface/trl/blob/830cadfc4c80bdced0d3753de392070e6760d1f5/trl/trainer/utils.py#L450)
```
'input_ids`, 'labels'
```
`_prepare_non_packed_dataloader` outputs some type of dataset, with [keys](https://github.com/huggingface/trl/blob/830cadfc4c80bdced0d3753de392070e6760d1f5/trl/trainer/sft_trainer.py#L393)
```
'input_ids`, 'attention_mask'
```
The lack of `'labels'` in the latter might be source of confusing bugs.

If this is intended behavior, I would at least suggest to post a warning in the documentation of the argument `packing` of `SFTTrainer` 


Discussion:
I agree it's a bit confusing but it shouldn't make a difference in practice: the `labels` in the first case are identical to `input_ids` which is the default value filled also when no `labels` are provided in the second case. Similarly for the `attention_mask` in the second case. 

If you indeed encountered weird behaviour let us know so we can fix. I think we can refactor this in the future to be a bit more coherent. cc @younesbelkada 
An example of undesired (or, at least, confusing) behavior is the following:
```

def do(packing, collator):
    # some code to prepare the dataset etc....

    trainer = SFTTrainer(
        model=my_model,
        train_dataset=my_dataset,
        eval_dataset=eval_dataset,
        packing=packing,
        dataset_text_field='text',
        data_collator=collator,
    )
    
    trainer.train()
    
do(True, None)                    # works
do(True, DefaultDataCollator())   # works
do(False, None)                   # works
do(False, DefaultDataCollator())  # does not work
```
Even just a warning not to use `DefaultDataCollator` I guess should be enough.

By `does not work` I mean that the model only gets `input_ids` and not `labels`.

Indeed, that's not ideal. I'll have a look at it hopefully soon, but if you find time in the meantime feel free to open a PR! 

Resolving patch:
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -402,7 +402,11 @@ def tokenize(element):
                 else:
                     self._dataset_sanity_checked = True
 
-            return {"input_ids": outputs["input_ids"], "attention_mask": outputs["attention_mask"]}
+            return {
+                "input_ids": outputs["input_ids"],
+                "labels": outputs["input_ids"],
+                "attention_mask": outputs["attention_mask"],
+            }
 
         tokenized_dataset = dataset.map(
             tokenize,
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -452,6 +452,7 @@ def __iter__(self):
                 yield {
                     "input_ids": torch.LongTensor(example),
                     "labels": torch.LongTensor(example),
+                    "attention_mask": torch.ones(len(example)),
                 }
 
 


Test patch:
Repository: huggingface/trl

Issue:
SFTTrainer removes dataset columns even when `remove_unused_columns = False`
I've noticed that SFTTrainer removes dataset columns before passing samples to the data collator, even when `remove_unused_columns` is set to `False` in the training arguments.

This happens here:
https://github.com/huggingface/trl/blob/main/trl/trainer/sft_trainer.py#L410

Any objections to adding a check for `args.remove_unused_columns is True` here?


Discussion:

Resolving patch:
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -258,6 +258,7 @@ def make_inputs_require_grad(module, input, output):
                 formatting_func,
                 num_of_sequences,
                 chars_per_token,
+                remove_unused_columns=args.remove_unused_columns if args is not None else True,
                 **dataset_kwargs,
             )
         if eval_dataset is not None:
@@ -273,6 +274,7 @@ def make_inputs_require_grad(module, input, output):
                     formatting_func,
                     num_of_sequences,
                     chars_per_token,
+                    remove_unused_columns=args.remove_unused_columns if args is not None else True,
                     **dataset_kwargs,
                 )
             if not _multiple:
@@ -348,6 +350,7 @@ def _prepare_dataset(
         formatting_func,
         num_of_sequences,
         chars_per_token,
+        remove_unused_columns=True,
         append_concat_token=True,
         add_special_tokens=True,
     ):
@@ -360,7 +363,13 @@ def _prepare_dataset(
 
         if not packing:
             return self._prepare_non_packed_dataloader(
-                tokenizer, dataset, dataset_text_field, max_seq_length, formatting_func, add_special_tokens
+                tokenizer,
+                dataset,
+                dataset_text_field,
+                max_seq_length,
+                formatting_func,
+                add_special_tokens,
+                remove_unused_columns,
             )
 
         else:
@@ -377,7 +386,14 @@ def _prepare_dataset(
             )
 
     def _prepare_non_packed_dataloader(
-        self, tokenizer, dataset, dataset_text_field, max_seq_length, formatting_func=None, add_special_tokens=True
+        self,
+        tokenizer,
+        dataset,
+        dataset_text_field,
+        max_seq_length,
+        formatting_func=None,
+        add_special_tokens=True,
+        remove_unused_columns=True,
     ):
         use_formatting_func = formatting_func is not None and dataset_text_field is None
         self._dataset_sanity_checked = False
@@ -407,7 +423,7 @@ def tokenize(element):
         tokenized_dataset = dataset.map(
             tokenize,
             batched=True,
-            remove_columns=dataset.column_names,
+            remove_columns=dataset.column_names if remove_unused_columns else None,
             num_proc=self.dataset_num_proc,
             batch_size=self.dataset_batch_size,
         )


Test patch:
Repository: huggingface/trl

Issue:
RewardTrainer fails with FSDP
I've just run into an odd issue with FSDP & RewardTrainer. It seems then when using FSDP, the output of the (sequence classification) model's `forward` function isn't as expected. 
Normally, it returns a `SequenceClassifierOutputWithPast` where `logits` contains a tensor with the logits, and `loss` is empty or contains some sort of generator object..
When using FSDP, I'm getting a `dict` inside the loss field (and oddly enough that dict again contains a single key `logits`, althouh that's not the issue).

Not sure why this happens, but the net effect is that when the RewardTrainer tries to get the logits through `model(...)[0]` (see [here](https://github.com/huggingface/trl/blob/3267be0fcd424c8d224d7f28a45a64358c857ed0/trl/trainer/reward_trainer.py#L220)), in the non-FSDP case it gets the logits, while in the FSDP case it gets the dict from the now non-empty`loss` field, and then fails a few lines later.

Two questions:
1. This is easily fixed by doing `model(...)["logits"]` instead. Any problem with doing that?
2. Purely out of curiosity, does anyone know why this behaves differently with FSDP?


To reproduce: Run `examples/scripts/reward_modeling.py` with accelerate + FSDP.

`forward` output in a single process:
```
SequenceClassifierOutputWithPast(loss=<generator object gather.<locals>.gather_map.<locals>.<genexpr> at 0x15360f993040>, logits=tensor([[...]], device='cuda:0', grad_fn=<GatherBackward>), past_key_values=None, hidden_states=None, attentions=None)
```

And in FSDP:
```
SequenceClassifierOutputWithPast(loss={'logits': tensor([[...]], device='cuda:1', grad_fn=<ToCopyBackward0>)}, logits=tensor([[...]], device='cuda:1', grad_fn=<ToCopyBackward0>), past_key_values=None, hidden_states=None, attentions=None)
```


Discussion:

Resolving patch:
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
--- a/trl/trainer/reward_trainer.py
+++ b/trl/trainer/reward_trainer.py
@@ -220,11 +220,13 @@ def compute_loss(
         rewards_chosen = model(
             input_ids=inputs["input_ids_chosen"],
             attention_mask=inputs["attention_mask_chosen"],
-        )[0]
+            return_dict=True,
+        )["logits"]
         rewards_rejected = model(
             input_ids=inputs["input_ids_rejected"],
             attention_mask=inputs["attention_mask_rejected"],
-        )[0]
+            return_dict=True,
+        )["logits"]
         # calculate loss, optionally modulate with margin
         if "margin" in inputs:
             loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean()


Test patch:
Repository: huggingface/trl

Issue:
`ORPOTrainer` fails with `flash-attention2`
## Description

The `ORPOTrainer` fails when training using `attn_implementation="flash_attention_2"`, since the cache is being used, and falls back to the default configuration i.e. `padding_side="right"` for the tokenizer in this case.

## Bug in code

Missing `use_cache=False` to prevent the model from using the cache, to avoid issues with Flash Attention 2.

https://github.com/huggingface/trl/blob/2ce8e45bb222db68fe9eb92dc8b2a2f1339a2d75/trl/trainer/orpo_trainer.py#L685-L689


Discussion:

Resolving patch:
diff --git a/trl/trainer/cpo_trainer.py b/trl/trainer/cpo_trainer.py
--- a/trl/trainer/cpo_trainer.py
+++ b/trl/trainer/cpo_trainer.py
@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import inspect
 import random
 import warnings
@@ -665,6 +666,7 @@ def concatenated_forward(
         outputs = model(
             concatenated_batch["concatenated_input_ids"],
             attention_mask=concatenated_batch["concatenated_attention_mask"],
+            use_cache=False,
             **model_kwargs,
         )
         all_logits = outputs.logits
diff --git a/trl/trainer/orpo_trainer.py b/trl/trainer/orpo_trainer.py
--- a/trl/trainer/orpo_trainer.py
+++ b/trl/trainer/orpo_trainer.py
@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import inspect
 import random
 import warnings
@@ -685,6 +686,7 @@ def concatenated_forward(
         outputs = model(
             concatenated_batch["concatenated_input_ids"],
             attention_mask=concatenated_batch["concatenated_attention_mask"],
+            use_cache=False,
             **model_kwargs,
         )
         all_logits = outputs.logits


Test patch:
Repository: huggingface/trl

Issue:
conversational data for SFTTrainer
For SFTTrainer, if we load the dataset using a conversational form (ChatML format), the function `apply_chat_template` is used (https://github.com/huggingface/trl/blob/v0.7.11/trl/extras/dataset_formatting.py#L55) with `tokenize=False`. Later in SFTTrainer, the data is tokenized again with `add_special_tokens=True`. In tokenizer like LLaMATokenizer, there will be two bos tokens at the very beginning: `

 ...`, which is not intended. Maybe we should modify `dataset_kwargs` at this line https://github.com/huggingface/trl/blob/v0.7.11/trl/trainer/sft_trainer.py#L246 so that `dataset_kwargs['add_special_tokens']=True`?


Discussion:
Yes that would make sense, would you like to open a PR for the fix?
cc @philschmid what do you think?
sure I will do that
Good idea @edixiong, thats what I currently do manually. https://www.philschmid.de/fine-tune-llms-in-2024-with-trl#4-fine-tune-llm-using-trl-and-the-sfttrainer

This probably should only be applied if the "chatml" or template is detected. 


Resolving patch:
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -264,6 +264,12 @@ def make_inputs_require_grad(module, input, output):
             # check if dataset has ChatML format or instruction format and is supported
             # if not stays #None
             formatting_func = get_formatting_func_from_dataset(train_dataset, tokenizer)
+            # if a template is detected, we don't need to add special tokens again
+            if formatting_func is not None:
+                if dataset_kwargs is None:
+                    dataset_kwargs = {"add_special_tokens": False}
+                else:
+                    dataset_kwargs["add_special_tokens"] = False
 
         if not packing:
             if dataset_text_field is None and formatting_func is None:


Test patch:
Repository: huggingface/trl

Issue:
a bug which leads to "Cuda: Out of memory" in CPOTrainer (cpo_trainer.py), trl 0.8.1 0.8.2, please fix this bug
there is a bug in CPOTrainer. when runing CPOTrainer after runing sevreal steps, the  usage of gpu memory increases and it raises the out-of-memory exception. we found that the exception is caused by missing the "detach" in line 741 of CPOTrainer (cpo_trainer.py), please fix this bug, be careful!

line 741 of CPOTrainer (cpo_trainer.py):  
metrics[f"{prefix}nll_loss"] = policy_nll_loss.cpu().mean()
->
metrics[f"{prefix}nll_loss"] = policy_nll_loss.detach().cpu().mean()



Discussion:

Resolving patch:
diff --git a/trl/trainer/cpo_trainer.py b/trl/trainer/cpo_trainer.py
--- a/trl/trainer/cpo_trainer.py
+++ b/trl/trainer/cpo_trainer.py
@@ -738,7 +738,7 @@ def get_batch_loss_metrics(
         metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.detach().mean().cpu()
         metrics[f"{prefix}logits/rejected"] = policy_rejected_logits.detach().mean().cpu()
         metrics[f"{prefix}logits/chosen"] = policy_chosen_logits.detach().mean().cpu()
-        metrics[f"{prefix}nll_loss"] = policy_nll_loss.cpu().mean()
+        metrics[f"{prefix}nll_loss"] = policy_nll_loss.detach().mean().cpu()
 
         return loss, metrics
 


Test patch:
Repository: huggingface/trl

Issue:
Speed up ZeRO-3 generation with DPO
Hi, a [recent PR](https://github.com/huggingface/trl/pull/1483) brought large improvements (x10) to PPO generation with ZeRO-3.
@lewtun, you mention on the PR that it can be adapted for other trainers. I gave it a quick shot and it seems that naive applying the context manager to trainers like DPO does not work:
```
in remove_hooks
    if model.optimizer is not None and hasattr(
       ^^^^^^^^^^^^^^^^^^^^
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'GPTNeoXForCausalLM' object has no attribute 'optimizer'
```
There seems to be an inconsistency between the base classes. Is there a reason why DPO is based on Trainer from transformers and PPO on BaseTrainer ? What would be the easy way to add this feature to other trainers ? Thanks !


Discussion:
Passing **self.model_wrapped** instead in **unwrap_model_for_generation** in gives:
```
deepspeed/runtime/zero/partitioned_param_coordinator.py", line 194, in record_parameters
    step_id = self.__step_id_module_fetched_for[sub_module.id].popleft()
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: pop from an empty deque
```
Is it related to the way the model removes/adds hooks ? 

Resolving patch:
diff --git a/trl/models/utils.py b/trl/models/utils.py
--- a/trl/models/utils.py
+++ b/trl/models/utils.py
@@ -1,3 +1,4 @@
+import itertools
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Literal, Optional, Tuple, Union
@@ -118,6 +119,9 @@ def remove_hooks(model: "DeepSpeedEngine") -> None:
     elif model.optimizer is not None:
         optimizer_offload = model.optimizer
 
+    for param in iter_params(optimizer_offload.module, recurse=True):
+        param.ds_active_sub_modules.clear()
+
     for hook in optimizer_offload.forward_hooks:
         hook.remove()
     for hook in optimizer_offload.backward_hooks:
@@ -127,6 +131,14 @@ def remove_hooks(model: "DeepSpeedEngine") -> None:
     optimizer_offload.backward_hooks = []
 
 
+def get_all_parameters(sub_module, recurse=False):
+    return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters())
+
+
+def iter_params(module, recurse=False):
+    return [param for _, param in get_all_parameters(module, recurse)]
+
+
 def add_hooks(model: "DeepSpeedEngine") -> None:
     """Adds the optimizer hooks from a DeepSpeed ZeRO-3 model."""
     if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"):
@@ -141,7 +153,6 @@ def unwrap_model_for_generation(
     model: Union["DistributedDataParallel", "DeepSpeedEngine"], accelerator: "Accelerator", is_peft_model: bool = False
 ) -> Union["PreTrainedModelWrapper", "DeepSpeedEngine"]:
     """Context manager to unwrap a model for generation.
-
     For ZeRO-3 models, we gather the weights once to speed up generation.
     """
     unwrapped_model = accelerator.unwrap_model(model)


Test patch: