training/train.py

#!/usr/bin/env python3
"""Cortex A 0.5 — SFT + 4-bit QAT (int8-int4) for Qwen3.5-0.8B on Kaggle 2x T4.

Full-parameter SFT with Unsloth QAT:
  * 4-bit weights
  * 8-bit dynamic activation quantization  (qat_scheme='int8-int4')

Resumes from Hugging Face, hard-stops at 11.5h, pushes checkpoints + metrics.
"""
from __future__ import annotations

import fcntl
import gc
import inspect
import json
import math
import os
import random
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")
os.environ.setdefault("OMP_NUM_THREADS", "4")

# ---------------------------------------------------------------------------
# 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"
# 4-bit weights + 8-bit dynamic activations (Unsloth official scheme)
QAT_SCHEME_PREF = os.environ.get("QAT_SCHEME", "int8-int4")
MAX_SAMPLES = os.environ.get("MAX_SAMPLES")
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


def dist_barrier():
    try:
        import torch.distributed as dist

        if dist.is_available() and dist.is_initialized():
            dist.barrier()
    except Exception:
        pass


def file_lock(path: str = "/tmp/cortex_data.lock"):
    """Exclusive lock so DDP ranks don't race HF dataset generation."""
    class _Lock:
        def __enter__(self):
            Path(path).parent.mkdir(parents=True, exist_ok=True)
            self.fh = open(path, "w")
            fcntl.flock(self.fh, fcntl.LOCK_EX)
            return self

        def __exit__(self, *exc):
            try:
                fcntl.flock(self.fh, fcntl.LOCK_UN)
            finally:
                self.fh.close()

    return _Lock()


