Fixed collate fn for bs>1. It will work fine for bs>1 for llama as well on single device.

Signed-off-by: Meet Patel <[email protected]>

Author: Meet Patel

QEfficient/cloud/finetune.py

@@ -133,13 +133,9 @@ def main(**kwargs):
     dataset_processer = tokenizer
 
     # Load and preprocess the dataset for training and validation
-    ctx_len = train_config.context_length
-    if ctx_len is None and hasattr(model.config, "max_position_embeddings"):
-        ctx_len = model.config.max_position_embeddings
+    dataset_train = get_preprocessed_dataset(dataset_processer, dataset_config, split="train", context_length=train_config.context_length)
 
-    dataset_train = get_preprocessed_dataset(dataset_processer, dataset_config, split="train", context_length=ctx_len)
-
-    dataset_val = get_preprocessed_dataset(dataset_processer, dataset_config, split="test", context_length=ctx_len)
+    dataset_val = get_preprocessed_dataset(dataset_processer, dataset_config, split="test", context_length=train_config.context_length)
 
     # TODO: vbaddi, check if its necessary to do this?
     # dataset_train = ConcatDataset(

QEfficient/finetune/dataset/imdb_dataset.py

@@ -28,11 +28,10 @@ def tokenize_add_label(sample):
         data = tokenizer(
             sample["text"],
             add_special_tokens=True,
-            max_length=context_length,
-            pad_to_max_length=True,
+            max_length=tokenizer.model_max_length,
         )
 
-        data["labels"] = sample["label"]
+        data["labels"] = [sample["label"]]
         return data
 
     dataset = dataset.map(tokenize_add_label, remove_columns=list(dataset.features))

QEfficient/finetune/utils/config_utils.py

@@ -88,19 +88,14 @@ def get_dataloader_kwargs(train_config, dataset, dataset_processer, mode):
                     num_replicas=dist.get_world_size(),
                     shuffle=False,
                 )
-                if train_config.task_type == "seq_classification":
-                    kwargs["collate_fn"] = default_data_collator
-                else:
-                    kwargs["collate_fn"] = DataCollatorForSeq2Seq(dataset_processer)
         else:
             kwargs["sampler"] = data_utils.DistributedSampler(
                 dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank(), shuffle=True
             )
             kwargs["batch_size"] = batch_size
             kwargs["drop_last"] = True
-            kwargs["collate_fn"] = default_data_collator
     else:
         kwargs["batch_size"] = batch_size
         kwargs["drop_last"] = True
-        kwargs["collate_fn"] = default_data_collator
+    kwargs["collate_fn"] = DataCollatorForSeq2Seq(dataset_processer)
     return kwargs

QEfficient/finetune/utils/train_utils.py

@@ -193,7 +193,7 @@ def train(
                         loss = model_outputs.loss  # Forward call
                         if train_config.task_type == "seq_classification":
                             logits = model_outputs.logits
-                            labels = batch["labels"]
+                            labels = batch["labels"][:, 0]
                             preds = torch.nn.functional.softmax(logits, dim=-1)
                             acc_helper.forward(preds, labels)
                     print("Mismatches detected:", verifier.get_perop_mismatch_count())
@@ -202,7 +202,7 @@ def train(
                     loss = model_outputs.loss  # Forward call
                     if train_config.task_type == "seq_classification":
                         logits = model_outputs.logits
-                        labels = batch["labels"]
+                        labels = batch["labels"][:, 0]
                         preds = torch.nn.functional.softmax(logits, dim=-1)
                         acc_helper.forward(preds, labels)
 
@@ -305,7 +305,7 @@ def train(
             dist.barrier()
             dist.all_reduce(train_epoch_loss, op=dist.ReduceOp.SUM)
             train_epoch_loss /= dist.get_world_size()
-            
+
         if train_config.task_type == "seq_classification":
             accuracy = acc_helper.compute()
             if train_config.enable_ddp:
@@ -515,7 +515,7 @@ def evaluation_acc(model, train_config, eval_dataloader, local_rank, tokenizer,
                 outputs = model(**batch)
             loss = outputs.loss
             logits = outputs.logits
-            labels = batch["labels"]
+            labels = batch["labels"][:, 0]
             if train_config.save_metrics:
                 val_step_loss.append(loss.detach().float().item())
                 preds = torch.nn.functional.softmax(logits, dim=-1)