File size: 4,362 Bytes
8f46582 | 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 118 119 120 121 122 123 124 125 126 127 | #!/usr/bin/env python3
"""Generate the L15 "make it work" sweep.
Diagnosis these configs are built from (see logs/diag_L15_*.log):
* Stage 0 is NOT the bottleneck. hop-1 sits at frontier 1.00 / ce_score 0.94-0.98
for the entire run and never degrades. The pinned stage-0 run saturates at
ce_score 0.945 by epoch 45 and only reaches 0.961 by epoch 150.
* The bottleneck is that ce_score decays monotonically with depth
(0.96 -> 0.44 by hop 15) and nothing in the live arms pushed back:
btNONE had no backtracking, btCE050 fires below 0.50 which only hop 15
ever reaches. Promotion on frontier@0.95 keeps advancing while balance rots.
* The L10 arm that worked used ce_score on BOTH gates at 0.95 and ended flat
(0.96-0.98 at every one of its 10 hops). No L15 arm has used that recipe.
Three differences from the L10 winner are therefore corrected here: both gates on
ce_score@0.95, bf16 off (the balance score compares near-equal log-probabilities,
where bf16's ~8-bit mantissa is a real noise source), and the stage-0 warm start
points at checkpoint_150 instead of the stale checkpoint_99.
The 4-layer arms warm-start by layer expansion: a strict=False load puts the
2-layer checkpoint's embeddings, lm_head, ln_f and blocks h.0/h.1 into the
4-layer model and leaves h.2/h.3 random. The only existing 4-layer checkpoint
(L12 stage-0) is unusable -- it never left chance at hop-1 frontier 0.027.
"""
import os
BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ARGS = os.path.join(BASE, "args")
STAGE0_CKPT = "ckpts/star-coconut-L15-bfs-stage0-warm/checkpoint_150"
M2L = "configs/symbol-2layer-8head-768dim-L20.json"
M4L = "configs/symbol-4layer-8head-768dim-L20.json"
TEMPLATE = """# {desc}
project: coconut
save_path: ckpts
name: "{name}"
only_eval: False
coconut: True
cot: False
no_thoughts: False
no_cot: False
c_thought: 1
max_latent_stage: {max_stage}
pad_latent_to_max: True
accuracy_staging: True
init_stage: {init_stage}
promote_metric: ce_score
promote_threshold: {thr}
promote_on_current_only: False
epochs_per_stage: 25
backtrack: {backtrack}
backtrack_metric: ce_score
backtrack_detect_threshold: {thr}
remember_rate: 0.3
revert_staging: False
eval_every: 5
log_every: 5
perhop_val_samples: 256
perhop_train_samples: 64
eval_print_full: False
backprop_depth: null
train_size: {train_size}
save_only_improve: False
save_every: 200
uniform_prob: 0.1
model_id: {model_id}
load_model_path: {ckpt}
seed: 0
resume: 0
bf16: False
train_path: data/star_2arm_L15_train_fo_bfs.json
val_path: data/star_2arm_L15_valid_fo_bfs.json
reset_optimizer: False
batch_size_training: 128
debug: False
gradient_accumulation_steps: 1
num_epochs: 4000
lr: !!float "1e-4"
grad_clip: !!float "1.0"
warmup_steps: 200
weight_decay: 0.01
bfs_variant: True
"""
# (name, model, ce threshold, train_size, desc)
ARMS = [
("L15_push_2L_ce95_50k", M2L, 0.95, 50000,
"Exact L10-winner recipe at L15: 2-layer, ce_score@0.95 on both gates, fp32, 50k."),
("L15_push_2L_ce95_100k", M2L, 0.95, 0,
"Data lever: same as the mirror arm but all 100k graphs."),
("L15_push_4L_ce95_50k", M4L, 0.95, 50000,
"Capacity lever: 4-layer via layer-expansion warm start, ce_score@0.95, 50k."),
("L15_push_4L_ce90_50k", M4L, 0.90, 50000,
"Capacity + looser gate: 4-layer, ce_score@0.90, 50k."),
]
written = []
for name, model_id, thr, train_size, desc in ARMS:
body = TEMPLATE.format(desc=desc, name=name, max_stage=15, init_stage=1,
thr=thr, backtrack="True", train_size=train_size,
model_id=model_id, ckpt=STAGE0_CKPT)
path = os.path.join(ARGS, name + ".yaml")
open(path, "w").write(body)
written.append(path)
# 4-layer stage-0 builder: insurance in case layer expansion does not hold.
# Pinned at stage 0 (max_latent_stage == init_stage == 0 so promote can never fire).
s0 = TEMPLATE.format(
desc="4-layer L15 stage-0, pinned. Produces a native 4-layer stage-0 checkpoint "
"in case layer-expansion warm start degrades hop-1.",
name="L15_push_s0_4L", max_stage=0, init_stage=0, thr=0.95,
backtrack="False", train_size=0, model_id=M4L, ckpt=STAGE0_CKPT)
path = os.path.join(ARGS, "L15_push_s0_4L.yaml")
open(path, "w").write(s0)
written.append(path)
for p in written:
print("wrote", os.path.relpath(p, BASE))
|