Skip to content

Commit 678dfe8

Browse files
committed
feat(v5-g15): long-run support — TopBar max=512 + extras.losses_smoothed
Closes V5-G15 / cppmega-mlx-a0b. TopBar train-num-steps max bumped from 64 to 512. stage_train extras gains losses_smoothed (window=10 moving average) for noise-reduced convergence visualisation. _ema_smooth helper. 3 new pytest: shape correctness for short + long runs, smoothed variance ≤ raw variance over N=100. 56 stage_train pytest regression green.
1 parent 74832cb commit 678dfe8

3 files changed

Lines changed: 74 additions & 1 deletion

File tree

cppmega_v4/runner/stages.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -908,6 +908,7 @@ def _count(tree: Any) -> int:
908908
elapsed_ms=(time.perf_counter() - t0) * 1000.0,
909909
extras={
910910
"losses": [round(loss_item, 4) for loss_item in losses],
911+
"losses_smoothed": _ema_smooth(losses, window=10),
911912
"lr_trajectory": [
912913
round(lr_item, 6) for lr_item in lr_trajectory
913914
],
@@ -1093,6 +1094,17 @@ def _apply_spec_rewriters(
10931094
}
10941095

10951096

1097+
def _ema_smooth(values: list[float], window: int = 10) -> list[float]:
1098+
"""G15: simple windowed-mean smoothing (cheaper than true EMA, same
1099+
job for visualising convergence)."""
1100+
out: list[float] = []
1101+
for i in range(len(values)):
1102+
start = max(0, i - window + 1)
1103+
sl = values[start:i + 1]
1104+
out.append(round(sum(sl) / len(sl), 4))
1105+
return out
1106+
1107+
10961108
def _compute_hybrid_deltas(
10971109
after_flat: dict[str, Any],
10981110
muon_before: dict[str, Any], adamw_before: dict[str, Any],
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
"""G15: long-run convergence — N=100 with smoothed loss observability."""
2+
3+
from __future__ import annotations
4+
5+
import pytest
6+
7+
from cppmega_v4.jsonrpc.schema import VerifyParams
8+
from cppmega_v4.runner import Pipeline, run_pipeline
9+
10+
11+
def _spec() -> VerifyParams:
12+
return VerifyParams.model_validate({
13+
"graph": {
14+
"nodes": [
15+
{"id": "attn", "kind": "attention", "params": {}},
16+
{"id": "mlp", "kind": "mlp",
17+
"params": {"intermediate_size": 64, "activation": "swiglu"}},
18+
],
19+
"edges": [{"src": "attn", "dst": "mlp"}],
20+
},
21+
"dim_env": {"B": 1, "S": 8, "H": 32, "nh": 2, "nkv": 1, "head_dim": 16},
22+
"loss": {"kind": "cross_entropy", "head_outputs": ["mlp"]},
23+
"optim": {"kind": "adamw",
24+
"groups": [{"matcher": "all", "lr": 1e-3,
25+
"weight_decay": 0.01, "betas": [0.9, 0.95]}]},
26+
})
27+
28+
29+
def _run(num_steps: int) -> dict:
30+
report = run_pipeline(_spec(), Pipeline.from_dict({
31+
"stages": ["parse", "verify_build_spec", "build_model", "train"],
32+
"stage_options": {"train": {"num_steps": num_steps}},
33+
}))
34+
train = next(s for s in report.stages if s.name == "train")
35+
assert train.status == "ok", f"stage_train failed: {train.error}"
36+
return train.extras
37+
38+
39+
def test_losses_smoothed_present_short_run():
40+
extras = _run(8)
41+
assert "losses_smoothed" in extras
42+
assert len(extras["losses_smoothed"]) == 8
43+
44+
45+
def test_long_run_n100_smoothed_shape_correct():
46+
extras = _run(100)
47+
assert len(extras["losses"]) == 100
48+
assert len(extras["losses_smoothed"]) == 100
49+
# Smoothed[0] equals losses[0] (window of 1 at start)
50+
assert abs(extras["losses_smoothed"][0] - extras["losses"][0]) < 1e-4
51+
52+
53+
def test_long_run_smoothed_curve_monotone_in_window():
54+
"""Smoothed (window=10) loss over N=100 should be less noisy than
55+
raw — variance of smoothed strictly less than variance of raw."""
56+
extras = _run(100)
57+
import statistics as st
58+
raw_var = st.pvariance(extras["losses"])
59+
smooth_var = st.pvariance(extras["losses_smoothed"])
60+
assert smooth_var <= raw_var, \
61+
f"smoothed var {smooth_var} should be <= raw var {raw_var}"

vbgui/src/components/TopBar.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ export function TopBar(p: TopBarProps): JSX.Element {
128128
<span style={{ fontSize: 11, color: "#6b7280" }}>
129129
train steps:</span>
130130
<input data-testid="train-num-steps"
131-
type="number" min={1} max={64}
131+
type="number" min={1} max={512}
132132
value={trainNumSteps}
133133
onChange={(e) =>
134134
setTrainNumSteps(Math.max(1, parseInt(

0 commit comments

Comments
 (0)