Skip to content

Commit 8798113

Browse files
committed
fix(path-c): seed direct-chain first brick with raw residual, not pre-normed (double-norm)
make_path_c_direct_chain_pre_step_runtime_owner seeded {first_brick}_hidden with layers[start].norm(prefix_hidden), but the fused chain's entry-RMSNorm re-applies the norm -> double-norm. The eager layer contract (hybrid_lm.py:687-697) is residual = hidden (raw), then delta = route_delta(hidden) which norms internally; the chain replays that, so the seed must be the RAW residual stream. Fix: boundary_hidden = prefix_hidden (raw) instead of layers[start].norm(prefix). Metal (chain start_layer=10): direct-chain runtime loss-vs-eager err 1.02e-2 -> 0.0 (exact), grad err 1.23e-2 -> 0.0. Flag-independent, target-agnostic (corrects CUDA equally). Tests 58 passed (direct-chain/pre-step-owner/value_and_grad/bridge). Found by the moe-span executor investigation; extracted minimal.
1 parent 299531d commit 8798113

1 file changed

Lines changed: 6 additions & 3 deletions

File tree

scripts/m04_train_step.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7626,10 +7626,13 @@ def make_path_c_direct_chain_pre_step_runtime_owner(
76267626
batch,
76277627
end_layer_index=start_layer_index,
76287628
)
7629-
layers = tuple(getattr(model, "layers", ()))
7629+
# The fused chain's first brick reads ``{first_brick}_hidden`` as the RAW
7630+
# residual baseline and applies its own entry-RMSNorm (a fused
7631+
# ``entry_rmsnorm`` surface bound to ``layers.{start}.norm.weight``) to
7632+
# derive the route input. Seeding the raw residual stream therefore matches
7633+
# the eager ``updated = hidden + route_delta(norm(hidden))`` contract;
7634+
# seeding the pre-normed hidden would double-apply the entry norm.
76307635
boundary_hidden = prefix_hidden
7631-
if start_layer_index < len(layers):
7632-
boundary_hidden = layers[start_layer_index].norm(prefix_hidden)
76337636
hidden_seed_names = {
76347637
"hidden",
76357638
f"{first_brick_name}_hidden",

0 commit comments

Comments
 (0)