# Qubit 7B (NeuroQuantum / QBNN) training: CLM pretraining on FineWeb2-edu-ja and SFT on
# Japanese conversation / code / math. Custom loop, FSDP (full shard) + BF16 + activation checkpointing.
#
#   torchrun --nproc_per_node=4 train_qubit7b.py --stage pretrain --ckpt-dir /ckpt/run1 ...
#
# Checkpoints (resumable):  <ckpt>/step_XXXXXXX/shard_rank{r}.pt  (FSDP flat-param shard + optimizer state, per rank)
#                           <ckpt>/step_XXXXXXX/trainer_state_rank{r}.json (step, sched, data state, rng)
# Best / final weights:     <ckpt>/best|final/model.safetensors (bf16, consolidated) + config + tokenizer
import argparse, json, math, os, random, re, shutil, subprocess, sys, tempfile, time, hashlib
from contextlib import nullcontext

import torch
import torch.distributed as dist
import torch.nn.functional as F

p = argparse.ArgumentParser()
p.add_argument("--stage", choices=["pretrain", "sft"], default="pretrain")
p.add_argument("--model-src", default="/qubit/7b")
p.add_argument("--size", choices=["7b", "tiny"], default="7b")
p.add_argument("--data-dir", default="/ckpt/data")
p.add_argument("--tokenizer", default="/ckpt/data/qubit7b_ja_spm32k.model")
p.add_argument("--ckpt-dir", required=True)
p.add_argument("--init-from", default="", help="consolidated safetensors to start from (e.g. pretrain best)")
p.add_argument("--resume", action="store_true")
p.add_argument("--seq-len", type=int, default=1024)
p.add_argument("--micro-bs", type=int, default=4)
p.add_argument("--grad-accum", type=int, default=8)
p.add_argument("--lr", type=float, default=3e-4)
p.add_argument("--min-lr-ratio", type=float, default=0.1)
p.add_argument("--warmup", type=int, default=500)
p.add_argument("--max-steps", type=int, default=1000)
p.add_argument("--weight-decay", type=float, default=0.1)
p.add_argument("--clip", type=float, default=1.0)
p.add_argument("--dropout", type=float, default=0.0)
p.add_argument("--save-every", type=int, default=500)
p.add_argument("--eval-every", type=int, default=250)
p.add_argument("--eval-batches", type=int, default=20)
p.add_argument("--gen-every", type=int, default=0, help="0 = only at start/end")
p.add_argument("--keep-ckpts", type=int, default=1)
p.add_argument("--max-hours", type=float, default=1e9, help="stop + save cleanly before job timeout")
p.add_argument("--sft-mix", default="conv:0.5,code:0.333,math:0.167")
p.add_argument("--fineweb-config", default="default")
p.add_argument("--pretrain-jsonl", default="", help="local jsonl with {text} instead of streaming (testing only)")
p.add_argument("--verify-restore", action="store_true", help="after the run: reload checkpoint and compare loss")
p.add_argument("--push-repo", default="", help="HF model repo to upload best/final weights (uses QUBIT_HUB_TOKEN)")
p.add_argument("--seed", type=int, default=1234)
args = p.parse_args()

# ---------------------------------------------------------------- distributed setup
dist_on = "RANK" in os.environ
if dist_on:
    from datetime import timedelta  # long timeout: rank0 writes multi-GB files to the bucket while others wait
    dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo", timeout=timedelta(hours=2))
rank = dist.get_rank() if dist_on else 0
world = dist.get_world_size() if dist_on else 1
local_rank = int(os.environ.get("LOCAL_RANK", 0))
cuda = torch.cuda.is_available()
device = torch.device("cuda", local_rank) if cuda else torch.device("cpu")
if cuda:
    torch.cuda.set_device(device)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
master = rank == 0
os.makedirs(args.ckpt_dir, exist_ok=True)
T0 = time.time()


def log(*a):
    if master:
        print(time.strftime("%H:%M:%S"), *a, flush=True)


def barrier():
    if dist_on:
        dist.barrier()


torch.manual_seed(args.seed + rank)
random.seed(args.seed + rank)

# ---------------------------------------------------------------- model
sys.path.insert(0, args.model_src)
from modeling_qubit_7b import NeuroQuantum, NeuroQuantumConfig, QBNNTransformerBlock  # noqa: E402

