RoPE 使用 bfloat16 時，8,192 個模型 token 位置只剩 897 個可區分

Cast the integers 8,190 through 8,200 to bfloat16 and you get eleven copies of 8,192.

That one line is why RoPE in bfloat16 is a trap. Rotary position embeddings tell attention where each token sits, and the number format most LLMs run in can't count past 256 without rounding. There's no NaN and no crash. A 128-token unit test comes back bit-identical. I kept seeing "force float32" around rotary code, so I built a small NumPy simulator to find out how bad the damage really is.

Across an 8,192-token sequence, only 897 distinct positions survived. A head whose entire job is "look at the previous token" dropped from 100% to 3.5%.

TL;DR

• bfloat16 has 8 bits of significand precision, so integers are exact only up to 256. Between 4,096 and 8,192, representable values sit 32 apart.
• Casting position_ids to bfloat16 turns 8,192 positions into 897 distinct values. 89.1% of adjacent token pairs end up at distance zero.
• In a synthetic previous-token head, top-1 accuracy fell from 8,191 of 8,191 queries to 286 (3.5%).
• Rounding the final cos/sin tables, q and k to bfloat16 cost nothing measurable: still 8,191 of 8,191. Only the steps before cos and sin need float32.
• Fix: keep inv_freq, positions and angles in float32, disable autocast around them, then cast cos/sin to the model dtype.

What does RoPE in bfloat16 actually break?

RoPE in bfloat16 breaks position information once a sequence passes 256 tokens, while every short test keeps passing. My setup used head_dim 128 and theta 10,000, with a synthetic attention head built so the correct answer is always token i - 1.

I tested all 8,191 query positions with 3 random heads. A hit means the previous token strictly beat every other key, so ties count as misses.

Your tests live in the first row. Then it falls off a cliff.

Why can't bfloat16 count past 256?

bfloat16 keeps float32's 8 exponent bits but only 7 stored mantissa bits, which gives 8 bits of precision with the implicit leading bit. The gap between neighboring values doubles at every power of two.

The first integer that can't be stored is 257, which rounds to 256. By 131,072 tokens you have 1,409 distinct positions and a worst-case error of 256.

Why does a rounded position wreck attention?

RoPE works because the attention score between two tokens depends only on their distance. It rotates q and k by angles proportional to position, and when you take the dot product, the absolute positions cancel and i - j remains.

Rounding destroys that cancellation, because round(i) - round(j) isn't i - j. Here's what happened to adjacent tokens in my 8,192-token run:

• 89.1% of neighbors had a distance of exactly 0
• 3.1% kept the correct distance of 1
• The rest jumped, up to a gap of 32

So "the previous token" and "this token" are usually the same place. Attention can't separate keys sitting on one rounded position, and a whole block of 32 tokens gets identical rotation.

Which bfloat16 cast in RoPE actually hurts?

The casts before the trig functions hurt and the casts after them don't. I put bfloat16 rounding at one stage at a time and averaged 3 random heads. Attention weight is on the correct token, using a softmax sharpness I picked, so compare those rows to each other and ignore the absolute value.

A cos value lives between -1 and 1, where bfloat16's relative error is harmless. An angle lives in the thousands of radians, and at 8,000 radians a gap of 32 is about five full rotations. The output of cos tolerates rounding. Its input doesn't.

A sampled run at Llama 3's rope_theta of 500,000 showed the same collapse, so a larger base doesn't rescue you.

Why is inv_freq in bfloat16 harder to catch?

An inv_freq table in bfloat16 keeps RoPE's relative property, so nearby attention looks perfect while long-range attention drifts. That 100% row above is misleading.

Each frequency is off by up to 0.37%, the same at every position. Distance still cancels cleanly, but the model now rotates at frequencies it wasn't trained on, and the phase error grows with distance: 0.078 radians at 64 tokens, 1.25 at 1,024, 5.01 at 4,096.

I tested heads aimed at a fixed offset, using 500 random heads per distance:

This is the nastiest variant. Local grammar stays fine and only the far end of your context gets quietly worse.

How does this bug get into real code?

It usually gets in through three ordinary lines, and none of them look like precision decisions.

A dtype borrowed from the activations:

# fine for x in float32, broken past 256 when x is bfloat16
t = torch.arange(seq_len, device=x.device, dtype=x.dtype)
freqs = torch.outer(t, self.inv_freq)

A buffer that gets downcast with the model. model.to(torch.bfloat16) casts floating-point buffers too, so a registered inv_freq is rounded on the spot. Calling .float() on it later changes the dtype and can't bring the lost bits back.

Autocast. If you build angles with a matmul of inv_freq and positions, autocast can run it in bfloat16. Hugging Face's Llama rotary code disables autocast around this exact computation, with a "Force float32" comment. Unsloth's early-2024 Gemma bug writeup also listed RoPE needing float32 instead of bfloat16.

How do you fix RoPE precision in PyTorch?

Compute everything up to cos and sin in float32 with autocast off, and cast only the result.

import torch

def rope_cos_sin(position_ids, head_dim, theta=10000.0, out_dtype=torch.bfloat16):
    device = position_ids.device
    with torch.autocast(device_type=device.type, enabled=False):
        exps = torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim
        inv_freq = 1.0 / (theta ** exps)
        angles = position_ids.to(torch.float32)[..., None] * inv_freq
        angles = torch.cat((angles, angles), dim=-1)
        return angles.cos().to(out_dtype), angles.sin().to(out_dtype)

inv_freq is rebuilt inside the function on purpose, so model.to() has no buffer to round.

Then add a test long enough to fail:

def test_rope_positions_stay_distinct():
    cos, _ = rope_cos_sin(torch.arange(8192), head_dim=128)
    assert torch.unique(cos.float(), dim=0).shape[0] == 8192

def test_inv_freq_stays_float32(model):
    for name, buf in model.named_buffers():
        if "inv_freq" in name:
            assert buf.dtype == torch.float32, name

In the simulator, that row count gave 8,192 for the correct path, 897 for bfloat16 positions and 7,533 for bfloat16 angles. It can't see the inv_freq drift, which is why the second test exists.

Is float16 safer for RoPE than bfloat16?

float16 is more precise and fails louder, and it's still the wrong choice. Its integers are exact up to 2,048, it keeps 4,097 distinct positions out of 8,192, and at 65,520 a position overflows to infinity. cos(inf) is NaN, so at least you'd notice.

bfloat16 has float32's range, so it never overflows. It just keeps rounding, and that's what makes it the quieter bug.

What did this test not measure?

This is a synthetic attention head in NumPy, with bfloat16 emulated by round-to-nearest-even on float32. I didn't run a trained model. Real networks spread position across many heads and lean on content too, so don't read 3.5% as a prediction for your model's accuracy. The PyTorch snippets above are the translation of that logic and weren't what produced these numbers.

To check a real model, run one long prompt through your RoPE path and through a float32 reference, then compare the logits at positions past 1,000.

So why does RoPE in bfloat16 collapse token positions?

RoPE in bfloat16 collapses token positions because bfloat16 has 8 bits of precision, so integers above 256 round to the nearest representable value, and the gap reaches 32 between 4,096 and 8,192. That left 897 distinct positions out of 8,192 in my test, and it broke the i - j cancellation RoPE depends on, taking a synthetic previous-token head from 100% to 3.5%. Keep inv_freq, positions and angles in float32 with autocast off, and cast only the final cos/sin tables to bfloat16.

Written by the developer behind Preterview, an interview prep platform.