Inference with model.generate( ) using a quantized model leads to assertion error

[Sandipan99]

System Info

Linux
transformers==4.52.4
bitsandbytes==0.46.1

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

m = "microsoft/phi-4"

bnb_config = BitsAndBytesConfig(
load_in_4bit=True
)

tokenizer = AutoTokenizer.from_pretrained(m)
model = AutoModelForCausalLM.from_pretrained(m, quantization_config=bnb_config, device_map='auto')
tokenizer.pad_token_id = tokenizer.eos_token_id

inputs = tokenizer.apply_chat_template(prompt, return_tensors="pt").to('cuda')
out = model.generate(inputs, max_new_tokens=50, synced_gpus=True)

Run with:

torchrun --nproc-per-node=2 script.py

Works perfectly fine with a single GPU setup, but produces assertion error when running on multiple GPUs

The error can be traced back to model.generate() function

error: Assertion error, python3.10/site-packages/bitsandbytes/nn/modules.py in fix_4bit_weight_quant_state_from_module
assert module.weight.shape[1] == 1

Expected behavior

Expect the model to execute generation without error

[imrrohitt]

Hi @maintainers 👋,

I'd like to work on this issue and would appreciate it if you could assign it to me.

From my understanding, the problem is that when loading AutoModelForCausalLM with torch_dtype=torch.float16, and the model is offloaded to disk using device_map="auto" with low_cpu_mem_usage=True, some modules (like the LM head) are unexpectedly loaded in float32 instead of float16. This results in dtype mismatches and inconsistencies during inference.

I plan to investigate the loading and offloading flow in transformers.modeling_utils.load_checkpoint_and_dispatch and related code paths to ensure torch_dtype is respected for all submodules, even when they are offloaded.

Let me know if there are any implementation preferences or edge cases you'd like me to consider.

[Sandipan99]

Two points in case it is helpful,

1. The above piece of code used to work with older versions of bitsandbytes and transformers.
2. I also experimented with low_cpu_mem_usage=False, it returned the same error.

[sameerajashyam]

Hi @maintainers 👋,

I'd like to work on this issue and work on a patch to handle the assertion error triggered during model.generate()` with quantized models across multiple GPUs.

My plan:

• Add a detection layer inside the Transformers `from_pretrained() or generation logic to gracefully warn or fallback when unsafe 4-bit quantized weight sharding is detected.
• Optionally provide a config switch to override device_map behavior when BitsAndBytesConfig(load_in_4bit=True)` is set.
• Keep the change backward-compatible with existing single-GPU use cases.

Please assign the issue to me if this direction aligns with your expectations.

[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.