caching_allocator_warmup crashes with AttributeError when loading bitsandbytes-quantized model with PEFT adapter

[LiaoBill]

System Info

• transformers version: 5.1.0
• Platform: Windows-11-10.0.26200-SP0
• Python version: 3.12.11
• Huggingface_hub version: 1.27.0
• Safetensors version: 0.8.0
• Accelerate version: 1.14.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (GPU?): 2.5.1+cu124 (True)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• bitsandbytes version: 0.50.0

Who can help?

@SunMarc @MekkCyber

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

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

quant_config = BitsAndBytesConfig(load_in_8bit=True)
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    quantization_config=quant_config,
    device_map="auto",
)
# Next call triggers _load_pretrained_model again via load_adapter
model = PeftModel.from_pretrained(base_model, "some-peft-adapter")

Traceback:

File "transformers/modeling_utils.py", line 4677, in caching_allocator_warmup
    total_byte_count = get_total_byte_count(model, accelerator_device_map, hf_quantizer)
File "transformers/modeling_utils.py", line 4634, in get_total_byte_count
    param = model.get_parameter_or_buffer(param_name)
File "transformers/modeling_utils.py", line 4550, in get_parameter_or_buffer
    raise AttributeError(f"`{target}` is neither a parameter, buffer, nor extra state.")
AttributeError: `layers.0.self_attn.q_proj.base_layer.SCB` is neither a parameter, buffer, nor extra state.

Expected behavior

A clear and concise description of what you would expect to happen:
Loading a PEFT adapter on top of a bitsandbytes-quantized model should succeed.

Now, caching_allocator_warmup (introduced in 5.x) crashes because expanded_device_map includes entries like layers.0.self_attn.q_proj.base_layer.SCB, which are bitsandbytes internal tensors stored as plain Python attributes on Int8Params (not registered via register_parameter()/register_buffer()), so get_parameter_or_buffer() cannot find them.

caching_allocator_warmup already exempts HQQ and Quark by skipping the warmup entirely for those quantization methods, but BITS_AND_BYTES is missing from that exemption list, I add that to the list fix the problem #47919

[SunMarc]

We shouldn't exclude bnb from caching_allocator_warmup. Can you check why it doesn't happen when loading

base_model = 
AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    quantization_config=quant_config,
    device_map="auto",
) 

but happens here:

model = PeftModel.from_pretrained(base_model, "some-peft-adapter")

[LiaoBill]

sure, I will dig deeper

[LiaoBill]

I think I find the reason for PeftModel.from_pretraied(xxx)'s error

The difference is when SCB gets materialized, in bnb SCB is a plain attribute on Int8Params that is only assigned during _quantize() (when the weights are moved to CUDA), it never registered as a parameter or buffer, it only appears in state_dict() because Linear8bitLt._save_to_state_dict explicitly writes it there. And caching_allocator_warmup assumes every key in state_dict() can be resolved with model.get_parameter_or_buffer(), for standard PyTorch models it works nicely however not for bnb 8bit once quantization has happened: state_dict() then contains unregistered keys (SCB, weight_format), so as caused the problem: AttributeError: layers.0.self_attn.q_proj.base_layer.SCB is neither a parameter, buffer, nor extra state.

happens like:

• During direct loading, warmup runs before weights are dispatched to GPU — SCB is still None, so state_dict() has no SCB key and nothing breaks.
• During PEFT loading (or the integrated from_pretrained(adapter_dir, ..., device_map) path), the base model is already quantized before the adapter load runs _load_pretrained_model a second time, so expected_keys includes SCB. expand_device_map maps it into the device map, and get_total_byte_count → get_parameter_or_buffer("...SCB") → AttributeError.

so the better way of solving this is to keep warmup for bnb but protect model.get_parameter_or_buffer by try-catch?
SCB is negligible in size (~1/1000th of the weights: a few MB vs tens of GB), so excluding it from the byte count has no practical impact. HQQ/QUARK also benefit since their non-registered keys are skipped the same way.
like:

for param_name, device in accelerator_device_map.items():
    if param_name in tied_param_names:
        continue
    try:
        param = model.get_parameter_or_buffer(param_name)
    except AttributeError:
        # Some quantization methods (e.g. bitsandbytes 8bit) include extra tensors
        # (SCB, weight_format) in state_dict() that are not registered as parameters
        # or buffers — they don't count toward the warmup pre-allocation.
        continue
    ...

what do you think? @SunMarc
If a protection seems good enough I will give a corresponding pr tomorrow

[LiaoBill]

I've changed pr #47919 correspondingly, waiting for furthur discussion~