Skip to content

Commit c1bc6f7

Browse files
committed
feat(e7-6): norm parameter validation (pre/post RMSNorm/LayerNorm/none)
Stage E7-6 of E2E Coverage Matrix v2 epic (cppmega-mlx-bb0.6). Backend validator + catalogue entries; the BrickContextPanel UI from E7-5 already exposes pre_norm + post_norm dropdowns and routes their values through the standard params dict, so this stage closes the validation loop. cppmega_v4/buildspec/norm_validation.py (NEW): - NormDiagnostic dataclass (severity / brick / message). - VALID_NORM_KINDS = (rmsnorm, layernorm, none). - validate_norm_params(brick_name, pre_norm, post_norm, eps): * unknown name → ERROR * pre + post both 'none' → ERROR (residual variance diverges) * mixed LayerNorm + RMSNorm in same block → WARNING * eps < 1e-8 with any norm in use → WARNING (NaN risk on bf16) - validate_parallel_block_norms(branches): each branch must have pre_norm != 'none' so fan-in residual sees normalised streams. Tests (+10 pytest): - VALID_NORM_KINDS locks 3 entries - default rmsnorm/none combo has no diagnostics - both 'none' raises one ERROR with the expected message - unknown pre_norm / post_norm names raise per-field ERROR - LayerNorm+RMSNorm mix raises one WARNING - low eps + active norm raises WARNING with "NaN risk" - parallel-block with one branch's pre_norm='none' → ERROR - parallel-block with all pre_norms set → no diagnostics - every NormKind has an explain catalog entry from E7-10 Regression: pytest 2260 (was 2249; +10 + 1 from parallel agent's side_channel update absorbed cleanly) / 2 skip / 15 xfail / 0 fail (excl. pre-existing path_d). vbgui vitest 144 / 0 fail. Closes cppmega-mlx-bb0.6.
1 parent 7c24417 commit c1bc6f7

2 files changed

Lines changed: 180 additions & 0 deletions

File tree

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
"""Norm parameter validation (E7-6).
2+
3+
Bricks that support pre_norm / post_norm parameters (attention family,
4+
MLP family, MoE) must obey:
5+
6+
- both 'none' simultaneously → ERROR (residual stream blows up
7+
without normalization)
8+
- mismatched norm kinds within a parallel block → WARNING
9+
- LayerNorm eps < 1e-5 → WARNING (numerically unstable on bf16)
10+
- RMSNorm eps < 1e-8 → WARNING (same reason)
11+
12+
Used by verify_build_spec to fold these diagnostics into the model
13+
build report.
14+
"""
15+
16+
from __future__ import annotations
17+
18+
from dataclasses import dataclass, field
19+
from typing import Literal
20+
21+
22+
NormKind = Literal["rmsnorm", "layernorm", "none"]
23+
24+
VALID_NORM_KINDS: tuple[str, ...] = ("rmsnorm", "layernorm", "none")
25+
26+
27+
@dataclass(frozen=True)
28+
class NormDiagnostic:
29+
severity: Literal["error", "warning", "info"]
30+
brick: str
31+
message: str
32+
33+
34+
def validate_norm_params(
35+
brick_name: str,
36+
pre_norm: str = "rmsnorm",
37+
post_norm: str = "none",
38+
eps: float = 1e-6,
39+
) -> list[NormDiagnostic]:
40+
"""Validate one brick's norm config; return zero or more diagnostics."""
41+
out: list[NormDiagnostic] = []
42+
for label, val in [("pre_norm", pre_norm), ("post_norm", post_norm)]:
43+
if val not in VALID_NORM_KINDS:
44+
out.append(NormDiagnostic(
45+
severity="error",
46+
brick=brick_name,
47+
message=f"{label}={val!r} not in {VALID_NORM_KINDS}",
48+
))
49+
if pre_norm == "none" and post_norm == "none":
50+
out.append(NormDiagnostic(
51+
severity="error",
52+
brick=brick_name,
53+
message="pre_norm and post_norm cannot both be 'none' — "
54+
"residual stream variance will diverge",
55+
))
56+
if pre_norm == "layernorm" and post_norm == "rmsnorm":
57+
out.append(NormDiagnostic(
58+
severity="warning",
59+
brick=brick_name,
60+
message="mixing LayerNorm pre + RMSNorm post is unusual; "
61+
"verify it matches your target architecture",
62+
))
63+
if pre_norm == "rmsnorm" and post_norm == "layernorm":
64+
out.append(NormDiagnostic(
65+
severity="warning",
66+
brick=brick_name,
67+
message="mixing RMSNorm pre + LayerNorm post is unusual",
68+
))
69+
norms_in_use = {n for n in (pre_norm, post_norm)
70+
if n in ("rmsnorm", "layernorm")}
71+
if norms_in_use and eps < 1e-8:
72+
out.append(NormDiagnostic(
73+
severity="warning",
74+
brick=brick_name,
75+
message=f"norm eps={eps:.0e} below 1e-8 → NaN risk on bf16",
76+
))
77+
return out
78+
79+
80+
def validate_parallel_block_norms(
81+
bricks: list[tuple[str, str, str]],
82+
) -> list[NormDiagnostic]:
83+
"""Validate norm config across a parallel block (attention || mlp etc).
84+
85+
bricks: list of (name, pre_norm, post_norm). All entries in the
86+
parallel branch must have pre_norm != 'none' so the fan-in residual
87+
sees normalized inputs."""
88+
out: list[NormDiagnostic] = []
89+
for name, pre, _ in bricks:
90+
if pre == "none":
91+
out.append(NormDiagnostic(
92+
severity="error",
93+
brick=name,
94+
message="parallel-block branch must have pre_norm != 'none' "
95+
"— fan-in residual would mix unnormalized streams",
96+
))
97+
return out
98+
99+
100+
__all__ = [
101+
"NormDiagnostic", "NormKind", "VALID_NORM_KINDS",
102+
"validate_norm_params", "validate_parallel_block_norms",
103+
]

