train: per-rank FSDP flat-param + optimizer shard checkpoints (replace DCP), local-first then bucket copy
518d1d7 verified | # 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 | |
| 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 | |
| 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() | |