Upload sft_qwen2_capybara.py with huggingface_hub

sft_qwen2_capybara.py

@@ -0,0 +1,203 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "trl==1.5.1",
+#     "peft==0.19.1",
+#     "transformers==5.9.0",
+#     "datasets==4.8.5",
+#     "accelerate==1.13.0",
+#     "trackio==0.26.0",
+#     "torch",
+# ]
+# ///
+"""
+LoRA SFT of Qwen/Qwen2-0.5B (BASE) on trl-lib/Capybara (conversational 'messages').
+
+Grounded in research (Principle 1), cross-checked against the canonical templates:
+- Skill template: huggingface-llm-trainer/scripts/train_sft_example.py
+  (SFTTrainer(model=str, train_dataset, peft_config=LoraConfig(...), args=SFTConfig(...)))
+- Canonical TRL SFT: https://github.com/huggingface/trl/blob/main/trl/scripts/sft.py
+- SFTTrainer consumes the 'messages' column directly and applies the chat template;
+  NO formatting_func needed (TRL "conversational" dataset support).
+
+KEY VERIFIED FACTS (by inspection, Principle 2):
+- Qwen/Qwen2-0.5B (BASE) ALREADY ships a ChatML chat template in its
+  tokenizer_config.json (327 chars, uses <|im_start|>/<|im_end|>). So we do NOT
+  inject a template; SFTTrainer applies the model's own template to 'messages'.
+  Subtlety: the base tokenizer's eos_token is <|endoftext|> while the template
+  ends turns with <|im_end|>. For SFT we set eos_token="<|im_end|>" via SFTConfig
+  so the model learns to stop at end-of-turn (and assistant_only_loss aligns).
+- trl-lib/Capybara: config 'default', splits {train:15806, test}, columns
+  ['source','messages','num_turns']; 'messages' is multi-turn alternating
+  user/assistant. Already in required schema (no mapping).
+- SFTConfig fields eos_token / assistant_only_loss / max_length / packing all
+  present in current trl (verified in trl/trainer/sft_config.py).
+
+LoRA recipe for a small instruct model (literature):
+- target ALL linear layers (q,k,v,o,gate,up,down) > attention-only
+  (QLoRA, Dettmers et al. 2023, arXiv:2305.14314; PEFT/TRL all-linear guidance).
+- r=16, alpha=32 (alpha=2r), dropout=0.05 (TRL ModelConfig LoRA defaults).
+- lr=2e-4, cosine schedule, warmup_ratio=0.03 (standard LoRA SFT LR for small models).
+- assistant_only_loss=True: loss only on assistant turns for multi-turn SFT.
+
+Monitoring: Trackio (report_to) + structured alerts at decision points (§5.6/§5.7).
+
+Resources (R14):
+- Model:   https://huggingface.co/Qwen/Qwen2-0.5B
+- Dataset: https://huggingface.co/datasets/trl-lib/Capybara
+"""
+
+import os
+
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import SFTConfig, SFTTrainer
+
+import trackio
+
+
+# --------------------------------------------------------------------------- #
+# Config (overridable via env so the SAME script serves smoke + full run)
+# --------------------------------------------------------------------------- #
+MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2-0.5B")
+DATASET_ID = os.environ.get("DATASET_ID", "trl-lib/Capybara")
+EOS_TOKEN = os.environ.get("EOS_TOKEN", "<|im_end|>")  # base eos is <|endoftext|>
+HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "krasserm/Qwen2-0.5B-Capybara-LoRA")
+
+SMOKE = os.environ.get("SMOKE", "0") == "1"
+LIMIT = int(os.environ.get("LIMIT", "0"))  # 0 = full split
+MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "2048"))
+
+TRACKIO_PROJECT = os.environ.get("TRACKIO_PROJECT", "qwen2-0.5b-capybara-lora")
+TRACKIO_SPACE = os.environ.get("TRACKIO_SPACE_ID", "")  # e.g. krasserm/trackio
+
+
+ALERT_WEBHOOK = os.environ.get("ALERT_WEBHOOK_URL") or None
+
+# Verified API (trackio 0.26.x, gradio-app/trackio __init__.py + skill alerts.md):
+#   trackio.alert(title, text, level=trackio.AlertLevel.{INFO,WARN,ERROR}, webhook_url)
+_LEVELS = {
+    "info": "INFO",
+    "success": "INFO",
+    "warn": "WARN",
+    "warning": "WARN",
+    "error": "ERROR",
+}
+
+
+def alert(level: str, title: str, message: str):
+    """Structured Trackio alert at a decision point (§5.7); never crash training."""
+    try:
+        lvl = getattr(trackio.AlertLevel, _LEVELS.get(level, "INFO"))
+        trackio.alert(title=title, text=message, level=lvl, webhook_url=ALERT_WEBHOOK)
+    except Exception as e:
+        print(f"[alert:{level}] {title} :: {message} (trackio.alert failed: {e})")
+
+
+def main():
+    # ------------------------------------------------------------------ #
+    # Monitoring init (live dashboard)
+    # ------------------------------------------------------------------ #
+    init_kwargs = {"project": TRACKIO_PROJECT}
+    if TRACKIO_SPACE:
+        init_kwargs["space_id"] = TRACKIO_SPACE
+    trackio.init(
+        config={
+            "model": MODEL_ID,
+            "dataset": DATASET_ID,
+            "method": "SFT+LoRA",
+            "eos_token": EOS_TOKEN,
+            "smoke": SMOKE,
+            "limit": LIMIT,
+            "max_length": MAX_LENGTH,
+        },
+        **init_kwargs,
+    )
+    alert("info", "Run started", f"SFT+LoRA {MODEL_ID} on {DATASET_ID} (smoke={SMOKE})")
+
+    # ------------------------------------------------------------------ #
+    # Data — consumed directly by SFTTrainer (no formatting_func)
+    # ------------------------------------------------------------------ #
+    ds = load_dataset(DATASET_ID, split="train")
+    if LIMIT > 0:
+        ds = ds.select(range(min(LIMIT, len(ds))))
+    if "messages" not in ds.column_names:
+        alert("error", "Schema mismatch", f"'messages' not in {ds.column_names}")
+        raise ValueError(f"Expected 'messages' column, got {ds.column_names}")
+    print(f"Loaded {len(ds)} examples; columns={ds.column_names}")
+
+    # ------------------------------------------------------------------ #
+    # LoRA config (all-linear, r16/alpha32) — research recipe
+    # ------------------------------------------------------------------ #
+    peft_config = LoraConfig(
+        r=16,
+        lora_alpha=32,
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=[
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+    )
+
+    # ------------------------------------------------------------------ #
+    # SFTConfig
+    # ------------------------------------------------------------------ #
+    sft_config = SFTConfig(
+        output_dir="/tmp/sft-out",
+        # BASE already ships a ChatML template; just set EOS to end-of-turn token
+        # so the model learns to stop (base eos is <|endoftext|>).
+        eos_token=EOS_TOKEN,
+        max_length=MAX_LENGTH,
+        packing=not SMOKE,            # packing on for the real run; off for tiny smoke
+        assistant_only_loss=True,     # loss only on assistant turns (multi-turn SFT)
+        # Optimization (LoRA SFT recipe)
+        learning_rate=2e-4,
+        lr_scheduler_type="cosine",
+        warmup_ratio=0.03,
+        weight_decay=0.0,
+        num_train_epochs=1 if SMOKE else 2,
+        max_steps=8 if SMOKE else -1,
+        per_device_train_batch_size=2 if SMOKE else 8,
+        gradient_accumulation_steps=1 if SMOKE else 4,
+        gradient_checkpointing=True,
+        bf16=True,
+        # Logging / saving / monitoring
+        logging_steps=1 if SMOKE else 10,
+        save_strategy="no" if SMOKE else "epoch",
+        report_to=["trackio"],
+        run_name=TRACKIO_PROJECT,
+        # Persistence (Principle 4) — disabled in smoke
+        push_to_hub=not SMOKE,
+        hub_model_id=HUB_MODEL_ID,
+        seed=42,
+    )
+
+    trainer = SFTTrainer(
+        model=MODEL_ID,
+        args=sft_config,
+        train_dataset=ds,
+        peft_config=peft_config,
+    )
+
+    alert("info", "Training start",
+          f"{len(ds)} ex | bs={sft_config.per_device_train_batch_size}"
+          f" x ga={sft_config.gradient_accumulation_steps} |"
+          f" epochs={sft_config.num_train_epochs} max_steps={sft_config.max_steps}")
+
+    result = trainer.train()
+    tr_loss = result.metrics.get("train_loss")
+    alert("success", "Training done", f"train_loss={tr_loss}")
+
+    if sft_config.push_to_hub:
+        trainer.push_to_hub(dataset_name=DATASET_ID)
+        alert("success", "Pushed to Hub", f"https://huggingface.co/{HUB_MODEL_ID}")
+    else:
+        trainer.save_model(sft_config.output_dir)
+
+    trackio.finish()
+
+
+if __name__ == "__main__":
+    main()