cfg_path = os.path.join(args.model_src, "config.json")
cfg_dict = json.load(open(cfg_path))
if args.size == "tiny":
    cfg_dict.update(embed_dim=256, hidden_dim=384, num_heads=4, num_layers=2)
cfg_dict["dropout"] = args.dropout
cfg_dict["gradient_checkpointing"] = True
cfg_dict["max_seq_len"] = max(cfg_dict["max_seq_len"], args.seq_len)
cfg = NeuroQuantumConfig(**cfg_dict)
assert args.seq_len <= cfg.max_seq_len


def build_model():
    # rank0 materialises real weights on CPU; other ranks start on meta and receive them via sync_module_states
    if master or not (dist_on and cuda):
        m = NeuroQuantum(cfg)
    else:
        with torch.device("meta"):
            m = NeuroQuantum(cfg)
    return m


def load_consolidated(m, path):
    from safetensors.torch import load_file
    sd = load_file(path)
    sd = {k: v.float() for k, v in sd.items()}
    missing, unexpected = m.load_state_dict(sd, strict=False)
    missing = [k for k in missing if not k.endswith(("last_correction_norm", "last_htilde_norm", "call_count")) and not k.startswith("text_embedding.")]
    assert not missing and not unexpected, (missing[:5], unexpected[:5])


def wrap(m):
    if not (dist_on and cuda):
        return m.to(device)
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy
    from torch.distributed.fsdp.wrap import ModuleWrapPolicy
    mp = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=None)
    return FSDP(
        m,
        auto_wrap_policy=ModuleWrapPolicy({QBNNTransformerBlock}),
        mixed_precision=mp,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        device_id=device,
        sync_module_states=True,
        param_init_fn=(lambda mod: mod.to_empty(device=device, recurse=False)) if not master else None,
        use_orig_params=True,
        limit_all_gathers=True,
    )


raw = build_model()
n_params = sum(p.numel() for p in raw.parameters())
log(f"model params: {n_params:,} (size={args.size}) cfg={cfg_dict}")
if args.init_from and master:
    load_consolidated(raw, args.init_from)
    log("initialised from", args.init_from)
# FSDP(1) rejects 0-dim parameters: store each QBNN lambda_base as shape [1] during training.
# Numerically identical (broadcasts the same); consolidated checkpoints are saved back as scalars.
for mod in raw.modules():
    if hasattr(mod, "lambda_base") and isinstance(mod.lambda_base, torch.nn.Parameter) and mod.lambda_base.dim() == 0:
        mod.lambda_base = torch.nn.Parameter(mod.lambda_base.data.reshape(1))
model = wrap(raw)
autocast = (lambda: torch.autocast("cuda", dtype=torch.bfloat16)) if cuda and not dist_on else nullcontext
# (under FSDP the MixedPrecision policy already computes in bf16)

decay, no_decay = [], []
for n, p_ in model.named_parameters():
    (no_decay if (p_.ndim < 2 or "norm" in n or "embedding" in n or "lambda_base" in n) else decay).append(p_)
opt = torch.optim.AdamW([{"params": decay, "weight_decay": args.weight_decay}, {"params": no_decay, "weight_decay": 0.0}],
                        lr=args.lr, betas=(0.9, 0.95), eps=1e-8, fused=cuda)


def lr_at(step):
    if step < args.warmup:
        return args.lr * (step + 1) / args.warmup
    prog = min(1.0, (step - args.warmup) / max(1, args.max_steps - args.warmup))
    return args.lr * (args.min_lr_ratio + (1 - args.min_lr_ratio) * 0.5 * (1 + math.cos(math.pi * prog)))


sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: lr_at(s) / args.lr)

# ---------------------------------------------------------------- tokenizer + data
import sentencepiece as spm
sp = spm.SentencePieceProcessor(model_file=args.tokenizer)
if args.size == "tiny" and sp.get_piece_size() != cfg.vocab_size:
    raise SystemExit("tiny test needs a 32k tokenizer too")
