| """Generate per-run training configs from a base yaml. |
| |
| The design covers 9 mixtures (5 monolingual + 4 EN-anchored bilingual |
| mixtures) x 2 tokenizer conditions (starved / destarved) for one chosen BPE |
| flavor = 18 runs at 30B tokens each. Four of them (EN-DE and EN-AR under both |
| conditions) are additionally extended to 100B via a WSD cooldown branch off a |
| long stable trunk -- the cheap-extension mechanism. |
| |
| A GPU-hour budget shortfall (2026-07-16: ~1200h available vs ~2615h needed |
| for the full 30-run all-pairs matrix) forced cutting the 12 non-EN-anchored |
| pairwise mixtures (de-fr, de-ar, de-zh, fr-ar, fr-zh, ar-zh) that a prior |
| expansion had added here; reverted to the original EN-anchored-only design. |
| Their partial checkpoints (0-12.4B tokens) are left on disk, unused. |
| |
| A run's config is this base block plus: name, langs, probs, tok_name, schedule |
| overrides, and (for cooldowns) a `branch` pointing at a trunk checkpoint. |
| |
| Run-name convention (kept in sync with eval.bts): |
| "<mix>__<tok>" e.g. en-ar__bl_destarved (30B deliverable) |
| "<mix>__<tok>__trunk" long stable trunk for extension |
| "<mix>__<tok>__100b" 100B cooldown deliverable |
| where <mix> is "en" or "en-ar" etc., <tok> is "<flavor>_<condition>". |
| """ |
| import copy |
| from pathlib import Path |
|
|
| from .langs import LANGS, ANCHOR, PARTNERS, tok_conditions, tok_name |
| from .paths import run_dir |
|
|
| BILINGUAL_MIXTURES = [f"{ANCHOR}-{p}" for p in PARTNERS] |
| MIXTURES = [*LANGS, *BILINGUAL_MIXTURES] |
| |
| |
| |
| EXTENDED_MIXTURES = [f"{ANCHOR}-de", f"{ANCHOR}-ar"] |
| TRUNK_TOKENS = 80e9 |
| |
| COOLDOWNS = {30e9: (24e9, 6e9), 100e9: (80e9, 20e9)} |
|
|
|
|
| def load_base(path) -> dict: |
| from . import _yaml |
| return _yaml.load(path) |
|
|
|
|
| def mix_langs(mix: str) -> list[str]: |
| return mix.split("-") |
|
|
|
|
| def _stamp(base: dict, name: str, mix: str, tok: str) -> dict: |
| cfg = copy.deepcopy(base) |
| cfg["name"] = name |
| cfg["langs"] = mix_langs(mix) |
| cfg["tok_name"] = tok |
| cfg.setdefault("probs", None) |
| return cfg |
|
|
|
|
| def wsd_run(base, mix, tok) -> dict: |
| """Self-contained warmup+stable+decay to 30B (base schedule as-is).""" |
| return _stamp(base, f"{mix}__{tok}", mix, tok) |
|
|
|
|
| def trunk_run(base, mix, tok) -> dict: |
| cfg = _stamp(base, f"{mix}__{tok}__trunk", mix, tok) |
| warm = cfg["schedule"]["warmup_tokens"] |
| cfg["schedule"] = {**cfg["schedule"], |
| "stable_tokens": TRUNK_TOKENS - warm, "decay_tokens": 0} |
| cfg["train"] = {**cfg["train"], |
| "stable_marks": sorted({b for b, _ in COOLDOWNS.values()})} |
| return cfg |
|
|
|
|
| def cooldown_run(base, mix, tok, total_tokens) -> dict: |
| branch_at, decay = COOLDOWNS[total_tokens] |
| suffix = "" if total_tokens == 30e9 else f"__{int(total_tokens/1e9)}b" |
| cfg = _stamp(base, f"{mix}__{tok}{suffix}", mix, tok) |
| cfg["schedule"] = {**cfg["schedule"], "warmup_tokens": 0, |
| "stable_tokens": 0, "decay_tokens": decay} |
| trunk_ckpt = (run_dir(f"{mix}__{tok}__trunk") / "checkpoints" |
| / f"stable_{int(branch_at/1e6)}M.pt") |
| cfg["branch"] = {"from": str(trunk_ckpt), "load_optim": True} |
| return cfg |
|
|
|
|
| def all_runs(base: dict, flavor: str) -> dict[str, dict]: |
| """Every run needed for the full experiment, keyed by run name.""" |
| runs: dict[str, dict] = {} |
| for mix in MIXTURES: |
| for cond in tok_conditions(flavor): |
| tok = tok_name(flavor, cond) |
| if mix in EXTENDED_MIXTURES: |
| runs[f"{mix}__{tok}__trunk"] = trunk_run(base, mix, tok) |
| for total in COOLDOWNS: |
| r = cooldown_run(base, mix, tok, total) |
| runs[r["name"]] = r |
| else: |
| r = wsd_run(base, mix, tok) |
| runs[r["name"]] = r |
| return runs |
|
|
|
|
| def all_30b_runs(base: dict, flavor: str) -> dict[str, dict]: |
| """The 30 independent 30B cells, without 100B-extension trunks. |
| |
| This mode is useful when the immediate experiment scope is only 30B. In |
| the full matrix, four of these names are cooldown branches from shared |
| trunks; here every name instead gets the base self-contained WSD schedule. |
| """ |
| runs: dict[str, dict] = {} |
| for mix in MIXTURES: |
| for cond in tok_conditions(flavor): |
| tok = tok_name(flavor, cond) |
| run = wsd_run(base, mix, tok) |
| runs[run["name"]] = run |
| return runs |
|
|
|
|
| def get_run(base_path, flavor: str, name: str, only_30b: bool = False) -> dict: |
| base = load_base(base_path) |
| runs = all_30b_runs(base, flavor) if only_30b else all_runs(base, flavor) |
| if name not in runs: |
| raise KeyError(f"unknown run '{name}'. Available:\n " + |
| "\n ".join(sorted(runs))) |
| return runs[name] |
|
|
|
|
| def list_runs(base_path, flavor: str, only_30b: bool = False) -> list[str]: |
| base = load_base(base_path) |
| runs = all_30b_runs(base, flavor) if only_30b else all_runs(base, flavor) |
| return sorted(runs) |
|
|