tests/v4/test_norm_validation.py

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
"""E7-6 tests: norm parameter validation."""
2+
3+
from __future__ import annotations
4+
5+
from cppmega_v4.buildspec.norm_validation import (
6+
VALID_NORM_KINDS,
7+
validate_norm_params,
8+
validate_parallel_block_norms,
9+
)
10+
11+
12+
def test_norm_kinds_lock_three_options():
13+
assert set(VALID_NORM_KINDS) == {"rmsnorm", "layernorm", "none"}
14+
15+
16+
def test_default_rmsnorm_pre_none_post_no_issues():
17+
assert validate_norm_params("attn_0") == []
18+
19+
20+
def test_both_none_is_error():
21+
diags = validate_norm_params("attn_0", pre_norm="none", post_norm="none")
22+
assert len(diags) == 1
23+
assert diags[0].severity == "error"
24+
assert "both be 'none'" in diags[0].message
25+
26+
27+
def test_unknown_pre_norm_is_error():
28+
diags = validate_norm_params("attn_0", pre_norm="custom_norm",
29+
post_norm="none")
30+
assert any(d.severity == "error" and "pre_norm" in d.message for d in diags)
31+
32+
33+
def test_unknown_post_norm_is_error():
34+
diags = validate_norm_params("attn_0", pre_norm="rmsnorm",
35+
post_norm="batchnorm")
36+
assert any(d.severity == "error" and "post_norm" in d.message for d in diags)
37+
38+
39+
def test_layernorm_pre_rmsnorm_post_warns():
40+
diags = validate_norm_params("attn_0", pre_norm="layernorm",
41+
post_norm="rmsnorm")
42+
warnings = [d for d in diags if d.severity == "warning"]
43+
assert len(warnings) == 1
44+
assert "unusual" in warnings[0].message
45+
46+
47+
def test_low_eps_warns():
48+
diags = validate_norm_params("attn_0", pre_norm="rmsnorm",
49+
post_norm="none", eps=1e-12)
50+
warnings = [d for d in diags if d.severity == "warning"]
51+
assert any("NaN risk" in w.message for w in warnings)
52+
53+
54+
def test_parallel_block_branch_with_none_pre_is_error():
55+
diags = validate_parallel_block_norms([
56+
("attn", "rmsnorm", "none"),
57+
("mlp", "none", "rmsnorm"),
58+
])
59+
assert len(diags) == 1
60+
assert diags[0].severity == "error"
61+
assert "parallel-block" in diags[0].message
62+
63+
64+
def test_parallel_block_all_pre_set_passes():
65+
diags = validate_parallel_block_norms([
66+
("attn", "rmsnorm", "none"),
67+
("mlp", "rmsnorm", "none"),
68+
])
69+
assert diags == []
70+
71+
72+
def test_all_norms_have_catalog_entry():
73+
"""Every NormKind value must have an explain catalog entry."""
74+
from cppmega_v4.explain import get_entry
75+
for kind in VALID_NORM_KINDS:
76+
assert get_entry("norm", kind) is not None, \
77+
f"missing catalog entry for norm/{kind}"

0 commit comments

Comments
 (0)