assert sp.get_piece_size() == cfg.vocab_size, (sp.get_piece_size(), cfg.vocab_size)
BOS, EOS, PAD = sp.bos_id(), sp.eos_id(), sp.pad_id()
SYS, USR, AST, END = (sp.piece_to_id(x) for x in ("<|system|>", "<|user|>", "<|assistant|>", "<|end|>"))
ROLE = {"system": SYS, "user": USR, "assistant": AST}
NL = sp.encode("\n")


def encode_chat(messages, add_generation=False):
    ids, labels = [BOS], [-100]
    for m in messages:
        body = sp.encode(m["content"])
        seg = [ROLE[m["role"]]] + NL + body + [END] + NL
        ids += seg
        if m["role"] == "assistant":
            labels += [-100] * (1 + len(NL)) + body + [END] + [-100] * len(NL)
        else:
            labels += [-100] * len(seg)
    if add_generation:
        ids += [AST] + NL
        labels += [-100] * (1 + len(NL))
    else:
        ids.append(EOS)
        labels.append(EOS)
    return ids, labels


JA_RE = re.compile(r"[぀-ゟ゠-ヿ一-鿿]")


def ja_ratio(s):
    s2 = re.sub(r"\s", "", s)
    return len(JA_RE.findall(s2)) / max(1, len(s2))


class PretrainStream:
    """Streams FineWeb2-edu-ja train split, filters low quality / dups, packs into seq_len+1 blocks."""

    def __init__(self):
        from datasets import load_dataset
        from datasets.distributed import split_dataset_by_node
        if args.pretrain_jsonl:
            ds = load_dataset("json", data_files=args.pretrain_jsonl, split="train", streaming=True)
        else:
            ds = load_dataset("hotchpotch/fineweb-2-edu-japanese", args.fineweb_config, split="train", streaming=True, token=False)
        ds = ds.shuffle(seed=args.seed, buffer_size=10_000)
        self.ds = split_dataset_by_node(ds, rank=rank, world_size=world)
        self.it = None
        self.buf = []
        self.seen = set()
        self.stats = {"docs": 0, "kept": 0, "short": 0, "low_ja": 0, "repetitive": 0, "dup": 0}

    def state_dict(self):
        return {"ds": self.ds.state_dict(), "buf": self.buf, "stats": self.stats}

    def load_state_dict(self, s):
        self.ds.load_state_dict(s["ds"])
        self.buf, self.stats = s["buf"], s["stats"]

    def _ok(self, t):
        if len(t) < 200:
            return "short"
        if ja_ratio(t) < 0.5:
            return "low_ja"
        lines = [l for l in t.splitlines() if l.strip()]
        if lines and len(set(lines)) / len(lines) < 0.6:
            return "repetitive"
        k = hashlib.sha1(t[:2000].encode()).digest()[:8]
        if k in self.seen:
            return "dup"
        if len(self.seen) > 5_000_000:
            self.seen.clear()
        self.seen.add(k)
        return None

    def next_block(self):
        L = args.seq_len + 1
        while len(self.buf) < L:
            if self.it is None:
                self.it = iter(self.ds)
            try:
                ex = next(self.it)
            except StopIteration:
                self.it = iter(self.ds)
                continue
            self.stats["docs"] += 1
            why = self._ok(ex["text"])
            if why:
                self.stats[why] += 1
                continue
            self.stats["kept"] += 1
            self.buf += [BOS] + sp.encode(ex["text"]) + [EOS]
        blk, self.buf = self.buf[:L], self.buf[L:]
        return torch.tensor(blk), None


