fix_eos_mask (#10588)

Author: gongel

llm/alignment/rl/run_rl.py

@@ -23,6 +23,7 @@
 from paddle.distributed import fleet
 
 from paddlenlp.datasets.rlhf_datasets import RLHFDataset, collate_fn
+from paddlenlp.generation import GenerationConfig
 from paddlenlp.rl.models.score_model import AutoModelForScore
 from paddlenlp.rl.trainer.ppo_trainer import PPOTrainer
 from paddlenlp.rl.utils.config_utils import (
@@ -358,6 +359,12 @@ def compute_metrics(eval_preds):
         accuracy = (eval_preds.predictions == 3).astype("float32").mean().item()
         return {"accuracy": accuracy}
 
+    try:
+        generation_config = GenerationConfig.from_pretrained(model_args.actor_model_name_or_path)
+    except:
+        logger.warning("Can't find generation config, so it will not use generation_config field in the model config")
+        generation_config = None
+
     trainer = PPOTrainer(
         actor_model=actor_model,
         reference_model=reference_model,
@@ -379,6 +386,7 @@ def compute_metrics(eval_preds):
             max_prompt_len=data_args.max_prompt_len if training_args.balance_batch else None,
         ),  # NOTE: enforce prompt padding to max_prompt_len when using balance_batch
         compute_metrics=compute_metrics,  # TODO: only used for grpo (kk datasets)
+        generation_config=generation_config,
     )
 
     # TODO(gongenlei) resume_from_checkpoint is not ready

paddlenlp/rl/trainer/ppo_trainer.py

@@ -33,6 +33,7 @@
 from rich.table import Table
 
 from ...data import DataCollator
+from ...generation import GenerationConfig
 from ...trainer.trainer import (
     EvalLoopOutput,
     EvalPrediction,
@@ -53,6 +54,7 @@
     PretrainedTokenizer,
 )
 from ...transformers.model_utils import _add_variant
+from ...trl import llm_utils
 from ...utils.env import PADDLE_WEIGHTS_NAME
 from ..algos.advantage import (
     add_kl_divergence_regularization,
@@ -71,6 +73,7 @@
     filter_valid_reward_groups,
     gather_and_pad,
     get_timer_label,
+    make_eos_mask,
     new_timer_log,
     pad_tensor,
     split_batch_by_rank,
@@ -228,6 +231,7 @@ def __init__(
         callbacks: Optional[List[TrainerCallback]] = None,
         optimizers: Tuple[paddle.optimizer.Optimizer, paddle.optimizer.lr.LRScheduler] = (None, None),
         preprocess_logits_for_metrics: Optional[Callable[[paddle.Tensor, paddle.Tensor], paddle.Tensor]] = None,
+        generation_config: Optional[GenerationConfig] = None,
     ):
         """
         Args:
@@ -359,6 +363,7 @@ def __init__(
         self.model = self.model_wrapped = self.DummyPPOModel()
         if self.timers:
             self.timers.log = types.MethodType(new_timer_log, self.timers)
+        self.generation_config = generation_config
 
     def create_actor_trainer(
         self,
@@ -1142,12 +1147,15 @@ def pad_batch_data(
             dtype=label_ids[0].dtype,
             padding_side="right",
         )
-        position_ids = make_position_ids_from_input_ids(input_ids)
+        position_ids = make_position_ids_from_input_ids(input_ids, pad_token_id=self.tokenizer.pad_token_id)
         return input_ids, label_ids, position_ids
 
     def distribute_gather_and_pad_data(self, batch):
         # group index for grpo
-        eos_mask = (batch["input_ids"] != self.tokenizer.pad_token_id)[:, batch["prompt"].shape[-1] :].to(
+        eos_mask = make_eos_mask(
+            batch["input_ids"][:, batch["prompt"].shape[-1] :],
+            eos_token_ids=llm_utils.get_eos_token_id(self.tokenizer, self.generation_config),
+        ).to(
             batch["log_probs"].dtype  # fix dtype
         )
         try:

paddlenlp/rl/utils/comm_utils.py

@@ -1015,7 +1015,7 @@ def process_prompt_and_response(micro_batch, pad_token_id=0):
     response = paddle.stack(padded_response_tensors, axis=0)
 
     micro_batch["input_ids"] = paddle.concat([micro_batch["prompt"], response], axis=1)
-    micro_batch["position_ids"] = make_position_ids_from_input_ids(micro_batch["input_ids"])
+    micro_batch["position_ids"] = make_position_ids_from_input_ids(micro_batch["input_ids"], pad_token_id=pad_token_id)
     key_to_slice = [
         "eos_mask",
         "kl_rewards",
@@ -1072,3 +1072,23 @@ def split_batch_into_micro_batches(total_batch, batch_size, pad_token_id=0):
         micro_batches.append(micro_batch)
 
     return micro_batches
+
+
+def make_eos_mask(response_id, eos_token_ids=0, dtype=paddle.int64):
+    """
+    end of sentence token can be int or list: 1 or [1, 2]
+    e.g. eos_token=1
+    response_id: [0, 0, 2, 42, 3, 5, 1, 0, 0]
+    eos_mask:     [1, 1, 1, 1,  1, 1, 1, 0, 0]
+    """
+    if isinstance(eos_token_ids, int):
+        eos_token_ids = [eos_token_ids]
+
+    eos_mask = paddle.zeros_like(response_id, dtype=paddle.bool)
+    for token_id in eos_token_ids:
+        eos_mask |= response_id == token_id
+
+    eos_mask = eos_mask.to("int64")
+    eos_mask = (paddle.cumsum(eos_mask, axis=1) - eos_mask).to("bool")
+    eos_mask = paddle.logical_not(eos_mask).to(dtype)
+    return eos_mask