Repository: huggingface/trl
Pull request: 🫷 Include stop token in policy model's generation_config
Another approach to fix #2387

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly. They may suggest changes to make the code even better.
-->

<!-- Remove if not applicable -->


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] 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:
eos_token config in PPOTrainer
### Feature request

It appears that the  generation config in PPOTrainer does not set an eos_token, resulting in each generation process continuing until it reaches the maximum length before stopping, which is quite time-consuming.

### Motivation

If the `eos_token` is set, it will significantly reduce the time spent during the generation phase.


### Your contribution

i'm sorry

Hints from issue discussion:
Thanks for reporting it. Would you like to open a PR to fix it?

@kechunFIVE, @qgallouedec, @dame-cell Seems that current code is inspired by https://iclr-blogposts.github.io/2024/blog/the-n-implementation-details-of-rlhf-with-ppo/, section *General implementation details*, 4.2. Authors tried to recreate results from early OpenAI work, but they say:
> Note that in a more recent codebase https://github.com/openai/summarize-from-feedback, OpenAI does stop sampling when encountering EOS token ([summarize_from_feedback/utils/experiment_helpers.py#L19](https://github.com/openai/summarize-from-feedback/blob/8af822a428c93432aa80ffbe5b065a8f93895669/summarize_from_feedback/utils/experiment_helpers.py#L19)). However in this work we aim to do a 1:1 replication, so we align the setting that could keep sampling even eos_token is encountered

Reading `PPOTrainer` implementation, it seems that `stop_token`/`stop_token_id` arguments should control when to stop generation: 
* When `stop_token` is set to `eos` value, everything that is after EOS is padded and therefore ignored in calculations (using masks). In that case, keeping generating after all sequences have EOS seems just wasting of time and resources. I've checked that between stopping on EOS and generating until max length, loss values are almost exact and minimal differences of the order of 1e-8 are caused by masked calculations.
* When `stop_token_id` is set, sequences get padded after this token, so generation should stop there, not model's EOS.
* The case when `stop_token != None` and `stop_token != 'eos'` are not supported, so they should probably raise an exception.

Commit message: Fix formatting

Patch to review:
diff --git a/trl/trainer/ppo_trainer.py b/trl/trainer/ppo_trainer.py
index 51897eeb44..723c3d8132 100644
--- a/trl/trainer/ppo_trainer.py
+++ b/trl/trainer/ppo_trainer.py
@@ -138,10 +138,18 @@ def __init__(
         if data_collator is None:
             data_collator = DataCollatorWithPadding(self.processing_class)
 
-        self.policy_model.generation_config.eos_token_id = (
-            None  # disable `pad_token_id` and `eos_token_id` because we just want to
-        )
-        self.policy_model.generation_config.pad_token_id = None  # generate tokens without truncation / padding
+        # Handle stop token settings
+        if args.stop_token and args.stop_token_id:
+            raise ValueError("You cannot set both `stop_token` and `stop_token_id`. ")
+        if args.stop_token:
+            if args.stop_token == "eos":
+                args.stop_token_id = processing_class.eos_token_id
+            else:
+                raise ValueError(
+                    f"Unknown `stop_token` {args.stop_token}. " f"Allowed values are: `eos`, None (no stop token)"
+                )
+        # Update policy model's generation_config to use provided stop token
+        self.policy_model.generation_config.eos_token_id = args.stop_token_id
 
         # peft support
         if not is_peft_available() and peft_config is not None:
@@ -220,8 +228,6 @@ def __init__(
         for module in [self.policy_model, self.ref_model, self.value_model, self.reward_model]:
             if module is not None:
                 disable_dropout_in_model(module)
-        if args.stop_token and args.stop_token == "eos":
-            args.stop_token_id = processing_class.eos_token_id
         self.model = PolicyAndValueWrapper(self.policy_model, self.value_model)
         self.model.config = self.policy_model.config  # needed for pushing to hub
         self.create_optimizer_and_scheduler(


Code review comments:
diff_hunk: @@ -138,10 +138,18 @@ def __init__(
         if data_collator is None:
             data_collator = DataCollatorWithPadding(self.processing_class)
 
-        self.policy_model.generation_config.eos_token_id = (
-            None  # disable `pad_token_id` and `eos_token_id` because we just want to
-        )
-        self.policy_model.generation_config.pad_token_id = None  # generate tokens without truncation / padding
+        # Handle stop token settings
+        if args.stop_token and args.stop_token_id:
+            raise ValueError("You cannot set both `stop_token` and `stop_token_id`. ")
+        if args.stop_token:
+            if args.stop_token == "eos":
+                args.stop_token_id = processing_class.eos_token_id
+            else:
+                raise ValueError(
+                    f"Unknown `stop_token` {args.stop_token}. " f"Allowed values are: `eos`, None (no stop token)"
+                )
+        # Update policy model's generation_config to use provided stop token
+        self.policy_model.generation_config.eos_token_id = args.stop_token_id
path: trl/trainer/ppo_trainer.py
text: @user1:
```suggestion
        elif args.stop_token:
            if args.stop_token == "eos":
                stop_token_id = processing_class.eos_token_id
            else:
                raise ValueError(
                    f"Unknown `stop_token` {args.stop_token}. " f"Allowed values are: `eos` and `None` (no stop token)"
                )
        elif args.stop_token_id:
             stop_token_id = processing_class.eos_token_id
        # Update policy model's generation_config to use provided stop token
        self.policy_model.generation_config.eos_token_id = args.stop_token_id
```

diff_hunk: @@ -138,10 +138,18 @@ def __init__(
         if data_collator is None:
             data_collator = DataCollatorWithPadding(self.processing_class)
 
-        self.policy_model.generation_config.eos_token_id = (
-            None  # disable `pad_token_id` and `eos_token_id` because we just want to
-        )
-        self.policy_model.generation_config.pad_token_id = None  # generate tokens without truncation / padding
+        # Handle stop token settings
+        if args.stop_token and args.stop_token_id:
+            raise ValueError("You cannot set both `stop_token` and `stop_token_id`. ")
path: trl/trainer/ppo_trainer.py
text: @user1:
```suggestion
            raise ValueError("You cannot set both `stop_token` and `stop_token_id`.")
```