Upload notebook.py with huggingface_hub

notebook.py

@@ -6,457 +6,127 @@ app = marimo.App(width="medium", auto_download=["html"])
 
 @app.cell
 def _():
-    # =============================================================== #
-    #  A) DragonCode 方案A — infra: imports, auth, helpers, data source
-    # =============================================================== #
-    # 規則：完全跳過 DragonCode-Tokenized-Pretrain（視為垃圾）；全部 token
-    # streaming FineWeb-Edu -> SlimPajama -> The-Pile；不下載完整檔案到本機。
-    import os, sys, time, json, math, shutil, tempfile, random, logging
-    from dataclasses import dataclass, field, asdict
-    from collections import deque
-    from typing import Optional
-    import numpy as np
-    import torch
-    import torch.nn.functional as F
-
-    # ---- auth (source from env; never hard-coded, to avoid leaking on a
-    #      public repo). Set HF_TOKEN in your Molab env before running: either
-    #      export it in the environment, or uncomment the next line and paste
-    #      your own token. ------------------------------------------------ #
-    DRAGON_HF_TOKEN = os.environ.get("HF_TOKEN", "") or ""
-    os.environ["HF_TOKEN"] = DRAGON_HF_TOKEN
-    os.environ["HF_HOME"] = "/home/marimo/.cache/huggingface"
-    if not DRAGON_HF_TOKEN:
-        print("⚠️  請先設定 HF_TOKEN 環境變數（例如 os.environ['HF_TOKEN']='hf_...'）再執行。")
-        print("    未設定 token 無法 load dataset / push checkpoint 至 HF。")
-
-    TOKENIZER_REPO = "bigcode/starcoder2-3b"
-    VOCAB_SIZE = 49152
-    MAX_SEQ_LEN = 2048
-
-    logger = logging.getLogger("dragoncode")
-    if not logger.handlers:
-        _h = logging.StreamHandler()
-        _h.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
-        logger.addHandler(_h)
-        logger.setLevel(logging.INFO)
-
-    # ---- HF helpers ------------------------------------------------------- #
-    def _hf_api():
-        from huggingface_hub import HfApi
-        return HfApi(token=DRAGON_HF_TOKEN)
-
-    def _run_with_retry(fn, tries=8, base_wait=4.0, max_wait=120.0, label="op"):
-        last = None
-        for attempt in range(1, tries + 1):
-            try:
-                return fn()
-            except Exception as e:
-                last = e
-                wait = min(base_wait * (2 ** (attempt - 1)), max_wait)
-                logger.warning("[%s] %d/%d failed: %s — retry in %.0fs", label, attempt, tries, e, wait)
-                time.sleep(wait)
-        raise RuntimeError(f"[{label}] failed after {tries} attempts: {last}")
-
-    def _seed_everything(seed):
-        random.seed(seed); np.random.seed(seed)
-        torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)
-
-    # ---- tokenizer -------------------------------------------------------- #
-    _tok = None
-    def _get_tokenizer():
-        global _tok
-        if _tok is not None:
-            return _tok
-        from transformers import AutoTokenizer
-        _tok = AutoTokenizer.from_pretrained(TOKENIZER_REPO, token=DRAGON_HF_TOKEN, trust_remote_code=True)
-        if _tok.pad_token is None:
-            _tok.pad_token = _tok.eos_token
-        if _tok.pad_token_id is None:
-            _tok.pad_token_id = _tok.eos_token_id
-        logger.info("Tokenizer %s (vocab=%d)", TOKENIZER_REPO, _tok.vocab_size)
-        return _tok
-
-    # ---- model / tier registry -------------------------------------------- #
-    def _build_model_config(hidden, layers, ffn, heads, kv_heads=None, name="DragonCode-150M"):
-        from transformers import LlamaConfig
-        return LlamaConfig(
-            vocab_size=VOCAB_SIZE, hidden_size=hidden, intermediate_size=ffn,
-            num_hidden_layers=layers, num_attention_heads=heads,
-            num_key_value_heads=kv_heads or heads, max_position_embeddings=MAX_SEQ_LEN,
-            rope_theta=10000.0, rms_norm_eps=1e-5, tie_word_embeddings=False,
-            hidden_act="silu", _name_or_path=f"DragonLimited/{name}",
-        )
-
-    @dataclass
-    class TierSpec:
-        name: str; size_str: str; hidden: int; layers: int; ffn: int
-        heads: int; kv_heads: Optional[int]; chinchilla_tokens: int
-
-    # Chinchilla 20x 鎖死: 150M->3.0B, 387M->7.74B, 787M->15.74B, 1.2B->24B, 2.4B->48B
-    def _tier_specs():
-        return {
-            "150M": TierSpec("DragonCode-150M", "150M", 768, 12, 3072, 12, None, 3_000_000_000),
-            "387M": TierSpec("DragonCode-387M", "387M", 1024, 22, 4096, 16, None, 7_740_000_000),
-            "787M": TierSpec("DragonCode-787M", "787M", 1280, 30, 5120, 16, 16, 15_740_000_000),
-            "1.2B": TierSpec("DragonCode-1.2B", "1.2B", 1536, 34, 6144, 24, 24, 24_000_000_000),
-            "2.4B": TierSpec("DragonCode-2.4B", "2.4B", 2048, 38, 8192, 32, 32, 48_000_000_000),
-        }
-
-    def _hf_repos():
-        return {
-            "150M": "DragonLimited/DragonCode-150M",
-            "387M": "DragonLimited/DragonCode-387M",
-            "787M": "DragonLimited/DragonCode-787M",
-            "1.2B": "DragonLimited/DragonCode-1.2B",
-            "2.4B": "DragonLimited/DragonCode-2.4B",
-            "family": "DragonLimited/DragonCode-Family",
-        }
-
-    # ---- 方案A data source: 純外部 streaming，絕不讀 Tokenized-Pretrain ---- #
-    DATA_SOURCES = [   # strict fallback 優先序
-        "HuggingFaceFW/fineweb-edu",
-        "cerebras/SlimPajama-627B",
-        "EleutherAI/the_pile",
-    ]
-
-    def _tokenize_string(text, tok):
-        ids = tok.encode(text)
-        if hasattr(ids, "ids"):
-            ids = ids.ids
-        return list(ids)
-
-    class StreamingCorpus:
-        """純外部 HF streaming 數據源（方案A）。單向前向 pass，唔讀 Tokenized repo。"""
-        def __init__(self, token, seq_len, budget_tokens, offset_tokens=0):
-            self.token = token; self.seq = seq_len; self.budget = budget_tokens
-            self.off = offset_tokens; self.data_used = 0
-
-        def __iter__(self):
-            from datasets import load_dataset
-            buf = deque(); cursor = self.off; need = self.budget
-            for src in DATA_SOURCES:
-                if cursor >= need:
-                    return
-                try:
-                    ds = load_dataset(src, split="train", streaming=True, token=self.token)
-                    logger.info("Streaming %s (cursor=%d budget=%d)", src, cursor, need)
-                except Exception as e:
-                    logger.warning("source %s failed (%s); next", src, repr(e)); continue
-                try:
-                    for row in ds:
-                        txt = row.get("text") or ""
-                        for v in _tokenize_string(txt, _get_tokenizer()):
-                            cursor += 1
-                            if cursor <= self.off:
-                                continue
-                            buf.append(int(v))
-                            while len(buf) >= self.seq:
-                                window = [buf.popleft() for _ in range(self.seq)]
-                                self.data_used = cursor
-                                yield torch.tensor(window, dtype=torch.long), cursor
-                                if cursor >= need:
-                                    return
-                except Exception as e:
-                    logger.warning("source %s mid-stream (%s); next", src, repr(e)); continue
-            logger.warning("All external sources exhausted at cursor %d (budget %d)", cursor, need)
-
-    def _make_batches(gen, batch_size, seq_len):
-        batch = []
-        for token_tensor, tok_cursor in gen:
-            batch.append(token_tensor)
-            if len(batch) == batch_size:
-                x = torch.stack(batch)[:, :-1].contiguous(); y = torch.stack(batch)[:, 1:].contiguous()
-                yield x, y, tok_cursor, True; batch = []
-        if batch:
-            x = torch.stack(batch)[:, :-1].contiguous(); y = torch.stack(batch)[:, 1:].contiguous()
-            yield x, y, tok_cursor, False
-
-
-    return (
-        DRAGON_HF_TOKEN,
-        MAX_SEQ_LEN,
-        StreamingCorpus,
-        json,
-        logger,
-        math,
-        os,
-        shutil,
-        tempfile,
-        time,
-        torch,
-    )
-
-
-@app.cell
-def _(DRAGON_HF_TOKEN, json, logger, os, shutil, tempfile, torch):
-    # =============================================================== #
-    #  B) DragonCode 方案A — HF checkpoint push / resume
-    # =============================================================== #
-    # checkpoint 直接推送至該 tier 嘅 HF public model repo；中斷後 resume
-    # 由 HF 拉取，完全唔依賴 Molab 本地儲存。
-
-    class _TrainState:
-        def __init__(self, step=0, tokens_seen=0, epoch=0, global_seed=0,
-                     best_val_loss=float("inf"), metadata=None):
-            self.step = step; self.tokens_seen = tokens_seen; self.epoch = epoch
-            self.global_seed = global_seed; self.best_val_loss = best_val_loss
-            self.metadata = metadata or {}
-        def to_dict(self):
-            return {"step": self.step, "tokens_seen": self.tokens_seen,
-                    "epoch": self.epoch, "global_seed": self.global_seed,
-                    "best_val_loss": self.best_val_loss, "metadata": self.metadata}
-
-
-    def _find_latest_hf_checkpoint(repo):
-        from huggingface_hub import list_repo_commits, hf_hub_download
-        try:
-            commits = _run_with_retry(lambda: list_repo_commits(repo, token=DRAGON_HF_TOKEN), label="list-ckpt")
-        except Exception as e:
-            logger.warning("list commits %s: %s", repo, e); return None
-        for c in commits:
-            sha = c.commit_id
-            try:
-                tree = _hf_api().repo_info(repo, revision=sha, repo_type="model").siblings
-                names = [t.rfilename for t in tree]
-                if "train_state.json" in names:
-                    st = hf_hub_download(repo, "train_state.json", revision=sha, token=DRAGON_HF_TOKEN)
-                    with open(st) as f: state = json.load(f)
-                    return sha, state
-            except Exception:
-                continue
-        return None
-
-
-    def _push_checkpoint_hf(repo, ckpt_paths, commit_message, size_cls):
-        api = _hf_api(); api.create_repo(repo, repo_type="model", exist_ok=True, private=False)
-        tmpdir = tempfile.mkdtemp(prefix="dc_ck_")
-        dst = os.path.join(tmpdir, "checkpoints", f"DragonCode-{size_cls}")
-        for name, src in ckpt_paths.items():
-            if os.path.isdir(src):
-                shutil.copytree(src, os.path.join(dst, name), dirs_exist_ok=True)
-            else:
-                os.makedirs(os.path.dirname(os.path.join(dst, name)), exist_ok=True)
-                shutil.copy(src, os.path.join(dst, name))
-        def _upload():
-            api.upload_folder(folder_path=os.path.join(tmpdir, "checkpoints"), repo_id=repo,
-                              repo_type="model", commit_message=commit_message,
-                              allow_duplicate_filename=True)
-        _run_with_retry(_upload, label="push-ckpt", tries=6)
-        shutil.rmtree(tmpdir, ignore_errors=True)
-        logger.info("Pushed checkpoint -> %s (%s)", repo, commit_message)
-
-
-    def _load_local_ckpt(path, model, optimizer, scheduler):
-        ck = torch.load(path, map_location="cpu", weights_only=False)
-        if model is not None: model.load_state_dict(ck["model"])
-        if optimizer is not None and ck.get("optimizer") is not None: optimizer.load_state_dict(ck["optimizer"])
-        if scheduler is not None and ck.get("lr_scheduler") is not None: scheduler.load_state_dict(ck["lr_scheduler"])
-        return ck
-
-
-    return
-
-
-@app.cell
-def _(
-    DRAGON_HF_TOKEN,
-    MAX_SEQ_LEN,
-    StreamingCorpus,
-    json,
-    logger,
-    math,
-    os,
-    shutil,
-    tempfile,
-    time,
-    torch,
-):
-    # =============================================================== #
-    #  C) DragonCode 方案A — pretrain 主迴圈（主執行緒，無後台thread）
-    # =============================================================== #
-    # 全部 CPU/GPU 繁重邏輯喺 Notebook Cell 主執行緒循序執行（Molab 合規）；
-    # 不開 daemon / 後台業務 thread。
-
-    def _lr_warmup(step, warm, total):
-        if step < warm: return step / max(1.0, warm)
-        return max(0.0, 1.0 - (step - warm) / max(1, total - warm))
-
-
-    def _try_resume(model, optimizer, scheduler, kwargs, repo, size_cls, step, tokens_seen, dev):
-        if not kwargs.get("resume", True):
-            return model, optimizer, scheduler, step, tokens_seen, 0
-        local_dir = kwargs.get("local_ckpt_dir", "/home/marimo/DragonCode/checkpoints")
-        ck = os.path.join(local_dir, f"DragonCode-{size_cls}", "checkpoint-latest.pt")
-        loaded = False
-        if os.path.exists(ck):
-            logger.info("Resume LOCAL %s", ck)
-            data = _load_local_ckpt(ck, model, optimizer, scheduler)
-            st = data["train_state"]; step = st["step"]; tokens_seen = st["tokens_seen"]
-            loaded = True
-        else:
-            hit = _find_latest_hf_checkpoint(repo)
-            if hit:
-                sha, st = hit
-                logger.info("Resume HF commit %s (step=%s tokens=%s)", sha, st.get("step"), st.get("tokens_seen"))
-                tmp = tempfile.mkdtemp(prefix="dc_rs_")
-                from huggingface_hub import snapshot_download
-                _run_with_retry(lambda: snapshot_download(repo, revision=sha, token=DRAGON_HF_TOKEN, local_dir=tmp),
-                                label="snap-resume")
-                ck = os.path.join(tmp, "checkpoints", f"DragonCode-{size_cls}", "checkpoint-latest.pt")
-                if os.path.exists(ck):
-                    data = _load_local_ckpt(ck, model, optimizer, scheduler)
-                    stt = data["train_state"]; step = stt["step"]; tokens_seen = stt["tokens_seen"]
-                    loaded = True
-                shutil.rmtree(tmp, ignore_errors=True)
-            else:
-                logger.info("No checkpoint found; starting fresh (%s)", kwargs.get("tier_name", "?"))
-        if loaded:
-            _seed_everything(kwargs.get("seed", 42) + step)
-        return model, optimizer, scheduler, step, tokens_seen, tokens_seen
-
-
-    def _save_and_push(model, optimizer, scheduler, step, tokens_seen, epoch, seed, kwargs, repo, size_cls, tier):
-        local_dir = kwargs.get("local_ckpt_dir", "/home/marimo/DragonCode/checkpoints")
-        cdir = os.path.join(local_dir, f"DragonCode-{size_cls}"); os.makedirs(cdir, exist_ok=True)
-        ckpt_path = os.path.join(cdir, "checkpoint-latest.pt")
-        st = _TrainState(step=step, tokens_seen=tokens_seen, epoch=epoch, global_seed=seed,
-                         best_val_loss=float("inf"), metadata={"tier": tier.size_cls, "name": tier.name})
-        torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(),
-                    "lr_scheduler": scheduler.state_dict(), "train_state": st.to_dict(),
-                    "format": "dragoncode-checkpoint-v1"}, ckpt_path)
-        sj = os.path.join(cdir, "train_state.json")
-        with open(sj, "w") as f: json.dump(st.to_dict(), f, indent=2)
-        _push_checkpoint_hf(repo, {"checkpoint-latest.pt": ckpt_path, "train_state.json": sj},
-                            f"[{tier.name}] checkpoint step={step} tokens={tokens_seen:,}", size_cls)
-        logger.info("Saved+push step=%d tokens=%d", step, tokens_seen)
-
-
-    def train_pretrain(tier, kwargs):
-        """主執行緒訓練單一 tier 至 Chinchilla 預算（純外部 streaming）。"""
-        torch.backends.cuda.matmul.allow_tf32 = False
-        torch.backends.cudnn.allow_tf32 = False
-        dev = kwargs["device"]; seed = kwargs.get("seed", 42)
-        seq_len = kwargs.get("seq_len", MAX_SEQ_LEN); batch_size = kwargs.get("batch_size", 8)
-        grad_acc = kwargs.get("grad_accum", 1); lr = kwargs.get("lr", 3e-4)
-        warmup_frac = kwargs.get("warmup_frac", 0.01); max_grad_norm = kwargs.get("max_grad_norm", 1.0)
-        save_every = kwargs.get("save_every_steps", 2000); log_every = kwargs.get("log_every", 50)
-        size_cls = tier.size_str; repo = kwargs["hf_repo"]; token_budget = tier.chinchilla_tokens
-        _seed_everything(seed); _get_tokenizer()
-
-        from transformers import LlamaConfig, LlamaForCausalLM
-        cfg = _build_model_config(tier.hidden, tier.layers, tier.ffn, tier.heads, tier.kv_heads, name=tier.name)
-        cfg.attn_implementation = "flash_attention_2"
-        model = LlamaForCausalLM(cfg).to(dev).to(torch.bfloat16)
-        n_params = sum(p.numel() for p in model.parameters())
-        logger.info("Pretrain %s on %s — %.1fM params, budget %d tokens", tier.name, dev, n_params/1e6, token_budget)
-
-        optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1, fused=True)
-        tokens_per_step_eff = batch_size * grad_acc * (seq_len - 1)
-        total_steps_est = math.ceil(token_budget / tokens_per_step_eff)
-        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda s: _lr_warmup(s, warmup_frac * total_steps_est, total_steps_est))
-
-        # 方案A：純外部 streaming
-        src = StreamingCorpus(DRAGON_HF_TOKEN, seq_len, token_budget, offset_tokens=0)
-        logger.info("數據源[方案A]: 跳過 DragonCode-Tokenized-Pretrain; 全部 token 流式讀取 FineWeb-Edu → SlimPajama → The-Pile. Budget=%d", token_budget)
-
-        step = 0; tokens_seen = 0; epoch = 0
-        model, optimizer, scheduler, step, tokens_seen, offset_toks = _try_resume(
-            model, optimizer, scheduler, kwargs, repo, size_cls, step, tokens_seen, dev)
-        if offset_toks:
-            src.off = offset_toks
-            logger.info("Resume offset=%d tokens（純外部流，不含舊shard）", offset_toks)
-
-        model.train(); gen = _make_batches(iter(src), batch_size, seq_len)
-        prog_start = time.time()
-        while tokens_seen < token_budget:
-            local_steps = 0
-            for x, y, tok_cursor, is_full in gen:
-                x = x.to(dev); y = y.to(dev)
-                out = model(x, labels=y); loss = out.loss / grad_acc; loss.backward()
-                local_steps += 1
-                tokens_seen = min(tok_cursor, token_budget)
-                cond = (is_full and (step + 1) % grad_acc == 0) or (not is_full)
-                if cond:
-                    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
-                    optimizer.step(); scheduler.step(); optimizer.zero_grad(set_to_none=True)
-                    step += 1
-                if step % log_every == 0 and step > 0:
-                    el = time.time() - prog_start; tps = tokens_seen / max(el, 1e-9)
-                    logger.info("[%s] step=%d tokens=%d/%d (%.2f%%) loss=%.4f lr=%.2e tps=%.0f",
-                                tier.name, step, tokens_seen, token_budget,
-                                100.0 * tokens_seen / token_budget, out.loss.item() * grad_acc,
-                                scheduler.get_last_lr()[0], tps)
-                if step % save_every == 0 and step > 0:
-                    _save_and_push(model, optimizer, scheduler, step, tokens_seen, epoch, seed, kwargs, repo, size_cls, tier)
-                if tokens_seen >= token_budget: break
-                if local_steps > 10_000_000: break
-            if tokens_seen < token_budget:
-                logger.warning("All 方案A external sources exhausted at token %d (budget %d) — stop %s.",
-                               tokens_seen, token_budget, tier.name)
-            break
-        _save_and_push(model, optimizer, scheduler, step, tokens_seen, epoch, seed, kwargs, repo, size_cls, tier)
-        logger.info("PRETRAIN DONE %s: %d/%d tokens", tier.name, tokens_seen, token_budget)
-        return {"model": model, "step": step, "tokens_seen": tokens_seen}
-
+    # ===================================================================== #
+    #  DragonCode LLM Family — production notebook (thin wrapper)
+    #  (c) 2026 Dragon Limited. All rights reserved.
+    #
+    #  This notebook is a THIN LAUNCHER over scripts/run_dragoncode.py.
+    #  The single source of truth for training logic lives in
+    #  ~/DragonCode/scripts/*.py  (NOT duplicated in cells).
+    #
+    #  Industrial-standard guarantees enforced by the scripts:
+    #   * 100% code-domain data — codeparrot/codeparrot-clean +
+    #     open-r1/codeforces-cots (permissive licenses only). NO generic web.
+    #   * 4 training bugs fixed: grad-accum dead-loop (local_steps counter),
+    #     lr=0 (token-progress schedule), bytes JSON serialization (base64),
+    #     allow_duplicate_filename (removed for hf_hub 1.24.0).
+    #   * resume-aware + idempotent stage markers (never restart from zero).
+    # ===================================================================== #
+    import os, sys, subprocess, json, time
+
+    SCRIPT_DIR = os.path.expanduser("~/DragonCode/scripts")
+    CONFIG_DIR = os.path.expanduser("~/DragonCode/configs")
+    LOG_DIR = os.path.expanduser("~/DragonCode/logs")
+
+    # Auth: the scripts force the "dragonlimited" account token internally, so
+    # we only need to make sure the box has *some* HF_TOKEN exported.
+    os.environ.setdefault("HF_TOKEN", os.environ.get("HF_TOKEN", ""))
+    os.environ["HF_HOME"] = os.path.expanduser("~/.cache/huggingface")
+
+    # 5-model tier order (Chinchilla-optimal, 20 tokens/param).
+    TIER_ORDER = ["150M", "387M", "787M", "1.2B", "2.4B"]
+    CHINCHILLA = {
+        "150M": 3_000_000_000,
+        "387M": 7_740_000_000,
+        "787M": 15_740_000_000,
+        "1.2B": 24_000_000_000,
+        "2.4B": 48_000_000_000,
+    }
 
-    return (train_pretrain,)
+    return (SCRIPT_DIR, CONFIG_DIR, LOG_DIR, TIER_ORDER, CHINCHILLA, os, subprocess, time)
 
 
 @app.cell
-def _(DRAGON_HF_TOKEN, torch, train_pretrain):
-    # =============================================================== #
-    #  D) DragonCode 方案A — runner：偵測下一個未完成 tier
-    # =============================================================== #
-
-    TIER_ORDER = ["150M", "387M", "787M", "1.2B", "2.4B"]
-
-
-    def hf_has_model(repo):
-        from huggingface_hub import list_repo_files
-        try:
-            files = list_repo_files(repo, repo_type="model", token=DRAGON_HF_TOKEN)
-            return any("config.json" in f or f.endswith(".safetensors") for f in files)
-        except Exception:
+def _(SCRIPT_DIR, CONFIG_DIR, LOG_DIR, os, subprocess, time):
+    # ===================================================================== #
+    #  Stage runner — delegates every stage to scripts/run_dragoncode.py
+    #  (the single source of truth). Stages per tier (domain-only 12-step):
+    #    pretrain → sft → dpo → golf → merge → verify → gguf
+    #  150M/387M: pretrain only. 2.4B: no DPO this cycle.
+    # ===================================================================== #
+    TIER_STAGES = {
+        "150M": ["pretrain"],
+        "387M": ["pretrain"],
+        "787M": ["pretrain", "sft", "dpo", "golf", "merge", "verify", "gguf"],
+        "1.2B": ["pretrain", "sft", "dpo", "golf", "merge", "verify", "gguf"],
+        "2.4B": ["pretrain", "sft", "golf", "merge", "verify", "gguf"],
+    }
+
+    def _run(args, logfile):
+        """Run a CLI stage, streaming stdout to its own log file (tail-friendly)."""
+        os.makedirs(LOG_DIR, exist_ok=True)
+        with open(logfile, "a") as lf:
+            lf.write(f"\n=== {time.strftime('%Y-%m-%dT%H:%M:%SZ')} {' '.join(args)} ===\n")
+            lf.flush()
+            proc = subprocess.Popen(
+                args, cwd=SCRIPT_DIR,
+                stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                text=True, bufsize=1,
+            )
+            assert proc.stdout is not None
+            for line in proc.stdout:
+                lf.write(line); lf.flush()
+                print(line, end="", flush=True)
+            rc = proc.wait()
+        return rc
+
+    def run_stage(tier, stage):
+        script = f"dragoncode_{stage}.py"
+        config = os.path.join(CONFIG_DIR, f"DragonCode-{tier}.yaml")
+        cmd = [sys.executable, script, "--tier", tier, "--config", config]
+        logfile = os.path.join(LOG_DIR, f"DragonCode-{tier}-{stage}.log")
+        print(f"\n[DRIVE] {tier}/{stage} -> {' '.join(cmd)}", flush=True)
+        rc = _run(cmd, logfile)
+        if rc != 0:
+            print(f"[DRIVE] {tier}/{stage} FAILED rc={rc} (see {logfile})", flush=True)
             return False
