SFT-only launcher mirroring working Colab

scripts/hf_sft_launcher.py

@@ -0,0 +1,105 @@
+"""HF Jobs SFT launcher — exact mirror of working Colab cells.
+
+No TRL OpenEnv, no env server, no GRPO. Pure SFT on 611 traces.
+Same code that works on Colab T4. H200 just runs it ~10× faster.
+"""
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+
+WORK = "/tmp/opsguard"
+HUB_REPO = "sai1906/opsguard-sft"
+
+
+def sh(cmd):
+    print(f"[sh] {cmd if isinstance(cmd, str) else ' '.join(cmd)}", flush=True)
+    return subprocess.run(cmd, shell=isinstance(cmd, str), check=True)
+
+
+def main():
+    if not os.environ.get("HF_TOKEN"):
+        raise SystemExit("HF_TOKEN required")
+
+    sh([sys.executable, "-m", "pip", "install", "-q", "-U",
+        "transformers>=4.46", "trl>=0.18", "peft>=0.13", "bitsandbytes>=0.44",
+        "datasets", "accelerate>=1.0", "huggingface_hub", "matplotlib", "networkx>=3"])
+
+    sh(f"rm -rf {WORK}")
+    sh(["git", "clone", "https://huggingface.co/spaces/sai1906/opsguard", WORK])
+    os.chdir(WORK)
+    sys.path.insert(0, WORK)
+
+    from huggingface_hub import login, HfApi
+    login(token=os.environ["HF_TOKEN"])
+    HfApi(token=os.environ["HF_TOKEN"]).create_repo(HUB_REPO, repo_type="model", exist_ok=True)
+    print(f"[preflight] repo {HUB_REPO} ready", flush=True)
+
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+
+    MODEL = "Qwen/Qwen2.5-7B-Instruct"
+    print(f"[model] loading {MODEL} 4bit...", flush=True)
+    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
+                              bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
+    tok = AutoTokenizer.from_pretrained(MODEL)
+    if tok.pad_token_id is None:
+        tok.pad_token_id = tok.eos_token_id
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL, quantization_config=bnb, torch_dtype=torch.bfloat16, device_map="auto",
+    )
+    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
+    lc = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0, bias="none",
+                    target_modules=["q_proj","k_proj","v_proj","o_proj",
+                                    "gate_proj","up_proj","down_proj"],
+                    task_type="CAUSAL_LM")
+    model = get_peft_model(model, lc)
+    model.print_trainable_parameters()
+
+    import json
+    from datasets import Dataset
+    rows = [json.loads(l) for l in open("data/sft_traces.jsonl")]
+    print(f"[data] {len(rows)} SFT traces loaded", flush=True)
+    texts = [r["prompt"] + "\n\nACTION:\n" + r["completion"] + tok.eos_token for r in rows]
+    ds = Dataset.from_list([{"text": t} for t in texts])
+
+    from trl import SFTConfig, SFTTrainer
+    import inspect
+
+    def safe_kwargs(cls, kw):
+        sig = inspect.signature(cls).parameters
+        return {k: v for k, v in kw.items() if k in sig}
+
+    raw = dict(
+        output_dir="/tmp/opsguard-sft",
+        per_device_train_batch_size=4,
+        gradient_accumulation_steps=4,
+        num_train_epochs=2,
+        learning_rate=1e-4,
+        warmup_ratio=0.03,
+        logging_steps=5,
+        save_strategy="epoch",
+        save_total_limit=1,
+        bf16=True,
+        max_seq_length=2048,
+        dataset_text_field="text",
+        report_to="none",
+        push_to_hub=False,
+    )
+    cfg = SFTConfig(**safe_kwargs(SFTConfig, raw))
+    trainer = SFTTrainer(model=model, train_dataset=ds, args=cfg, processing_class=tok)
+    trainer.train()
+
+    print("[push] saving + pushing to Hub...", flush=True)
+    model.save_pretrained("/tmp/opsguard-sft-lora")
+    tok.save_pretrained("/tmp/opsguard-sft-lora")
+    HfApi(token=os.environ["HF_TOKEN"]).upload_folder(
+        folder_path="/tmp/opsguard-sft-lora", repo_id=HUB_REPO, repo_type="model",
+    )
+    print(f"DONE: pushed to https://huggingface.co/{HUB_REPO}", flush=True)
+
+
+if __name__ == "__main__":
+    main()