Made fixes to training script based on recent findings.

Signed-off-by: meetkuma <[email protected]>

Author: Meet Patel

QEfficient/cloud/finetune.py

@@ -5,6 +5,7 @@
 #
 # -----------------------------------------------------------------------------
 
+import os
 import random
 import warnings
 from typing import Any, Dict, Optional, Union
@@ -139,7 +140,7 @@ def load_model_and_tokenizer(
         train_config.model_name if train_config.tokenizer_name is None else train_config.tokenizer_name
     )
     if not tokenizer.pad_token_id:
-        tokenizer.pad_token_id = tokenizer.eos_token_id
+        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
 
     # If there is a mismatch between tokenizer vocab size and embedding matrix,
     # throw a warning and then expand the embedding matrix
@@ -195,7 +196,9 @@ def apply_peft(
     else:
         peft_config = generate_peft_config(train_config, peft_config_file, **kwargs)
         model = get_peft_model(model, peft_config)
-    model.print_trainable_parameters()
+
+    if os.getenv("LOCAL_RANK", 0) == 0:
+        model.print_trainable_parameters()
 
     return model
 

QEfficient/finetune/data/sampler.py

@@ -49,14 +49,18 @@ def __init__(
     ) -> None:
         random.seed(seed)
         self.batch_sampler = LengthBasedBatchSampler(
-            data_source, batch_size=batch_size, drop_last=True, shuffle=shuffle
+            data_source, batch_size=batch_size, drop_last=False, shuffle=shuffle
         )
         self.num_replicas = num_replicas
         self.rank = rank
+        assert len(self.batch_sampler) % self.num_replicas == 0, (
+            "Length of batch samples should be divisible by number to processes in DDP."
+        )
+        self.sampler_len = len(self.batch_sampler) // self.num_replicas
+        self.max_length = len(self.batch_sampler)
 
     def __iter__(self):
-        max_length = len(self.batch_sampler) // self.num_replicas * self.num_replicas
-        return islice(self.batch_sampler, self.rank, max_length, self.num_replicas)
+        return islice(self.batch_sampler, self.rank, self.max_length, self.num_replicas)
 
     def __len__(self):
-        return len(self.batch_sampler) // self.num_replicas
+        return self.sampler_len

QEfficient/finetune/dataset/alpaca_dataset.py

@@ -11,6 +11,8 @@
 import torch
 from torch.utils.data import Dataset
 
+from QEfficient.finetune.dataset.helper import IGNORE_INDEX
+
 PROMPT_DICT = {
     "prompt_input": (
         "Below is an instruction that describes a task, paired with an input that provides further context. "
@@ -42,8 +44,6 @@ def __len__(self):
         return len(self.ann)
 
     def __getitem__(self, index):
-        IGNORE_INDEX = -100  # The default setting
-
         ann = self.ann[index]
         if ann.get("input", "") == "":
             prompt = PROMPT_DICT["prompt_no_input"].format_map(ann)

QEfficient/finetune/dataset/grammar_dataset.py

@@ -10,6 +10,8 @@
 from datasets import load_dataset
 from torch.utils.data import Dataset
 
+from QEfficient.finetune.dataset.helper import IGNORE_INDEX
+
 
 class grammar(Dataset):
     def __init__(self, tokenizer, csv_name=None, context_length=None):
@@ -58,7 +60,7 @@ def convert_to_features(self, example_batch):
         sample = {
             "input_ids": prompt_ids + label_ids,
             "attention_mask": [1] * len(prompt_ids + label_ids),
-            "labels": [-100] * len(prompt_ids) + label_ids,
+            "labels": [IGNORE_INDEX] * len(prompt_ids) + label_ids,
         }
 
         return sample

QEfficient/finetune/dataset/gsm8k_dataset.py

@@ -9,6 +9,8 @@
 
 from datasets import Dataset, load_dataset
 
+from QEfficient.finetune.dataset.helper import IGNORE_INDEX
+
 default_instruction = """### Instruction: Solve the math question using a basic calculator.
 Calculator can be invoked using the format: <<expression=answer>>.
 "expression" can be one of the 4 arithmetic operations, and "answer" will be filled in for you.
@@ -26,9 +28,8 @@ def tokenize_and_mask(row: Dict[str, str], *, tokenizer, instruction) -> Dict[st
 
     input_str = tokenizer.bos_token + instruction.format(**row)
     ques_ids = tokenizer(input_str, add_special_tokens=False, return_attention_mask=False)["input_ids"]
-    ans_ids = tokenizer(row["answer"] + tokenizer.eos_token, add_special_tokens=False, return_attention_mask=False)[
-        "input_ids"
-    ]
+    ans_str = row["answer"] + tokenizer.eos_token
+    ans_ids = tokenizer(ans_str, add_special_tokens=False, return_attention_mask=False)["input_ids"]
     input_ids = ques_ids + ans_ids
 
     # State machine to recognize <<expression=answer>> and mask answer
@@ -39,11 +40,11 @@ def tokenize_and_mask(row: Dict[str, str], *, tokenizer, instruction) -> Dict[st
         elif mode == 1 and token in equal_tokens:
             mode = 2
         elif mode == 2:
-            ans_ids[i] = -100
+            ans_ids[i] = IGNORE_INDEX
             if token in end_tokens:
                 mode = 0
 
-    labels = [-100] * len(ques_ids) + ans_ids
+    labels = [IGNORE_INDEX] * len(ques_ids) + ans_ids
 
     inputs = {"input_ids": input_ids, "labels": labels}
     return inputs
@@ -54,7 +55,7 @@ def pad_to_max_length(row: Dict[str, list], *, tokenizer, max_length: int) -> Di
     return {
         "input_ids": row["input_ids"] + [tokenizer.pad_token_id] * (max_length - length),
         "attention_mask": [1] * length + [0] * (max_length - length),
-        "labels": row["labels"] + [-100] * (max_length - length),
+        "labels": row["labels"] + [IGNORE_INDEX] * (max_length - length),
     }
 
 

QEfficient/finetune/dataset/helper.py

@@ -0,0 +1,8 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# -----------------------------------------------------------------------------
+
+IGNORE_INDEX = -100

QEfficient/finetune/dataset/samsum_dataset.py

@@ -7,9 +7,11 @@
 
 import datasets
 
+from QEfficient.finetune.dataset.helper import IGNORE_INDEX
+
 
 def get_preprocessed_samsum(dataset_config, tokenizer, split, context_length=None):
-    dataset = datasets.load_dataset("Samsung/samsum", split=split, trust_remote_code=True)
+    dataset = datasets.load_dataset("knkarthick/samsum", split=split, trust_remote_code=True)
 
     prompt = "Summarize this dialog:\n{dialog}\n---\nSummary:\n"
 
@@ -35,10 +37,15 @@ def tokenize_add_label(sample):
             pad_to_max_length=True,
         )
 
+        labels = [IGNORE_INDEX] * len(prompt) + summary
+        # labels = [l if l != tokenizer.pad_token_id else -100 for l in labels]
+        # sentence: <bos> <prompt> <summary> <eos> <pad>
+        # labels  : -100  -100     <summary> <eos> -100
+
         sample = {
             "input_ids": prompt + summary,
             "attention_mask": [1] * (len(prompt) + len(summary)),
-            "labels": [-100] * len(prompt) + summary,
+            "labels": labels,
         }
 
         return sample

QEfficient/finetune/utils/config_utils.py

@@ -115,10 +115,26 @@ def generate_dataset_config(dataset_name: str) -> Any:
     return dataset_config
 
 
+def pad_dataset(dataset, batch_size, num_replicas):
+    reminder = len(dataset) % (batch_size * num_replicas)
+    if reminder == 0:
+        return dataset
+
+    sample_input = dataset[0]
+    sample_input["labels"] = [-100] * len(sample_input["labels"])
+    num_pads = (batch_size * num_replicas) - reminder
+    for _ in range(num_pads):
+        dataset = dataset.add_item(sample_input)
+    return dataset
+
+
 def get_dataloader_kwargs(train_config, dataset, dataset_processer, mode):
     kwargs = {}
     batch_size = train_config.batch_size_training if mode == "train" else train_config.val_batch_size
     if train_config.enable_ddp:
+        print("Length of dataset before: ", len(dataset))
+        dataset = pad_dataset(dataset, batch_size, 2)
+        print("Length of dataset after: ", len(dataset))
         if train_config.enable_sorting_for_ddp:
             if train_config.context_length:
                 raise ValueError(
@@ -134,13 +150,12 @@ def get_dataloader_kwargs(train_config, dataset, dataset_processer, mode):
                 )
         else:
             kwargs["sampler"] = data_utils.DistributedSampler(
-                dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True
+                dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True, drop_last=False
             )
             kwargs["batch_size"] = batch_size
-            kwargs["drop_last"] = True
     else:
         kwargs["batch_size"] = batch_size
-        kwargs["drop_last"] = True
+    kwargs["drop_last"] = False
     kwargs["collate_fn"] = DataCollatorForSeq2Seq(dataset_processer)
     return kwargs
 

QEfficient/finetune/utils/helper.py

@@ -0,0 +1,13 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# -----------------------------------------------------------------------------
+
+import os
+
+
+def print_rank_0(msg):
+    if os.getenv("LOCAL_RANK", None) in [None, 0]:
+        print(msg)