Add training script with QLoRA

train.py

@@ -1,25 +1,19 @@
-#!/usr/bin/env python3
-"""
-Hebrew-first SFT training for Gemma 4 E2B/E4B.
-DictaLM-style recipe: 50/50 Hebrew/English mix, LoRA r=64, 3 epochs.
-"""
-
 import os
-import sys
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
 import random
 import json
 from datasets import load_dataset, concatenate_datasets
-from transformers import AutoModelForImageTextToText, AutoTokenizer
+from transformers import AutoModelForImageTextToText, AutoTokenizer, BitsAndBytesConfig
 from peft import LoraConfig
 from trl import SFTConfig, SFTTrainer
 import trackio
 import torch
 from transformers import TrainerCallback
 
-# Trackio init
 trackio.init(
-    project=os.environ.get("TRACKIO_PROJECT", "hebrew-gemma4"),
-    space_id=os.environ.get("TRACKIO_SPACE_ID", "ssdataanalysis/mlintern-heb4"),
+    project="hebrew-gemma4",
+    space_id="ssdataanalysis/mlintern-heb4",
 )
 
 class TrackioAlertCallback(TrainerCallback):
@@ -28,43 +22,23 @@ class TrackioAlertCallback(TrainerCallback):
             loss = logs["loss"]
             step = state.global_step
             if loss > 5.0 and step > 50:
-                trackio.alert(
-                    title="High Loss Warning",
-                    text=f"loss={loss:.3f} at step {step} — lr may be too high, consider reducing",
-                    level="WARN"
-                )
+                trackio.alert(title="High Loss Warning", text=f"loss={loss:.3f} at step {step} lr too high", level="WARN")
             elif step % 100 == 0:
-                trackio.alert(
-                    title="Training Progress",
-                    text=f"loss={loss:.3f} at step {step}",
-                    level="INFO"
-                )
-
+                trackio.alert(title="Training Progress", text=f"loss={loss:.3f} at step {step}", level="INFO")
     def on_evaluate(self, args, state, control, metrics=None, **kwargs):
         if metrics and "eval_loss" in metrics:
-            trackio.alert(
-                title="Eval Complete",
-                text=f"eval_loss={metrics['eval_loss']:.3f} at step {state.global_step}",
-                level="INFO"
-            )
+            trackio.alert(title="Eval Complete", text=f"eval_loss={metrics['eval_loss']:.3f} at step {state.global_step}", level="INFO")
 
 def convert_hebrew_qa_to_messages(example):
-    """Convert yuvalav/hebrew-qa to messages format."""
     instruction = example.get("instruction", "")
     input_text = example.get("input", "")
     output = example.get("output", "")
     user_content = instruction
     if input_text and str(input_text).strip():
         user_content += "\n" + str(input_text)
-    return {
-        "messages": [
-            {"role": "user", "content": user_content},
-            {"role": "assistant", "content": output},
-        ]
-    }
+    return {"messages": [{"role": "user", "content": user_content}, {"role": "assistant", "content": output}]}
 
 def convert_hebrew_chatml_to_messages(example):
-    """Convert itayl/hebrewQA-chatml to messages format."""
     conversations = example.get("conversations", [])
     messages = []
     for turn in conversations:
@@ -79,164 +53,103 @@ def convert_hebrew_chatml_to_messages(example):
     return {"messages": messages}
 
 def prepare_dataset(hebrew_ratio=0.5, max_total=120000, seed=42):
-    """Prepare mixed Hebrew-English instruction dataset."""
     random.seed(seed)
-
-    hebrew_samples_target = int(max_total * hebrew_ratio)
-    english_samples_target = max_total - hebrew_samples_target
-
     datasets_list = []
-
-    # Hebrew datasets
     print("Loading Hebrew datasets...")
-
-    # 1. yuvalav/hebrew-qa (~30K)
     ds_he1 = load_dataset("yuvalav/hebrew-qa", split="train")
     ds_he1 = ds_he1.map(convert_hebrew_qa_to_messages, remove_columns=ds_he1.column_names)
     datasets_list.append(("hebrew-qa", ds_he1))
     print(f"  hebrew-qa: {len(ds_he1)}")
