CRITICAL FIX: mHC base/scale ordering was wrong
Checkpoint order is [pre, post, res] not [pre, res, post]: - base[0:4] = S_pre, base[4:8] = S_post, base[8:24] = S_res - scale[0] = alpha_pre, scale[1] = alpha_post, scale[2] = alpha_res - W_stacked rows: [W_pre(4), W_post(4), W_res(16)] - Projection split: A_raw=proj[:,0:4], C_raw=proj[:,4:8], B_raw=proj[:,8:24] This was causing B_l to be near-identity and C_l to be near-2.0, leading to exponential residual stream growth.
This commit is contained in:
@@ -138,17 +138,20 @@ class mHCBlock:
|
||||
"""Load from checkpoint tensors.
|
||||
|
||||
fn: (24, 28672) FP32
|
||||
base: (24,)
|
||||
scale: (3,) — [alpha_pre, alpha_res, alpha_post]
|
||||
base: (24,) — [pre(4), post(4), res(16)]
|
||||
scale: (3,) — [alpha_pre, alpha_post, alpha_res]
|
||||
"""
|
||||
n = self.n_hc
|
||||
# CRITICAL: checkpoint base order is [pre, post, res], not [pre, res, post]
|
||||
# This matches the old working code: base[:4]=pre, base[4:8]=post, base[8:24]=res
|
||||
self.W_stacked = fn.to(device=self.device, dtype=torch.float32).contiguous()
|
||||
self.S_pre = base[0:n].reshape(1, n).to(device=self.device, dtype=torch.float32)
|
||||
self.S_res = base[n:n + n*n].reshape(n, n).to(device=self.device, dtype=torch.float32)
|
||||
self.S_post = base[n + n*n:].reshape(n, 1).to(device=self.device, dtype=torch.float32)
|
||||
self.S_post = base[n:2*n].reshape(n, 1).to(device=self.device, dtype=torch.float32)
|
||||
self.S_res = base[2*n:2*n + n*n].reshape(n, n).to(device=self.device, dtype=torch.float32)
|
||||
# CRITICAL: checkpoint scale order is [alpha_pre, alpha_post, alpha_res]
|
||||
self.alpha_pre = scale[0].item() if scale.numel() > 0 else 0.01
|
||||
self.alpha_res = scale[1].item() if scale.numel() > 1 else 0.01
|
||||
self.alpha_post = scale[2].item() if scale.numel() > 2 else 0.01
|
||||
self.alpha_post = scale[1].item() if scale.numel() > 1 else 0.01
|
||||
self.alpha_res = scale[2].item() if scale.numel() > 2 else 0.01
|
||||
|
||||
def _project_and_rms(self, X_flat):
|
||||
"""Compute RMSNorm(X_flat) @ W_stacked.T → (T, N_proj) BF16.
|
||||
@@ -173,11 +176,11 @@ class mHCBlock:
|
||||
X_flat = X_l.reshape(T, self.K_proj) # (T, K_proj)
|
||||
proj = self._project_and_rms(X_flat).float() # (T, N_proj) FP32
|
||||
|
||||
# Split projection
|
||||
i0, i1, i2, i3 = 0, n, n + n*n, self.N_proj
|
||||
A_raw = proj[:, i0:i1]
|
||||
B_raw = proj[:, i1:i2]
|
||||
C_raw = proj[:, i2:i3]
|
||||
# Split projection — order matches W_stacked rows: [pre(4), post(4), res(16)]
|
||||
i0, i1, i2, i3 = 0, n, 2*n, self.N_proj
|
||||
A_raw = proj[:, i0:i1] # (T, n_hc) — pre
|
||||
C_raw = proj[:, i1:i2] # (T, n_hc) — post
|
||||
B_raw = proj[:, i2:i3] # (T, n_hc²) — res
|
||||
|
||||
# Add biases and scale by gating alphas (paper eqs. 3-5)
|
||||
A_tilde = self.alpha_pre * A_raw + self.S_pre
|
||||
|
||||
Reference in New Issue
Block a user