[QEff Finetune]: Added support for gradient checkpointing in the finetuning script. (quic#338)

Added --gradient_checkpointing new CLI flag to enable this feature.
Currently this is enabled for all the HF models which has
"supports_gradient_checkpointing" attribute set to True.

---------

Signed-off-by: Meet Patel <quic_meetkuma@quicinc.com>

Signed-off-by: eplatero <quic_eplatero@quicinc.com>

Author: quic-meetkuma

QEfficient/finetune/dataset/dataset_config.py

@@ -24,6 +24,9 @@
 from QEfficient.finetune.dataset.samsum_dataset import (
     get_preprocessed_samsum as get_samsum_dataset,
 )
+from QEfficient.finetune.dataset.samsum_dataset import (
+    get_samsum_collate_fn,
+)
 
 DATASET_PREPROC = {
     "alpaca_dataset": partial(get_alpaca_dataset),
@@ -36,3 +39,7 @@
 DATALOADER_COLLATE_FUNC = {
     "custom_dataset": get_data_collator,
 }
+DATALOADER_COLLATE_FUNC = {
+    "custom_dataset": get_data_collator,
+    "samsum_dataset": get_samsum_collate_fn,
+}

QEfficient/finetune/dataset/samsum_dataset.py

@@ -6,6 +6,8 @@
 # -----------------------------------------------------------------------------
 
 import datasets
+import torch
+from torch.nn.utils.rnn import pad_sequence
 
 
 def get_preprocessed_samsum(dataset_config, tokenizer, split, context_length=None):
@@ -46,3 +48,22 @@ def tokenize_add_label(sample):
     dataset = dataset.map(tokenize_add_label, remove_columns=list(dataset.features))
 
     return dataset
+
+
+def collate_fn(batch):
+    eos_token = batch[0]["input_ids"][-1]
+
+    input_ids = pad_sequence(
+        [torch.tensor(b["input_ids"], dtype=torch.int32) for b in batch], batch_first=True, padding_value=eos_token
+    )
+    attn_mask = pad_sequence(
+        [torch.tensor(b["attention_mask"], dtype=torch.int32) for b in batch], batch_first=True, padding_value=0
+    )
+    labels = pad_sequence(
+        [torch.tensor(b["labels"], dtype=torch.long) for b in batch], batch_first=True, padding_value=eos_token
+    )
+    return {"input_ids": input_ids, "attention_mask": attn_mask, "labels": labels}
+
+
+def get_samsum_collate_fn(dataset_processer, dataset_config):
+    return collate_fn