p4_stop_probe: becomes the no-checkpointing soak at the real 262,144-token step shape (180 steps, 3 pushes, a forced cold resume, the 2M-token validation pass, peak memory asserted) as well as the stop-after-steps probe -- the measurement that answers D-017 with evidence
Browse files- kernels/p4_stop_probe.py +203 -177
kernels/p4_stop_probe.py
CHANGED
|
@@ -1,177 +1,203 @@
|
|
| 1 |
-
# P4 probe: does --stop-after-steps really produce a resumable, Hub-verified segment boundary?
|
| 2 |
-
#
|
| 3 |
-
# Why this exists and what it buys with ~0.4 GPU-hours: `--stop-after-steps` is the mechanism every Phase 4
|
| 4 |
-
# session ends on, and Gate 3 never exercised it -- P3's legs used `--max-steps`, so each leg *was* a whole
|
| 5 |
-
# run. The launcher's plan, the trainer's stop, HubPush's push-verify-pointer-prune order, the resume scan
|
| 6 |
-
# that refuses a stale pointer, and the final-push-plus-terminal-pointer path are therefore untested
|
| 7 |
-
# together. E-035/4 is precisely a defect in that untested seam, and the rehearsal kernel cannot reach it
|
| 8 |
-
# because it needs two T4s. This does, against the published mix and the frozen 22L geometry, in one
|
| 9 |
-
# session: leg 1 trains to a mid-run stop step, leg 2 wipes the disk, resumes from the Hub and runs to the
|
| 10 |
-
# horizon.
|
| 11 |
-
#
|
| 12 |
-
# The probe writes to its own checkpoint repo, never to Cion-lab/ounce100m-ckpt: the first real session has
|
| 13 |
-
# to find that repo absent, which is the RepoMissing branch E-034 was about.
|
| 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 |
-
rc
|
| 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 |
-
|
| 128 |
-
|
| 129 |
-
|
| 130 |
-
|
| 131 |
-
|
| 132 |
-
|
| 133 |
-
|
| 134 |
-
|
| 135 |
-
|
| 136 |
-
|
| 137 |
-
|
| 138 |
-
|
| 139 |
-
|
| 140 |
-
|
| 141 |
-
|
| 142 |
-
|
| 143 |
-
|
| 144 |
-
|
| 145 |
-
|
| 146 |
-
|
| 147 |
-
|
| 148 |
-
|
| 149 |
-
|
| 150 |
-
|
| 151 |
-
|
| 152 |
-
|
| 153 |
-
|
| 154 |
-
|
| 155 |
-
|
| 156 |
-
|
| 157 |
-
|
| 158 |
-
"
|
| 159 |
-
|
| 160 |
-
|
| 161 |
-
|
| 162 |
-
|
| 163 |
-
"
|
| 164 |
-
|
| 165 |
-
|
| 166 |
-
|
| 167 |
-
|
| 168 |
-
|
| 169 |
-
"
|
| 170 |
-
}
|
| 171 |
-
|
| 172 |
-
|
| 173 |
-
|
| 174 |
-
|
| 175 |
-
|
| 176 |
-
|
| 177 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# P4 probe: does --stop-after-steps really produce a resumable, Hub-verified segment boundary?
|
| 2 |
+
#
|
| 3 |
+
# Why this exists and what it buys with ~0.4 GPU-hours: `--stop-after-steps` is the mechanism every Phase 4
|
| 4 |
+
# session ends on, and Gate 3 never exercised it -- P3's legs used `--max-steps`, so each leg *was* a whole
|
| 5 |
+
# run. The launcher's plan, the trainer's stop, HubPush's push-verify-pointer-prune order, the resume scan
|
| 6 |
+
# that refuses a stale pointer, and the final-push-plus-terminal-pointer path are therefore untested
|
| 7 |
+
# together. E-035/4 is precisely a defect in that untested seam, and the rehearsal kernel cannot reach it
|
| 8 |
+
# because it needs two T4s. This does, against the published mix and the frozen 22L geometry, in one
|
| 9 |
+
# session: leg 1 trains to a mid-run stop step, leg 2 wipes the disk, resumes from the Hub and runs to the
|
| 10 |
+
# horizon.
|
| 11 |
+
#
|
| 12 |
+
# The probe writes to its own checkpoint repo, never to Cion-lab/ounce100m-ckpt: the first real session has
|
| 13 |
+
# to find that repo absent, which is the RepoMissing branch E-034 was about.
|
| 14 |
+
#
|
| 15 |
+
# It carries a second question, the one the user pushed back on (D-017): gradient checkpointing costs 32 %
|
| 16 |
+
# of the throughput (9,696 -> 12,792 tok/s, 29.7 h -> 22.5 h) and gives up the memory headroom, sitting at
|
| 17 |
+
# 12.25 GB of ~14.56. So this cell runs the *real* geometry at --accum 32, micro 4, WITHOUT checkpointing,
|
| 18 |
+
# for 180 steps, through three push/verify/prune cycles and one forced cold resume and the end-of-run
|
| 19 |
+
# validation pass -- the three things a 30-step throughput cell cannot show: allocator drift across a few
|
| 20 |
+
# hundred steps, the save path's host/GPU copies while the card is 84 % full, and the eval forward pass.
|
| 21 |
+
# Peak memory is asserted, not eyeballed. If it holds, the run adopts it as D-018 with a finer push cadence
|
| 22 |
+
# as the bounded blast radius; if it does not, D-017 stands and this is the measurement that says so.
|
| 23 |
+
import hashlib, json, os, shutil, signal, subprocess, sys, threading, time
|
| 24 |
+
|
| 25 |
+
os.chdir("/kaggle/working")
|
| 26 |
+
sys.path.insert(0, "/kaggle/working")
|
| 27 |
+
REV = "91fb319faee801a74fb3199e63111db84673715f"
|
| 28 |
+
WANT = {
|
| 29 |
+
"ounce100m_credentials.py": ("ounce100m_credentials.py",
|
| 30 |
+
"6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
|
| 31 |
+
"shard_dataset.py": ("train/shard_dataset.py",
|
| 32 |
+
"f35653bf4c8f2cfe7bb0c2c7835e308505fb5a84b4eb002d0f82f2cada768ca6"),
|
| 33 |
+
"hubckpt.py": ("train/hubckpt.py",
|
| 34 |
+
"568c31b300906cb8d78a59ef890baa9731b18e064e3d69b0eaea4b45c546f23e"),
|
| 35 |
+
"train_ounce100m.py": ("train/train_ounce100m.py",
|
| 36 |
+
"280f7ddc4632de30565426f1417f59c528a9ac0a59cfcfe18bc209216ede7c5e"),
|
| 37 |
+
}
|
| 38 |
+
BASE = "https://huggingface.co/Cion-lab/ounce100m-code/resolve/" + REV
|
| 39 |
+
for p, (rp, want) in sorted(WANT.items()):
|
| 40 |
+
assert subprocess.run(["curl", "-sfL", f"{BASE}/{rp}", "-o", p]).returncode == 0, ("fetch", rp)
|
| 41 |
+
got = hashlib.sha256(open(p, "rb").read()).hexdigest()
|
| 42 |
+
assert got == want, ("SHA MISMATCH", rp, got[:16], want[:16])
|
| 43 |
+
print("OK", p, got[:12], flush=True)
|
| 44 |
+
|
| 45 |
+
import ounce100m_credentials as C
|
| 46 |
+
print("creds:", json.dumps(C.install(verify=True)), flush=True)
|
| 47 |
+
import hubckpt
|
| 48 |
+
from huggingface_hub import HfApi
|
| 49 |
+
|
| 50 |
+
MIX = "Cion-lab/ounce100m-mix-v1"
|
| 51 |
+
PROBE = "Cion-lab/ounce100m-ckpt-probe" # never the real run's repo, see the header
|
| 52 |
+
ROOT, RUN = "/kaggle/working/mixroot", "/kaggle/working/run"
|
| 53 |
+
# 180 steps at the run's real step shape (262,144 tokens/step) = 47,185,920 tokens of the published mix:
|
| 54 |
+
# 139,842,880 at the mid-run stop, 61,118,720 after the cold resume. Pushes every 60 steps, which is also
|
| 55 |
+
# what a no-checkpointing main run would need as its recovery interval.
|
| 56 |
+
TOKENS, PUSH_EVERY, STOP1, STEPS = 47185920, 60, 120, 180
|
| 57 |
+
TPS = 262144
|
| 58 |
+
T0 = time.time()
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def run(argv, label, timeout):
|
| 62 |
+
print("=== " + label, flush=True)
|
| 63 |
+
t0 = time.time()
|
| 64 |
+
e = dict(os.environ); e["PYTHONPATH"] = "/kaggle/working"
|
| 65 |
+
p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
|
| 66 |
+
env=e, bufsize=1, start_new_session=True)
|
| 67 |
+
killed = []
|
| 68 |
+
|
| 69 |
+
def _kill():
|
| 70 |
+
killed.append(True)
|
| 71 |
+
try:
|
| 72 |
+
os.killpg(os.getpgid(p.pid), signal.SIGTERM)
|
| 73 |
+
except Exception:
|
| 74 |
+
p.kill()
|
| 75 |
+
|
| 76 |
+
timer = threading.Timer(timeout, _kill)
|
| 77 |
+
timer.daemon = True
|
| 78 |
+
timer.start()
|
| 79 |
+
keep, lines = [], []
|
| 80 |
+
try:
|
| 81 |
+
for line in p.stdout:
|
| 82 |
+
line = line.rstrip("\n")
|
| 83 |
+
lines.append(line)
|
| 84 |
+
if line.startswith(("CKPT ", "RUN_JSON ", "resume from", "auto-resume", "precision:",
|
| 85 |
+
"params:", "mix:", "checkpoint hub target", "segment boundary",
|
| 86 |
+
"latest.json", "TRAIN DONE", "validation loss", "Traceback", "Error")):
|
| 87 |
+
print(" KEY>", line[:300], flush=True)
|
| 88 |
+
keep.append(line)
|
| 89 |
+
del keep[:-40]
|
| 90 |
+
finally:
|
| 91 |
+
timer.cancel()
|
| 92 |
+
rc = p.wait()
|
| 93 |
+
if killed:
|
| 94 |
+
print(" TIMEOUT after %d s" % timeout, flush=True)
|
| 95 |
+
rc = -9
|
| 96 |
+
if rc != 0:
|
| 97 |
+
print(" TAIL:\n" + "\n".join(keep)[-2500:], flush=True)
|
| 98 |
+
print("%s_RC %s seconds %.1f elapsed %.0f" % (label, rc, time.time() - t0, time.time() - T0),
|
| 99 |
+
flush=True)
|
| 100 |
+
return rc, "\n".join(lines)
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
rc, out = run([sys.executable, "-c",
|
| 104 |
+
"import sys; sys.path.insert(0, '/kaggle/working')\n"
|
| 105 |
+
"from huggingface_hub import snapshot_download\n"
|
| 106 |
+
"p = snapshot_download(repo_id='%s', repo_type='dataset',\n"
|
| 107 |
+
" local_dir='/kaggle/working/mixroot', max_workers=4)\n"
|
| 108 |
+
"print('mix at', p)\n" % MIX], "FETCH_MIX", 2400)
|
| 109 |
+
if rc != 0:
|
| 110 |
+
raise SystemExit("VERDICT P4PROBE_STOP could not fetch the published mix")
|
| 111 |
+
man = json.load(open(os.path.join(ROOT, "manifest.json")))
|
| 112 |
+
print("mix", man["n_shards"], "shards", format(int(man["total_tokens"]), ","), "tokens", flush=True)
|
| 113 |
+
|
| 114 |
+
common = [sys.executable, "-u", "train_ounce100m.py", "--root", ROOT, "--out", RUN,
|
| 115 |
+
"--hub-repo", PROBE, "--prune", "--seq-len", "1024", "--attn", "eager", "--grad-ckpt",
|
| 116 |
+
"--micro-batch", "4", "--accum", "32", "--no-grad-ckpt",
|
| 117 |
+
"--tokens", str(TOKENS), "--lr", "6e-4",
|
| 118 |
+
"--push-every-steps", str(PUSH_EVERY), "--val-tokens", "2000000", "--log-every", "10",
|
| 119 |
+
"--resume", "auto"]
|
| 120 |
+
|
| 121 |
+
shutil.rmtree(RUN, ignore_errors=True)
|
| 122 |
+
rc1, o1 = run(["torchrun", "--nproc_per_node=2"] + common + ["--stop-after-steps", str(STOP1)],
|
| 123 |
+
"LEG1_STOP_EARLY", 4200)
|
| 124 |
+
api = HfApi(token=os.environ["HF_TOKEN"])
|
| 125 |
+
try:
|
| 126 |
+
listed = sorted(hubckpt.hub_listing(PROBE, "dataset", token=os.environ["HF_TOKEN"]))
|
| 127 |
+
except Exception as e:
|
| 128 |
+
listed = ["<listing failed: %s>" % type(e).__name__]
|
| 129 |
+
pushed = sorted({k.split("/")[1] for k in listed if k.startswith("ckpt/checkpoint-")
|
| 130 |
+
and len(k.split("/")) > 1})
|
| 131 |
+
ptr = hubckpt.latest_pointer(PROBE, token=os.environ["HF_TOKEN"])
|
| 132 |
+
print("LEG1 pushed:", pushed, "pointer:", {k: ptr.get(k) for k in ("step", "path_in_repo", "error")},
|
| 133 |
+
flush=True)
|
| 134 |
+
|
| 135 |
+
shutil.rmtree(RUN, ignore_errors=True) # force the cold-resume path (E-029's shape)
|
| 136 |
+
rc2, o2 = run(["torchrun", "--nproc_per_node=2"] + common + ["--stop-after-steps", "0"],
|
| 137 |
+
"LEG2_TO_HORIZON", 3000)
|
| 138 |
+
try:
|
| 139 |
+
listed2 = sorted(hubckpt.hub_listing(PROBE, "dataset", token=os.environ["HF_TOKEN"]))
|
| 140 |
+
except Exception as e:
|
| 141 |
+
listed2 = ["<listing failed: %s>" % type(e).__name__]
|
| 142 |
+
ptr2 = hubckpt.latest_pointer(PROBE, token=os.environ["HF_TOKEN"])
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def harvest(text, tag):
|
| 146 |
+
for line in (text or "").splitlines():
|
| 147 |
+
if line.startswith(tag):
|
| 148 |
+
try:
|
| 149 |
+
return json.loads(line[len(tag):].strip())
|
| 150 |
+
except Exception:
|
| 151 |
+
return {"unparsed": line[:200]}
|
| 152 |
+
return {}
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
rj1, rj2 = harvest(o1, "RUN_JSON "), harvest(o2, "RUN_JSON ")
|
| 156 |
+
res = {
|
| 157 |
+
"leg1_rc": rc1, "leg2_rc": rc2,
|
| 158 |
+
"leg1": {"final_step": rj1.get("final_step"), "segment_stop": rj1.get("segment_stop"),
|
| 159 |
+
"tokens_consumed": rj1.get("tokens_consumed"), "tok_per_s": rj1.get("tok_per_s"),
|
| 160 |
+
"params": rj1.get("params"), "peak_gpu_gb": rj1.get("peak_gpu_gb"),
|
| 161 |
+
"tok_per_s": rj1.get("tok_per_s"), "final_loss": rj1.get("final_loss"),
|
| 162 |
+
"pushed": pushed, "pointer": {k: ptr.get(k) for k in ("step", "path_in_repo")}},
|
| 163 |
+
"leg2": {"final_step": rj2.get("final_step"), "segment_stop": rj2.get("segment_stop"),
|
| 164 |
+
"tokens_consumed": rj2.get("tokens_consumed"), "val_ppl": rj2.get("val_ppl"),
|
| 165 |
+
"peak_gpu_gb": rj2.get("peak_gpu_gb"), "tok_per_s": rj2.get("tok_per_s"),
|
| 166 |
+
"final_loss": rj2.get("final_loss"),
|
| 167 |
+
"pointer_after": {k: ptr2.get(k) for k in ("step", "path_in_repo")},
|
| 168 |
+
"has_final": any(k.startswith("final/") for k in listed2)},
|
| 169 |
+
"expect": {"steps_planned": STEPS, "push_every": PUSH_EVERY, "stop1": STOP1,
|
| 170 |
+
"tokens_at_stop1": STOP1 * TPS, "tokens_at_horizon": STEPS * TPS, "grad_ckpt": False},
|
| 171 |
+
}
|
| 172 |
+
res["checks"] = {
|
| 173 |
+
"leg1_stopped_at_the_stop_step": rj1.get("final_step") == STOP1,
|
| 174 |
+
"leg1_reported_a_segment": rj1.get("segment_stop") is True,
|
| 175 |
+
# Integer compare: sorted() over names would order "checkpoint-120" before "checkpoint-60".
|
| 176 |
+
"leg1_pushed_every_interval": sorted(int(x.split("-")[1]) for x in pushed) ==
|
| 177 |
+
list(range(PUSH_EVERY, STOP1 + 1, PUSH_EVERY)),
|
| 178 |
+
"leg1_pointer_at_stop": ptr.get("step") == STOP1,
|
| 179 |
+
"leg2_resumed_and_finished": rj2.get("final_step") == STEPS,
|
| 180 |
+
"leg2_was_not_a_segment": rj2.get("segment_stop") is False,
|
| 181 |
+
"leg2_pushed_final": any(k.startswith("final/") for k in listed2),
|
| 182 |
+
"all_steps_present": sorted(int(x.split("-")[1]) for x in
|
| 183 |
+
{k.split("/")[1] for k in listed2 if k.startswith("ckpt/checkpoint-")})
|
| 184 |
+
== list(range(PUSH_EVERY, STEPS + 1, PUSH_EVERY)),
|
| 185 |
+
"leg2_pointer_terminal": ptr2.get("step") == STEPS and ptr2.get("path_in_repo") == "final",
|
| 186 |
+
"tokens_match_the_arithmetic": rj1.get("tokens_consumed") == STOP1 * TPS
|
| 187 |
+
and rj2.get("tokens_consumed") == STEPS * TPS,
|
| 188 |
+
"params_are_the_frozen_model": rj1.get("params") == 106194240,
|
| 189 |
+
"checkpointing_really_off": rj1.get("grad_ckpt") is False,
|
| 190 |
+
# 13.6 GB of ~14.56 usable leaves the run room for a fragmentation spike and for the two CUDA contexts;
|
| 191 |
+
# above that the answer to the user's question is "not on this hardware", not "keep going".
|
| 192 |
+
"peak_memory_below_13_6_gb": (rj1.get("peak_gpu_gb") or 99) <= 13.6
|
| 193 |
+
and (rj2.get("peak_gpu_gb") or 99) <= 13.6,
|
| 194 |
+
"no_nan_and_loss_moved": (rj1.get("final_loss") or 1e9) < 11.0
|
| 195 |
+
and (rj2.get("final_loss") or 1e9) < 11.0,
|
| 196 |
+
}
|
| 197 |
+
res["PROBE_PASSED"] = all(res["checks"].values()) and rc1 == 0 and rc2 == 0
|
| 198 |
+
print("PROBE_JSON_BEGIN")
|
| 199 |
+
print(json.dumps(res, indent=1, default=str))
|
| 200 |
+
print("PROBE_JSON_END")
|
| 201 |
+
print("VERDICT P4PROBE", "PASS" if res["PROBE_PASSED"] else "FAIL",
|
| 202 |
+
[k for k, v in res["checks"].items() if not v], "seconds", round(time.time() - T0, 1), flush=True)
|
| 203 |
+
raise SystemExit(0 if res["PROBE_PASSED"] else 5)
|