Repository: huggingface/trl
Pull request: 🔀 Rename `get_batch_sample` and add `num_items_in_batch` to `compute_loss`
# What does this PR do?

follows https://github.com/huggingface/transformers/pull/34198


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a GitHub issue? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Problem statement:
TypeError: XPOTrainer.training_step() takes 3 positional arguments but 4 were given  
### System Info

- Platform: Linux-5.4.0-42-generic-x86_64-with-glibc2.35
- Python version: 3.10.15
- PyTorch version: 2.5.0
- CUDA device(s): NVIDIA A800-SXM4-80GB, NVIDIA A800-SXM4-80GB, NVIDIA A800-SXM4-80GB, NVIDIA A800-SXM4-80GB, NVIDIA A800-SXM4-80GB, NVIDIA A800-SXM4-80GB, NVIDIA A800-SXM4-80GB, NVIDIA A800-SXM4-80GB
- Transformers version: 4.46.0.dev0
- Accelerate version: 1.0.1
- Accelerate config: not found
- Datasets version: 3.0.1
- HF Hub version: 0.25.2
- TRL version: 0.12.0.dev0+a67f214
- bitsandbytes version: 0.44.1
- DeepSpeed version: 0.15.2
- Diffusers version: not installed
- Liger-Kernel version: 0.3.1
- LLM-Blender version: not installed
- OpenAI version: not installed
- PEFT version: 0.13.2

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder
- [ ] My own task or dataset (give details below)

### Reproduction

After updating the transformers to support processing_class=tokenizer, a script that previously worked started producing the following traceback:
[rank7]: Traceback (most recent call last):
[rank7]:   File "/llm-align/trl/xpo.py", line 108, in <module>
[rank7]:     trainer.train()
[rank7]:   File "/llm-align/miniconda3/envs/trl/lib/python3.10/site-packages/transformers/trainer.py", line 2112, in train
[rank7]:     return inner_training_loop(
[rank7]:   File "/llm-align/miniconda3/envs/trl/lib/python3.10/site-packages/transformers/trainer.py", line 2458, in _inner_training_loop
[rank7]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank7]: TypeError: XPOTrainer.training_step() takes 3 positional arguments but 4 were given
Could you please advise on how to resolve this issue?

Here is my script:
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
deepspeed  --num_gpus 8 --master_port=29501  xpo.py \
    --deepspeed ds_config.json \
    --do_train \
    --model_name_or_path  /llm-align/qwen2.5-14B-chat \
    --reward_model_path /llm-align/weizhe/qwen2-0.5B-reward \
    --dataset_name  /llm-align/ultrafeedback-prompt \
    --learning_rate 5.0e-6 \
    --beta 0.1 \
    --torch_dtype bfloat16 \
    --output_dir /llm-align/qwen2.5-14B-xpo-lora \
    --num_train_epochs 3 \
    --max_new_tokens 256 \
    --warmup_ratio 0.1 \
    --missing_eos_penalty 1.0 \
    --overwrite_output_dir \
    --logging_steps 50 \
    --optim paged_adamw_32bit \
    --save_steps 200 \
    --save_total_limit 5 \
    --bf16 \
    --lr_scheduler_type 'cosine' \
    --load_in_4bit \
    --use_bnb_nested_quant \
    --use_peft  \
    --lora_r 16 \
    --lora_alpha 16 \
    --attn_implementation flash_attention_2 \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 2 \
    --lora_target_modules all-linear \
    --ddp_timeout 180000000

and  the train_dataset and test_dataset are json like:
```json
[
    {
        "prompt": [
            {
                "content": "create a table with 5 meals per day for 2 days, this is prepared for a 20 year old female. \nit sould be vegan, i should not contain nuts.\nshow a table with the meal, description, calorie count \nshow it in this style:\nDay n\nMeal n: meal name\n\nn ingredient\nn ingredient\nn ingredient\nn calories",
                "role": "user"
            }
        ]
    },
    {
        "prompt": [
            {
                "content": "In this task you will be given a list of integers. You should find the maximum absolute difference between 2 integers in the list. The absolute difference is the absolute value of one integer subtracted by another. The output should be a single integer which is the largest possible absolute distance.\nQ: [31, 28, -27]\nA: ",
                "role": "user"
            }
        ]
    },
]
```
I’m not sure where the error occurred. Could you please advise? thanks

### Expected behavior

NO

Hints from issue discussion:

Commit message: `num_items_in_batch` in `training_step`

