From 48cdc65c0eafa08fc39b5bfa49ef7c42943ac91d Mon Sep 17 00:00:00 2001
From: vin <git@vineetk.net>
Date: Sat, 23 Aug 2025 17:23:59 -0400
Subject: add finetuning scripts

---
 finetune.py | 99 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 length.py   | 23 ++++++++++++++
 2 files changed, 122 insertions(+)
 create mode 100644 finetune.py
 create mode 100644 length.py

diff --git a/finetune.py b/finetune.py
new file mode 100644
index 0000000..bca7ac1
--- /dev/null
+++ b/finetune.py
@@ -0,0 +1,99 @@
+import torch
+from datasets import load_dataset
+from peft import LoraConfig
+from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+from trl import SFTTrainer, SFTConfig
+import os
+
+#MODEL_ID = "google/gemma-3-270m-it"
+MODEL_ID = "google/gemma-3-1b-it"
+DATASET_TRAIN_PATH = "train.jsonl"
+DATASET_VAL_PATH = "validation.jsonl"
+OUTPUT_DIR = "./gemma3-1b-it-sponsors"
+LORA_R = 16
+LORA_ALPHA = 32
+LORA_DROPOUT = 0.05
+LORA_TARGET_MODULES = [
+    "q_proj", "k_proj", "v_proj", "o_proj",
+    "gate_proj", "up_proj", "down_proj",
+]
+
+def main():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    # A pad token is required for training
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        attn_implementation='eager',
+        offload_folder="./offload",
+        offload_state_dict=True,
+    )
+    model.config.use_cache = False         # disable kv-cache during training
+    model.gradient_checkpointing_enable()  # enable gradient checkpointing
+
+    train_dataset = load_dataset("json", data_files=DATASET_TRAIN_PATH, split="train")
+    validation_dataset = load_dataset("json", data_files=DATASET_VAL_PATH, split="train")
+
+    peft_config = LoraConfig(
+        r=LORA_R,
+        lora_alpha=LORA_ALPHA,
+        target_modules=LORA_TARGET_MODULES,
+        lora_dropout=LORA_DROPOUT,
+        bias="none",
+        task_type="CAUSAL_LM"
+    )
+
+    print("Model and LoRA config loaded. Setting up trainer...")
+
+    training_args = SFTConfig(
+        output_dir=OUTPUT_DIR,
+        
+        # Training parameters
+        num_train_epochs=3,
+        per_device_eval_batch_size=1,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=8,
+        gradient_checkpointing=True,
+        learning_rate=1e-4,
+        lr_scheduler_type="linear",
+        optim="adamw_torch",
+        
+        # Model and data parameters
+        max_length=5120,
+        dataset_text_field="text",
+        
+        # Technical parameters
+        bf16=True,
+        max_grad_norm=1.0,
+        warmup_ratio=0.03,
+        
+        # Logging, saving, and evaluation
+        logging_steps=25,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+    )
+    training_args.eval_accumulation_steps = 4
+
+    trainer = SFTTrainer(
+        model=model,
+        processing_class=tokenizer,
+        args=training_args,
+        peft_config=peft_config,
+        train_dataset=train_dataset,
+        eval_dataset=validation_dataset,
+    )
+
+    print("Trainer initialized. Starting fine-tuning...")
+    trainer.train()
+
+    print("Training complete. Saving the best model...")
+    trainer.save_model(os.path.join(OUTPUT_DIR, "final_checkpoint"))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/length.py b/length.py
new file mode 100644
index 0000000..06f8055
--- /dev/null
+++ b/length.py
@@ -0,0 +1,23 @@
+import json
+from transformers import AutoTokenizer
+
+MODEL_ID = "google/gemma-3-270m-it"
+DATASET_PATH = "train.jsonl"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+token_lengths = []
+
+print(f"Loading dataset from {DATASET_PATH}...")
+with open(DATASET_PATH, 'r', encoding='utf-8') as f:
+    for line in f:
+        entry = json.loads(line)
+        text = entry.get("text", "")
+        # Tokenize the text and get the number of tokens
+        length = len(tokenizer(text).input_ids)
+        token_lengths.append(length)
+
+print(f"Number of examples: {len(token_lengths)}")
+if token_lengths:
+    print(f"Min length:    {min(token_lengths)}")
+    print(f"Max length:    {max(token_lengths)}")
+    print(f"Average length:  {sum(token_lengths) / len(token_lengths):.2f}")
-- 
cgit v1.2.3