upgrade support to trl 0.27+

Signed-off-by: Yash Mehan <yashmehan@gmail.com>

Author: yash4242

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
 "sentencepiece>=0.1.99,<0.3",
 "tokenizers<=0.23.0",
 "tqdm>=4.66.2,<5.0",
-"trl>=0.19.1,<0.20.0",
+"trl>=0.27.0,<0.29.0",
 "peft>=0.18.1,<0.19.0",
 "datasets>=4.0.0,<5.0.0",
 "simpleeval>=0.9.13,<2.0",

tests/data/test_data_preprocessing.py

@@ -22,7 +22,6 @@
 from datasets import Dataset, DatasetDict, IterableDataset
 from PIL import Image
 from transformers import AutoProcessor, AutoTokenizer, DataCollatorForSeq2Seq
-from trl import DataCollatorForCompletionOnlyLM
 import datasets
 import numpy as np
 import pyarrow
@@ -69,7 +68,7 @@
 # Local
 from tuning.config import configs
 from tuning.config.acceleration_configs import AttentionAndDistributedPackingConfig
-from tuning.data.collators import VisionDataCollator
+from tuning.data.collators import DataCollatorForCompletionOnlyLM, VisionDataCollator
 from tuning.data.data_config import (
     DataHandlerConfig,
     DataPreProcessorConfig,

tuning/data/collators.py

@@ -12,6 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Standard
+from typing import Any, Optional, Union
+import logging
+
+# Third Party
+from transformers import DataCollatorForLanguageModeling
+import numpy as np
+import torch
+
 # Local
 from tuning.data.utils import try_convert_bytes_dict_to_pil
 
@@ -91,3 +100,239 @@ def __call__(self, features):
         batch["labels"] = labels
 
         return batch
+
+
+class DataCollatorForCompletionOnlyLM(DataCollatorForLanguageModeling):
+    """
+    Data collator used for completion tasks.
+    It ensures that all the tokens of the labels
+    are set to an 'ignore_index'
+    when they do not come from the assistant.
+    This ensure that the loss is only calculated on the completion made by
+    the assistant.
+
+    Args:
+        response_template (`Union[str, list[int]]`):
+            the template form that indicates the
+            start of the response, typically
+            something like '### Response:\n'. It
+            can also be passed as tokenized ids,
+            which can be useful when using a tokenizer
+            that encodes the response
+            differently if it does not have proper context.
+        instruction_template (`Union[str, list[int]]`):
+            the template form that indicates the start
+            of the human instruction, typically
+            something like '###
+            Human:\n'. Useful for assistant-style
+            conversation datasets. It can also be passed
+            as tokenized ids.
+        mlm (`bool`, *optional*, defaults to `False`): Whether
+        to use masked language modeling in the underlying
+            `DataCollatorForLanguageModeling` class.
+            Note that this option currently has no effect but is present
+             for flexibility and backwards-compatibility.
+        ignore_index (`int`, *optional*, defaults to `-100`):
+            The index to use to ignore the initial tokens with
+    """
+
+    def __init__(
+        self,
+        *args,
+        response_template: Union[str, list[int]],
+        instruction_template: Optional[Union[str, list[int]]] = None,
+        mlm: bool = False,
+        ignore_index: int = -100,
+        padding_free: bool = False,
+        **kwargs,
+    ):
+        super().__init__(*args, mlm=mlm, **kwargs)
+
+        self.instruction_template = instruction_template
+        if isinstance(instruction_template, str):
+            # The user provides a string, must tokenize
+            self.instruction_token_ids = self.tokenizer.encode(
+                self.instruction_template, add_special_tokens=False
+            )
+        else:
+            # The user already provides the token ids
+            self.instruction_token_ids = instruction_template
+
+        self.response_template = response_template
+        if isinstance(response_template, str):
+            # The user provides a string, must tokenize
+            self.response_token_ids = self.tokenizer.encode(
+                self.response_template, add_special_tokens=False
+            )
+        else:
+            # The user already provides the token ids
+            self.response_token_ids = response_template
+
+        if (
+            not self.mlm
+            and self.instruction_template
+            and self.tokenizer.pad_token_id == self.tokenizer.eos_token_id
+        ):
+            logging.warning(
+                "The pad_token_id and eos_token_id values "
+                "of this tokenizer are identical. "
+                "If you are planning for multi-turn training, "
+                "it can result in the model continuously generating "
+                "questions and answers without eos token. "
+                "To avoid this, set the pad_token_id to a different value.",
+            )
+
+        self.ignore_index = ignore_index
+        self.padding_free = padding_free
+
+    def torch_call(
+        self, examples: list[Union[list[int], Any, dict[str, Any]]]
+    ) -> dict[str, Any]:
+        batch = super().torch_call(examples)
+
+        if self.instruction_template is None:
+            for i in range(len(examples)):
+                response_token_ids_start_idx = None
+
+                for idx in np.where(batch["labels"][i] == self.response_token_ids[0])[
+                    0
+                ]:
+                    # `response_token_ids` is
+                    # `'### Response:\n'`, here we are just making sure
+                    # that the token IDs match
+                    if (
+                        self.response_token_ids
+                        == batch["labels"][i][
+                            idx : idx + len(self.response_token_ids)
+                        ].tolist()
+                    ):
+                        response_token_ids_start_idx = idx
+
+                if response_token_ids_start_idx is None:
+                    logging.warning(
+                        "Could not find response key %s in the following instance: "
+                        "%s. This instance will be ignored in loss "
+                        "calculation. Note, if this happens often, "
+                        "consider increasing the `max_length`.",
+                        self.response_template,
+                        self.tokenizer.decode(batch["input_ids"][i]),
+                    )
+                    batch["labels"][i, :] = self.ignore_index
+                else:
+                    response_token_ids_end_idx = response_token_ids_start_idx + len(
+                        self.response_token_ids
+                    )
+
+                    # Make pytorch loss function ignore all
+                    # tokens up through the end of the response key
+                    batch["labels"][i, :response_token_ids_end_idx] = self.ignore_index
+
+        else:
+            for i in range(len(examples)):
+                response_token_ids_idxs = []
+                human_token_ids_idxs = []
+
+                for assistant_idx in np.where(
+                    batch["labels"][i] == self.response_token_ids[0]
+                )[0]:
+                    # find the indexes of the start of a response.
+                    if (
+                        self.response_token_ids
+                        == batch["labels"][i][
+                            assistant_idx : assistant_idx + len(self.response_token_ids)
+                        ].tolist()
+                    ):
+                        response_token_ids_idxs.append(
+                            assistant_idx + len(self.response_token_ids)
+                        )
+
+                if len(response_token_ids_idxs) == 0:
+                    logging.warning(
+                        "Could not find response key %s in the following instance: "
+                        "%s. This instance will be ignored in loss "
+                        "calculation. Note, if this happens often, "
+                        "consider increasing the `max_length`.",
+                        self.response_template,
+                        self.tokenizer.decode(batch["input_ids"][i]),
+                    )
+                    batch["labels"][i, :] = self.ignore_index
+
+                human_token_ids = self.instruction_token_ids
+                for human_idx in np.where(batch["labels"][i] == human_token_ids[0])[0]:
+                    # find the indexes of the start of a human answer.
+                    if (
+                        human_token_ids
+                        == batch["labels"][i][
+                            human_idx : human_idx + len(human_token_ids)
+                        ].tolist()
+                    ):
+                        human_token_ids_idxs.append(human_idx)
+
+                if len(human_token_ids_idxs) == 0:
+                    logging.warning(
+                        "Could not find instruction key `%s` in the following instance: "
+                        "%s. This instance will be ignored in loss "
+                        "calculation. Note, if this happens often, "
+                        "consider increasing the `max_length`.",
+                        self.instruction_template,
+                        self.tokenizer.decode(batch["input_ids"][i]),
+                    )
+                    batch["labels"][i, :] = self.ignore_index
+
+                if (
+                    len(human_token_ids_idxs) > 0
+                    and len(response_token_ids_idxs) > 0
+                    and human_token_ids_idxs[0] > response_token_ids_idxs[0]
+                ):
+                    human_token_ids_idxs = [0] + human_token_ids_idxs
+
+                for idx, (start, end) in enumerate(
+                    zip(human_token_ids_idxs, response_token_ids_idxs)
+                ):
+                    # Make pytorch loss function ignore all non response tokens
+                    if idx != 0:
+                        batch["labels"][i, start:end] = self.ignore_index
+                    else:
+                        batch["labels"][i, :end] = self.ignore_index
+
+                if len(response_token_ids_idxs) < len(human_token_ids_idxs):
+                    batch["labels"][i, human_token_ids_idxs[-1] :] = self.ignore_index
+
+        if self.padding_free:
+            # remove padding, `attention_mask` and add `position_ids`
+            attn_mask = batch.pop("attention_mask")
+            batch["input_ids"] = batch["input_ids"][attn_mask.bool()].unsqueeze(0)
+            batch["position_ids"] = (
+                attn_mask.cumsum(1)[attn_mask.bool()].unsqueeze(0) - 1
+            )
+            batch["labels"] = batch["labels"][attn_mask.bool()].unsqueeze(0)
+            batch["labels"][batch["position_ids"] == 0] = self.ignore_index
+
+            # Calculate cumulative sequence lengths for queries and
+            # keys to prevent graph breaks during further computations.
+            flattened_position_ids = batch["position_ids"].flatten()
+            indices_q = torch.arange(
+                flattened_position_ids.size(0),
+                device=flattened_position_ids.device,
+                dtype=torch.int32,
+            )
+            batch["cu_seq_lens_q"] = torch.cat(
+                (
+                    indices_q[flattened_position_ids == 0],
+                    torch.tensor(
+                        flattened_position_ids.size(),
+                        device=flattened_position_ids.device,
+                        dtype=torch.int32,
+                    ),
+                )
+            ).unsqueeze(0)
+            batch["cu_seq_lens_k"] = batch["cu_seq_lens_q"]
+
+            # Determine maximum sequence lengths to
+            # prevent graph breaks during further computations.
+            batch["max_length_k"] = torch.tensor(
+                [flattened_position_ids.max().item() + 1]
+            )
+            batch["max_length_q"] = batch["max_length_k"]
+
+        return batch

tuning/data/data_preprocessing_utils.py

@@ -22,11 +22,10 @@
     DataCollatorForSeq2Seq,
     LlavaProcessor,
 )
-from trl import DataCollatorForCompletionOnlyLM
 
 # Local
 from tuning.config import configs
-from tuning.data.collators import VisionDataCollator
+from tuning.data.collators import DataCollatorForCompletionOnlyLM, VisionDataCollator
 
 logger = logging.getLogger(__name__)