SecureCoder trainer v1

train_securecoder.py

@@ -0,0 +1,642 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "unsloth",
+#     "datasets",
+#     "trl>=0.22",
+#     "transformers>=4.57",
+#     "trackio",
+#     "huggingface_hub",
+# ]
+# ///
+"""SecureCoder: QLoRA fine-tune for code + tool calling + cybersecurity.
+
+Default base: Qwen/Qwen3-Coder-30B-A3B-Instruct (Apache-2.0, 30.5B MoE, ~3B
+active) - a MoE that trains like a small model and runs like a useful one.
+
+Runs anywhere; same file for local validation, a GPU smoke test, and the real
+run:
+
+  uv run train_securecoder.py --validate-only          # no GPU needed
+  uv run train_securecoder.py --smoke --output-repo you/securecoder-smoke
+  uv run train_securecoder.py --num-epochs 1 --output-repo you/securecoder-30b-pro
+
+Launch on Hugging Face Jobs (see README.md for why the URL form is used):
+
+  hf jobs run -d --flavor l40sx1 --timeout 12h --secrets HF_TOKEN \\
+      ghcr.io/astral-sh/uv:python3.12-bookworm \\
+      uv run --no-project https://huggingface.co/USER/securecoder-scripts/resolve/main/train_securecoder.py \\
+      -- --num-epochs 1 --output-repo USER/securecoder-30b-pro
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+import random
+import sys
+import time
+from dataclasses import dataclass
+from typing import Any
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+log = logging.getLogger("securecoder")
+
+
+# --------------------------------------------------------------------------
+# Data mix
+# --------------------------------------------------------------------------
+@dataclass
+class Source:
+    """One dataset feeding the mix.
+
+    ``kind`` selects the converter ('auto' sniffs columns). ``limit`` is how
+    many rows are taken - the sources differ wildly in size, so the cap *is*
+    the recipe. Adjust the numbers, not the code.
+    """
+
+    repo: str
+    limit: int
+    kind: str = "auto"
+    config: str | None = None
+    split: str = "train"
+    note: str = ""
+
+
+MIX: list[Source] = [
+    # ---- tool calling ----------------------------------------------------
+    Source("NousResearch/hermes-function-calling-v1", 9000, "tools", "func_calling",
+           note="Hermes FC: conversations + JSON tool schemas"),
+    Source("NousResearch/hermes-function-calling-v1", 3000, "tools", "func_calling_singleturn",
+           note="single-turn tool selection"),
+    Source("lockon/xlam-function-calling-60k", 10000, "xlam", "dataset",
+           note="xLAM: query/answers/tools API-call pairs"),
+    # ---- coding ----------------------------------------------------------
+    Source("ise-uiuc/Magicoder-OSS-Instruct-75K", 10000, "magicoder",
+           note="self-instruct code problems + solutions"),
+    # ---- cybersecurity ---------------------------------------------------
+    Source("Trendyol/Trendyol-Cybersecurity-Instruction-Tuning-Dataset", 8000, "sua",
+           note="security instruction tuning"),
+    Source("AlicanKiraz0/Cybersecurity-Dataset-Fenrir-v2.1", 5000, "sua",
+           note="broad security Q&A"),
+    Source("Humanlearning/CyberSecurity_OWASP-sft-dataset", 3000, "messages",
+           note="OWASP / secure-coding SFT"),
+    Source("dpevzner/Cybersecurity_Reasoning_Dataset", 3000, "secops", "default", "test",
+           note="command interpretation reasoning (goal -> unified_interpretation)"),
+    Source("MrClipperz134/CTF-Instruct", 3000, "io",
+           note="CTF instruction/output"),
+    Source("TrueNix/ctf-solver-dataset", 3000, "messages",
+           note="CTF solving trajectories"),
+    # ---- capability replay ----------------------------------------------
+    Source("mlabonne/FineTome-100k", 3000, "messages",
+           note="general instruct replay so chat ability does not drift"),
+]
+
+
+def source_name(src: Source) -> str:
+    return src.repo + (f" [{src.config}]" if src.config else "")
+
+# --------------------------------------------------------------------------
+# Schema sniffing -> OpenAI-style chat messages
+# --------------------------------------------------------------------------
+ROLE_ALIASES = {
+    "human": "user", "user": "user", "gpt": "assistant", "assistant": "assistant",
+    "system": "system", "tool": "tool", "function": "tool", "function_call": "tool",
+    "observation": "tool", "chatgpt": "assistant",
+}
+
+
+def _as_list(value: Any) -> list:
+    """Accept a JSON string or an already-parsed list."""
+    if value is None:
+        return []
+    if isinstance(value, list):
+        return value
+    if isinstance(value, str):
+        try:
+            parsed = json.loads(value)
+        except json.JSONDecodeError:
+            return []
+        return parsed if isinstance(parsed, list) else [parsed]
+    return []
+
+
+def _normalise_tool_schema(tool: Any) -> dict | None:
+    """Force a tool schema into the shape Qwen's chat template expects:
+    {"type": "function", "function": {"name", "description", "parameters"}}."""
+    if not isinstance(tool, dict):
+        return None
+    fn = tool.get("function") if "function" in tool else tool
+    if not isinstance(fn, dict) or not fn.get("name"):
+        return None
+    return {
+        "type": "function",
+        "function": {
+            "name": fn["name"],
+            "description": fn.get("description", ""),
+            "parameters": fn.get("parameters") or {"type": "object", "properties": {}},
+        },
+    }
+
+
+def _parse_calls(value: Any) -> list[dict] | None:
+    """Return OpenAI tool_calls if ``value`` is one or more JSON function calls."""
+    if isinstance(value, str):
+        text = value.strip()
+        if not text.startswith(("{", "[")):
+            return None
+        try:
+            value = json.loads(text)
+        except json.JSONDecodeError:
+            return None
+    items = value if isinstance(value, list) else [value]
+    if not items or not all(isinstance(i, dict) and "name" in i for i in items):
+        return None
+    calls = []
+    for i, item in enumerate(items):
+        arguments = item.get("arguments", item.get("parameters", {}))
+        if not isinstance(arguments, str):
+            arguments = json.dumps(arguments)
+        calls.append({
+            "id": f"call_{i}",
+            "type": "function",
+            "function": {"name": item["name"], "arguments": arguments},
+        })
+    return calls
+
+
+def _tools_from_row(row: dict) -> list[dict]:
+    raw = row.get("tools")
+    candidates = [raw] if isinstance(raw, dict) else _as_list(raw)
+    tools = []
+    for candidate in candidates:
+        norm = _normalise_tool_schema(candidate)
+        if norm:
+            tools.append(norm)
+    return tools
+
+def _messages_from_any(row: dict, kind: str) -> tuple[list[dict], list[dict]]:
+    """Convert one dataset row into (messages, tools).
+
+    Returns empty messages when a row cannot be converted confidently; the
+    caller counts those, so a silent schema change shows up in the logs instead
+    of quietly training on nothing.
+    """
+    tools = _tools_from_row(row)
+    messages: list[dict] = []
+
+    # --- explicit chat formats (Hermes, OWASP SFT, ctf-solver, FineTome) ---
+    if isinstance(row.get("conversations"), list) or isinstance(row.get("messages"), list):
+        for turn in row.get("conversations") or row.get("messages") or []:
+            if not isinstance(turn, dict):
+                continue
+            role = ROLE_ALIASES.get(str(turn.get("role") or turn.get("from") or "").lower())
+            if role is None:
+                continue
+            content = turn.get("content", turn.get("value", ""))
+            if role == "assistant":
+                calls = _parse_calls(content)
+                if calls:
+                    messages.append({"role": "assistant", "content": None, "tool_calls": calls})
+                    continue
+            if role == "tool":
+                if not isinstance(content, str):
+                    content = json.dumps(content)
+                messages.append({"role": "tool", "content": content,
+                                 "tool_call_id": turn.get("tool_call_id", "call_0")})
+                continue
+            if not isinstance(content, str):
+                content = json.dumps(content) if content is not None else ""
+            if content.strip():
+                messages.append({"role": role, "content": content})
+        return messages, tools
+
+    # --- xLAM: query + answers + tools ------------------------------------
+    if kind == "xlam" or (row.get("query") and row.get("answers")):
+        calls = _parse_calls(_as_list(row.get("answers")))
+        if calls and row.get("query"):
+            messages.append({"role": "user", "content": str(row["query"])})
+            messages.append({"role": "assistant", "content": None, "tool_calls": calls})
+        return messages, tools
+
+    # --- system / user / assistant ----------------------------------------
+    if row.get("user") and row.get("assistant"):
+        if row.get("system"):
+            messages.append({"role": "system", "content": str(row["system"])})
+        messages.append({"role": "user", "content": str(row["user"])})
+        calls = _parse_calls(row["assistant"])
+        if calls:
+            messages.append({"role": "assistant", "content": None, "tool_calls": calls})
+        else:
+            messages.append({"role": "assistant", "content": str(row["assistant"])})
+        return messages, tools
+
+    # --- instruction / output (CTF-Instruct) ------------------------------
+    if row.get("instruction") and (row.get("output") or row.get("response")):
+        user = str(row["instruction"])
+        if row.get("input"):
+            user = f"{user}\n\n{row['input']}"
+        messages.append({"role": "user", "content": user})
+        messages.append({"role": "assistant",
+                         "content": str(row.get("output") or row.get("response"))})
+        return messages, tools
+
+    # --- Magicoder: problem / solution ------------------------------------
+    if row.get("problem") and row.get("solution"):
+        messages.append({"role": "user", "content":
+            "You are an exceptionally intelligent coding assistant that consistently "
+            "delivers reliable and accurate responses.\n\n" + str(row["problem"])})
+        messages.append({"role": "assistant", "content": str(row["solution"])})
+        return messages, tools
+
+    # --- SecOps reasoning: goal/command -> unified_interpretation ---------
+    if kind == "secops" and row.get("unified_interpretation"):
+        ask = [f"Tool: {row.get('tool', 'shell')}", f"Goal: {row.get('goal', '')}"]
+        for key in ("command", "command_sequence", "nmap_context"):
+            if row.get(key):
+                ask.append(f"{key}: {row[key]}")
+        ask.append("Explain what the output means, what it tells you about the target, "
+                   "and what the next step should be.")
+        messages.append({"role": "user", "content": "\n".join(str(a) for a in ask)})
+        messages.append({"role": "assistant", "content": str(row["unified_interpretation"])})
+        return messages, tools
+
+    # --- generic prompt/completion fallback -------------------------------
+    for pkey, ckey in (("prompt", "completion"), ("question", "answer"), ("input", "output")):
+        if row.get(pkey) and row.get(ckey):
+            messages.append({"role": "user", "content": str(row[pkey])})
+            messages.append({"role": "assistant", "content": str(row[ckey])})
+            return messages, tools
+
+    return [], tools
+
+# --------------------------------------------------------------------------
+# Loading, rendering, dataset construction
+# --------------------------------------------------------------------------
+def load_source(src: Source, token: str | None, progress: bool = False) -> list[dict]:
+    """Pull up to ``limit`` rows from one Hub dataset, streaming so we never
+    download more than we need."""
+    from datasets import load_dataset
+
+    kwargs: dict[str, Any] = {"split": src.split, "streaming": True}
+    if src.config:
+        kwargs["name"] = src.config
+    if token:
+        kwargs["token"] = token
+
+    ds = load_dataset(src.repo, **kwargs)
+    rows = []
+    for i, row in enumerate(ds):
+        if i >= src.limit:
+            break
+        rows.append(dict(row))
+        if progress and i and i % 2500 == 0:
+            log.info("    %s: %d rows...", source_name(src), i)
+    return rows
+
+
+def render_record(tokenizer, messages: list[dict], tools: list[dict] | None = None) -> str:
+    """Render to the model's native chat format (Qwen3 emits <tool_call> blocks)."""
+    return tokenizer.apply_chat_template(
+        messages,
+        tools=tools or None,
+        tokenize=False,
+        add_generation_prompt=False,
+    )
+
+
+def build_dataset(tokenizer, sources: list[Source], token: str | None, validate: bool):
+    """Returns (records, stats). ``records`` are {"text", "source"} dicts."""
+    records: list[dict] = []
+    stats: list[dict] = []
+
+    for src in sources:
+        entry = {"source": source_name(src), "note": src.note, "kept": 0, "skipped": 0,
+                 "tool_samples": 0, "chars": 0, "error": None}
+        try:
+            rows = load_source(src, token, progress=validate)
+            for row in rows:
+                messages, tools = _messages_from_any(row, src.kind)
+                has_answer = any(
+                    (m.get("content") or m.get("tool_calls")) for m in messages
+                    if m["role"] == "assistant"
+                )
+                if not messages or not has_answer:
+                    entry["skipped"] += 1
+                    continue
+                try:
+                    text = render_record(tokenizer, messages, tools)
+                except Exception as exc:  # noqa: BLE001 - bad template input, skip row
+                    if entry["skipped"] < 3:
+                        log.warning("    render failed (%s): %s", source_name(src), exc)
+                    entry["skipped"] += 1
+                    continue
+                if len(text) < 40 or len(text) > 120_000:
+                    entry["skipped"] += 1
+                    continue
+                records.append({"text": text, "source": source_name(src)})
+                entry["kept"] += 1
+                entry["chars"] += len(text)
+                if tools:
+                    entry["tool_samples"] += 1
+        except Exception as exc:  # noqa: BLE001 - one bad dataset must not kill the run
+            entry["error"] = repr(exc)
+            log.error("  %s failed: %s", source_name(src), exc)
+
+        stats.append(entry)
+        log.info("  %-58s kept=%-6d skipped=%-5d tools=%-5d",
+                 entry["source"], entry["kept"], entry["skipped"], entry["tool_samples"])
+
+    random.shuffle(records)
+    return records, stats
+
+
+def print_stats(stats: list[dict], records: list[dict], tokenizer=None) -> None:
+    total_chars = sum(r["chars"] for r in stats if not r["error"])
+    print("\n" + "=" * 78)
+    print("DATA MIX")
+    print("=" * 78)
+    print(f"{'source':<58}{'kept':>7}{'skip':>7}{'tools':>7}")
+    for s in stats:
+        print(f"{s['source']:<58}{s['kept']:>7}{s['skipped']:>7}{s['tool_samples']:>7}")
+        if s["error"]:
+            print(f"    !! {s['error'][:120]}")
+    tool_rows = sum(s["tool_samples"] for s in stats)
+    print("-" * 78)
+    print(f"total rows : {len(records):,}")
+    print(f"tool rows  : {tool_rows:,} ({100 * tool_rows / max(len(records), 1):.1f}%)")
+    print(f"total chars: {total_chars:,}  (~{total_chars // 4:,} tokens)")
+
+# --------------------------------------------------------------------------
+# Model + training
+# --------------------------------------------------------------------------
+# Attention + router only by default: on a 128-expert MoE, adapting every expert
+# MLP means ~800M trainable parameters, which dominates VRAM and step time.
+# Pass --target-modules all-linear when you want the MLP/expert capacity too.
+ATTENTION_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate"]
+
+
+def load_model_and_tokenizer(args):
+    from unsloth import FastLanguageModel
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=args.base_model,
+        max_seq_length=args.max_seq_length,
+        dtype=None,
+        load_in_4bit=not args.no_4bit,
+    )
+
+    targets: Any = args.target_modules
+    if isinstance(targets, str) and targets != "all-linear":
+        targets = [t.strip() for t in targets.split(",") if t.strip()]
+
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=args.lora_r,
+        target_modules=targets,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=0.0,
+        bias="none",
+        use_gradient_checkpointing="unsloth",
+        random_state=args.seed,
+        use_rslora=False,
+    )
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    log.info("trainable parameters: %s (%.2f%% of the model)",
+             f"{trainable:,}", 100 * trainable / max(sum(p.numel() for p in model.parameters()), 1))
+    return model, tokenizer
+
+
+def make_sft_config(**kwargs):
+    """TRL renamed max_seq_length -> max_length; support both."""
+    from trl import SFTConfig
+
+    try:
+        return SFTConfig(max_length=kwargs.pop("max_seq_length"), **kwargs)
+    except TypeError:
+        kwargs["max_seq_length"] = kwargs.get("max_seq_length")
+        return SFTConfig(**kwargs)
+
+
+def build_sft_config(args, has_eval: bool, steps_per_epoch: int | None):
+    import torch
+
+    bf16 = torch.cuda.is_bf16_supported()
+    cfg: dict[str, Any] = dict(
+        output_dir=args.output_dir,
+        per_device_train_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.grad_accum,
+        warmup_ratio=0.03,
+        learning_rate=args.learning_rate,
+        max_grad_norm=1.0,
+        weight_decay=0.01,
+        lr_scheduler_type=args.lr_scheduler,
+        optim="adamw_8bit",
+        logging_steps=args.logging_steps,
+        save_steps=args.save_steps,
+        save_total_limit=2,
+        seed=args.seed,
+        bf16=bf16,
+        fp16=not bf16,
+        max_seq_length=args.max_seq_length,
+        dataset_text_field="text",
+        packing=args.packing,
+        report_to=args.report_to,
+        run_name=args.run_name or "securecoder",
+        remove_unused_columns=False,
+    )
+    if args.max_steps > 0:
+        cfg["max_steps"] = args.max_steps
+    else:
+        cfg["num_train_epochs"] = args.num_epochs
+
+    if has_eval:
+        cfg["eval_strategy"] = "steps"
+        cfg["eval_steps"] = args.save_steps
+        cfg["per_device_eval_batch_size"] = 1
+        cfg["do_eval"] = True
+    return make_sft_config(**cfg)
+
+
+def init_trackio(args):
+    if args.report_to == "none":
+        return
+    try:
+        import trackio
+
+        if args.trackio_space:
+            trackio.init(project=args.trackio_project, space_id=args.trackio_space)
+        else:
+            trackio.init(project=args.trackio_project)
+        log.info("trackio initialised (project=%s space=%s)", args.trackio_project, args.trackio_space)
+    except Exception as exc:  # noqa: BLE001 - monitoring must never kill training
+        log.warning("trackio init failed (%s); continuing without it", exc)
+
+
+def train(args, model, tokenizer, records: list[dict]):
+    from datasets import Dataset
+    from trl import SFTTrainer
+
+    random.Random(args.seed).shuffle(records)
+    split_at = len(records) - args.eval_samples if args.eval_samples > 0 else len(records)
+    train_ds = Dataset.from_list(records[:split_at])
+    eval_ds = Dataset.from_list(records[split_at:]) if args.eval_samples > 0 else None
+    log.info("train rows=%d eval rows=%d", len(train_ds), len(eval_ds) if eval_ds else 0)
+
+    steps_per_epoch = len(train_ds) // max(args.batch_size * args.grad_accum, 1)
+    cfg = build_sft_config(args, eval_ds is not None, steps_per_epoch)
+    init_trackio(args)
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        args=cfg,
+    )
+
+    started = time.time()
+    stats = trainer.train()
+    elapsed = time.time() - started
+    log.info("training finished in %.1f min (final loss %.4f)",
+             elapsed / 60, stats.metrics.get("train_loss", float("nan")))
+
+    if eval_ds is not None:
+        try:
+            metrics = trainer.evaluate()
+            log.info("eval_loss %.4f (train %.4f)",
+                     metrics.get("eval_loss", float("nan")),
+                     stats.metrics.get("train_loss", float("nan")))
+        except Exception as exc:  # noqa: BLE001
+            log.warning("eval failed: %s", exc)
+
+    return trainer, stats, elapsed
+
+
+def save_and_push(args, model, tokenizer):
+    from huggingface_hub import HfApi
+
+    api = HfApi()
+    api.create_repo(args.output_repo, repo_type="model", exist_ok=True, private=args.private)
+
+    model.save_pretrained(args.output_dir)
+    tokenizer.save_pretrained(args.output_dir)
+    log.info("pushing LoRA adapter to %s", args.output_repo)
+    model.push_to_hub(args.output_repo, tokenizer=tokenizer)
+
+    if args.merge_repo:
+        api.create_repo(args.merge_repo, repo_type="model", exist_ok=True, private=args.private)
+        log.info("merging to 16-bit and pushing to %s (large upload)", args.merge_repo)
+        model.push_to_hub_merged(args.merge_repo, tokenizer=tokenizer, save_method="merged_16bit")
+
+# --------------------------------------------------------------------------
+# CLI
+# --------------------------------------------------------------------------
+def parse_args(argv=None):
+    p = argparse.ArgumentParser(description="SecureCoder QLoRA fine-tune")
+
+    p.add_argument("--base-model", default="Qwen/Qwen3-Coder-30B-A3B-Instruct")
+    p.add_argument("--output-repo", default=None, help="Hub repo for the LoRA adapter")
+    p.add_argument("--merge-repo", default=None, help="optional Hub repo for a 16-bit merge")
+    p.add_argument("--output-dir", default="securecoder-out")
+    p.add_argument("--private", action="store_true", help="create Hub repos as private")
+
+    p.add_argument("--max-seq-length", type=int, default=4096)
+    p.add_argument("--batch-size", type=int, default=2)
+    p.add_argument("--grad-accum", type=int, default=8)
+    p.add_argument("--learning-rate", type=float, default=2e-4)
+    p.add_argument("--lr-scheduler", default="cosine")
+    p.add_argument("--num-epochs", type=float, default=1.0)
+    p.add_argument("--max-steps", type=int, default=0, help="overrides --num-epochs when > 0")
+    p.add_argument("--eval-samples", type=int, default=200, help="0 disables evaluation")
+    p.add_argument("--logging-steps", type=int, default=10)
+    p.add_argument("--save-steps", type=int, default=250)
+    p.add_argument("--packing", action="store_true", default=True)
+    p.add_argument("--no-packing", dest="packing", action="store_false")
+    p.add_argument("--seed", type=int, default=3407)
+
+    p.add_argument("--lora-r", type=int, default=32)
+    p.add_argument("--lora-alpha", type=int, default=32)
+    p.add_argument("--no-4bit", action="store_true")
+    p.add_argument("--target-modules", default=",".join(ATTENTION_TARGETS),
+                   help="comma-separated suffixes, or 'all-linear' to include expert MLPs")
+
+    p.add_argument("--report-to", default="trackio", choices=["trackio", "none"])
+    p.add_argument("--trackio-project", default="securecoder")
+    p.add_argument("--trackio-space", default=None, help="e.g. Taimwe/securecoder-trackio")
+    p.add_argument("--run-name", default=None)
+
+    p.add_argument("--validate-only", action="store_true",
+                   help="load a small sample of each source, print the mix, exit (no GPU)")
+    p.add_argument("--validate-per-source", type=int, default=40)
+    p.add_argument("--show-samples", type=int, default=3)
+    p.add_argument("--smoke", action="store_true",
+                   help="tiny end-to-end run: 200 rows/source, 20 steps")
+    return p.parse_args(argv)
+
+
+def apply_smoke(args) -> None:
+    args.max_steps = args.max_steps or 20
+    args.eval_samples = min(args.eval_samples, 20)
+    args.save_steps = 20
+    args.max_seq_length = min(args.max_seq_length, 2048)
+    global MIX
+    MIX = [Source(s.repo, 200, s.kind, s.config, s.split, s.note) for s in MIX]
+
+
+def main(argv=None) -> int:
+    args = parse_args(argv)
+    if args.smoke:
+        apply_smoke(args)
+    token = os.environ.get("HF_TOKEN")
+
+    if args.validate_only:
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(args.base_model)
+        sources = [Source(s.repo, args.validate_per_source, s.kind, s.config, s.split, s.note)
+                   for s in MIX]
+        records, stats = build_dataset(tokenizer, sources, token, validate=True)
+        print_stats(stats, records, tokenizer)
+        for i, rec in enumerate(records[: args.show_samples], 1):
+            print("\n" + "-" * 78)
+            print(f"SAMPLE {i}  [{rec['source']}]  {len(rec['text'])} chars")
+            print("-" * 78)
+            print(rec["text"][:1500])
+        return 0
+
+    import torch
+
+    if not torch.cuda.is_available():
+        log.error("no CUDA device - use --validate-only locally, or run on HF Jobs / Colab")
+        return 1
+    log.info("GPU: %s", torch.cuda.get_device_name(0))
+
+    if not args.output_repo:
+        log.error("--output-repo is required (the container/VM is ephemeral)")
+        return 1
+
+    model, tokenizer = load_model_and_tokenizer(args)
+    records, stats = build_dataset(tokenizer, MIX, token, validate=False)
+    print_stats(stats, records, tokenizer)
+    if len(records) < 100:
+        log.error("only %d usable rows - refusing to train", len(records))
+        return 1
+
+    trainer, stats_train, elapsed = train(args, model, tokenizer, records)
+    save_and_push(args, model, tokenizer)
+
+    print("\n" + "=" * 78)
+    print(f"DONE  rows={len(records):,}  time={elapsed / 60:.1f} min  "
+          f"loss={stats_train.metrics.get('train_loss', float('nan')):.4f}")
+    print(f"adapter: https://huggingface.co/{args.output_repo}")
+    if args.merge_repo:
+        print(f"merged : https://huggingface.co/{args.merge_repo}")
+    print("=" * 78)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())