InternLM2RotaryEmbedding.inv_freq is never re-initialized by from_pretrained β€” reads as uninitialized memory (denormals, sometimes NaN), not the RoPE frequency table

#2
by francistm - opened

Environment

  • transformers: 5.15.0
  • torch: 2.13.0 (CPU)
  • Python: 3.14.7
  • Platform: Linux, x86_64
  • Model: internlm/internlm2_5-1_8b-chat, revision 57dcbc5455bf050d27a80f0ebd7808387caa714d
  • trust_remote_code=True (this model ships its own modeling_internlm2.py)

Expected behavior

After AutoModelForCausalLM.from_pretrained(...), every layer's
attention.rotary_emb.inv_freq should hold the RoPE frequency table
1 / (rope_theta ** (arange(0, head_dim, 2) / head_dim)), as computed by
InternLM2RotaryEmbedding.__init__.

Actual behavior

inv_freq holds whatever bytes happened to be in the memory allocated for it β€”
denormals and zeros in this reproduction, literal NaN in some runs β€” not the computed
frequency table. Downstream, rotary_emb(...) returns NaN cos/sin, and the whole
forward pass is NaN from the first attention layer.

Reproduction

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "internlm/internlm2_5-1_8b-chat", dtype=torch.float32, trust_remote_code=True
)
inv_freq = model.model.layers[0].attention.rotary_emb.inv_freq
expected = 1.0 / (1_000_000.0 ** (torch.arange(0, 128, 2, dtype=torch.float32) / 128))
print("inv_freq[:8] =", inv_freq[:8].tolist())
print("expected[:8] =", expected[:8].tolist())
print("torch.allclose(inv_freq, expected) =", torch.allclose(inv_freq, expected))

Output, this session, this environment (loaded from a local snapshot at the revision
above β€” confirmed byte-identical to the file raw from the Hub repo's main branch, so
this is not a local-only artifact):

inv_freq[:8] = [1.1749077385612814e-25, 3.0644996116319425e-41, 0.0, 0.0, 7.006492321624085e-45, 0.0, 7.006492321624085e-45, 0.0]
expected[:8] = [1.0, 0.8058422207832336, 0.6493816375732422, 0.5232991576194763, 0.4216965138912201, 0.33982083201408386, 0.2738419771194458, 0.22067339718341827]
torch.allclose(inv_freq, expected) = False

Note expected[0] is exactly 1.0 (base**0); inv_freq[0] is a denormal nowhere near
that. This is not a rounding difference β€” the buffer was never written.

Root cause

InternLM2RotaryEmbedding.__init__ computes inv_freq and registers it
persistent=False:

inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)

persistent=False is correct on its own β€” the value is derived from config, not learned,
so it has no business in the checkpoint's state dict. The problem is downstream:
from_pretrained's standard loading path constructs the model, then calls
InternLM2PreTrainedModel._init_weights to fill in whatever the checkpoint didn't
supply. That method only handles two cases:

def _init_weights(self, module):
    std = self.config.initializer_range
    if isinstance(module, nn.Linear):
        module.weight.data.normal_(mean=0.0, std=std)
        if module.bias is not None:
            module.bias.data.zero_()
    elif isinstance(module, nn.Embedding):
        module.weight.data.normal_(mean=0.0, std=std)
        if module.padding_idx is not None:
            module.weight.data[module.padding_idx].zero_()

Nothing here (or anywhere else in the file) recomputes inv_freq for an
InternLM2RotaryEmbedding module. Since the buffer is persistent=False, it is never in
the checkpoint either, so nothing ever writes it after the module's storage is allocated
β€” it is simply left holding whatever the allocator handed back.

Suggested fix

Add a branch to _init_weights that recomputes inv_freq for
InternLM2RotaryEmbedding (this also covers its two subclasses,
InternLM2LinearScalingRotaryEmbedding and InternLM2DynamicNTKScalingRotaryEmbedding,
since both call the same __init__ and neither overrides the base inv_freq formula):

elif isinstance(module, InternLM2RotaryEmbedding):
    module.inv_freq.data.copy_(
        1.0 / (module.base ** (
            torch.arange(0, module.dim, 2, dtype=torch.int64, device=module.inv_freq.device).float()
            / module.dim
        ))
    )

Sign up or log in to comment