class SFTStream:
    """Packs chat examples (loss on assistant tokens only) from the prepared jsonl files, mixed by weight."""

    def __init__(self, split="train"):
        self.sources, self.weights = {}, {}
        for kv in args.sft_mix.split(","):
            k, w = kv.split(":")
            rows = [json.loads(l) for l in open(os.path.join(args.data_dir, f"sft_{k}_{split}.jsonl"), encoding="utf-8")]
            rows = rows[rank::world] if split == "train" else rows
            self.sources[k] = rows
            self.weights[k] = float(w)
        self.pos = {k: 0 for k in self.sources}
        self.rng = random.Random(args.seed + rank)
        self.buf_i, self.buf_l = [], []
        self.used = {k: 0 for k in self.sources}

    def state_dict(self):
        return {"pos": self.pos, "rng": self.rng.getstate(), "buf_i": self.buf_i, "buf_l": self.buf_l, "used": self.used}

    def load_state_dict(self, s):
        self.pos, self.buf_i, self.buf_l, self.used = s["pos"], s["buf_i"], s["buf_l"], s["used"]
        st = s["rng"]
        self.rng.setstate((st[0], tuple(st[1]), st[2]))

    def next_block(self):
        L = args.seq_len + 1
        while len(self.buf_i) < L:
            k = self.rng.choices(list(self.weights), weights=list(self.weights.values()))[0]
            rows = self.sources[k]
            r = rows[self.pos[k] % len(rows)]
            self.pos[k] += 1
            self.used[k] += 1
            ids, lab = encode_chat(r["messages"])
            ids, lab = ids[:L], lab[:L]
            self.buf_i += ids
            self.buf_l += lab
        bi, bl = self.buf_i[:L], self.buf_l[:L]
        self.buf_i, self.buf_l = self.buf_i[L:], self.buf_l[L:]
        return torch.tensor(bi), torch.tensor(bl)


def get_batch(stream):
    xs, ls = [], []
    for _ in range(args.micro_bs):
        x, l = stream.next_block()
        xs.append(x)
        ls.append(x if l is None else l)
    x = torch.stack(xs).to(device, non_blocking=True)
    l = torch.stack(ls).to(device, non_blocking=True)
    return x[:, :-1].contiguous(), l[:, 1:].contiguous()


def loss_fn(inp, tgt):
    with autocast():
        logits = model(inp)
    return F.cross_entropy(logits.float().view(-1, logits.size(-1)), tgt.reshape(-1), ignore_index=-100)


# fixed evaluation sets (held out, never trained on)
def build_eval_sets():
    sets = {}
    L = args.seq_len + 1
    # pretrain val: FineWeb2 *test* split docs
    toks = []
    for line in open(os.path.join(args.data_dir, "pretrain_val.jsonl"), encoding="utf-8"):
        toks += [BOS] + sp.encode(json.loads(line)["text"]) + [EOS]
        if len(toks) > L * args.eval_batches * args.micro_bs * world + L:
            break
    blocks = [torch.tensor(toks[i:i + L]) for i in range(0, len(toks) - L, L)]
    sets["web"] = [(b, b) for b in blocks]
    for k in ("conv", "code", "math"):
        rows = [json.loads(l) for l in open(os.path.join(args.data_dir, f"sft_{k}_eval.jsonl"), encoding="utf-8")]
        out = []
        for r in rows:
            ids, lab = encode_chat(r["messages"])
            ids, lab = ids[:L], lab[:L]
            pad = L - len(ids)
            out.append((torch.tensor(ids + [PAD] * pad), torch.tensor(lab + [-100] * pad)))
        sets[k] = out
    return sets


@torch.no_grad()
def evaluate(eval_sets):
    model.eval()
    res = {}
    for name, items in eval_sets.items():
        mine = items[rank::world][: args.eval_batches * args.micro_bs]
        n_min = min(len(items[r::world][: args.eval_batches * args.micro_bs]) for r in range(world))
        mine = mine[:n_min]
        tot, cnt = torch.zeros((), device=device), torch.zeros((), device=device)
        for i in range(0, len(mine), args.micro_bs):
            chunk = mine[i:i + args.micro_bs]
            x = torch.stack([c[0] for c in chunk]).to(device)
            l = torch.stack([c[1] for c in chunk]).to(device)
            with autocast():
                logits = model(x[:, :-1])
            t = l[:, 1:]
            ce = F.cross_entropy(logits.float().reshape(-1, logits.size(-1)), t.reshape(-1), ignore_index=-100, reduction="sum")
            tot += ce
            cnt += (t != -100).sum()
        if dist_on:
            dist.all_reduce(tot)
            dist.all_reduce(cnt)
        v = (tot / cnt.clamp(min=1)).item()
        res[name] = {"loss": round(v, 4), "ppl": round(math.exp(min(v, 50)), 2), "tokens": int(cnt.item())}
    model.train()
    return res


