Muonclip

Fireworks AI

GLM 5.2 is live! Opus-level intelligence at open-source rates. Pay per token on serverless. Try it today.

Blog

Muonclip Deep-dive into MuonClip: Fixing Attention Score Explosions in Transformer Training

PUBLISHED 7/15/2025

Table of Contents The Attention Mechanism: A Quick Refresher The Scaling Challenge: Why Do QK Scores Explode? MuonClip to the Rescue: Rescaling at the Source Visualizing qk-clip in Action: A Toy Example

Table of Contents

Interactive visualization for MuonClip, brought to you from Fireworks.ai With the release of Kimi-K2 , a state of the art tool calling and instruction following model, Kimi team also talked about how they scaled up their pre-training, with a new optimizer, MuonClip. Honestly we don’t see new optimizers that often, so let’s dive into this a little more to understand how this helped the Kimi team scale their training. Specifically, this was the part of the blog https://moonshotai.github.io/Kimi-K2/ related to MuonClip. So for people who are bad at math like me, what are they talking and how exactly does it solves their scaling problem. The Attention Mechanism: A Quick Refresher

Before we hit the problem, let's recall how attention works in transformers (the backbone of most LLMs like GPT or Llama). Attention lets the model "focus" on relevant parts of the input sequence. It does this by computing query (Q) , key (K) , and value (V) projections from the input embeddings. The magic happens in the attention scores (often called "logits" in this context, but we'll call them "QK scores" to avoid confusion with output probabilities). These are dot products between queries and keys, scaled by the square root of the dimension for stability: High scores mean the model pays more attention to that key when aggregating values. But if these scores blow up to extreme values during training, things go haywire—leading to NaNs, gradients vanishing or exploding, and your entire run crashing. The Scaling Challenge: Why Do QK Scores Explode?

As you scale LLMs to billions of parameters and trillions of tokens (like Kimi K2's 15.5T-token pretraining), instabilities creep in. Moonshot AI noticed this especially when using the Muon optimizer —a high-efficiency alternative to the trusty AdamW that's great for speeding up training but a bit more aggressive. If you are interested in learning more about Muon, you can read more about it Kimi's paper around Muon , and this blog from Keller Jordan around this topic. Existing fixes for the QK score explosion problem? Things like logit soft-capping (clamping scores to a max value) or query-key normalization (normalizing Q and K vectors) sound promising, but Moonshot found them lacking. Soft-capping can distort the attention distribution unnaturally, while normalization might not address the root cause in the weights themselves. Enter MuonClip , Moonshot's upgrade to Muon that tackles this head-on with a technique called qk-clip . MuonClip to the Rescue: Rescaling at the Source

MuonClip keeps Muon's speed advantages but adds a post-update safety net. After each Muon step (which orthogonalizes updates for balance—more on that in a sec), qk-clip checks the potential QK scores. If the max score exceeds a threshold t (say, 1.0 in our demo), it rescales W_q and W_k directly: • Compute η = t / max_score (so η G . size ( 1 ) if transpose : X = X . T for _ in range ( steps ) : A = X @ X . T B = b * A + c * A @ A X = a * X + B @ X if transpose : X = X . T return X

def muon_update ( grad , momentum , beta = 0.95 , ns_steps = 5 , nesterov = True ) : momentum . lerp_ ( grad , 1 - beta ) if nesterov : update = grad . clone ( ) . lerp_ ( momentum , beta ) else : update = momentum . clone ( ) update = zeropower_via_newtonschulz5 ( update , steps = ns_steps ) scale_factor = max ( 1 , grad . size ( - 2 ) / grad . size ( - 1 ) ) ** 0.5 update *= scale_factor return update

def apply_clip ( W_q , W_k , alpha , t , x , eps = 1e-7 ) : q = x @ W_q . T # (batch, seq, fan_out) k = x @ W_k . T scores = torch . einsum ( 'bid,bjd->bij' , q , k ) / np . sqrt ( W_q . size ( 0 ) ) # / sqrt(fan_out) max_score = scores . max ( ) if max_score > t : eta = t / ( max_score + eps ) scale_q = eta alpha scale_k = eta ( 1 - alpha ) W_q *= scale_q W_k *= scale_k return True , max_score . item ( ) return False , max_score . item ( )

st . title ( 'MuonClip Clipping Visualization - Early vs Late Training' )

fan_out = st . slider ( 'Output Dim (fan_out, e.g., d_head)' , min_value = 2 , max_value = 10 , value = 4 , step = 1 ) fan_in = st . slider ( 'Input Dim (fan_in, e.g., d_model)' , min_value = 2 , max_value = 10 , value = 4 , step = 1 ) seq_len = st . slider ( 'Sequence Length for Simulation' , min_value = 2 , max_value = 10 , value = 4 , step = 1 ) seed = st . slider ( 'Random Seed' , min_value = 0 , max_value = 100 , value = 42 , step = 1 ) beta = st . slider ( 'Momentum Beta' , min_value = 0.5 , max_value = 0.99 , value = 0.95 , step = 0.01 ) nesterov = st . checkbox ( 'Use Nesterov Momentum' , value = True ) alpha = st . slider ( 'Alpha for Clip' , min_value = 0.0 , max_value = 1.0 , value = 0.5 , step = 0.05 ) lr_early = st . slider ( 'LR for Early Training (high to simulate explosion)' , min_value = 0.1 , max_value = 200.0 , value = 50.0 , step = 1.0 ) lr_late = st . slider ( 'LR for Late Training (low for stable)' , min_value = 0.1 , max_value = 200.0 , value = 5.0 , step = 1.0 ) t = st . slider ( 'Clip Threshold t' , min_value = 10.0 , max_value = 200.0 , value = 100.0 , step = 10.0 )

torch . manual_seed ( seed )

Initial weights small

init_scale = 0.1 W_q = torch . randn ( fan_out , fan_in ) * init_scale W_k = torch . randn ( fan_out , fan_in ) * init_scale

Simulate forward for scores

x = torch . randn ( 1 , seq_len , fan_in )

Function to compute for a scenario

def compute_scenario ( lr , t , scenario_name ) :

Fix grad_scale to 1, as normalization makes scale irrelevant for magnitude

grad_scale = 1.0

Simulate grads

grad_q = torch . randn ( fan_out , fan_in ) * grad_scale grad_k = torch . randn ( fan_out , fan_in ) * grad_scale

Momentum buffers

momentum_q = torch . zeros_like ( grad_q ) momentum_k = torch . zeros_like ( grad_k )

Compute Muon updates

update_q = muon_update ( grad_q , momentum_q , beta = beta , ns_steps = 5 , nesterov = nesterov ) update_k =...