-
-    # 2. itayl/hebrewQA-chatml (~30K)
     ds_he2 = load_dataset("itayl/hebrewQA-chatml", split="train")
     ds_he2 = ds_he2.map(convert_hebrew_chatml_to_messages, remove_columns=ds_he2.column_names)
     datasets_list.append(("hebrewQA-chatml", ds_he2))
     print(f"  hebrewQA-chatml: {len(ds_he2)}")
-
-    total_hebrew = len(ds_he1) + len(ds_he2)
-    print(f"Total Hebrew: {total_hebrew}")
-
-    # English datasets
     print("Loading English datasets...")
-
-    # 3. OpenHermes 2.5 H4 (~950K, take subset)
     ds_en1 = load_dataset("HuggingFaceTB/OpenHermes-2.5-H4", split="train_sft")
     ds_en1 = ds_en1.remove_columns([c for c in ds_en1.column_names if c != "messages"])
-    # Filter to only user/assistant/system roles
     def filter_messages(example):
         msgs = example.get("messages", [])
         return all(m.get("role") in ["user", "assistant", "system"] for m in msgs)
     ds_en1 = ds_en1.filter(filter_messages)
-    # Sample
-    if len(ds_en1) > english_samples_target:
-        ds_en1 = ds_en1.shuffle(seed=seed).select(range(english_samples_target))
+    english_target = max_total - (len(ds_he1) + len(ds_he2))
+    if len(ds_en1) > english_target:
+        ds_en1 = ds_en1.shuffle(seed=seed).select(range(english_target))
     datasets_list.append(("OpenHermes", ds_en1))
     print(f"  OpenHermes: {len(ds_en1)}")
-
-    # Combine and shuffle
     all_datasets = [d for _, d in datasets_list]
     combined = concatenate_datasets(all_datasets)
     combined = combined.shuffle(seed=seed)
-
-    # Verify format
-    sample = combined[0]
-    print(f"\nSample messages: {json.dumps(sample['messages'][:2], ensure_ascii=False)}")
-
-    total = len(combined)
-    hebrew_count = len(ds_he1) + len(ds_he2)
-    print(f"\nFinal dataset: {total} samples ({hebrew_count} Hebrew, {len(ds_en1)} English)")
-    print(f"Hebrew ratio: {hebrew_count/total:.2%}")
-
+    print(f"Final dataset: {len(combined)} samples")
     return combined
 