PROMPTS = {
    "japanese": "日本の首都はどこですか？",
    "conversation": "こんにちは。あなたは何ができますか？",
    "math": "125 × 48 を計算してください。",
    "coding": "Pythonでクイックソートを書いてください。",
    "reasoning": "なぜ飛行機は空を飛べるのですか？",
}


def check_code(text):
    m = re.search(r"```(?:python)?\n(.*?)```", text, re.S)
    code = m.group(1) if m else text
    test = code + "\n\nimport random\n_a=[random.randint(0,99) for _ in range(30)]\n_f=[v for k,v in list(globals().items()) if callable(v) and 'sort' in k.lower()]\nprint(bool(_f) and _f[0](list(_a))==sorted(_a))\n"
    try:
        r = subprocess.run([sys.executable, "-I", "-c", test], capture_output=True, text=True, timeout=5, cwd=tempfile.gettempdir())
        return r.stdout.strip().endswith("True")
    except Exception:
        return False


@torch.no_grad()
def generate_all(max_new=96):
    model.eval()
    out = {}
    for key, q in PROMPTS.items():
        ids, _ = encode_chat([{"role": "user", "content": q}], add_generation=True)
        x = torch.tensor([ids], device=device)
        gen = []
        for _ in range(max_new):
            with autocast():
                logits = model(x[:, -args.seq_len:])[0, -1].float()
            for t in set(gen[-64:]):  # mild repetition penalty
                logits[t] = logits[t] / 1.3 if logits[t] > 0 else logits[t] * 1.3
            nxt = int(torch.argmax(logits))
            if dist_on:  # keep ranks in lockstep
                tt = torch.tensor([nxt], device=device)
                dist.broadcast(tt, 0)
                nxt = int(tt.item())
            if nxt in (EOS, END):
                break
            gen.append(nxt)
            x = torch.cat([x, torch.tensor([[nxt]], device=device)], 1)
        text = sp.decode(gen)
        rec = {"prompt": q, "output": text}
        if key == "math":
            rec["correct_6000"] = "6000" in text.replace(",", "")
        if key == "coding" and master:
            rec["code_runs_and_sorts"] = check_code(text)
        if key == "japanese":
            rec["mentions_tokyo"] = "東京" in text
        out[key] = rec
    model.train()
    return out


# ---------------------------------------------------------------- checkpointing
# Resumable checkpoints = each rank's own FSDP flat-parameter shards + its local optimizer state,
# saved with torch.save (one file per rank). torch.distributed.checkpoint's FSDP1/use_orig_params
# optimizer handling failed to reload in the 7B smoke tests on torch 2.14, so it is not used here.
# Requirement: resume with the same number of GPUs (asserted). Consolidated bf16 weights are saved separately.
from torch.distributed.checkpoint.state_dict import get_model_state_dict, StateDictOptions

LOCAL_CKPT = os.environ.get("LOCAL_CKPT_DIR", os.path.join(tempfile.gettempdir(), "qubit_ckpt_local"))


def _flat_params():
    if not (dist_on and cuda):
        return None
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
    return [m._flat_param for m in FSDP.fsdp_modules(model) if getattr(m, "_flat_param", None) is not None]


def _copy_shared(src_dir, dst_dir, names):
    # every rank copies a disjoint subset of files (single node: all ranks share the local disk);
    # large sequential copies are fast on the bucket mount, many small random reads are not
    os.makedirs(dst_dir, exist_ok=True)
    for i, n in enumerate(sorted(names)):
        if i % world == rank:
            shutil.copyfile(os.path.join(src_dir, n), os.path.join(dst_dir, n) + ".part")
            os.replace(os.path.join(dst_dir, n) + ".part", os.path.join(dst_dir, n))
    barrier()


