Non-persistent buffers materialized uninitialized by from_pretrained (_move_missing_keys_from_meta_to_device)
System Info
transformersversion: 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
examplesfolder (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_dictgives correct buffer values; onlyfrom_pretrainedcorrupts them. (low_cpu_mem_usage=Falsedoes not help; it appears to be ignored in recent versions.) - In a real model this caused RoPE
inv_freqbuffers to be polluted with ~1e38 garbage ->freqsoverflow 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.
Source: huggingface/transformers