add gsm8k + finetuning + more pretrain data

configs/gsm8k.yaml

@@ -0,0 +1,48 @@
+output_dir: "" # just a placeholder
+bf16: true
+model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
+label_names: ["labels"]
+# eval_on_start: True
+# eval_strategy: "steps"
+# eval_steps: 500
+# save_strategy: "no"
+# # save_steps: 500
+# logging_strategy: "steps"
+# logging_steps: 100
+# use_liger_kernel: true
+# remove_unused_columns: false
+
+# needed to avoid OOM by compute the metrics batch by batch
+# w/o this the trainer stores logits of all sample in memory...
+# batch_eval_metrics: true
+
+per_device_train_batch_size: 8
+per_device_eval_batch_size: 8
+max_val_samples_per_ds: 1000
+# optim: schedule_free_adamw
+
+learning_rate: 0.00002
+# lr_scheduler_type: "constant_with_warmup"
+neftune_noise_alpha: 5
+weight_decay: 0.01
+# warmup_ratio: 0.1
+warmup_steps: 100
+
+dataloader_prefetch_factor: 8
+dataloader_num_workers: 8
+# LoRA
+lora_r: 8
+lora_dropout: 0.05
+target_modules:
+  - down_proj
+  - up_proj
+
+# data
+train_ds_names:
+- gsm8k
+
+val_ds_names:
+- gsm8k
+
+load_best_model_at_end: true
+metric_for_best_model: eval_gsm8k_loss

src/ctx_to_lora/configs.py

@@ -260,6 +260,10 @@ class CtxTrainingArguments:
         default=ExperimentSetup.LORA,
         metadata={"help": "Experiment setup - LoRA, HyperLoRA, or full finetuning"},
     )
+    from_pretrained_checkpoint: str = field(
+        default=None,
+        metadata={"help": "Path to the pretrained checkpoint."},
+    )
     max_base_len: Optional[int] = field(
         default=2**13,
         metadata={"help": "Maximum base length for training."},
@@ -281,11 +285,11 @@ class CtxTrainingArguments:
         metadata={"help": "Wandb notes for the experiment."},
     )
     add_repeat_prompt: bool = field(
-        default=True,
+        default=False,
         metadata={"help": "Whether to add repeat prompt to the dataset."},
     )
     add_negative_prompt: bool = field(
-        default=True,
+        default=False,
         metadata={"help": "Whether to add negative prompt to the dataset."},
     )
     use_kl_loss: bool = field(

src/ctx_to_lora/data_utils.py

@@ -85,6 +85,62 @@ DS_KWARGS = {
             split="train",
         ),
     ),
+    "drop": dict(
+        train=dict(
+            path="ucinlp/drop",
+            split="train[900:]",
+        ),
+        validation=dict(
+            path="ucinlp/drop",
+            split="train[:900]",
+        ),
+    ),
+    "narrativeqa": dict(
+        train=dict(
+            path="nvidia/ChatQA-Training-Data",
+            name="narrativeqa",
+            split="train",
+        ),
+    ),
+    "quoref": dict(
+        train=dict(
+            path="nvidia/ChatQA-Training-Data",
+            name="quoref",
+            split="train",
+        ),
+    ),
+    "ropes": dict(
+        train=dict(
+            path="nvidia/ChatQA-Training-Data",
+            name="ropes",
+            split="train",
+        ),
+    ),
+    "synthetic_convqa": dict(
+        train=dict(
+            path="nvidia/ChatQA-Training-Data",
+            name="synthetic_convqa",
+            split="train",
+        ),
+    ),
+    "gsm8k": dict(
+        train=dict(
+            path="openai/gsm8k",
+            name="main",
+            split="train[100:]",
+        ),
+        validation=dict(
+            path="openai/gsm8k",
+            name="main",
+            split="train[:100]",
+        ),
+    ),
+    "openmathintx-2": dict(
+        train=dict(
+            path="nvidia/OpenMathInstruct-2",
+            split="train_1M[:100000]",
+        ),
+    ),
 }
 
 
@@ -473,6 +529,45 @@ def get_preprocessing_fn(ds_name: str) -> Callable[[dict[str, Any]], dict[str, A
                 "response": sample["answers"]["text"][0],
             }
 
+    elif ds_name == "drop":
+
+        def f(sample):
+            return {
+                "context": sample["passage"],
+                "prompt": sample["question"],
+                "response": ", ".join(sample["answers_spans"]["spans"]),
+            }
+
+    elif ds_name in ["narrativeqa", "quoref", "ropes", "synthetic_convqa"]:
+
+        def f(sample):
+            response = sample["answers"][0]
+            if isinstance(response, list):
+                response = response[0]
+            return {
+                "context": sample["document"],
+                "prompt": sample["messages"][0]["content"],
+                "response": response,
+            }
+
+    elif ds_name == "openmathintx-2":
+
+        def f(sample):
+            return {
+                "context": sample["problem"],
+                "prompt": sample["problem"],
+                "response": sample["generated_solution"],
+            }
+
+    elif "gsm8k" in ds_name:
+
+        def f(sample):
+            return {
+                "context": sample["question"],
+                "prompt": sample["question"],
+                "response": sample["answer"],
+            }
+
     return f