#!/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 "" in (text or ""): return text return f"\n{r}\n\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("" 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()