The confusion surrounding bnb_4bit_compute_dtype, torch_dtype, and prepare_model_for_kbit_training.

[xiaobingbuhuitou]

I want to implement QLoRA fine-tuning for a model with dtype=Float32 based on PEFT. When I load the base model using from_pretrained("PATH", BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16 )), without setting torch_dtype, the model's dtype changes to Float16, and the obtained last_hidden_state also becomes Float16. However, when I set torch_dtype=Float32, both the model's dtype and last_hidden_state remain Float32. But when I wrap the quantized model with prepare_model_for_kbit_training(), everything changes back to Float32. I would like to know if using the prepare_model_for_kbit_training() function causes bnb_4bit_compute_dtype and torch_dtype to become ineffective. Additionally, I would like to ask when it is necessary to set prepare_model_for_kbit_training(). Furthermore, what determines the impact on the base model's and last_hidden_state's data types? Thank you.

[ErenAta16]

prepare_model_for_kbit_training() is exactly why everything ends up back at Float32, and it happens unconditionally, independent of torch_dtype or bnb_4bit_compute_dtype.

Looking at the current implementation in peft/utils/other.py:

for param in model.parameters():
    if (
        (param.dtype == torch.float16) or (param.dtype == torch.bfloat16)
    ) and param.__class__.__name__ != "Params4bit":
        param.data = param.data.to(torch.float32)

It walks every parameter in the model and upcasts any fp16/bf16 tensor to fp32, except ones whose class is Params4bit, which is the packed 4-bit weight format bitsandbytes uses. That's the standard QLoRA setup: keep the frozen 4-bit base weights quantized, but put everything else, layernorms, embeddings, the lm head, any adapters, into fp32 for training stability, since mixed fp16/bf16 optimizer states get numerically fragile fast, especially for layernorm.

That explains both things you're seeing:

Without torch_dtype, from_pretrained keeps the non-quantized layers in whatever dtype the checkpoint was saved in, often fp16, so before calling prepare_model_for_kbit_training your hidden states come out fp16.

With torch_dtype=torch.float32, those same layers are created in fp32 from the start, so the later call to prepare_model_for_kbit_training is effectively a no-op for you, it was already fp32.

After calling prepare_model_for_kbit_training, regardless of what you started with, everything except Params4bit gets force-upcast to fp32. That's the "changes back to Float32" behavior in your first case.

On the other two questions:

bnb_4bit_compute_dtype doesn't touch the dtype of the model's parameters or hidden states at all. It only controls what dtype bitsandbytes dequantizes into internally, on the fly, for the matmul inside a Linear4bit layer (dequantize 4-bit, cast to compute_dtype, matmul, cast back). It's a compute-precision setting scoped to the quantized layers themselves.

prepare_model_for_kbit_training is needed whenever you're doing QLoRA (load_in_4bit or load_in_8bit) and want stable training. Beyond the dtype casting, it also freezes the base model's parameters and registers the input-embedding hook gradient checkpointing needs to work through frozen, non-differentiable-input layers. Skipping it on a quantized base model is a common cause of loss not decreasing or gradient-checkpointing errors.

So yes, calling it will override whatever torch_dtype/bnb_4bit_compute_dtype combination you set, for every parameter except the frozen 4-bit weights. That's by design, not a bug.