PreTrainedModel._load_pretrained_model doesn't move non-persistent buffers to cpu with `low_cpu_mem_usage`

[shingjan]

System Info

With env variable INIT_INCLUDE_BUFFERS=True in accelerate, big LLM models will be loaded in meta tensors first and use actual weights during inference. However, PreTrainedModel._load_pretrained_model doesn't init non-persistent buffer tensors with low_cpu_mem_usage and accelerate, leading to a ValueError: inv_freq is on the meta device, we need a value to put in on 0. while accelerate tries to add hook to module with named_module_tensors.

Who can help?

cc: @sgugger

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

My repro, after this PR from accelerate:

import torch
from transformers import (
    AutoModelForCausalLM,
)

import os
os.environ["ACCELERATE_INIT_INCLUDE_BUFFERS"] = "1"

with torch.no_grad():
    model_real = AutoModelForCausalLM.from_pretrained(
        "lmsys/vicuna-7b-v1.3", device_map='auto', torch_dtype=torch.float16, trust_remote_code=True
    )

stacktrace:

Traceback (most recent call last):
  File "repro/vicuna.py", line 34, in <module>
    model_real = AutoModelForCausalLM.from_pretrained(
  File "/lib/transformers/src/transformers/models/auto/auto_factory.py", line 511, in from_pretrained
    return model_class.from_pretrained(
  File "/lib/transformers/src/transformers/modeling_utils.py", line 2996, in from_pretrained
    dispatch_model(model, **kwargs)
  File "/lib/accelerate/src/accelerate/big_modeling.py", line 385, in dispatch_model
    attach_align_device_hook_on_blocks(
  File "/lib/accelerate/src/accelerate/hooks.py", line 536, in attach_align_device_hook_on_blocks
    attach_align_device_hook_on_blocks(
  File "/lib/accelerate/src/accelerate/hooks.py", line 536, in attach_align_device_hook_on_blocks
    attach_align_device_hook_on_blocks(
  File "/lib/accelerate/src/accelerate/hooks.py", line 536, in attach_align_device_hook_on_blocks
    attach_align_device_hook_on_blocks(
  File "/lib/accelerate/src/accelerate/hooks.py", line 506, in attach_align_device_hook_on_blocks
    add_hook_to_module(module, hook)
  File "/lib/accelerate/src/accelerate/hooks.py", line 155, in add_hook_to_module
    module = hook.init_hook(module)
  File "/lib/accelerate/src/accelerate/hooks.py", line 253, in init_hook
    set_module_tensor_to_device(module, name, self.execution_device)
  File "/lib/accelerate/src/accelerate/utils/modeling.py", line 277, in set_module_tensor_to_device
    raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.")
ValueError: inv_freq is on the meta device, we need a `value` to put in on 0.

Expected behavior

The above loading should work

[sgugger]

As I said before, this is not something the Transformers library supports, so I'm not surprised this doesn't work. This tensor inv_freq is created at init in the model and doesn't have a key in the state dict since it is a non-persistent buffer so there is nothing we can do for this bug.

[shingjan]

I was thinking that maybe we can put a fix to work around it. We can either:

1. Ignore non-persistent buffer in accelerate.utils.modeling.named_module_tensors. However non-persistent buffers could be used for computation at inference. or
2. Init meta tensor to empty in accelerate.utils.modeling.set_module_tensor_to_device if moving meta tensor to non-meta tensors. This carries the assumption that meta tensor initialization to empty tensor is correct. or
3. Init non-persistent buffers in PreTrainedModel._load_pretrained_model.

By

This tensor inv_freq is created at init in the model and doesn't have a key in the state dict since it is a non-persistent buffer so there is nothing we can do for this bug.

I take that you think no.3 is not a viable solution. If so how do the first two options sound to you? Thanks!

[sgugger]

1. won't work as we need them as you stated
2. won't work either since the tensor needs to have the value set at the init
3. won't work either since we don't know to which value initialize that tensor

It's just not possible to fully initialize on the meta device models that have non-persistent buffers that are initialized at init since we lose that value and have no way to recover it. I'm afraid you will need to manually set back those tensors to the proper value.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.