def save_ckpt(step, stream, metrics_tail):
    name = f"step_{step:07d}"
    local, remote = os.path.join(LOCAL_CKPT, name), os.path.join(args.ckpt_dir, name)
    os.makedirs(local, exist_ok=True)
    t = time.time()
    flats = _flat_params()
    payload = {
        "model": [f.detach().cpu() for f in flats] if flats is not None else {k: v.cpu() for k, v in model.state_dict().items()},
        "optim": opt.state_dict(),
        "step": step, "world": world,
    }
    torch.save(payload, os.path.join(local, f"shard_rank{rank}.pt"))
    del payload
    st = {"step": step, "sched": sched.state_dict(), "rng_cpu": torch.get_rng_state().tolist(),
          "data": stream.state_dict(), "metrics_tail": metrics_tail, "args": vars(args), "world": world}
    with open(os.path.join(local, f"trainer_state_rank{rank}.json"), "w") as f:
        json.dump(st, f)
    barrier()
    t_local = time.time() - t
    _copy_shared(local, remote, os.listdir(local))   # persist to the bucket
    if master:
        # completion marker written last -> only complete checkpoints are ever resumed
        json.dump({"step": step, "time": time.time(), "secs": round(time.time() - t, 1), "world": world}, open(os.path.join(remote, "COMPLETE"), "w"))
        for root in (args.ckpt_dir, LOCAL_CKPT):
            olds = sorted(x for x in os.listdir(root) if x.startswith("step_") and x != name)
            for old in olds[: max(0, len(olds) - (args.keep_ckpts - 1))]:
                shutil.rmtree(os.path.join(root, old), ignore_errors=True)
    barrier()
    log(f"checkpoint saved {remote} (local {t_local:.0f}s, total {time.time() - t:.0f}s)")
    return remote


def latest_ckpt():
    if not os.path.isdir(args.ckpt_dir):
        return None
    c = sorted(x for x in os.listdir(args.ckpt_dir) if x.startswith("step_") and os.path.exists(os.path.join(args.ckpt_dir, x, "COMPLETE")))
    return os.path.join(args.ckpt_dir, c[-1]) if c else None


def load_ckpt(d, stream):
    t = time.time()
    local = os.path.join(LOCAL_CKPT, os.path.basename(d))
    if not os.path.exists(os.path.join(local, f"shard_rank{world - 1}.pt")):
        _copy_shared(d, local, [n for n in os.listdir(d) if n != "COMPLETE"])
        log(f"copied checkpoint {d} -> local in {time.time() - t:.0f}s")
    st = json.load(open(os.path.join(local, f"trainer_state_rank{rank}.json")))
    assert st["world"] == world, f"checkpoint was written with world={st['world']}, now {world}"
    payload = torch.load(os.path.join(local, f"shard_rank{rank}.pt"), map_location=device, weights_only=False)
    flats = _flat_params()
    with torch.no_grad():
        if flats is not None:
            assert len(flats) == len(payload["model"])
            for f, saved in zip(flats, payload["model"]):
                assert f.shape == saved.shape, (f.shape, saved.shape)
                f.data.copy_(saved)
        else:
            model.load_state_dict(payload["model"])
    opt.load_state_dict(payload["optim"])
    del payload
    sched.load_state_dict(st["sched"])
    torch.set_rng_state(torch.tensor(st["rng_cpu"], dtype=torch.uint8))
    if stream is not None:
        stream.load_state_dict(st["data"])
    barrier()
    log(f"checkpoint loaded in {time.time() - t:.0f}s")
    return st["step"]


def save_consolidated(tag, extra):
    d = os.path.join(args.ckpt_dir, tag)
    t = time.time()
    sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
    if master:
        from safetensors.torch import save_file
        os.makedirs(d, exist_ok=True)
        sd = {k: v.to(torch.bfloat16).contiguous() for k, v in sd.items() if not k.endswith(("last_correction_norm", "last_htilde_norm", "call_count")) and not k.startswith("text_embedding.")}  # text_embedding is an alias of token_embedding
        sd = {k: (v.reshape(()) if k.endswith("lambda_base") else v) for k, v in sd.items()}  # original scalar shape
        for fn in ("SAVED_OK", "model.safetensors"):  # marker removed first; rewritten only after verification
            if os.path.exists(os.path.join(d, fn)):
                os.remove(os.path.join(d, fn))
        save_file(sd, os.path.join(d, "model.safetensors"), metadata={"format": "pt", "arch": "NeuroQuantum-QBNN", "params": str(n_params)})
        c = dict(cfg_dict)
        c.update({"architecture": "NeuroQuantum (Qubit QBNN)", "num_parameters": n_params, "torch_dtype": "bfloat16"})
        json.dump(c, open(os.path.join(d, "config.json"), "w"), indent=2)
        shutil.copy(args.tokenizer, os.path.join(d, "tokenizer.model"))
        shutil.copy(os.path.join(args.model_src, "modeling_qubit_7b.py"), d)
        json.dump(extra, open(os.path.join(d, "training_info.json"), "w"), ensure_ascii=False, indent=2)
        # verify: re-read file, compare tensor count + checksum of a few tensors
        from safetensors import safe_open
        with safe_open(os.path.join(d, "model.safetensors"), "pt") as f:
            keys = list(f.keys())
            assert len(keys) == len(sd)
            for k in keys[:3] + keys[-3:]:
                assert torch.equal(f.get_tensor(k), sd[k]), k
            nparam = sum(math.prod(f.get_slice(k).get_shape()) for k in keys)
        json.dump({"tensors": len(keys), "params": nparam, "time": time.time()}, open(os.path.join(d, "SAVED_OK"), "w"))
        log(f"consolidated {tag} saved+verified: {len(keys)} tensors, {nparam:,} params, {time.time() - t:.0f}s")
    barrier()
    return d