-def train(model_id, output_dir, hebrew_ratio=0.5, max_total=120000):
-    print(f"=== Training {model_id} -> {output_dir} ===")
-
-    # Dataset
-    train_dataset = prepare_dataset(hebrew_ratio=hebrew_ratio, max_total=max_total)
-
-    # Create a small eval set (first 1000 samples)
-    eval_dataset = train_dataset.select(range(min(1000, len(train_dataset))))
-    train_dataset = train_dataset.select(range(min(1000, len(train_dataset)), len(train_dataset)))
-
-    # Tokenizer
-    print("Loading tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-
-    # Model
-    print("Loading model...")
-    model = AutoModelForImageTextToText.from_pretrained(
-        model_id,
-        attn_implementation="sdpa",
-        dtype="bfloat16",
-        device_map="auto",
-    )
-
-    # LoRA config - Dicta style: high rank, all linear layers
-    peft_config = LoraConfig(
-        r=64,
-        lora_alpha=16,
-        lora_dropout=0.05,
-        bias="none",
-        task_type="CAUSAL_LM",
-        target_modules="all-linear",
-        exclude_modules=["vision_tower", "multi_modal_projector"],
-    )
-
-    # Training args
-    training_args = SFTConfig(
-        output_dir=output_dir,
-        num_train_epochs=3,
-        per_device_train_batch_size=1,
-        gradient_accumulation_steps=8,
-        learning_rate=2e-4,
-        lr_scheduler_type="cosine",
-        warmup_ratio=0.03,
-        weight_decay=0.01,
-        max_length=4096,
-        packing=True,
-        bf16=True,
-        use_liger_kernel=True,
-        logging_strategy="steps",
-        logging_steps=10,
-        logging_first_step=True,
-        eval_strategy="steps",
-        eval_steps=100,
-        save_strategy="epoch",
-        save_total_limit=2,
-        push_to_hub=True,
-        hub_model_id=output_dir,
-        report_to="trackio",
-        run_name=output_dir,
-        remove_unused_columns=False,
-        disable_tqdm=True,
-        dataset_num_proc=8,
-        gradient_checkpointing=True,
-    )
-
-    # Trainer
-    trainer = SFTTrainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        eval_dataset=eval_dataset,
-        peft_config=peft_config,
-        processing_class=tokenizer,
-        callbacks=[TrackioAlertCallback()],
-    )
-
-    # Train
-    print("Starting training...")
-    trainer.train()
-
-    # Save
-    trainer.save_model(output_dir)
-    trainer.push_to_hub()
-
-    trackio.alert(
-        title="Training Complete",
-        text=f"Model {output_dir} training completed successfully",
-        level="INFO"
-    )
+model_id = os.environ.get("MODEL_ID", "google/gemma-4-E2B-it")
+output_dir = os.environ.get("OUTPUT_DIR", "ssdataanalysis/gemma-4-E2B-hebrew-first")
+print(f"=== Training {model_id} -> {output_dir} ===")
+
+train_dataset = prepare_dataset(hebrew_ratio=0.5, max_total=120000)
+eval_dataset = train_dataset.select(range(min(1000, len(train_dataset))))
+train_dataset = train_dataset.select(range(min(1000, len(train_dataset)), len(train_dataset)))
+
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+print("Loading model with 4-bit quantization...")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+model = AutoModelForImageTextToText.from_pretrained(
+    model_id,
+    attn_implementation="sdpa",
+    quantization_config=bnb_config,
+    device_map="auto",
+)
 
-    print(f"Done! Model saved to {output_dir}")
+peft_config = LoraConfig(
+    r=64, lora_alpha=16, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
+    target_modules="all-linear",
+    exclude_modules=["vision_tower", "multi_modal_projector"],
+)
 
-if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_id", type=str, required=True)
-    parser.add_argument("--output_dir", type=str, required=True)
-    parser.add_argument("--hebrew_ratio", type=float, default=0.5)
-    parser.add_argument("--max_total", type=int, default=120000)
-    args = parser.parse_args()
+training_args = SFTConfig(
+    output_dir=output_dir,
+    num_train_epochs=3,
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=8,
+    learning_rate=2e-4,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.03,
+    weight_decay=0.01,
+    max_length=4096,
+    packing=True,
+    bf16=True,
+    use_liger_kernel=True,
+    logging_strategy="steps",
+    logging_steps=10,
+    logging_first_step=True,
+    eval_strategy="steps",
+    eval_steps=100,
+    save_strategy="epoch",
+    save_total_limit=2,
+    push_to_hub=True,
+    hub_model_id=output_dir,
+    report_to="trackio",
+    run_name=output_dir,
+    remove_unused_columns=False,
+    disable_tqdm=True,
+    dataset_num_proc=8,
+    gradient_checkpointing=True,
+)
 
-    train(args.model_id, args.output_dir, args.hebrew_ratio, args.max_total)
+trainer = SFTTrainer(
+    model=model, args=training_args, train_dataset=train_dataset,
+    eval_dataset=eval_dataset, peft_config=peft_config,
+    processing_class=tokenizer, callbacks=[TrackioAlertCallback()],
+)
+print("Starting training...")
+trainer.train()
+trainer.save_model(output_dir)
+trainer.push_to_hub()
+trackio.alert(title="Training Complete", text=f"Model {output_dir} training completed successfully", level="INFO")
+print(f"Done! Model saved to {output_dir}")