diff options
| -rw-r--r-- | docs/campaign/FW135M_BP_BASELINE.md | 89 | ||||
| -rw-r--r-- | ep_run/baseline_configs/fw135m_matched_bp.json | 42 | ||||
| -rw-r--r-- | ep_run/fw135m_baseline.py | 117 |
3 files changed, 248 insertions, 0 deletions
diff --git a/docs/campaign/FW135M_BP_BASELINE.md b/docs/campaign/FW135M_BP_BASELINE.md new file mode 100644 index 0000000..3f81ba8 --- /dev/null +++ b/docs/campaign/FW135M_BP_BASELINE.md @@ -0,0 +1,89 @@ +# FW135M BP Twin Run Sheet + +Source of truth: `docs/BASELINE_SPEC.md`. + +Status: commands configured on branch `xiang`; no training has started. + +## First width-only rung + +- Shape: `L12/C768/H12/T256` +- Head dimension: `64` +- Exact local parameter count: `135,303,936` +- Data: existing FineWeb-Edu 32k bins +- Batch: `B24` +- Target: `20N = 2,706,078,720` tokens +- Complete-batch exposure: `2,706,081,792` tokens +- Trainer argument: `--steps 440442` (440,443 inclusive updates) + +## What changes from 72M + +- Width: `512 -> 768` +- Heads: `8 -> 12` +- Parameters: `72.11M -> 135.30M` +- Run length: increases to preserve approximately 20 tokens/parameter + +## What stays identical + +- L12 depth and 64-dimensional heads +- T256 and B24 +- FineWeb-Edu data and 32k tokenizer +- OLMo2-style model implementation +- Muon hybrid optimizer and cosine schedule +- Weight decay 0.1 and BF16 autocast +- Seed list and validation cadence +- Existing `casc_bp_train.py` code path + +## LR sweep + +At this smallest new rung only: + +- `7e-4` +- `1e-3` +- `1.4e-3` + +Only Adam-side LR changes. Carry the winner to larger rungs. + +No result may be quoted with fewer than two seeds. Record best validation CE, final-step CE, and tail median. Save every 5,000 steps and log validation every 100 steps. + +## Dry-run launcher + +From `ep_run/`: + +```bash +python -m py_compile fw135m_baseline.py +python fw135m_baseline.py --mode smoke +python fw135m_baseline.py --mode sweep +python fw135m_baseline.py --mode full +``` + +Commands print by default. Add `--execute` only when ready. + +Smoke: + +```bash +python fw135m_baseline.py --mode smoke --execute +``` + +LR sweep: + +```bash +python fw135m_baseline.py --mode sweep --execute +``` + +Selected full run: + +```bash +python fw135m_baseline.py --mode full --lr <selected-lr> --seed 1 --execute +python fw135m_baseline.py --mode full --lr <selected-lr> --seed 2 --execute +``` + +W&B remains off until configured. Later: + +```bash +python fw135m_baseline.py --mode full --lr <selected-lr> \ + --wandb_project <project> --execute +``` + +## Final EP/BP comparison + +The EP arm must use the same shape, data, B/T, steps, optimizer family, selected shared settings, seeds, and evaluation cadence. Only the training rule and EP-specific `β` controls differ. diff --git a/ep_run/baseline_configs/fw135m_matched_bp.json b/ep_run/baseline_configs/fw135m_matched_bp.json new file mode 100644 index 0000000..0e22686 --- /dev/null +++ b/ep_run/baseline_configs/fw135m_matched_bp.json @@ -0,0 +1,42 @@ +{ + "source_of_truth": "docs/BASELINE_SPEC.md", + "purpose": "First approximately 2x BP twin rung after the completed 72M FineWeb model", + "changed_from_72m": { + "width": "512 -> 768", + "heads": "8 -> 12", + "parameters": "72,114,688 -> 135,303,936", + "training_updates": "440,443 updates to preserve approximately 20 tokens/parameter" + }, + "held_fixed_from_72m": { + "layers": 12, + "head_dim": 64, + "context": 256, + "global_sequence_batch": 24, + "data": "fineweb_edu", + "vocab": 32768, + "architecture": "local OLMo2-style", + "optimizer": "Muon on hidden matrices plus AdamW on remaining parameters", + "muon_lr": 0.02, + "adam_side_lr_center": 0.001, + "weight_decay": 0.1, + "warmup_steps": 500, + "schedule": "cosine to 0.1 of peak", + "precision": "bf16 autocast with fp32 parameters/states" + }, + "locked_run": { + "tag": "fw135m_bp", + "layers": 12, + "width": 768, + "heads": 12, + "head_dim": 64, + "context": 256, + "batch": 24, + "parameters": 135303936, + "target_tokens_20N": 2706078720, + "trainer_steps_argument": 440442, + "actual_updates": 440443, + "actual_tokens": 2706081792 + }, + "bp_lr_sweep_smallest_rung_only": [0.0007, 0.001, 0.0014], + "minimum_seeds_before_reporting": 2 +} diff --git a/ep_run/fw135m_baseline.py b/ep_run/fw135m_baseline.py new file mode 100644 index 0000000..ddcc5c0 --- /dev/null +++ b/ep_run/fw135m_baseline.py @@ -0,0 +1,117 @@ +"""Dry-run-by-default launcher for BASELINE_SPEC.md's first BP twin rung.""" +from __future__ import annotations + +import argparse +import subprocess +import sys + + +VOCAB = 32768 +LAYERS = 12 +WIDTH = 768 +HEADS = 12 +CONTEXT = 256 +BATCH = 24 +PARAMETERS = 135303936 +FULL_STEPS_ARG = 440442 # casc_bp_train.py loops inclusively: 440,443 updates +LR_SWEEP = ("7e-4", "1e-3", "1.4e-3") + + +def parameter_count(vocab=VOCAB, layers=LAYERS, width=WIDTH): + hidden = ((8 * width // 3) + 63) // 64 * 64 + return 2 * vocab * width + width + layers * ( + 4 * width * width + 3 * width * hidden + 4 * width + ) + + +def command(tag, lr, steps, warmup, seed, wandb_project): + return [ + sys.executable, + "casc_bp_train.py", + "--tag", + tag, + "--L", + str(LAYERS), + "--C", + str(WIDTH), + "--H", + str(HEADS), + "--T", + str(CONTEXT), + "--B", + str(BATCH), + "--steps", + str(steps), + "--lr", + lr, + "--warmup", + str(warmup), + "--amp", + "--olmo2", + "--wd", + "0.1", + "--opt", + "muon", + "--muon_lr", + "0.02", + "--cosine", + "--lr_min_ratio", + "0.1", + "--data", + "fineweb_edu", + "--seed", + str(seed), + "--save_every", + "5000", + "--log", + "100", + "--wandb", + wandb_project, + "--wandb_run", + tag if wandb_project else "", + ] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--mode", choices=("smoke", "sweep", "full"), default="smoke") + ap.add_argument("--execute", action="store_true", help="run commands; default only prints") + ap.add_argument("--lr", default="1e-3", help="selected LR for --mode full") + ap.add_argument("--seed", type=int, default=1) + ap.add_argument("--sweep_steps", type=int, default=FULL_STEPS_ARG) + ap.add_argument("--wandb_project", default="") + args = ap.parse_args() + + count = parameter_count() + if count != PARAMETERS: + raise RuntimeError(f"parameter formula returned {count}, expected {PARAMETERS}") + + if args.mode == "smoke": + runs = [(f"fw135m_bp_smoke_s{args.seed}", args.lr, 400, 50)] + elif args.mode == "sweep": + runs = [ + ( + f"fw135m_bp_lr{lr.replace('-', 'm').replace('.', 'p')}_s{args.seed}", + lr, + args.sweep_steps, + 500, + ) + for lr in LR_SWEEP + ] + else: + runs = [(f"fw135m_bp_s{args.seed}", args.lr, FULL_STEPS_ARG, 500)] + + print( + f"# L{LAYERS} C{WIDTH} H{HEADS} T{CONTEXT} B{BATCH} " + f"| {PARAMETERS:,} params | target 20N={20 * PARAMETERS:,} tokens", + flush=True, + ) + for tag, lr, steps, warmup in runs: + cmd = command(tag, lr, steps, warmup, args.seed, args.wandb_project) + print(" ".join(f'"{item}"' if item == "" else item for item in cmd), flush=True) + if args.execute: + subprocess.run(cmd, check=True) + + +if __name__ == "__main__": + main() |