def push(d, msg):
    tok = os.environ.get("QUBIT_HUB_TOKEN")
    if not (master and args.push_repo and tok):
        return
    try:
        from huggingface_hub import HfApi
        HfApi(token=tok).upload_folder(folder_path=d, repo_id=args.push_repo, path_in_repo=f"{args.stage}/{os.path.basename(d)}", commit_message=msg)
        log("pushed", d, "->", args.push_repo)
    except Exception as e:  # token may have expired; bucket copy is the source of truth
        log("push failed (bucket copy is kept):", type(e).__name__, str(e)[:200])


# ---------------------------------------------------------------- main loop
stream = PretrainStream() if args.stage == "pretrain" else SFTStream("train")
eval_sets = build_eval_sets()
metrics_path = os.path.join(args.ckpt_dir, "metrics.jsonl")


def mlog(rec):
    if master:
        with open(metrics_path, "a") as f:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")


step = 0
if args.resume and latest_ckpt():
    d = latest_ckpt()
    step = load_ckpt(d, stream)
    log(f"resumed from {d} at step {step}")
best = float("inf")
best_file = os.path.join(args.ckpt_dir, "best", "training_info.json")
if os.path.exists(best_file):
    best = json.load(open(best_file)).get("val_loss", best)

tok_per_step = args.micro_bs * args.seq_len * args.grad_accum * world
log(f"world={world} tokens/step={tok_per_step:,} max_steps={args.max_steps} stage={args.stage}")
if step == 0:
    ev = evaluate(eval_sets)
    gen = generate_all()
    log("pre-train eval:", json.dumps(ev, ensure_ascii=False))
    log("pre-train generations:", json.dumps(gen, ensure_ascii=False))
    mlog({"step": 0, "type": "eval", "eval": ev, "gen": gen})

