sysnote
Research

Fine-tuning Nemotron, Falcon-H1 or Kimi-Linear on new languages can silently skip a fifth of your training steps, a new paper finds

If you fine-tune a hybrid or Mamba-style model on Telugu, Mongolian or chat data, the loss can look normal while a large share of your optimizer steps do nothing. That is the finding of "Near-Zero Token Embeddings Silently Stall Training in Gated Sequence Models," a paper posted to arXiv on November 3 by Xin Liu, Boyuan Sun, Yuxin Wang, Haoran Song, Mingjian Zhao and Mengyu Liu (Washington University in St. Louis and Nankai University). There's no NaN and no error message. The only sign is grad_norm: inf in the logs.

Where the dead rows come from

The authors scanned 291 Hugging Face repositories covering 280 distinct embedding tables. They looked for input-embedding rows with a norm below 1e-6 that a real tokenizer can actually produce. None of the 101 models with tied input and output embeddings had any. Among untied models, 47 of 152 Transformers had them, including Llama-3.1-8B and OLMoE-1B-7B. So did 20 of 38 gated or hybrid models, among them Nemotron-3-Nano-4B, Falcon-H1-1.5B and 7B, Mamba-Codestral-7B and Kimi-Linear-48B-Base.

The mechanism is mundane. A token that never shows up in pretraining gets no gradient on its input row. AdamW's weight decay keeps shrinking that row anyway, and in an untied model nothing pushes back. In OLMoE-1B-7B, 250 of 251 such rows sit at norms around 3.2–3.5e-12, and none of the 251 tokens appears in a 364-million-token sample of the training corpus. The affected tokens are telling: Telugu and Basque subwords in Nemotron, Cyrillic Mongolian in Falcon-H1, CRLF code fragments in Mamba-Codestral, Nemotron-H-4B-Base's <s>, and Kimi-Linear's own chat tokens <|im_user|> and <|im_end|>.

Why gated models blow up

In an ordinary Transformer a near-zero embedding is mostly harmless once attention mixes it with earlier tokens. Gated and linear-recurrent layers are different. A SiLU output gate or a linear readout query that gets a near-zero input outputs nearly zero. Pre-RMSNorm then rescales a tiny input back to unit size, which multiplies the backward gradient by up to about 1,000 per layer. The paper measures embedding gradients of 1e30 to 1e35: a peak of 4.5e34 for Nemotron-3-Nano-4B and 1.0e32 for xLSTM-7B.

The silent part is in gradient clipping. PyTorch's clip_grad_norm_, the Hugging Face Trainer and FSDP2 compute one global norm in FP32 by squaring the entries. Any entry above about 1.84e19 squares past FP32's maximum. The norm becomes inf, the clip coefficient becomes zero, and every gradient in the model is multiplied by zero for that step. Megatron-Core only raises an error if check_for_nan_in_grad is enabled. Without it, its FP32 clipping zeroes the gradients silently too.

We checked the core claim ourselves in PyTorch 2.15.0 on CPU. Two parameters got gradients, one with a single entry of 2e19 and one with ordinary values of 0.5. clip_grad_norm_(..., 1.0) returned inf and set both gradients to all zeros, with no warning. At 1e19 the norm stays finite, but the normal parameter's gradient gets scaled down to about 5e-20, which is effectively zero too.

How much it costs

In Telugu continual pretraining of Nemotron-3-Nano-4B, just 62 near-zero rows (0.047% of the vocabulary) were hit in 22.5% of batches. Across 3 seeds of a 1,000-step run, 198 steps (19.8%) got no fresh update at all. English web, SFT, code and math data triggered nothing. Multilingual corpora triggered 10.6% to 100% of batches, depending on model and language. For Kimi-Linear, <|im_user|> appears at position 0 of every chat-formatted example, so the trigger rate in chat SFT is 100%.

LoRA fine-tuning doesn't touch the input embeddings, so the paper reports it as unaffected. So is inference.

The fixes

For released checkpoints the authors propose GroupClip: clip the backbone and the embedding table separately, and compute the embedding norm in FP64.

clip_grad_norm_(backbone_params, max_norm=1.0)
gn = embed_tokens.weight.grad.double().norm().item()
embed_tokens.weight.grad.mul_(min(1.0, 1.0 / (gn + 1e-12)))

On the Nemotron Telugu run, GroupClip brought stalled steps to zero and reached the baseline's 1,000-step perplexity (13.04 on FLORES-200) by step 800. In the paper's words, the dead rows "self-heal on first contact": they grow to normal size after their first real update. GroupClip also matched or beat a 10-method ablation that included skip-on-inf, ZClip, freezing the embeddings and FP64 global norms. On Falcon-H1-1.5B Mongolian it lowered step-100 perplexity from 40.40 to 37.47.

For pretraining from scratch they propose GateFloor, which replaces the gate with SiLU(z) + δ, for δ between 0.01 and 0.1. That way a gate never fully closes. In a 370M-parameter 3:1 hybrid trained on 1.57B tokens, it cut gradients from 1e-30 inputs by more than 28 orders of magnitude, at roughly the same validation perplexity (22.28 vs 22.45). The authors say plainly that GateFloor was tested only up to 370M parameters, with one or two seeds per run. Whether it holds at production scale is open.

The paper lists no code repository. If you're about to run full fine-tuning of an untied hybrid model on non-English or chat-template data, there's a cheap check: scan the embedding table for rows with a norm below 1e-6, and watch for inf in your logged gradient norm.

Comments

No comments yet. Be the first.

Write a comment

Comments are published after review.