Upload training/train.py with huggingface_hub

training/train.py

@@ -0,0 +1,1027 @@
+#!/usr/bin/env python3
+"""Cortex A 0.5 — SFT + 4-bit QAT for Qwen3.5-0.8B on Kaggle 2x T4.
+
+Full-parameter SFT with Unsloth QAT (4-bit weights + 4-bit dynamic activations
+when TorchAO exposes it, else official Unsloth int4 / int8-int4). Designed to
+resume from Hugging Face, stop cleanly at 11.5h, and push checkpoints + metrics.
+"""
+from __future__ import annotations
+
+import gc
+import inspect
+import json
+import math
+import os
+import random
+import shutil
+import time
+import traceback
+from pathlib import Path
+from typing import Any, Optional
+
+# Each torchrun process must see a single GPU so Unsloth's 1-GPU assert is happy
+# while HF Trainer still DDP-coordinates via RANK/WORLD_SIZE.
+if "LOCAL_RANK" in os.environ:
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(os.environ["LOCAL_RANK"])
+
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("UNSLOTH_DISABLE_STATISTICS", "1")
+os.environ.setdefault("NCCL_P2P_DISABLE", "1")
+os.environ.setdefault("NCCL_IB_DISABLE", "1")
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+HF_REPO = os.environ.get("HF_REPO", "PRATYUSH-BHARDWAJ/Cortex_A_0.5")
+MODEL_NAME = os.environ.get("MODEL_NAME", "unsloth/Qwen3.5-0.8B")
+OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/kaggle/working/cortex_sft")
+MAX_TRAIN_SECONDS = int(float(os.environ.get("MAX_TRAIN_HOURS", "11.5")) * 3600)
+SEED = int(os.environ.get("SEED", "3407"))
+MAX_SEQ = int(os.environ.get("MAX_SEQ", "8192"))
+BATCH = int(os.environ.get("BATCH", "2"))
+GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "4"))
+LR = float(os.environ.get("LR", "2e-5"))
+EPOCHS = float(os.environ.get("EPOCHS", "3"))
+SAVE_STEPS = int(os.environ.get("SAVE_STEPS", "80"))
+EVAL_STEPS = int(os.environ.get("EVAL_STEPS", "80"))
+LOG_STEPS = int(os.environ.get("LOG_STEPS", "5"))
+WARMUP_RATIO = float(os.environ.get("WARMUP_RATIO", "0.03"))
+WEIGHT_DECAY = float(os.environ.get("WEIGHT_DECAY", "0.01"))
+MAX_GRAD_NORM = float(os.environ.get("MAX_GRAD_NORM", "1.0"))
+VAL_SIZE = int(os.environ.get("VAL_SIZE", "256"))
+PACKING = os.environ.get("PACKING", "1") != "0"
+QAT_SCHEME_PREF = os.environ.get("QAT_SCHEME", "int4-dyn")
+MAX_SAMPLES = os.environ.get("MAX_SAMPLES")  # smoke test cap
+HER_CAP = int(os.environ.get("HER_CAP", "8000"))
+
+START_TS = time.time()
+IS_MAIN = int(os.environ.get("RANK", "0")) == 0
+
+
+def log(*a):
+    if IS_MAIN:
+        print(*a, flush=True)
+
+
+def is_rank0() -> bool:
+    return int(os.environ.get("RANK", "0")) == 0
+
+
+# ---------------------------------------------------------------------------
+# Dataset catalog (frontier distillation + roleplay)
+# ---------------------------------------------------------------------------
+# Core mix is r0b0tlab (Qwen3.8-Max + GLM-5.2 + Kimi K3) which already has a
+# held-out val split. Complements fill Claude Opus 4.6/4.7/4.8 reasoning,
+# DeepSeek V4 coding, and character roleplay (the core set is only 1.1% dialogue).
+DATASETS = [
+    {
+        "id": "r0b0tlab/qwen3.8-max-glm5.2-kimi-k3-distillation",
+        "config": "sft_balanced",
+        "split": "train",
+        "domain": "mixed-sota",
+        "weight": 1.0,
+        "required": True,
+    },
+    {
+        "id": "Jackrong/DeepSeek-V4-Distill-8000x",
+        "domain": "coding",
+        "weight": 1.0,
+    },
+    {
+        "id": "Jackrong/Claude-opus-4.7-TraceInversion-5000x",
+        "domain": "reasoning",
+        "weight": 0.9,
+    },
+    {
+        "id": "Jackrong/Claude-opus-4.6-TraceInversion-9000x",
+        "domain": "reasoning",
+        "weight": 0.7,
+    },
+    {
+        "id": "Roman1111111/claude-opus-4.6-10000x",
+        "domain": "reasoning",
+        "weight": 0.6,
+    },
+    {
+        "id": "nohurry/Opus-4.6-Reasoning-3000x-filtered",
+        "domain": "reasoning",
+        "weight": 1.0,
+    },
+    {
+        "id": "ansulev/claude-opus-4.8-distill-5k",
+        "domain": "reasoning",
+        "weight": 1.0,
+    },
+    {
+        "id": "TeichAI/Claude-Opus-4.6-Reasoning-887x",
+        "domain": "reasoning-long",
+        "weight": 1.2,
+    },
+    {
+        "id": "lordx64/reasoning-distill-claude-opus-4-7-max",
+        "domain": "reasoning",
+        "weight": 0.8,
+    },
+    {
+        "id": "Jackrong/Qwen3.5-reasoning-700x",
+        "domain": "reasoning",
+        "weight": 1.0,
+    },
+    {
+        "id": "beyoru/Aesir-Character-CoT-roleplay",
+        "domain": "roleplay",
+        "weight": 1.5,
+    },
+    {
+        "id": "ChengyuDu0123/HER-Dataset",
+        "config": "sft_multi_turn",
+        "domain": "roleplay",
+        "weight": 1.0,
+        "cap": HER_CAP,
+    },
+]
+
+
+# ---------------------------------------------------------------------------
+# Message normalization
+# ---------------------------------------------------------------------------
+def _as_text(content) -> Optional[str]:
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts = []
+        for p in content:
+            if isinstance(p, str):
+                parts.append(p)
+            elif isinstance(p, dict):
+                if p.get("type") in ("image", "image_url", "video", "video_url"):
+                    return None
+                parts.append(str(p.get("text") or p.get("content") or ""))
+        return "\n".join(x for x in parts if x)
+    if isinstance(content, dict):
+        if content.get("type") in ("image", "image_url"):
+            return None
+        return str(content.get("text") or content.get("content") or "")
+    return str(content)
+
+
+def _role_of(turn: dict) -> str:
+    r = (turn.get("role") or turn.get("from") or turn.get("speaker") or "user")
+    r = str(r).lower().strip()
+    mapping = {
+        "human": "user",
+        "gpt": "assistant",
+        "bot": "assistant",
+        "model": "assistant",
+        "ai": "assistant",
+        "assistant": "assistant",
+        "user": "user",
+        "system": "system",
+        "tool": "tool",
+        "function": "tool",
+        "observation": "tool",
+    }
+    return mapping.get(r, "user" if r not in ("assistant", "system", "tool") else r)
+
+
+def extract_messages(ex: dict) -> Optional[list]:
+    for key in ("messages", "conversations", "conversation", "chat"):
+        raw = ex.get(key)
+        if not raw:
+            continue
+        if isinstance(raw, str):
+            try:
+                raw = json.loads(raw)
+            except Exception:
+                continue
+        if not isinstance(raw, list) or not raw:
+            continue
+        out = []
+        for t in raw:
+            if not isinstance(t, dict):
+                continue
+            role = _role_of(t)
+            text = _as_text(t.get("content") or t.get("value") or t.get("text") or t.get("message"))
+            if text is None:
+                return None
+            if not str(text).strip() and role != "tool":
+                continue
+            msg = {"role": role, "content": str(text)}
+            if t.get("tool_calls"):
+                msg["tool_calls"] = t["tool_calls"]
+            if t.get("name"):
+                msg["name"] = t["name"]
+            out.append(msg)
+        if out:
+            return out
+
+    # instruction / input / output
+    instr = ex.get("instruction") or ex.get("question") or ex.get("prompt")
+    outp = ex.get("output") or ex.get("response") or ex.get("completion") or ex.get("answer")
+    if instr and outp:
+        msgs = []
+        sys = ex.get("system") or ex.get("system_prompt")
+        if sys:
+            msgs.append({"role": "system", "content": str(sys)})
+        user = str(instr)
+        inp = ex.get("input")
+        if inp and str(inp).strip() and str(inp).strip() not in user:
+            user = user + "\n\n" + str(inp)
+        msgs.append({"role": "user", "content": user})
+        msgs.append({"role": "assistant", "content": str(outp)})
+        return msgs
+
+    text = ex.get("text")
+    if isinstance(text, str) and len(text.strip()) > 32:
+        return [{"role": "user", "content": "Continue."}, {"role": "assistant", "content": text}]
+    return None
+
+
+def apply_template(tokenizer, messages, tools=None) -> Optional[str]:
+    kwargs = {"tokenize": False, "add_generation_prompt": False}
+    if tools:
+        kwargs["tools"] = tools
+    attempts = [
+        dict(kwargs, chat_template_kwargs={"enable_thinking": False}),
+        dict(kwargs, enable_thinking=False),
+        dict(kwargs),
+    ]
+    for kw in attempts:
+        try:
+            text = tokenizer.apply_chat_template(messages, **kw)
+            if isinstance(text, str) and text.strip():
+                return text
+        except TypeError:
+            continue
+        except Exception:
+            continue
+    return None
+
+
+# ---------------------------------------------------------------------------
+# QAT helpers — 4-bit weights + 4-bit dynamic activations when available
+# ---------------------------------------------------------------------------
+def _linear_filter(mod, fqn: str) -> bool:
+    import torch.nn as nn
+
+    if not isinstance(mod, nn.Linear):
+        return False
+    if getattr(mod, "in_features", 0) < 128:
+        return False
+    skip = ("embed", "lm_head", "norm", "visual", "vision", "patch")
+    f = (fqn or "").lower()
+    if any(s in f for s in skip):
+        return False
+    return True
+
+
+def resolve_qat_scheme():
+    """Return (scheme_obj_or_str, label). Prefer 4-bit dynamic activations."""
+    pref = QAT_SCHEME_PREF.lower()
+    if pref in ("int4", "int8-int4", "fp8-int4", "fp8-fp8", "int8", "cactus"):
+        return pref, pref
+
+    # Try TorchAO 4-bit dynamic activation + 4-bit weight configs
+    import_paths = [
+        ("torchao.quantization", "Int4DynamicActivationInt4WeightConfig"),
+        ("torchao.quantization.quant_api", "Int4DynamicActivationInt4WeightConfig"),
+        ("torchao.prototype.quantization", "Int4DynamicActivationInt4WeightConfig"),
+        ("torchao.quantization", "IntxDynamicActivationIntxWeightConfig"),
+    ]
+    for mod_name, cls_name in import_paths:
+        try:
+            import importlib
+            import torch
+
+            mod = importlib.import_module(mod_name)
+            cls = getattr(mod, cls_name)
+            if cls_name == "IntxDynamicActivationIntxWeightConfig":
+                try:
+                    from torchao.quantization.granularity import PerGroup, PerToken
+                except Exception:
+                    from torchao.quantization import PerGroup, PerToken  # type: ignore
+                cfg = cls(
+                    activation_dtype=torch.int4,
+                    weight_dtype=torch.int4,
+                    weight_granularity=PerGroup(32),
+                )
+            else:
+                try:
+                    cfg = cls(group_size=32)
+                except TypeError:
+                    cfg = cls()
+            log(f"[QAT] using TorchAO {cls_name} (4-bit weights + 4-bit dynamic activations)")
+            return cfg, f"torchao:{cls_name}"
+        except Exception as e:
+            log(f"[QAT] {mod_name}.{cls_name} unavailable: {e}")
+
+    # Official Unsloth schemes. int8-int4 = dynamic int8 activations + int4 weights.
+    # True int4 activations are not in Unsloth's public qat_scheme list; int4 is
+    # 4-bit weight QAT (the accuracy-recovery path Unsloth + TorchAO ship).
+    if pref in ("int4-dyn", "int4-int4", "auto"):
+        log("[QAT] falling back to Unsloth qat_scheme='int4' (4-bit weight QAT)")
+        return "int4", "int4"
+    return "int4", "int4"
+
+
+def apply_qat(model, scheme) -> Any:
+    if scheme is None:
+        return model
+    # Unsloth internal prepare
+    try:
+        from unsloth.models._utils import _prepare_model_for_qat
+
+        log("[QAT] _prepare_model_for_qat")
+        return _prepare_model_for_qat(model, scheme)
+    except Exception as e:
+        log(f"[QAT] unsloth prepare failed: {e}")
+
+    try:
+        from torchao.quantization import quantize_
+        from torchao.quantization.qat import QATConfig
+        from torchao.quantization import Int4WeightOnlyConfig
+
+        base = scheme if not isinstance(scheme, str) else Int4WeightOnlyConfig(group_size=128)
+        quantize_(model, QATConfig(base, step="prepare"), filter_fn=_linear_filter)
+        log("[QAT] torchao QATConfig(prepare) applied")
+        return model
+    except Exception as e:
+        log(f"[QAT] torchao prepare failed: {e}")
+        return model
+
+
+def freeze_vision(model):
+    n_freeze = 0
+    keys = ("visual", "vision_tower", "vision_model", "merger", "patch_embed", "vision_encoder")
+    for name, p in model.named_parameters():
+        nl = name.lower()
+        if any(k in nl for k in keys):
+            p.requires_grad = False
+            n_freeze += p.numel()
+    log(f"[model] froze vision params: {n_freeze/1e6:.2f}M")
+
+
+def count_trainable(model) -> tuple[int, int]:
+    t = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    a = sum(p.numel() for p in model.parameters())
+    return t, a
+
+
+# ---------------------------------------------------------------------------
+# Callbacks
+# ---------------------------------------------------------------------------
+def make_callbacks(tokenizer=None):
+    from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
+
+    class TimeLimitCallback(TrainerCallback):
+        def on_step_end(self, args, state: TrainerState, control: TrainerControl, **kwargs):
+            elapsed = time.time() - START_TS
+            remaining = MAX_TRAIN_SECONDS - elapsed
+            if remaining <= 180:  # 3 min buffer inside the 11.5h cap for the final save
+                log(f"[time] stopping at {elapsed/3600:.2f}h (limit {MAX_TRAIN_SECONDS/3600:.2f}h)")
+                control.should_training_stop = True
+                control.should_save = True
+            return control
+
+    class MetricsCallback(TrainerCallback):
+        def __init__(self):
+            self.t0 = time.time()
+            self.last_t = self.t0
+            self.last_tokens = 0
+            self.history = []
+
+        def on_log(self, args, state: TrainerState, control: TrainerControl, logs=None, **kwargs):
+            if not is_rank0() or not logs:
+                return
+            logs = dict(logs)
+            loss = logs.get("loss") or logs.get("train_loss")
+            eval_loss = logs.get("eval_loss")
+            ppl = math.exp(min(float(loss), 20)) if loss is not None else None
+            val_ppl = math.exp(min(float(eval_loss), 20)) if eval_loss is not None else None
+            now = time.time()
+            tokens = int(getattr(state, "num_input_tokens_seen", 0) or 0)
+            dt = max(now - self.last_t, 1e-6)
+            tok_s = (tokens - self.last_tokens) / dt if tokens else logs.get("train_tokens_per_second")
+            self.last_t, self.last_tokens = now, tokens
+            trainer = kwargs.get("model")
+            mtp = None
+            # pull last mtp from trainer if present
+            tr = kwargs.get("trainer") or None
+            rec = {
+                "step": int(state.global_step),
+                "epoch": float(state.epoch or 0),
+                "loss": None if loss is None else round(float(loss), 6),
+                "mtp_loss": logs.get("mtp_loss") or logs.get("aux_loss"),
+                "ppl": None if ppl is None else round(float(ppl), 4),
+                "val_loss": None if eval_loss is None else round(float(eval_loss), 6),
+                "val_ppl": None if val_ppl is None else round(float(val_ppl), 4),
+                "tok_s": None if tok_s is None else round(float(tok_s), 1),
+                "grad_norm": logs.get("grad_norm"),
+                "lr": logs.get("learning_rate"),
+                "tokens_seen": tokens,
+                "elapsed_h": round((now - START_TS) / 3600, 4),
+                "remaining_h": round(max(MAX_TRAIN_SECONDS - (now - START_TS), 0) / 3600, 4),
+                "gpu_mem_gb": None,
+            }
+            try:
+                import torch
+
+                if torch.cuda.is_available():
+                    rec["gpu_mem_gb"] = round(torch.cuda.max_memory_allocated() / 1024**3, 2)
+            except Exception:
+                pass
+            self.history.append(rec)
+            log(
+                f"[metrics] step={rec['step']} loss={rec['loss']} mtp={rec['mtp_loss']} "
+                f"ppl={rec['ppl']} val_loss={rec['val_loss']} val_ppl={rec['val_ppl']} "
+                f"tok/s={rec['tok_s']} gnorm={rec['grad_norm']} lr={rec['lr']} "
+                f"mem={rec['gpu_mem_gb']}G t={rec['elapsed_h']}h"
+            )
+            try:
+                Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+                live = Path(OUTPUT_DIR) / "live_metrics.json"
+                live.write_text(json.dumps({"latest": rec, "history": self.history[-400:]}, indent=2))
+                if rec["step"] % max(LOG_STEPS * 4, 20) == 0 or rec["step"] <= 3:
+                    from huggingface_hub import HfApi
+
+                    HfApi(token=HF_TOKEN).upload_file(
+                        path_or_fileobj=str(live),
+                        path_in_repo="training/live_metrics.json",
+                        repo_id=HF_REPO,
+                        repo_type="model",
+                        token=HF_TOKEN,
+                    )
+            except Exception as e:
+                log(f"[metrics] hub upload skipped: {e}")
+
+        def on_save(self, args, state, control, **kwargs):
+            if not is_rank0():
+                return
+            pointer = {
+                "global_step": int(state.global_step),
+                "checkpoint": f"checkpoint-{state.global_step}",
+                "output_dir": args.output_dir,
+                "updated_unix": int(time.time()),
+            }
+            p = Path(OUTPUT_DIR) / "RESUME_POINTER.json"
+            p.write_text(json.dumps(pointer, indent=2))
+            try:
+                from huggingface_hub import HfApi
+
+                HfApi(token=HF_TOKEN).upload_file(
+                    path_or_fileobj=str(p),
+                    path_in_repo="training/RESUME_POINTER.json",
+                    repo_id=HF_REPO,
+                    repo_type="model",
+                    token=HF_TOKEN,
+                )
+            except Exception as e:
+                log(f"[save] pointer upload skipped: {e}")
+
+    return [TimeLimitCallback(), MetricsCallback()]
+
+
+# ---------------------------------------------------------------------------
+# Resume
+# ---------------------------------------------------------------------------
+def find_resume() -> Optional[str]:
+    out = Path(OUTPUT_DIR)
+    local = sorted(out.glob("checkpoint-*"), key=lambda p: int(p.name.split("-")[-1]) if p.name.split("-")[-1].isdigit() else -1)
+    if local:
+        log(f"[resume] local {local[-1]}")
+        return str(local[-1])
+    try:
+        from huggingface_hub import HfApi, snapshot_download
+
+        api = HfApi(token=HF_TOKEN)
+        files = api.list_repo_files(HF_REPO, repo_type="model")
+        ckpts = []
+        for f in files:
+            if "checkpoint-" in f and f.endswith("trainer_state.json"):
+                try:
+                    step = int(f.split("checkpoint-")[1].split("/")[0])
+                    ckpts.append((step, f))
+                except Exception:
+                    pass
+        if not ckpts:
+            # pointer file
+            if "training/RESUME_POINTER.json" in files:
+                log("[resume] pointer exists but no checkpoint files listed yet")
+            return None
+        step, _ = max(ckpts)
+        dest = out / f"checkpoint-{step}"
+        log(f"[resume] downloading checkpoint-{step} from hub")
+        snapshot_download(
+            HF_REPO,
+            repo_type="model",
+            allow_patterns=[f"**/*checkpoint-{step}/**", f"checkpoint-{step}/**"],
+            local_dir=str(out),
+            token=HF_TOKEN,
+        )
+        if dest.exists():
+            return str(dest)
+        # hub may nest
+        found = list(out.rglob(f"checkpoint-{step}/trainer_state.json"))
+        if found:
+            return str(found[0].parent)
+    except Exception as e:
+        log(f"[resume] hub lookup failed: {e}")
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Data
+# ---------------------------------------------------------------------------
+def load_and_mix(tokenizer):
+    from datasets import Dataset, concatenate_datasets, load_dataset
+
+    pieces = []
+    val_ds = None
+    report = []
+
+    def take_split(dsobj, split_name="train"):
+        if hasattr(dsobj, "keys"):
+            keys = list(dsobj.keys())
+            if split_name in dsobj:
+                return dsobj[split_name]
+            if "train" in dsobj:
+                return dsobj["train"]
+            return dsobj[keys[0]]
+        return dsobj
+
+    for spec in DATASETS:
+        ds_id = spec["id"]
+        cfg = spec.get("config")
+        try:
+            kwargs = {"path": ds_id}
+            if cfg:
+                kwargs["name"] = cfg
+            log(f"[data] loading {ds_id}" + (f" ({cfg})" if cfg else ""))
+            raw = load_dataset(**kwargs)
+            split = take_split(raw, spec.get("split", "train"))
+
+            # r0b0tlab: prefer native validation
+            if ds_id.startswith("r0b0tlab/") and val_ds is None:
+                if hasattr(raw, "keys") and "validation" in raw:
+                    val_ds = raw["validation"]
+                elif "split" in split.column_names:
+                    val_ds = split.filter(lambda x: str(x.get("split", "")).lower() in ("validation", "val"))
+                    split = split.filter(lambda x: str(x.get("split", "train")).lower() in ("train", ""))
+
+            cap = spec.get("cap")
+            if cap and len(split) > cap:
+                split = split.shuffle(seed=SEED).select(range(cap))
+
+            n_before = len(split)
+            # upsample via repetition for weight > 1
+            w = spec.get("weight", 1.0)
+            copies = max(int(round(w)), 1)
+            if copies > 1:
+                split = concatenate_datasets([split] * copies)
+
+            pieces.append(split)
+            report.append({"id": ds_id, "rows": n_before, "used": len(split), "domain": spec.get("domain")})
+            log(f"[data]   {n_before} rows -> {len(split)} used ({spec.get('domain')})")
+        except Exception as e:
+            msg = f"[data] FAILED {ds_id}: {e}"
+            log(msg)
+            traceback.print_exc()
+            if spec.get("required"):
+                raise
+            report.append({"id": ds_id, "error": str(e)})
+
+    if not pieces:
+        raise RuntimeError("No datasets loaded")
+
+    train = concatenate_datasets(pieces).shuffle(seed=SEED)
+    if MAX_SAMPLES:
+        n = min(int(MAX_SAMPLES), len(train))
+        train = train.select(range(n))
+        if val_ds is not None:
+            val_ds = val_ds.select(range(min(64, len(val_ds))))
+
+    def to_text(ex):
+        msgs = extract_messages(ex)
+        if not msgs:
+            return {"text": ""}
+        tools = ex.get("tools")
+        if isinstance(tools, str):
+            try:
+                tools = json.loads(tools)
+            except Exception:
+                tools = None
+        if tools == []:
+            tools = None
+        text = apply_template(tokenizer, msgs, tools=tools)
+        return {"text": text or ""}
+
+    log("[data] applying chat template…")
+    num_proc = min(4, os.cpu_count() or 2)
+    cols = train.column_names
+    train = train.map(to_text, remove_columns=cols, num_proc=num_proc, desc="format-train")
+    train = train.filter(lambda x: isinstance(x.get("text"), str) and len(x["text"]) > 48)
+
+    if val_ds is not None:
+        vcols = val_ds.column_names
+        val_ds = val_ds.map(to_text, remove_columns=vcols, num_proc=num_proc, desc="format-val")
+        val_ds = val_ds.filter(lambda x: isinstance(x.get("text"), str) and len(x["text"]) > 48)
+        if len(val_ds) > VAL_SIZE:
+            val_ds = val_ds.shuffle(seed=SEED).select(range(VAL_SIZE))
+    else:
+        # last-resort holdout
+        n = min(VAL_SIZE, max(1, len(train) // 40))
+        val_ds = train.select(range(n))
+        train = train.select(range(n, len(train)))
+
+    log(f"[data] train={len(train)} val={len(val_ds)}")
+    if is_rank0():
+        Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+        (Path(OUTPUT_DIR) / "dataset_report.json").write_text(json.dumps({"report": report, "train": len(train), "val": len(val_ds)}, indent=2))
+    return train, val_ds, report
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+def load_model(max_seq: int):
+    import torch
+
+    scheme, scheme_label = resolve_qat_scheme()
+    loader = None
+    loader_name = None
+    for name in ("FastVisionModel", "FastModel", "FastLanguageModel"):
+        try:
+            import unsloth
+
+            loader = getattr(unsloth, name)
+            loader_name = name
+            break
+        except Exception:
+            continue
+    if loader is None:
+        from unsloth import FastLanguageModel as loader
+
+        loader_name = "FastLanguageModel"
+    log(f"[model] loader={loader_name} name={MODEL_NAME} seq={max_seq} qat={scheme_label}")
+
+    kwargs = dict(
+        model_name=MODEL_NAME,
+        max_seq_length=max_seq,
+        load_in_4bit=False,
+        load_in_8bit=False,
+        full_finetuning=True,
+        use_gradient_checkpointing="unsloth",
+        dtype=None,
+        token=HF_TOKEN,
+    )
+    model = tokenizer = None
+    # Prefer passing qat_scheme when it's an official string
+    if isinstance(scheme, str):
+        try:
+            model, tokenizer = loader.from_pretrained(**kwargs, qat_scheme=scheme)
+        except TypeError:
+            model, tokenizer = loader.from_pretrained(**kwargs)
+            model = apply_qat(model, scheme)
+    else:
+        try:
+            model, tokenizer = loader.from_pretrained(**kwargs, qat_scheme=scheme)
+        except TypeError:
+            model, tokenizer = loader.from_pretrained(**kwargs)
+            model = apply_qat(model, scheme)
+
+    freeze_vision(model)
+    t, a = count_trainable(model)
+    log(f"[model] trainable {t/1e6:.2f}M / {a/1e6:.2f}M")
+
+    # Verify fake-quant modules landed
+    n_fq = 0
+    for m in model.modules():
+        if "FakeQuant" in m.__class__.__name__ or "QAT" in m.__class__.__name__:
+            n_fq += 1
+    log(f"[QAT] fake-quant / QAT modules detected: {n_fq}")
+    if n_fq == 0:
+        log("[QAT] WARNING: no fake-quant modules found — applying prepare again")
+        model = apply_qat(model, scheme if not isinstance(scheme, str) else scheme)
+
+    return model, tokenizer, loader, scheme_label
+
+
+def build_sft_config(max_seq, batch, ga, resume_dir=None):
+    from trl import SFTConfig
+
+    params = set(inspect.signature(SFTConfig.__init__).parameters)
+    use_bf16 = False
+    use_fp16 = True
+    try:
+        import torch
+
+        # T4 has no bf16 tensor cores
+        major, _ = torch.cuda.get_device_capability(0) if torch.cuda.is_available() else (0, 0)
+        use_bf16 = major >= 8
+        use_fp16 = not use_bf16
+    except Exception:
+        pass
+
+    cfg = dict(
+        output_dir=OUTPUT_DIR,
+        per_device_train_batch_size=batch,
+        per_device_eval_batch_size=1,
+        gradient_accumulation_steps=ga,
+        num_train_epochs=EPOCHS,
+        learning_rate=LR,
+        warmup_ratio=WARMUP_RATIO,
+        weight_decay=WEIGHT_DECAY,
+        max_grad_norm=MAX_GRAD_NORM,
+        logging_steps=LOG_STEPS,
+        save_steps=SAVE_STEPS,
+        eval_steps=EVAL_STEPS,
+        save_total_limit=2,
+        lr_scheduler_type="cosine",
+        optim="adamw_8bit",
+        seed=SEED,
+        report_to="none",
+        fp16=use_fp16,
+        bf16=use_bf16,
+        dataloader_num_workers=2,
+        dataloader_pin_memory=True,
+        remove_unused_columns=False,
+        hub_model_id=HF_REPO,
+        hub_strategy="every_save",
+        push_to_hub=bool(HF_TOKEN),
+        hub_private_repo=True,
+        hub_token=HF_TOKEN,
+        save_safetensors=True,
+        logging_first_step=True,
+        load_best_model_at_end=False,
+        greater_is_better=False,
+        metric_for_best_model="eval_loss",
+    )
+    # eval strategy names differ
+    if "eval_strategy" in params:
+        cfg["eval_strategy"] = "steps"
+        cfg["save_strategy"] = "steps"
+    elif "evaluation_strategy" in params:
+        cfg["evaluation_strategy"] = "steps"
+        cfg["save_strategy"] = "steps"
+
+    if "max_length" in params:
+        cfg["max_length"] = max_seq
+    elif "max_seq_length" in params:
+        cfg["max_seq_length"] = max_seq
+
+    if "dataset_text_field" in params:
+        cfg["dataset_text_field"] = "text"
+    if "packing" in params and PACKING:
+        cfg["packing"] = True
+    if "padding_free" in params and PACKING:
+        cfg["padding_free"] = True
+    if "assistant_only_loss" in params:
+        cfg["assistant_only_loss"] = True
+    if "completion_only_loss" in params:
+        cfg["completion_only_loss"] = True
+    if "include_num_input_tokens_seen" in params:
+        cfg["include_num_input_tokens_seen"] = True
+    if "dataset_kwargs" in params:
+        cfg["dataset_kwargs"] = {"skip_prepare_dataset": False}
+    if "ddp_find_unused_parameters" in params:
+        cfg["ddp_find_unused_parameters"] = False
+    if "gradient_checkpointing" in params:
+        cfg["gradient_checkpointing"] = True
+    if "optim_args" in params:
+        pass
+    # drop unknown later via filter
+    cfg = {k: v for k, v in cfg.items() if k in params or k in ("output_dir",)}
+    # SFTConfig inherits TrainingArguments — keep keys it accepts
+    try:
+        return SFTConfig(**{k: v for k, v in cfg.items() if k in params})
+    except TypeError as e:
+        log(f"[cfg] SFTConfig retry after {e}")
+        # drop extras
+        ok = {}
+        for k, v in cfg.items():
+            try:
+                SFTConfig(**{**ok, k: v})
+                ok[k] = v
+            except TypeError:
+                log(f"[cfg] drop {k}")
+        return SFTConfig(**ok)
+
+
+class CortexTrainer:
+    """Factory wrapping TRL SFTTrainer with MTP-aware compute_loss."""
+
+    @staticmethod
+    def build(model, tokenizer, train_ds, val_ds, args):
+        from trl import SFTTrainer
+
+        class _T(SFTTrainer):
+            def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+                outputs = model(**inputs)
+                loss = outputs.loss if hasattr(outputs, "loss") else outputs[0]
+                extra = {}
+                for key in ("mtp_loss", "aux_loss"):
+                    val = getattr(outputs, key, None)
+                    if val is not None:
+                        try:
+                            extra[key] = float(val.detach().float().mean().item())
+                        except Exception:
+                            pass
+                if extra:
+                    self._last_mtp = extra
+                    try:
+                        self.log(extra)
+                    except Exception:
+                        pass
+                return (loss, outputs) if return_outputs else loss
+
+        kw = dict(model=model, args=args, train_dataset=train_ds, eval_dataset=val_ds)
+        sig = inspect.signature(SFTTrainer.__init__)
+        if "processing_class" in sig.parameters:
+            kw["processing_class"] = tokenizer
+        elif "tokenizer" in sig.parameters:
+            kw["tokenizer"] = tokenizer
+        if "callbacks" in sig.parameters:
+            kw["callbacks"] = make_callbacks(tokenizer)
+        try:
+            return _T(**kw)
+        except TypeError:
+            kw.pop("eval_dataset", None)
+            return _T(**kw)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+    random.seed(SEED)
+    Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+
+    if HF_TOKEN:
+        from huggingface_hub import login, HfApi
+
+        login(token=HF_TOKEN, add_to_git_credential=False)
+        if is_rank0():
+            try:
+                HfApi(token=HF_TOKEN).create_repo(HF_REPO, repo_type="model", private=True, exist_ok=True)
+            except Exception as e:
+                log(f"[hub] create_repo: {e}")
+
+    import torch
+
+    log(f"[hw] cuda={torch.cuda.is_available()} n={torch.cuda.device_count()}")
+    if torch.cuda.is_available():
+        for i in range(torch.cuda.device_count()):
+            p = torch.cuda.get_device_properties(i)
+            log(f"[hw] gpu{i} {p.name} {p.total_mem/1024**3:.1f}GB" if hasattr(p, "total_mem") else f"[hw] gpu{i} {p.name} {p.total_memory/1024**3:.1f}GB")
+
+    resume = find_resume()
+    seq, batch, ga = MAX_SEQ, BATCH, GRAD_ACCUM
+    backoff = [
+        (MAX_SEQ, BATCH, GRAD_ACCUM),
+        (MAX_SEQ, max(1, BATCH // 2), GRAD_ACCUM * 2),
+        (4096, 2, 4),
+        (4096, 1, 8),
+        (2048, 2, 4),
+        (2048, 1, 8),
+    ]
+    # de-dup
+    seen = set()
+    plans = []
+    for t in backoff:
+        if t not in seen:
+            seen.add(t)
+            plans.append(t)
+
+    last_err = None
+    for seq, batch, ga in plans:
+        model = tokenizer = trainer = None
+        try:
+            log(f"[run] seq={seq} batch={batch} ga={ga} packing={PACKING}")
+            model, tokenizer, loader, scheme_label = load_model(seq)
+            train_ds, val_ds, report = load_and_mix(tokenizer)
+            args = build_sft_config(seq, batch, ga)
+            trainer = CortexTrainer.build(model, tokenizer, train_ds, val_ds, args)
+            if is_rank0():
+                try:
+                    if hasattr(trainer, "model"):
+                        pass
+                    (Path(OUTPUT_DIR) / "run_config.json").write_text(
+                        json.dumps(
+                            {
+                                "model": MODEL_NAME,
+                                "repo": HF_REPO,
+                                "seq": seq,
+                                "batch": batch,
+                                "ga": ga,
+                                "lr": LR,
+                                "epochs": EPOCHS,
+                                "qat": scheme_label,
+                                "packing": PACKING,
+                                "max_hours": MAX_TRAIN_SECONDS / 3600,
+                                "world_size": int(os.environ.get("WORLD_SIZE", "1")),
+                                "resume": resume,
+                                "data": report,
+                            },
+                            indent=2,
+                        )
+                    )
+                except Exception:
+                    pass
+            log("[train] starting")
+            trainer.train(resume_from_checkpoint=resume)
+            last_err = None
+            break
+        except torch.cuda.OutOfMemoryError as e:
+            last_err = e
+            log(f"[OOM] seq={seq} batch={batch}: {e}")
+            try:
+                del trainer, model, tokenizer
+            except Exception:
+                pass
+            gc.collect()
+            torch.cuda.empty_cache()
+            resume = None  # fresh after OOM
+            continue
+        except Exception as e:
+            last_err = e
+            log(f"[run] failed: {e}")
+            traceback.print_exc()
+            # OOM-like
+            if "out of memory" in str(e).lower():
+                try:
+                    del trainer, model, tokenizer
+                except Exception:
+                    pass
+                gc.collect()
+                torch.cuda.empty_cache()
+                resume = None
+                continue
+            raise
+
+    if last_err is not None and trainer is None:
+        raise last_err
+
+    # Final save (still in QAT-prepare / fake-quant state so the next session can resume)
+    if is_rank0() and trainer is not None:
+        log("[save] final checkpoint")
+        trainer.save_model(OUTPUT_DIR)
+        try:
+            tokenizer.save_pretrained(OUTPUT_DIR)
+        except Exception:
+            pass
+        try:
+            trainer.push_to_hub(commit_message=f"Cortex A 0.5 SFT step {trainer.state.global_step}")
+        except Exception as e:
+            log(f"[save] push_to_hub: {e}")
+            try:
+                from huggingface_hub import HfApi
+
+                HfApi(token=HF_TOKEN).upload_folder(
+                    folder_path=OUTPUT_DIR,
+                    repo_id=HF_REPO,
+                    repo_type="model",
+                    token=HF_TOKEN,
+                    ignore_patterns=["*.tmp", "checkpoint-*/*.pt"],
+                )
+            except Exception as e2:
+                log(f"[save] upload_folder: {e2}")
+
+        # If we actually finished epochs (not time-killed), convert QAT fake-quant → real 4-bit
+        elapsed = time.time() - START_TS
+        finished = elapsed < (MAX_TRAIN_SECONDS - 300) and trainer.state.global_step > 0
+        if finished:
+            log("[QAT] converting fake-quant → real quantized weights")
+            try:
+                from torchao.quantization import quantize_
+                from torchao.quantization.qat import QATConfig
+
+                quantize_(trainer.model, QATConfig(step="convert"))
+                qdir = str(Path(OUTPUT_DIR) / "qat_converted")
+                Path(qdir).mkdir(exist_ok=True)
+                try:
+                    trainer.model.save_pretrained_torchao(qdir, tokenizer)
+                except TypeError:
+                    trainer.model.save_pretrained_torchao(qdir)
+                except Exception:
+                    trainer.model.save_pretrained(qdir)
+                    tokenizer.save_pretrained(qdir)
+                from huggingface_hub import HfApi
+
+                HfApi(token=HF_TOKEN).upload_folder(
+                    folder_path=qdir,
+                    path_in_repo="qat_converted",
+                    repo_id=HF_REPO,
+                    repo_type="model",
+                    token=HF_TOKEN,
+                )
+            except Exception as e:
+                log(f"[QAT] convert skipped: {e}")
+
+        log(f"[done] step={trainer.state.global_step} elapsed={elapsed/3600:.2f}h")
+
+
+if __name__ == "__main__":
+    main()