[QEff Finetune] Adding dataset padding changes

QEfficient/finetune/data/sampler.py

@@ -4,11 +4,9 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 # -----------------------------------------------------------------------------
-
 import random
 from itertools import islice
 
-import numpy as np
 import torch
 
 
@@ -22,17 +20,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 +45,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
@@ -62,19 +62,41 @@ def get_dataloader_kwargs(train_config, dataset, dataset_processer, split):
     return kwargs
 
 
+def padding_dataset(train_config, dataset):
+    dataset = dataset.map(lambda x: {"input_length": len(x["input_ids"])})
+    if train_config.enable_sorting_for_ddp:
+        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 * train_config.train_batch_size) > 0:
+        padding_size = num_replicas - (len(dataset) % (num_replicas * train_config.train_batch_size))
+
+    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])
+    return combined_dataset
+
+
 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)
+    if train_config.enable_ddp:
+        dataset = padding_dataset(train_config, dataset)
+
     dl_kwargs = get_dataloader_kwargs(train_config, 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,
         num_workers=train_config.num_workers_dataloader,

QEfficient/finetune/utils/train_utils.py

@@ -151,7 +151,7 @@ def train(
 
         # enable profile for qaic
         qaic_profile.start_profiling(device, 1) if train_config.use_profiler else None
-
+        num_dummy_samples = 0
         for step, batch in enumerate(train_dataloader):
             # resume training from a particular checkpoint, assuming the dataset is not shuffled
             if train_config.use_peft and train_config.from_peft_checkpoint:
@@ -192,6 +192,10 @@ 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)
+                            num_dummy_samples += 1
+
                         if train_config.task_type == "seq_classification":
                             logits = model_outputs.logits
                             labels = batch["labels"][:, 0]
@@ -201,15 +205,18 @@ 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)
+                        num_dummy_samples += 1
+
                     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 +244,17 @@ def train(
                     step_metric_val = float(torch.exp(loss.detach().float()))
                 train_step_metric.append(step_metric_val)
 
+            # Accumalate gradients
+            complete_accum_steps = (
+                len(train_dataloader) - len(train_dataloader) % train_config.gradient_accumulation_steps
+            )
+            if step < complete_accum_steps:
+                num_samples_in_cur_update = train_config.gradient_accumulation_steps
+            else:
+                num_samples_in_cur_update = len(train_dataloader) % train_config.gradient_accumulation_steps
+
+            loss = loss / num_samples_in_cur_update
+
             if train_config.grad_scaler:
                 scaler.scale(loss).backward()  # backward pass
             else:
@@ -296,14 +314,22 @@ def train(
 
         if loss_0_counter.item() == train_config.convergence_counter:
             if train_config.use_peft and train_config.from_peft_checkpoint and epoch == intermediate_epoch:
-                train_epoch_loss = total_loss / (step - intermediate_step)
+                train_epoch_loss = (
+                    0.0 if total_loss == 0.0 else total_loss / (step - intermediate_step - num_dummy_samples)
+                )
             else:
-                train_epoch_loss = total_loss / step
+                train_epoch_loss = 0.0 if total_loss == 0.0 else total_loss / (step - num_dummy_samples)
         else:
             if train_config.use_peft and train_config.from_peft_checkpoint and epoch == intermediate_epoch:
-                train_epoch_loss = total_loss / (len(train_dataloader) - intermediate_step)
+                train_epoch_loss = (
+                    0.0
+                    if total_loss == 0.0
+                    else total_loss / (len(train_dataloader) - intermediate_step - num_dummy_samples)
+                )
             else:
-                train_epoch_loss = total_loss / len(train_dataloader)
+                train_epoch_loss = (
+                    0.0 if total_loss == 0.0 else total_loss / (len(train_dataloader) - num_dummy_samples)
+                )
 
         if train_config.task_type == "seq_classification":
             metric_val = acc_helper.compute()
@@ -421,6 +447,7 @@ def evaluation_helper(model, train_config, eval_dataloader, device):
     eval_loss = 0.0  # Initialize evaluation loss
     device_type = torch.device(device).type
 
+    num_dummy_samples = 0
     for step, batch in enumerate(tqdm(eval_dataloader, colour="green", desc="evaluating Epoch", dynamic_ncols=True)):
         #  stop when the maximum number of eval steps is reached
         if train_config.max_eval_step > 0 and step > train_config.max_eval_step:
@@ -439,6 +466,10 @@ 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)
+                num_dummy_samples += 1
+
             if train_config.task_type == "seq_classification":
                 logits = outputs.logits
                 labels = batch["labels"][:, 0]
@@ -455,7 +486,7 @@ def evaluation_helper(model, train_config, eval_dataloader, device):
             eval_loss += loss.detach().float()
 
     # Compute average loss and metric
-    eval_epoch_loss = eval_loss / len(eval_dataloader)
+    eval_epoch_loss = 0.0 if eval_loss == 0.0 else eval_loss / (len(eval_dataloader) - num_dummy_samples)
     if train_config.task_type == "seq_classification":
         eval_metric = acc_helper.compute()
     else: