Add DragonCode family pretraining marimo notebook (方案A data source; token from env)

notebook.py

@@ -0,0 +1,465 @@
+import marimo
+
+__generated_with = "0.24.0"
+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}
+
+
+    return (train_pretrain,)
+
+
+@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:
+            return False
+
+
+    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"]}
+
+
+    return (run_pipeline,)
+
+
+@app.cell
+def _(run_pipeline):
+    # =============================================================== #
+    #  E) DragonCode 方案A — 手動啟動 cell
+    # =============================================================== #
+    # 請手動點擊呢個 cell 嘅 Run (▶) 開始 150M 預訓練。
+    # 若中斷，重新 Run 即由 HF checkpoint resume（唔依賴 Molab 本地）。
+    run_pipeline(device="cuda")
+
+    return
+
+
+if __name__ == "__main__":
+    app.run()