Skip to content

Commit 6ff8362

Browse files
committed
feat(v7-f53): dimension-scaling sweep on llama3_8b preset
Closes V7-F53 (cppmega-mlx-2yhv): pins the constructor flow the user described — pick preset, instantiate at small H, run through the system, scale H up, still works. Tests (tests/v4/test_dim_scaling_sweep.py): 5/5 — * llama3_8b at H in {64, 128, 256, 512}: train status=ok, losses finite, weight_delta_norm > 0. * params_bytes strictly monotone with H across the sweep.
1 parent 86b3a37 commit 6ff8362

1 file changed

Lines changed: 84 additions & 0 deletions

File tree

tests/v4/test_dim_scaling_sweep.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
"""V7-F53: dimension-scaling sweep on llama3_8b preset.
2+
3+
Proves the constructor flow the user described:
4+
pick preset → instantiate at small H → run through system →
5+
expand H → still works.
6+
7+
Pytest parametrises H in {64, 128, 256, 512} and asserts:
8+
(a) all four configurations land train.status='ok' with finite
9+
losses.
10+
(b) per-H weight_delta_norm > 0 (training really happened).
11+
(c) param footprint scales monotonically with H.
12+
"""
13+
14+
from __future__ import annotations
15+
16+
import pytest
17+
18+
from cppmega_v4.architectures.presets import build_preset_specs
19+
from cppmega_v4.jsonrpc.schema import VerifyParams
20+
from cppmega_v4.runner import Pipeline, run_pipeline
21+
22+
23+
def _spec_at(H: int) -> VerifyParams:
24+
specs = build_preset_specs("llama3_8b", hidden_size=H)
25+
return VerifyParams.model_validate({
26+
"graph": {
27+
"nodes": [
28+
{"id": f"n{i}", "kind": s["kind"],
29+
"params": s.get("params", {})}
30+
for i, s in enumerate(specs)
31+
],
32+
"edges": [
33+
{"src": f"n{i}", "dst": f"n{i + 1}"}
34+
for i in range(len(specs) - 1)
35+
],
36+
},
37+
"dim_env": {"B": 1, "S": 8, "H": H,
38+
"nh": max(2, H // 64), "nkv": max(1, H // 128),
39+
"head_dim": 64, "num_experts": 4, "top_k": 2},
40+
"loss": {"kind": "cross_entropy",
41+
"head_outputs": [f"n{len(specs) - 1}"]},
42+
"optim": {"kind": "adamw",
43+
"groups": [{"matcher": "all", "lr": 1e-3,
44+
"weight_decay": 0.01,
45+
"betas": [0.9, 0.95]}]},
46+
})
47+
48+
49+
def _train_at(H: int, num_steps: int = 4) -> dict:
50+
rep = run_pipeline(_spec_at(H), Pipeline.from_dict({
51+
"stages": ["parse", "verify_build_spec", "build_model",
52+
"estimate_memory", "train"],
53+
"stage_options": {"train": {"num_steps": num_steps}},
54+
}))
55+
return {
56+
s.name: s for s in rep.stages
57+
}
58+
59+
60+
@pytest.mark.parametrize("H", [64, 128, 256, 512])
61+
def test_v7_f53_each_dim_trains_with_finite_loss(H):
62+
stages = _train_at(H, num_steps=4)
63+
tr = stages["train"]
64+
assert tr.status == "ok", f"H={H} train failed: {tr.error}"
65+
for L in tr.extras["losses"]:
66+
assert -1e10 < L < 1e10, f"H={H} loss non-finite: {L}"
67+
# Weights moved.
68+
assert tr.extras["weight_delta_norm"] > 0, (
69+
f"H={H} weight_delta_norm zero — model didn't train"
70+
)
71+
72+
73+
def test_v7_f53_param_count_monotone_with_H():
74+
"""estimate_memory.params_bytes should grow with hidden size."""
75+
sizes: dict[int, int] = {}
76+
for H in (64, 128, 256, 512):
77+
stages = _train_at(H, num_steps=2)
78+
sizes[H] = int(stages["estimate_memory"].extras["params_bytes"])
79+
# Monotone increase across H.
80+
for a, b in zip(sorted(sizes.keys()), sorted(sizes.keys())[1:]):
81+
assert sizes[a] < sizes[b], (
82+
f"params_bytes not monotone: H={a}{sizes[a]} vs "
83+
f"H={b}{sizes[b]}"
84+
)

0 commit comments

Comments
 (0)