Quantized model loading failing with custom device mesh
Author: radulescupetru
Pull request: no
Labels: name: Tensor Parallel

name: bug

### System Info

```
- `transformers` version: 4.55.3
- Platform: Linux-5.10.228-219.884.amzn2.x86_64-x86_64-with-glibc2.35
- Python version: 3.11.13
- Huggingface_hub version: 0.34.4
- Safetensors version: 0.6.2
- Accelerate version: 1.5.2
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.8.0+cu128 (CUDA)
- 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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H200
```

### Who can help?

@ArthurZucker @SunMarc @MekkCyber 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

I'm having this issue when trying to load a quantized model when passing a custom TP mesh. Here's a quick repro below. Running the same script, but with the bf16 model goes through with no error.
Command: `torchrun --standalone --nnodes=1 --nproc-per-node=8 script.py`
```
import os

import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from transformers import AutoModelForCausalLM

if __name__ == "__main__":
    world_size = os.environ["WORLD_SIZE"]
    tp = 4
    dp = 2
    device_mesh = init_device_mesh("cuda", (tp, dp), mesh_dim_names=("tp", "dp"))
    model = AutoModelForCausalLM.from_pretrained(
        "Qwen/Qwen3-0.6B-FP8", torch_dtype="auto", device_map="auto", device_mesh=device_mesh["tp"]
    )
    dist.destroy_process_group()

-----------------
[rank2]:     model = AutoModelForCausalLM.from_pretrained(
[rank2]:             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/conda/lib/python3.11/site-packages/transformers/models/auto/auto_factory.py", line 600, in from_pretrained
[rank2]:     return model_class.from_pretrained(
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/conda/lib/python3.11/site-packages/transformers/modeling_utils.py", line 317, in _wrapper
[rank2]:     return func(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/conda/lib/python3.11/site-packages/transformers/modeling_utils.py", line 5074, in from_pretrained
[rank2]:     ) = cls._load_pretrained_model(
[rank2]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/conda/lib/python3.11/site-packages/transformers/modeling_utils.py", line 5537, in _load_pretrained_model
[rank2]:     _error_msgs, disk_offload_index, cpu_offload_index = load_shard_file(args)
[rank2]:                                                          ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/conda/lib/python3.11/site-packages/transformers/modeling_utils.py", line 975, in load_shard_file
[rank2]:     disk_offload_index, cpu_offload_index = _load_state_dict_into_meta_model(
[rank2]:                                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/conda/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank2]:     return func(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/conda/lib/python3.11/site-packages/transformers/modeling_utils.py", line 806, in _load_state_dict_into_meta_model
[rank2]:     not hf_quantizer.check_quantized_param(
[rank2]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/conda/lib/python3.11/site-packages/transformers/quantizers/quantizer_finegrained_fp8.py", line 151, in check_quantized_param
[rank2]:     if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
[rank2]:                                    ^^^^^^^^^^^^^^^^^
[rank2]: AttributeError: 'builtins.PySafeSlice' object has no attribute 'dtype'
```

### Expected behavior

Quantized model is loaded correctly.

Comments:
cc @SunMarc and @S1ro1 

cc @MekkCyber as it touches fp8 with tp 

hey @radulescupetru, did you try to pass `tp_plan="auto"` and not `device_mesh` ? For loading quantized model with tp,  I wouldn't recommend using `device_mesh` for now as this is not indented for and we added this quite recently for training purposes. 

I think this will fail anyway, even with only tp_plan. Will check what's the issue

@SunMarc It fails with the same error when removing `device_mesh` and adding `tp_plan=auto` (Need to also remote `device_map=auto` as they are mutually exclusive).
My use case involves having the model loaded in a submesh so it I can scale the dp dimension. Afaik `tp_plan=auto` or `device_map=auto` will use all the available devices to shard the model, right?

@S1ro1 The script at https://github.com/huggingface/accelerate/blob/5dd3d0b6901983bcb6de1b69687333a324382524/examples/fsdp2/nd_parallel.py fails for this model as well, same error, when using `accelerate launch ...` too.
The odd thing is that if running eleuther_evaluation or evaluation using lighteval this doesn't show../

So the issue for this is that the `check_quantized_param` method expects params as tensors, we pass them as PySlice objects, which do not have `dtype` attribute, therefore this [check](https://github.com/huggingface/transformers/blob/4f9b4e62bc52a52b19a6a4a1a6bfc61a3f5b65b1/src/transformers/quantizers/quantizer_finegrained_fp8.py#L149C1-L159C1) fails. I assume it's gonna be as simple as converting the slice to param with slice[...] and then checking dtype, will check that out!

@S1ro1 Changing that line to `tensor_name == "weight" and param_value.get_dtype() != 'F8_E4M3'` also works. I guess it can be change to detect whether a regular param is passed or a PySlice object.

@radulescupetru Perfect! Do you want to open the PR with the fix yourself or should I? I think it would have to be fixed in all the quantizers we have

@S1ro1 I could give it a go. I took a quick look, but it appears that `PySafeSlice` class is not exposed from rust so it makes it a bit harder to add that check. Do you have any preferences on how that should be handled?
My initial thought was to change the annotation to `param_value: "torch.Tensor" | "PySafeSlice"` in all implementations and add an extra check based on the types here: https://github.com/huggingface/safetensors/blob/54b35101e4fead20d5104468e1093250f267ead1/bindings/python/py_src/safetensors/torch.py#L400
In the absence of that I'm not sure what would be the cleanest thing, could do `try...catch` and add a comment saying that during TP param_value comes as PySafeSlice...
Let me know your thoughts.


I think the cleanest would be to use `hasattr(slice, "dtype")` and use the mapping you sent if it doesn't have it.

Sorry for being a bit late to the party. I think the simplest way is to use the method `slice_dtype = slice.get_dtype()` if the param is a `PySafeSlice` or don't have a `dtype` attribute

@MekkCyber See this PR: https://github.com/huggingface/transformers/pull/40425
It checks for the dtype attribute. I couldn't find a way to import `PySafeSlice` to check for that type, if you have some pointers in that regard I can modify the code.


Thanks for the fast reply and the PR @radulescupetru, I think checking `param_value.__class__.__name__=='PySafeSlice'` should be enough, no need to import the `PySafeSlice` class

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.