RoPE corruption due to TF32 on B200 (possibly others) when using `torch.set_float32_matmul_precision("high")`

System Info

• transformers version: 4.51.3
• Platform: Linux
• Python version: 3.10.19
• PyTorch version (GPU?): torch 2.7.0a0+79aa174 (torch.version.cuda 12.9)
• GPU: reproduced on NVIDIA B200 (Blackwell, sm_100, capability 10.0); not reproduced on A100-SXM4-80GB (capability 8.0)
• Using GPU in script?: yes

Description

The rotary-embedding angle computation

freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)

in Qwen3RotaryEmbedding.forward (and derived functions, Llama, Qwen2/3, Mistral, etc.) is a float32 matrix multiply. The surrounding torch.autocast(..., enabled=False) block correctly forces fp32 dtype, but does not control the multiplication backend. When torch.set_float32_matmul_precision("high") is set, the runtime may dispatch this matmul to a GEMM tensor-core kernel which may be run in TF32 (depending on architecture). This is for example the case when ran on B200 with batch size > 1.

TF32 truncates the fp32 mantissa from 23 to 10 bits. This is problematic for the position ids: TF32 can only represent every integer exactly up to 2^11 = 2048. Beyond that the grid spacing exceeds 1, so position_ids such as 2049 get truncated (to 2048). This causes large errors in RoPE computations, see reproduction below.

Possible fixes

1. Locally force the fp32 backend around the angle computation:

   with torch.backends.cuda.matmul.allow_tf32_ctx(False):   # or set allow_tf32=False around this block
       freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)

1. Avoid the matmul:

   freqs = (position_ids_expanded.float().transpose(1, 2) * inv_freq_expanded.float().transpose(1, 2))

I tested out both on B200/A100 and confirmed that this fixes the issue.

Possibly related reports

• thinking-machines-lab/batch_invariant_ops#23, reproduces the identical RoPE corruption and error magnitude
• ByteDance-Seed/VeOmni#959

Who can help?

@ArthurZucker

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Run the following on B200:

import torch
from transformers import Qwen3Config
from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding

cfg = Qwen3Config(hidden_size=1024, num_attention_heads=8, num_key_value_heads=8,
                  head_dim=128, rope_theta=1e7, max_position_embeddings=8192)
re = Qwen3RotaryEmbedding(cfg).cuda()
pos = torch.arange(0, 4096, device="cuda")
x = torch.zeros(1, 4096, 1, dtype=torch.bfloat16, device="cuda")

def cos_for(prec, bs):
    torch.set_float32_matmul_precision(prec)
    pid = pos[None, :].expand(bs, -1)          # matmul batch = bs
    with torch.no_grad():
        c, _ = re(x, pid)
    return c.float()

for bs in (1, 2, 4):
    d = (cos_for("high", bs) - cos_for("highest", bs)).abs()[:, 2048:].max().item()
    print(f"Qwen3RotaryEmbedding  bs={bs}  | high-vs-highest @2048-4096 = {d:.5f}")

Outputs (depending on GPU used):

# A100-SXM4-80GB
Qwen3RotaryEmbedding  bs=1  | high-vs-highest @2048-4096 = 0.00000
Qwen3RotaryEmbedding  bs=2  | high-vs-highest @2048-4096 = 0.00000
Qwen3RotaryEmbedding  bs=4  | high-vs-highest @2048-4096 = 0.00000

# NVIDIA B200
Qwen3RotaryEmbedding  bs=1  | high-vs-highest @2048-4096 = 0.00000
Qwen3RotaryEmbedding  bs=2  | high-vs-highest @2048-4096 = 1.28516
Qwen3RotaryEmbedding  bs=4  | high-vs-highest @2048-4096 = 1.28516

Expected behavior

The RoPE angle matmul must be computed in true IEEE fp32 regardless of the global float32_matmul_precision setting, matching the intent of the existing autocast(enabled=False) guard. RoPE output numbers between GPU architectures should match (up to corresponding floating point precision).