Patch to review:
diff --git a/trl/trainer/bco_trainer.py b/trl/trainer/bco_trainer.py
index 91461a9b0d..c6ce2d4902 100644
--- a/trl/trainer/bco_trainer.py
+++ b/trl/trainer/bco_trainer.py
@@ -1260,6 +1260,7 @@ def compute_loss(
         model: Union[PreTrainedModel, nn.Module],
         inputs: Dict[str, Union[torch.Tensor, Any]],
         return_outputs=False,
+        num_items_in_batch=None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
         if not self.use_dpo_data_collator:
             warnings.warn(
@@ -1290,7 +1291,7 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
             return None
         return SequentialSampler(self.train_dataset)
 
-    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
+    def generate_from_model_and_ref(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
         """Generate samples from the model and reference model for the given batch of inputs."""
 
         # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with
@@ -1407,7 +1408,7 @@ def evaluation_loop(
                 "prompt_attention_mask": itemgetter(*target_indicies)(random_batch["prompt_attention_mask"]),
                 "prompt": itemgetter(*target_indicies)(random_batch["prompt"]),
             }
-            policy_output_decoded, ref_output_decoded = self.get_batch_samples(self.model, target_batch)
+            policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, target_batch)
 
             self.log(
                 {
diff --git a/trl/trainer/cpo_trainer.py b/trl/trainer/cpo_trainer.py
index 5847cb182b..c85ffb0a50 100644
--- a/trl/trainer/cpo_trainer.py
+++ b/trl/trainer/cpo_trainer.py
@@ -828,6 +828,7 @@ def compute_loss(
         model: Union[PreTrainedModel, nn.Module],
         inputs: Dict[str, Union[torch.Tensor, Any]],
         return_outputs=False,
+        num_items_in_batch=None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
         if not self.use_dpo_data_collator:
             warnings.warn(
@@ -847,7 +848,7 @@ def compute_loss(
             return (loss, metrics)
         return loss
 
-    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
+    def generate_from_model(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
         """Generate samples from the model and reference model for the given batch of inputs."""
 
         # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with
@@ -938,7 +939,7 @@ def evaluation_loop(
             random_batch = self.data_collator(random_batch_dataset)
             random_batch = self._prepare_inputs(random_batch)
 
-            policy_output_decoded = self.get_batch_samples(self.model, random_batch)
+            policy_output_decoded = self.generate_from_model(self.model, random_batch)
 
             self.log(
                 {
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
index 8b9843cd91..082a627ce0 100644
--- a/trl/trainer/dpo_trainer.py
+++ b/trl/trainer/dpo_trainer.py
@@ -1547,6 +1547,7 @@ def compute_loss(
         model: Union[PreTrainedModel, nn.Module],
         inputs: Dict[str, Union[torch.Tensor, Any]],
         return_outputs=False,
+        num_items_in_batch=None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
         compute_loss_context_manager = amp.autocast("cuda") if self._peft_has_been_casted_to_bf16 else nullcontext()
         with compute_loss_context_manager:
@@ -1561,7 +1562,7 @@ def compute_loss(
             return (loss, metrics)
         return loss
 
-    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
+    def generate_from_model_and_ref(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
         """Generate samples from the model and reference model for the given batch of inputs."""
 
         # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with
@@ -1672,7 +1673,7 @@ def evaluation_loop(
             random_batch = self.data_collator(random_batch_dataset)
             random_batch = self._prepare_inputs(random_batch)
 
-            policy_output_decoded, ref_output_decoded = self.get_batch_samples(self.model, random_batch)
+            policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, random_batch)
 
             self.log(
                 {
diff --git a/trl/trainer/gkd_trainer.py b/trl/trainer/gkd_trainer.py
index 1b7c77557d..49e93e269b 100644
--- a/trl/trainer/gkd_trainer.py
+++ b/trl/trainer/gkd_trainer.py
@@ -215,7 +215,7 @@ def generalized_jsd_loss(
         else:
             return jsd
 
-    def compute_loss(self, model, inputs, return_outputs=False):
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
         # compute student output
         outputs_student = model(
             input_ids=inputs["input_ids"],
@@ -273,7 +273,9 @@ def generate_on_policy_outputs(model, inputs, generation_config, pad_token_id=No
 
         return generated_tokens, new_attention_mask, new_labels
 
-    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
+    def training_step(
+        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None
+    ) -> torch.Tensor:
         """
         Perform a training step for the Generalized Knowledge Distillation (GKD) model.
 
@@ -298,7 +300,7 @@ def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor,
             inputs["attention_mask"] = new_attention_mask
             inputs["labels"] = new_labels
 
-        loss = super().training_step(model, inputs)
+        loss = super().training_step(model, inputs, num_items_in_batch)
         return loss
 
     def _prepare_deepspeed(self, model: PreTrainedModelWrapper):
diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py
index ab9ba87e41..7f32424812 100644
--- a/trl/trainer/kto_trainer.py
+++ b/trl/trainer/kto_trainer.py
@@ -1234,6 +1234,7 @@ def compute_loss(
         model: Union[PreTrainedModel, nn.Module],
         inputs: Dict[str, Union[torch.Tensor, Any]],
         return_outputs=False,
+        num_items_in_batch=None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
         if not self.use_dpo_data_collator:
             warnings.warn(
@@ -1264,7 +1265,7 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
             return None
         return SequentialSampler(self.train_dataset)
 
-    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
+    def generate_from_model_and_ref(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
         """Generate samples from the model and reference model for the given batch of inputs."""
 
         # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with
@@ -1383,7 +1384,7 @@ def evaluation_loop(
                 "prompt_attention_mask": random_batch["prompt_attention_mask"][target_indicies],
                 "prompt": itemgetter(*target_indicies)(random_batch["prompt"]),
             }
-            policy_output_decoded, ref_output_decoded = self.get_batch_samples(self.model, target_batch)
+            policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, target_batch)
 
             self.log(
                 {
diff --git a/trl/trainer/nash_md_trainer.py b/trl/trainer/nash_md_trainer.py
index db0c3046b3..73aab7899a 100644
--- a/trl/trainer/nash_md_trainer.py
+++ b/trl/trainer/nash_md_trainer.py
@@ -328,7 +328,9 @@ def gather_mean(tensor):
         self.stats["beta"].append(self.beta)
         self.stats["mixture_coef"].append(self.mixture_coef)
 
-    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
+    def training_step(
+        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None
+    ) -> torch.Tensor:
         model.train()
 
         # Apply chat template and tokenize the input
diff --git a/trl/trainer/online_dpo_trainer.py b/trl/trainer/online_dpo_trainer.py
index ffc407b57d..c480c61fc5 100644
--- a/trl/trainer/online_dpo_trainer.py
+++ b/trl/trainer/online_dpo_trainer.py
@@ -366,7 +366,9 @@ def get_eval_dataloader(self, eval_dataset: Optional[Union[str, Dataset]] = None
 
         return self.accelerator.prepare(eval_dataloader)
 
-    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
+    def training_step(
+        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None
+    ) -> torch.Tensor:
         model.train()
 
         # Apply chat template and tokenize the input.
diff --git a/trl/trainer/orpo_trainer.py b/trl/trainer/orpo_trainer.py
index 4edbf9b1a5..54b97cd6d1 100644
--- a/trl/trainer/orpo_trainer.py
+++ b/trl/trainer/orpo_trainer.py
@@ -844,6 +844,7 @@ def compute_loss(
         model: Union[PreTrainedModel, nn.Module],
         inputs: Dict[str, Union[torch.Tensor, Any]],
         return_outputs=False,
+        num_items_in_batch=None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
         if not self.use_dpo_data_collator:
             warnings.warn(
@@ -866,7 +867,7 @@ def compute_loss(
             return (loss, metrics)
         return loss
 
-    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
+    def generate_from_model(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
         """Generate samples from the model and reference model for the given batch of inputs."""
 
         # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with
@@ -957,7 +958,7 @@ def evaluation_loop(
             random_batch = self.data_collator(random_batch_dataset)
             random_batch = self._prepare_inputs(random_batch)
 
-            policy_output_decoded = self.get_batch_samples(self.model, random_batch)
+            policy_output_decoded = self.generate_from_model(self.model, random_batch)
 
             self.log(
                 {
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
index 787c6cbd54..0ebdee68b4 100644
--- a/trl/trainer/reward_trainer.py
+++ b/trl/trainer/reward_trainer.py
@@ -266,6 +266,7 @@ def compute_loss(
         model: Union[PreTrainedModel, nn.Module],
         inputs: Dict[str, Union[torch.Tensor, Any]],
         return_outputs=False,
+        num_items_in_batch=None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
         if not self.use_reward_data_collator:
             warnings.warn(
diff --git a/trl/trainer/xpo_trainer.py b/trl/trainer/xpo_trainer.py
index 0255e6206f..a154875821 100644
--- a/trl/trainer/xpo_trainer.py
+++ b/trl/trainer/xpo_trainer.py
@@ -377,7 +377,9 @@ def gather_mean(tensor):
         self.stats["alpha"].append(self.alpha)
         self.stats["beta"].append(self.beta)
 
-    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]) -> torch.Tensor:
+    def training_step(
+        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None
+    ) -> torch.Tensor:
         model.train()
 
         # Apply chat template and tokenize the input


Code review comments:
diff_hunk: @@ -866,7 +867,7 @@ def compute_loss(
             return (loss, metrics)
         return loss
 
-    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
+    def generate_from_model(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
path: trl/trainer/orpo_trainer.py
text: @author:
```suggestion
    def generate_from_model(self, model, batch: Dict[str, torch.LongTensor]) -> str:
```

diff_hunk: @@ -847,7 +848,7 @@ def compute_loss(
             return (loss, metrics)
         return loss
 
-    def get_batch_samples(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
+    def generate_from_model(self, model, batch: Dict[str, torch.LongTensor]) -> Tuple[str, str]:
path: trl/trainer/cpo_trainer.py
text: @author:
```suggestion
    def generate_from_model(self, model, batch: Dict[str, torch.LongTensor]) -> str:
```