Non-persistent buffers materialized uninitialized by from_pretrained (_move_missing_keys_from_meta_to_device)

[taoshaoyu]

System Info

• transformers version: 5.17.0
• Platform: Linux-6.8.0-124-generic-x86_64-with-glibc2.35
• Python version: 3.12.0
• Huggingface_hub version: 1.31.0
• Safetensors version: 0.8.0
• Accelerate version: 1.14.0
• Accelerate config: not found
• DeepSpeed version: 0.19.6
• PyTorch version (accelerator?): 2.12.1+cu130 (CUDA)
• Using distributed or parallel set-up in script?:
• Using GPU in script?:
• GPU type: NVIDIA GeForce RTX 3060 Laptop GPU

Who can help?

No response

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

import torch
from transformers import PretrainedConfig, PreTrainedModel

class ToyConfig(PretrainedConfig):
    model_type = "toy"
    def __init__(self, hidden=8, **kw):
        self.hidden = hidden
        super().__init__(**kw)

class ToyModel(PreTrainedModel):
    config_class = ToyConfig
    _tied_weights_keys = []
    all_tied_weights_keys = {}
    base_model_prefix = ""

    def __init__(self, config):
        super().__init__(config)
        self.lin = torch.nn.Linear(config.hidden, config.hidden)
        self.register_buffer("magic", torch.arange(8, dtype=torch.float32),
                             persistent=False)

import tempfile
d = tempfile.mkdtemp()
ToyModel(ToyConfig()).save_pretrained(d)

for trial in range(3):
    m = ToyModel.from_pretrained(d)
    print(f"trial#{trial} expected [0,1,2,3,4,5,6,7], got:", m.magic.tolist())

Actual output (values vary across processes; repeated tiny values like
~4e-41 look like pointer fragments, i.e. uninitialized heap memory):

trial#0 expected [0,1,2,3,4,5,6,7], got: [-1.681428329902701e-06, 4.123320731275774e-41, 1.8150632398209942e+21, 3.397448126755519e-41, 7.087288612715111e+22, 1.774807231375223e+28, 6.977648479786776e+22, 1.3563127992606555e-19]
trial#1 expected [0,1,2,3,4,5,6,7], got: [1.8482483472156955e+31, 7.867477360614283e+34, 1.3563156426940112e-19, 3.0330088035509495e+32, 4.6538183192164926e+33, 1.9943137160545966e-19, 4.462551399272987e+30, 1.3587065824052016e-19]
trial#2 expected [0,1,2,3,4,5,6,7], got: [3.548079985193908e-05, -4.914141955897349e-25, -4.4084102609044554e+30, 2.1205226318872167e-27, -1.7764443341638393e+29, 1.0588361175850368e-22, 2.965495775668143e-26, 6.60211260661292e-11]

Additional notes:

• Loading the same checkpoint via from_config + torch.load +
  load_state_dict gives correct buffer values; only from_pretrained
  corrupts them. (low_cpu_mem_usage=False does not help; it appears to
  be ignored in recent versions.)
• In a real model this caused RoPE inv_freq buffers to be polluted with
  ~1e38 garbage -> freqs overflow to inf -> cos/sin NaN -> the whole
  forward produced NaN, intermittently across processes depending on the
  garbage values.

Expected behavior

magic keeps the value assigned in __init__ ([0.0, 1.0, ..., 7.0]).
Non-persistent buffers should be re-initialized (or re-computed) after
from_pretrained, not left as uninitialized memory.

Suspected root cause: _move_missing_keys_from_meta_to_device
(modeling_utils.py, ~line 4730) moves checkpoint-missing keys — which
include non-persistent buffers — from the meta device to real memory with
allocation only, and the low-cpu-mem loading path skips init_weights for
them.

Workaround on the model side: do not store constants/state in
non-persistent buffers; recompute them inside forward instead.

[zucchini-nlp]

You need to override self._init_weights to init any buffers, since the weight init has changed in v5 release

[KashmirAwana]

It looks like the non-persistent buffer "magic" isn't being initialized correctly when loading from a pretrained model. I'd check how _move_missing_keys_from_meta_to_device handles non-persistent buffers.