diff --git a/single_shot_inference.py b/single_shot_inference.py index 09d48998..9dff6bcc 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -107,12 +107,7 @@ class mHCBlock: mixed_residual = torch.einsum('tij,tjh->tjh', comb_mix, residual.float()) post_term = post_mix.unsqueeze(-1) * F_out.unsqueeze(1).float() residual_next = mixed_residual + post_term - # Emergency RMSNorm (remove once MoE provides balance) - _T = residual_next.shape[0] - rn_f = residual_next.reshape(_T, self.n_hc, -1) - rms = rn_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt() - return (rn_f * rms).bfloat16() - + return residual_next.bfloat16() # ===================================================================== # RoPE