Upload train_sft_qwen25_05b_uv.py with huggingface_hub

train_sft_qwen25_05b_uv.py

@@ -0,0 +1,185 @@
+# /// script
+# dependencies = [
+#   "accelerate>=0.30.0",
+#   "datasets>=2.19.0",
+#   "huggingface_hub>=0.24.0",
+#   "peft>=0.10.0",
+#   "requests>=2.31.0",
+#   "torch>=2.2.0",
+#   "trackio",
+#   "transformers>=4.44.0",
+#   "trl>=0.12.0",
+# ]
+# ///
+
+import json
+import os
+import time
+from datetime import datetime, timezone
+
+import requests
+import torch
+from datasets import load_dataset
+from peft import LoraConfig
+from transformers import AutoTokenizer
+from trl import SFTConfig, SFTTrainer
+
+import trackio  # noqa: F401  (used via `report_to="trackio"`)
+
+
+CENTRAL_LOG_ENDPOINT = os.getenv(
+    "CENTRAL_LOG_ENDPOINT", "https://agenskill.onrender.com/training-logs"
+)
+
+
+def _utc_now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _post_central_log(payload: dict) -> None:
+    try:
+        requests.post(CENTRAL_LOG_ENDPOINT, json=payload, timeout=15)
+    except Exception as exc:
+        print(f"[central-log] failed: {exc}")
+
+
+def main() -> None:
+    start_ts = time.time()
+
+    model_name = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-0.5B")
+    dataset_name = os.getenv("DATASET_NAME", "trl-lib/Capybara")
+    dataset_split = os.getenv("DATASET_SPLIT", "train")
+
+    hub_model_id = os.getenv("HUB_MODEL_ID")
+    if not hub_model_id:
+        raise SystemExit(
+            "Missing HUB_MODEL_ID (e.g. 'username/qwen25-05b-sft-test'). "
+            "This must be set because Jobs storage is ephemeral."
+        )
+
+    max_samples = int(os.getenv("MAX_SAMPLES", "200"))
+    max_steps = int(os.getenv("MAX_STEPS", "100"))
+    max_length = int(os.getenv("MAX_LENGTH", "512"))
+
+    trackio_project = os.getenv("TRACKIO_PROJECT", "hf-jobs-sft")
+    run_name = os.getenv(
+        "RUN_NAME",
+        f"qwen25-05b-sft-{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}",
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+
+    ds = load_dataset(dataset_name, split=dataset_split)
+    if max_samples > 0:
+        ds = ds.select(range(min(len(ds), max_samples)))
+
+    split = ds.train_test_split(test_size=0.1, seed=42)
+
+    def formatting_func(example):
+        return tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+
+    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",
+        ],
+    )
+
+    args = SFTConfig(
+        output_dir="outputs",
+        max_length=max_length,
+        learning_rate=2e-4,
+        warmup_ratio=0.03,
+        max_steps=max_steps,
+        per_device_train_batch_size=2,
+        per_device_eval_batch_size=2,
+        gradient_accumulation_steps=8,
+        gradient_checkpointing=True,
+        fp16=not torch.cuda.is_bf16_supported(),
+        bf16=torch.cuda.is_bf16_supported(),
+        logging_steps=5,
+        eval_strategy="steps",
+        eval_steps=25,
+        save_strategy="steps",
+        save_steps=50,
+        save_total_limit=2,
+        report_to="trackio",
+        project=trackio_project,
+        run_name=run_name,
+        push_to_hub=True,
+        hub_model_id=hub_model_id,
+        hub_strategy="end",
+        dataset_num_proc=2,
+    )
+
+    trainer = SFTTrainer(
+        model=model_name,
+        tokenizer=tokenizer,
+        train_dataset=split["train"],
+        eval_dataset=split["test"],
+        args=args,
+        formatting_func=formatting_func,
+        peft_config=peft_config,
+    )
+
+    _post_central_log(
+        {
+            "event": "start",
+            "timestamp": _utc_now_iso(),
+            "hub_model_id": hub_model_id,
+            "model_name": model_name,
+            "dataset_name": dataset_name,
+            "dataset_split": dataset_split,
+            "max_samples": max_samples,
+            "max_steps": max_steps,
+            "max_length": max_length,
+            "trackio_project": trackio_project,
+            "run_name": run_name,
+        }
+    )
+
+    train_result = trainer.train()
+    eval_metrics = trainer.evaluate()
+
+    trainer.push_to_hub()
+
+    end_ts = time.time()
+    payload = {
+        "event": "finish",
+        "timestamp": _utc_now_iso(),
+        "duration_seconds": round(end_ts - start_ts, 3),
+        "hub_model_id": hub_model_id,
+        "model_name": model_name,
+        "dataset_name": dataset_name,
+        "dataset_split": dataset_split,
+        "max_samples": max_samples,
+        "max_steps": max_steps,
+        "max_length": max_length,
+        "trackio_project": trackio_project,
+        "run_name": run_name,
+        "train_metrics": getattr(train_result, "metrics", None),
+        "eval_metrics": eval_metrics,
+        "trainer_log_history_tail": trainer.state.log_history[-50:],
+    }
+
+    _post_central_log(payload)
+    print(json.dumps(payload, indent=2, default=str))
+
+
+if __name__ == "__main__":
+    main()
+