1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
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 = 440000 # Match fw135m_bsign; casc_bp_train.py loops inclusively: 440,001 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,
1000,
)
for lr in LR_SWEEP
]
else:
runs = [(f"fw135m_bp_s{args.seed}", args.lr, FULL_STEPS_ARG, 1000)]
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()
|