+        print(f"[DRIVE] {tier}/{stage} OK", flush=True)
+        return True
 
+    return (TIER_STAGES, run_stage)
 
-    def next_tier():
-        for t in TIER_ORDER:
-            if not hf_has_model(_hf_repos()[t]):
-                return t
-        return None
 
-
-    def run_pipeline(tier=None, device="cuda", dry_run=False):
-        """喺主執行緒訓練下一個未完成 tier（或指定 tier）。"""
-        if tier is None:
-            tier = next_tier()
-        if tier is None:
-            print("[RUN] 所有 tier 都已有公開 weights。無需再做。")
-            return {"tier": None, "tokens_seen": None}
-        spec = _tier_specs()[tier]
-        kwargs = dict(tier_name=spec.name, hf_repo=_hf_repos()[tier], device=device, seed=42)
-        print(f"[RUN] PRETRAIN {spec.name} — budget {spec.chinchilla_tokens/1e9:.2f}B tokens (方案A外部streaming)")
-        if dry_run:
-            print("[RUN] dry_run:", spec.name); return {"tier": tier, "tokens_seen": 0}
-        out = train_pretrain(spec, kwargs)
-        print(f"[RUN] PRETRAIN {tier} DONE: tokens_seen={out['tokens_seen']}")
-        torch.cuda.empty_cache()
-        return {"tier": tier, "tokens_seen": out["tokens_seen"]}
+@app.cell
+def _(TIER_ORDER, TIER_STAGES, run_stage, time):
+    # ===================================================================== #
+    #  Full pipeline — sequential, resume-safe, single epoch per tier.
+    #  Unique stop condition: user interrupt. No early-exit, no auto-exit.
+    # ===================================================================== #
+    def run_all():
+        for tier in TIER_ORDER:
+            for stage in TIER_STAGES[tier]:
+                ok = run_stage(tier, stage)
+                attempt = 0
+                while not ok and attempt < 3:
+                    attempt += 1
+                    time.sleep(8)
+                    print(f"[DRIVE] {tier}/{stage} retry {attempt}/3", flush=True)
+                    ok = run_stage(tier, stage)
+                if not ok:
+                    print(f"[DRIVE] {tier}/{stage} failed after retries — stopping pipeline", flush=True)
+                    return
+        print("[DRIVE] All 5 models complete.", flush=True)
 
 
-    return (run_pipeline,)
+    return (run_all,)
 
 
 @app.cell
-def _(run_pipeline):
-    # =============================================================== #
-    #  E) DragonCode 方案A — 手動啟動 cell
-    # =============================================================== #
-    # 請手動點擊呢個 cell 嘅 Run (▶) 開始 150M 預訓練。
-    # 若中斷，重新 Run 即由 HF checkpoint resume（唔依賴 Molab 本地）。
-    run_pipeline(device="cuda")
+def _(run_all):
+    # ===================================================================== #
+    #  LAUNCH — press Run (▶) on this cell to train all 5 models end to end.
+    # ===================================================================== #
+    run_all()
 
     return