model.train()
flops_per_tok = 6 * n_params + 12 * cfg.num_layers * cfg.embed_dim * min(args.seq_len, cfg.attention_window)
t_last, tokens_since, losses = time.time(), 0, []
stop_reason = "max_steps"
while step < args.max_steps:
    for g in opt.param_groups:
        g["lr"] = lr_at(step)
    tot = torch.zeros((), device=device)
    t_data = 0.0
    for micro in range(args.grad_accum):
        td = time.time()
        x, y = get_batch(stream)
        t_data += time.time() - td
        # no FSDP.no_sync(): it would keep full unsharded fp32 grads (~28 GB/GPU for 7B)
        loss = loss_fn(x, y) / args.grad_accum
        loss.backward()
        tot += loss.detach()
    if dist_on and cuda:
        gn = model.clip_grad_norm_(args.clip)
    else:
        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
    opt.step()
    opt.zero_grad(set_to_none=True)
    sched.step()
    step += 1
    if dist_on:
        dist.all_reduce(tot)
        tot /= world
    lv = tot.item()
    if not math.isfinite(lv):
        log(f"non-finite loss at step {step}; aborting")
        stop_reason = "nan"
        break
    losses.append(lv)
    tokens_since += tok_per_step
    if step % 10 == 0 or step <= 5:
        dt = time.time() - t_last
        tps = tokens_since / dt
        mem = torch.cuda.max_memory_allocated() / 2**30 if cuda else 0
        mfu = tps * flops_per_tok / (world * 312e12) if cuda else 0
        rec = {"step": step, "type": "train", "loss": round(lv, 4), "lr": lr_at(step), "grad_norm": round(float(gn), 3),
               "tok_per_s": round(tps), "mfu": round(mfu, 3), "max_mem_gib": round(mem, 1), "data_s": round(t_data, 2),
               "elapsed_h": round((time.time() - T0) / 3600, 3)}
        if args.stage == "pretrain":
            rec["filter"] = stream.stats
        log(json.dumps(rec, ensure_ascii=False))
        mlog(rec)
        t_last, tokens_since = time.time(), 0
    out_of_time = (time.time() - T0) / 3600 > args.max_hours
    if step % args.eval_every == 0 or step == args.max_steps or out_of_time:
        ev = evaluate(eval_sets)
        key = "web" if args.stage == "pretrain" else None
        vl = ev[key]["loss"] if key else sum(ev[k]["loss"] for k in ("conv", "code", "math")) / 3
        rec = {"step": step, "type": "eval", "train_loss_avg": round(sum(losses[-50:]) / len(losses[-50:]), 4), "val_loss": round(vl, 4), "eval": ev}
        if args.gen_every and step % args.gen_every == 0:
            rec["gen"] = generate_all()
        log("eval:", json.dumps(rec, ensure_ascii=False))
        mlog(rec)
        if vl < best:
            best = vl
            d = save_consolidated("best", {"step": step, "val_loss": vl, "eval": ev, "stage": args.stage})
            push(d, f"best @ step {step} val_loss={vl:.4f}")
    if step % args.save_every == 0 or out_of_time:
        save_ckpt(step, stream, losses[-20:])
    if out_of_time:
        stop_reason = "max_hours"
        break

# ---------------------------------------------------------------- end of run
if step % args.save_every != 0 and stop_reason != "max_hours":
    save_ckpt(step, stream, losses[-20:])
gen = generate_all()
log("final generations:", json.dumps(gen, ensure_ascii=False))
mlog({"step": step, "type": "final", "stop_reason": stop_reason, "gen": gen})
d = save_consolidated("final", {"step": step, "stop_reason": stop_reason, "gen": gen, "stage": args.stage})
push(d, f"final @ step {step}")

if args.verify_restore:
    # reference loss on a fixed held-out batch with current weights
    fixed = eval_sets["web"][: args.micro_bs] if args.stage == "pretrain" else eval_sets["conv"][: args.micro_bs]
    fx = torch.stack([c[0] for c in fixed]).to(device)
    fl = torch.stack([c[1] for c in fixed]).to(device)
    model.eval()
    with torch.no_grad():
        ref = loss_fn(fx[:, :-1], fl[:, 1:]).item()
    # scramble weights, then restore from the latest checkpoint
    with torch.no_grad():
        for p_ in model.parameters():
            p_.mul_(0.0)
    with torch.no_grad():
        broken = loss_fn(fx[:, :-1], fl[:, 1:]).item()
    d = latest_ckpt()
    rstep = load_ckpt(d, None)
    with torch.no_grad():
        got = loss_fn(fx[:, :-1], fl[:, 1:]).item()
    ok = abs(ref - got) < 1e-3
    res = {"type": "restore_test", "ckpt": d, "ckpt_step": rstep, "ref_loss": ref, "zeroed_loss": broken, "restored_loss": got, "ok": ok}
    log("restore test:", json.dumps(res))
    mlog(res)
    # continue training 2 steps from restored state to prove optimizer/scheduler resume works
    model.train()
    for _ in range(2):
        x, y = get_batch(stream)
        l = loss_fn(x, y)
        l.backward()
        opt.step()
        opt.zero_grad(set_to_none=True)
        sched.step()
    log("resume-continue ok, lr now", sched.get_last_lr()[0])
    mlog({"type": "resume_continue_ok", "lr": sched.get_last_lr()[0]})

log(f"DONE stage={args.stage} step={step} stop={stop_reason} elapsed_h={(time.time() - T0) / 3600:.2f}")
if dist_on:
    dist.destroy_process_group()