Tensor Parallelism with Quantized Models
Author: HuangBugWei
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.52.0.dev0
- Platform: Linux-5.4.17-2136.339.5.el8uek.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.13
- Huggingface_hub version: 0.30.1
- Safetensors version: 0.5.2
- Accelerate version: 1.6.0
- Accelerate config:    not found
- DeepSpeed version: 0.16.3
- PyTorch version (GPU?): 2.6.0+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: yes
- Using GPU in script?: yes
- GPU type: NVIDIA GeForce RTX 3090

### Who can help?

@zucchini-nlp @gante 

### Information

- [ ] The official example scripts
- [x] 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

```python
import torch
from transformers import AutoTokenizer, Gemma3ForCausalLM

model_name = "google/gemma-3-12b-it-qat-int4-unquantized"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,      
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = Gemma3ForCausalLM.from_pretrained(
    model_name, 
    torch_dtype=torch.bfloat16,
    # device_map="auto",
    quantization_config=bnb_config,
    tp_plan="auto",
    ).eval()

print(model.model.layers[0].mlp.down_proj)
# Linear4bit(in_features=15360, out_features=3840, bias=False)
print(model.model.layers[0].mlp.down_proj.weight)
# DTensor(local_tensor=tensor([[-0.0032, -0.0145,  0.0042,  ...,  0.0006, -0.0073,  0.0051],
#     [ 0.0015,  0.0056, -0.0059,  ...,  0.0033, -0.0056, -0.0012],
#     [ 0.0074,  0.0083, -0.0046,  ...,  0.0042,  0.0012, -0.0051],
#     ...,
#     [ 0.0076, -0.0005, -0.0044,  ..., -0.0013,  0.0123,  0.0025],
#     [-0.0014, -0.0096,  0.0026,  ..., -0.0032, -0.0060, -0.0024],
#     [ 0.0034,  0.0026, -0.0089,  ..., -0.0012, -0.0013,  0.0035]],
#    device='cuda:0', dtype=torch.bfloat16), device_mesh=DeviceMesh('cuda', [0, 1]), placements=(Shard(dim=1),))
```

then ```torchrun --nproc-per-node 2 script.py```

This takes about 20GB of GPU memory on each of the two, while the 4-bit model only takes 12GB in total.

### Expected behavior

I'm trying to use tensor parallelism with int4 quantized models. 
However, I've noticed that the model weights are being loaded in bf16 format instead of int4 as expected.

I understand that representing int4 can be complex, potentially involving concatenation and storage in int8. 
However, I also came across this PR: [https://github.com/huggingface/transformers/pull/37790](https://github.com/huggingface/transformers/pull/37790) which mentions "Fix tensor parallel with non-floating dtypes."

This leads me to wonder:
Is tensor parallelism currently supported for int4 quantized models in Transformers, and the bf16 loading indicates a problem in my script?
Or, is using tensor parallelism with int4 quantized models not currently a supported feature in Transformers due to the complexities of int4 representation?

Any clarification on this would be greatly appreciated.

Thanks!


Comments:
@HuangBugWei I think you are using unquantized version of Gemma 3. Can you please check your model name once? 

```
import torch
from transformers import AutoTokenizer, Gemma3ForCausalLM
model_name = "google/gemma-3-12b-it-qat-int4-unquantized"`

@Tanuj-rai 
Thank you for your reply! 
The 'unquantized' in the model name indeed indicates that it hasn't been quantized. 
My understanding is that because 4-bit models cannot currently be stored in the `safetensors` format, they convert the model to floating-point format after training. 
We can then quantize it ourselves when using it (which is what my subsequent quantization configuration, `BitsAndBytesConfig`, is for).


cc @SunMarc for quantization

cc @MekkCyber as he added support for tp + quant for some methods 

Hi @HuangBugWei, for now we don't support quantizing a model while using tp, we can only quantize using `device_map`

Hi @MekkCyber ,
Thank you for the clarification.
I was wondering if I had made a mistake in my script.🫠😅