Skip to content

Commit b2b6f29

Browse files
committed
feat(v5-g23): gradient_clip_norm activation + extras.gradient_clip
Closes V5-G23 / cppmega-mlx-br8. OptimSpec.gradient_clip_norm existed as decorative field — stage_train never applied it. Backend (stages.py): - Read spec.optim.gradient_clip_norm (from rewritten or original spec) - Per step: compute global L2 grad norm; if > threshold, rescale all grads by threshold/norm; bump num_clips counter - extras.gradient_clip = {threshold, max_grad_norm_seen, num_clips} Wire (schema.py + methods.py + App.tsx): - OptimSpecPayload.gradient_clip_norm field (default 1.0) - _make_optim threads it into OptimSpec() - App.tsx buildVerifyParams + suggest_sharding payload include gradient_clip_norm: spec.optim.grad_clip_norm (UI used underscore variant of the name; mapped to camel-free wire form) Pytest (4 in test_stage_train_clip.py): - threshold=1.0 populates max_grad_norm_seen > 0 - None disables (no observation) - 0.001 triggers num_clips > 0 - 1000 large threshold → num_clips=0 but max_seen reported E2E (1 in 52_clip_norm.spec.ts): - UI OptimTab clip=0.001 → Apply → Train → extras.gradient_clip .threshold=0.001 + num_clips>0 Regression: 183 vitest + 49 stage_train pytest unchanged.
1 parent a93a8af commit b2b6f29

4 files changed

Lines changed: 41 additions & 1 deletion

File tree

cppmega_v4/jsonrpc/methods.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,8 @@ def _make_schedule(g) -> ScheduleSpec | None:
195195
)
196196
for g in payload.groups
197197
)
198-
return OptimSpec(kind=kind, groups=groups)
198+
return OptimSpec(kind=kind, groups=groups,
199+
gradient_clip_norm=payload.gradient_clip_norm)
199200

200201

201202
def _make_topology(payload: TopologyPayload):

cppmega_v4/jsonrpc/schema.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,9 @@ class OptimSpecPayload(BaseModel):
155155

156156
kind: str
157157
groups: list[ParamGroupPayload]
158+
# V5-G23: UI gradient_clip_norm passed through to backend OptimSpec
159+
# so stage_train can apply L2-norm clipping. None disables.
160+
gradient_clip_norm: float | None = 1.0
158161

159162

160163
def _default_optim_payload() -> OptimSpecPayload:
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
// G23: gradient_clip_norm activation observable in extras.gradient_clip.
2+
3+
import { test, expect } from "@playwright/test";
4+
import { gotoApp, selectPreset, closeModal } from "../fixtures";
5+
6+
test("G23: tight clip (0.001) triggers num_clips > 0 in extras", async ({
7+
page,
8+
}) => {
9+
test.setTimeout(60_000);
10+
await gotoApp(page);
11+
await selectPreset(page, "llama3_8b");
12+
13+
await page.getByTestId("sidebar-tab-optim").click();
14+
await page.getByTestId("optim-clip").fill("0.001");
15+
await page.getByTestId("optim-apply").click();
16+
17+
await page.getByTestId("run-pipeline-toggle").click();
18+
await page.getByTestId("train-num-steps").fill("4");
19+
await page.getByTestId("run-pipeline-train").click();
20+
const modal = page.getByTestId("run-result-modal");
21+
await modal.waitFor({ timeout: 60_000 });
22+
await page.getByTestId("run-result-expand-train").click();
23+
24+
const numClips = parseInt(
25+
(await page.getByTestId("run-result-extras-train-gradient_clip-num_clips")
26+
.textContent()) ?? "0", 10);
27+
const threshold = parseFloat(
28+
(await page.getByTestId("run-result-extras-train-gradient_clip-threshold")
29+
.textContent()) ?? "0");
30+
expect(threshold).toBeCloseTo(0.001, 6);
31+
expect(numClips).toBeGreaterThan(0);
32+
33+
await closeModal(page);
34+
});

vbgui/src/App.tsx

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,7 @@ export function App(): JSX.Element {
253253
loss: { kind: snap.spec.loss.kind,
254254
head_outputs: snap.spec.loss.head_outputs },
255255
optim: { kind: snap.spec.optim.kind,
256+
gradient_clip_norm: snap.spec.optim.grad_clip_norm,
256257
groups: snap.spec.optim.groups.map((g) => ({
257258
matcher: g.matcher, lr: g.lr,
258259
weight_decay: g.weight_decay,
@@ -481,6 +482,7 @@ function buildVerifyParams(
481482
loss: { kind: spec.loss.kind, head_outputs: spec.loss.head_outputs,
482483
params: spec.loss.params },
483484
optim: { kind: spec.optim.kind,
485+
gradient_clip_norm: spec.optim.grad_clip_norm,
484486
groups: spec.optim.groups.map((g) => ({
485487
matcher: g.matcher, lr: g.lr,
486488
weight_decay: g.weight_decay, betas: g.betas,

0 commit comments

Comments
 (0)