# ---------------------------------------------------------------------------
# Dataset catalog (frontier distillation + roleplay)
# ---------------------------------------------------------------------------
# r0b0tlab README configs omit data_files, so load_dataset(id, name=cfg) globs
# EVERY parquet (glm47_native schema clash). We load via data_dir instead.
DATASETS = [
    {
        "id": "r0b0tlab/qwen3.8-max-glm5.2-kimi-k3-distillation",
        "config": "sft_balanced",
        "data_dir": "data/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 _with_think(text: str, reasoning) -> str:
    if not reasoning:
        return text
    r = str(reasoning).strip()
    if not r:
        return text
    if "<think>" in (text or ""):
        return text
    return f"<think>\n{r}\n</think>\n{text}"


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
            reasoning = t.get("reasoning_content") or t.get("reasoning") or t.get("thought")
            if role == "assistant":
                text = _with_think(str(text), reasoning)
            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"]
            if t.get("tool_call_id"):
                msg["tool_call_id"] = t["tool_call_id"]
            out.append(msg)
        if out:
            return out

    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)
        assistant = _with_think(str(outp), ex.get("reasoning_content") or ex.get("reasoning"))
        msgs.append({"role": "user", "content": user})
        msgs.append({"role": "assistant", "content": assistant})
        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
    has_think = any("<think>" in str(m.get("content", "")) for m in messages)
    attempts = []
    if has_think:
        attempts.append(dict(kwargs, chat_template_kwargs={"enable_thinking": True}))
        attempts.append(dict(kwargs, enable_thinking=True))
    attempts.extend(
        [
            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 — 4-bit weights + 8-bit dynamic activations
# ---------------------------------------------------------------------------
_QAT_ALIASES = {
    "int8-int4": "int8-int4",
    "int8int4": "int8-int4",
    "int4-int8": "int8-int4",
    "int8_int4": "int8-int4",
    "dyn-int8": "int8-int4",
    "int4-dyn": "int8-int4",  # previous default; map to the requested scheme
    "int4-dyn-act8": "int8-int4",
    "auto": "int8-int4",
    "int4": "int4",
    "int4-wo": "int4",
    "fp8-int4": "fp8-int4",
    "fp8-fp8": "fp8-fp8",
    "int8": "int8",
    "cactus": "cactus",
}


def resolve_qat_scheme():
    pref = (QAT_SCHEME_PREF or "int8-int4").lower().strip()
    scheme = _QAT_ALIASES.get(pref, "int8-int4")
    if scheme == "int8-int4":
        log("[QAT] Unsloth qat_scheme='int8-int4'  (INT4 weights + INT8 dynamic activations)")
    else:
        log(f"[QAT] Unsloth qat_scheme='{scheme}'")
    return scheme, scheme


def apply_qat(model, scheme) -> Any:
    if scheme is None:
        return model
    try:
        from unsloth.models._utils import _prepare_model_for_qat

        log(f"[QAT] _prepare_model_for_qat({scheme!r})")
        return _prepare_model_for_qat(model, scheme)
    except Exception as e:
        log(f"[QAT] unsloth prepare failed: {e}")

    try:
        from torchao.quantization import quantize_, Int8DynamicActivationIntxWeightConfig, Int4WeightOnlyConfig
        from torchao.quantization.qat import QATConfig
        from torchao.quantization.granularity import PerGroup
        import torch

        if scheme == "int8-int4":
            base = Int8DynamicActivationIntxWeightConfig(
                weight_dtype=torch.int4, weight_granularity=PerGroup(32)
            )
        else:
            base = Int4WeightOnlyConfig(group_size=128)
        quantize_(model, QATConfig(base, step="prepare"))
        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


def count_qat_modules(model) -> int:
    n = 0
    names = []
    for m in model.modules():
        cn = m.__class__.__name__
        if any(s in cn for s in ("FakeQuant", "QAT", "FakeQuantize", "Int8DType", "ChosenFakeQuant")):
            n += 1
            if len(names) < 8:
                names.append(cn)
    log(f"[QAT] fake-quant / QAT modules detected: {n}  e.g. {names}")
    return n


# ---------------------------------------------------------------------------
# Callbacks
# ---------------------------------------------------------------------------
def make_callbacks(tokenizer=None):
    from transformers import TrainerCallback, TrainerControl, TrainerState

    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:
                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
            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,
                "qat_scheme": QAT_SCHEME_PREF,
            }
            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,
                "qat_scheme": "int8-int4",
                "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:
            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)
        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_raw_dataset(spec: dict):
    """Load one HF dataset, never globbing sibling configs with different schemas."""
    from datasets import load_dataset

    ds_id = spec["id"]
    cfg = spec.get("config")
    folder = spec.get("data_dir")
    if folder is None and ds_id.startswith("r0b0tlab/") and cfg:
        folder = f"data/{cfg}"

    attempts = []
    if folder:
        attempts.append(dict(path=ds_id, data_dir=folder))
        attempts.append(
            dict(
                path=ds_id,
                data_files={
                    "train": f"{folder}/train-*.parquet",
                    "validation": f"{folder}/validation-*.parquet",
                    "test": f"{folder}/test-*.parquet",
                },
            )
        )
        attempts.append(
            dict(
                path="parquet",
                data_files={
                    "train": f"hf://datasets/{ds_id}/{folder}/train-*.parquet",
                    "validation": f"hf://datasets/{ds_id}/{folder}/validation-*.parquet",
                },
            )
        )
    std = {"path": ds_id}
    if cfg:
        std["name"] = cfg
    attempts.append(std)

    last = None
    for kw in attempts:
        try:
            log(f"[data] load {ds_id} via { {k: v for k, v in kw.items() if k != 'path'} or 'default' }")
            return load_dataset(**kw)
        except Exception as e:
            last = e
            log(f"[data]   failed: {type(e).__name__}: {e}")
    raise last


def load_and_mix(tokenizer):
    from datasets import concatenate_datasets

    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

    with file_lock():
        for spec in DATASETS:
            ds_id = spec["id"]
            try:
                log(f"[data] loading {ds_id}" + (f" ({spec.get('config')})" if spec.get("config") else ""))
                raw = load_raw_dataset(spec)
                split = take_split(raw, spec.get("split", "train"))

                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 getattr(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)
                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:
                log(f"[data] FAILED {ds_id}: {e}")
                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(2, 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:
        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,
        qat_scheme=scheme,
    )
    try:
        model, tokenizer = loader.from_pretrained(**kwargs)
    except TypeError as e:
        log(f"[model] from_pretrained qat_scheme TypeError: {e} — retrying then apply_qat")
        kwargs.pop("qat_scheme", None)
        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")

    n_fq = count_qat_modules(model)
    if n_fq == 0:
        log("[QAT] WARNING: no fake-quant modules — applying prepare again")
        model = apply_qat(model, scheme)
        n_fq = count_qat_modules(model)
    if n_fq == 0:
        raise RuntimeError(
            "QAT did not attach fake-quant modules. int8-int4 (4-bit weights + "
            "8-bit dynamic activations) is required. Check Unsloth/TorchAO install."
        )
    log(f"[QAT] ENABLED scheme={scheme_label} modules={n_fq}")
    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

        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",
    )
    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

    cfg = {k: v for k, v in cfg.items() if k in params or k in ("output_dir",)}
    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}")
        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

    try:
        torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", "4")))
    except Exception:
        pass

    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)
            mem = getattr(p, "total_memory", None) or getattr(p, "total_mem", 0)
            log(f"[hw] gpu{i} {p.name} {mem/1024**3:.1f}GB cap={torch.cuda.get_device_capability(i)}")

    resume = find_resume()
    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),
    ]
    seen = set()
    plans = []
    for t in backoff:
        if t not in seen:
            seen.add(t)
            plans.append(t)

    last_err = None
    trainer = None
    for seq, batch, ga in plans:
        model = tokenizer = trainer = None
        try:
            log(f"[run] seq={seq} batch={batch} ga={ga} packing={PACKING} qat={QAT_SCHEME_PREF}")
            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:
                    (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,
                                "qat_detail": "int4 weights + int8 dynamic activations",
                                "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
            continue
        except Exception as e:
            last_err = e
            log(f"[run] failed: {e}")
            traceback.print_exc()
            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

    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 int8-int4 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}")

        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 int8-int4 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()