Adding dataset padding changes

Signed-off-by: Swati Allabadi <[email protected]>

Author: Swati Allabadi

QEfficient/finetune/data/sampler.py

@@ -4,7 +4,6 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 # -----------------------------------------------------------------------------
-
 import random
 from itertools import islice
 
@@ -22,17 +21,15 @@ def __init__(self, data_source, batch_size: int, drop_last: bool, shuffle: bool
         self.batch_size = batch_size
         self.drop_last = drop_last
         self.shuffle = shuffle
+        self.data_source = data_source
 
     def __iter__(self):
-        ids = np.argsort(self.lengths, kind="mergesort")
+        ids = [i for i in range(len(self.data_source))]
         if self.drop_last:
             ids = ids[: len(ids) // self.batch_size * self.batch_size]
 
         batches = [ids[i : i + self.batch_size] for i in range(0, len(ids), self.batch_size)]
 
-        if self.shuffle:
-            random.shuffle(batches)
-
         for b in batches:
             yield b
 
@@ -49,7 +46,7 @@ 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

QEfficient/finetune/utils/dataset_utils.py

@@ -4,7 +4,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 # -----------------------------------------------------------------------------
-
+import datasets
 import torch
 import torch.distributed as dist
 from transformers.data import DataCollatorForSeq2Seq
@@ -64,19 +64,35 @@ def get_dataloader_kwargs(train_config, dataset, dataset_processer, split):
 
 def get_dataloader(tokenizer, dataset_config, train_config, split: str = "train"):
     dataset = get_preprocessed_dataset(tokenizer, dataset_config, split, context_length=train_config.context_length)
-    dl_kwargs = get_dataloader_kwargs(train_config, dataset, tokenizer, split)
+    dataset = dataset.select(range(0, 10))
+    dataset = dataset.map(lambda x: {"input_length": len(x["input_ids"])})
+    dataset = dataset.sort("input_length")
+    dataset = dataset.remove_columns("input_length")
+    dummy_row = next(iter(dataset))
+    dummy_row["labels"] = [-100] * len(dummy_row["labels"])
+    padding_size = 0
+    num_replicas = dist.get_world_size()
+    if len(dataset) % num_replicas > 0:
+        padding_size = num_replicas - len(dataset) % num_replicas
+
+    dummy_data = [dummy_row.copy() for _ in range(padding_size)]
+    dummy_dataset = datasets.Dataset.from_list(dummy_data)
+    combined_dataset = datasets.concatenate_datasets([dataset, dummy_dataset])
+
+    dl_kwargs = get_dataloader_kwargs(train_config, combined_dataset, tokenizer, split)
 
     # FIXME (Meet): Add custom data collator registration from the outside by the user.
     custom_data_collator = get_custom_data_collator(tokenizer, dataset_config)
+
     if custom_data_collator:
         print("custom_data_collator is used")
         dl_kwargs["collate_fn"] = custom_data_collator
 
     print(f"length of dataset_{split}", len(dataset))
-
     # Create data loader
+
     dataloader = torch.utils.data.DataLoader(
-        dataset,
+        combined_dataset,
         num_workers=train_config.num_workers_dataloader,
         pin_memory=True,
         **dl_kwargs,

QEfficient/finetune/utils/train_utils.py

@@ -192,6 +192,9 @@ def train(
                     ) as verifier:
                         model_outputs = model(**batch)
                         loss = model_outputs.loss  # Forward call
+                        if (batch["labels"] != -100).sum() == 0:
+                            loss = loss.nan_to_num(nan=0.0)
+
                         if train_config.task_type == "seq_classification":
                             logits = model_outputs.logits
                             labels = batch["labels"][:, 0]
@@ -201,15 +204,17 @@ def train(
                 else:
                     model_outputs = model(**batch)
                     loss = model_outputs.loss  # Forward call
+                    if (batch["labels"] != -100).sum() == 0:
+                        loss = loss.nan_to_num(nan=0.0)
+
                     if train_config.task_type == "seq_classification":
                         logits = model_outputs.logits
                         labels = batch["labels"][:, 0]
                         preds = torch.nn.functional.softmax(logits, dim=-1)
                         acc_helper.forward(preds, labels)
 
             total_loss += loss.detach().float()
-            # Accumalate gradients
-            loss = loss / train_config.gradient_accumulation_steps
+
             if train_config.enable_ddp:
                 if local_rank == 0:
                     if loss <= train_config.convergence_loss:
@@ -237,6 +242,9 @@ def train(
                     step_metric_val = float(torch.exp(loss.detach().float()))
                 train_step_metric.append(step_metric_val)
 
+            # Accumalate gradients
+            loss = loss / train_config.gradient_accumulation_steps
+
             if train_config.grad_scaler:
                 scaler.scale(loss).backward()  # backward pass
             else:
@@ -439,6 +447,9 @@ def evaluation_helper(model, train_config, eval_dataloader, device):
                 outputs = model(**batch)
             loss = outputs.loss
 
+            if (batch["labels"] != -100).sum() == 0:
+                loss = loss.nan_to_num(nan=0.0)
+
             if train_config.task_type == "seq_classification":
                 logits = outputs.logits
                 labels = batch["labels"][:, 0]