Repository: huggingface/transformers
Pull request: Tensor Parallelism and `mps` device

# What does this PR do?

The current implementation does not work with the `mps` device and TP.

## Example script

script.py
```
import os
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

world_size = int(os.environ.get("WORLD_SIZE", "1"))
device = "cpu" # or "mps"

# model_id = "meta-llama/Llama-3.2-3B-instruct"
model_id = "Qwen/Qwen3-0.6B"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id , dtype=torch.bfloat16, tp_plan="auto",
)
model = model.to(device)

inputs = tokenizer("Hello, what's your name? My name is Michael.", return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
model.train()

rank = torch.distributed.get_rank() if world_size > 1 else 0

for step in range(10):
    outputs = model(**inputs, labels=inputs["input_ids"])
    loss = outputs.loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    print(f"Step {step}, TP Rank: {rank}, Loss: {loss.item()}")
```

**Command**: `uv run torchrun --nproc_per_node=2 ./script.py`

### With the `mps` device
```
W0306 12:57:46.508000 11409 torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs.
W0306 12:57:46.519000 11409 torch/distributed/run.py:852]
W0306 12:57:46.519000 11409 torch/distributed/run.py:852] *****************************************
W0306 12:57:46.519000 11409 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W0306 12:57:46.519000 11409 torch/distributed/run.py:852] *****************************************
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 311/311 [00:00<00:00, 12702.59it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 311/311 [00:00<00:00, 12944.87it/s]
The tied weights mapping and config for this model specifies to tie model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning
The tied weights mapping and config for this model specifies to tie model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning
/Users/michaelbenayoun/projects/peft/.venv/lib/python3.13/site-packages/torch/autograd/function.py:583: UserWarning: The operator 'c10d::allreduce_' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:15.)
  return super().apply(*args, **kwargs)  # type: ignore[misc]
/Users/michaelbenayoun/projects/peft/.venv/lib/python3.13/site-packages/torch/autograd/function.py:583: UserWarning: The operator 'c10d::allreduce_' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:15.)
  return super().apply(*args, **kwargs)  # type: ignore[misc]
Step 0, TP Rank: 0, Loss: 8.390717506408691
Step 0, TP Rank: 1, Loss: 13.627806663513184
Step 1, TP Rank: 0, Loss: 7.506132125854492
Step 1, TP Rank: 1, Loss: 7.839929103851318
Step 2, TP Rank: 0, Loss: 6.150381088256836
Step 2, TP Rank: 1, Loss: 6.395097732543945
Step 3, TP Rank: 1, Loss: 5.616857051849365
Step 3, TP Rank: 0, Loss: 5.9035115242004395
Step 4, TP Rank: 0, Loss: 4.873610973358154
Step 4, TP Rank: 1, Loss: 9.064188003540039
Step 5, TP Rank: 1, Loss: 7.211321830749512
Step 5, TP Rank: 0, Loss: 2.953089714050293
Step 6, TP Rank: 1, Loss: 5.9867024421691895
Step 6, TP Rank: 0, Loss: 3.9124104976654053
Step 7, TP Rank: 0, Loss: 3.087076425552368
Step 7, TP Rank: 1, Loss: 4.329252243041992
Step 8, TP Rank: 0, Loss: 2.5245583057403564
Step 8, TP Rank: 1, Loss: 2.348144054412842
Step 9, TP Rank: 1, Loss: 1.581630825996399
Step 9, TP Rank: 0, Loss: 4.455328941345215
```

### With the `cpu` device

```
W0306 12:59:49.657000 12843 torch/distributed/elastic/multiprocessing/redirects.py:29] NOTE: Redirects are currently not supported in Windows or MacOs.
W0306 12:59:49.667000 12843 torch/distributed/run.py:852]
W0306 12:59:49.667000 12843 torch/distributed/run.py:852] *****************************************
W0306 12:59:49.667000 12843 torch/distributed/run.py:852] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W0306 12:59:49.667000 12843 torch/distributed/run.py:852] *****************************************
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 311/311 [00:00<00:00, 14302.16it/s]
Loading weights: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 311/311 [00:00<00:00, 12377.74it/s]
The tied weights mapping and config for this model specifies to tie model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning
The tied weights mapping and config for this model specifies to tie model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with `tie_word_embeddings=False` to silence this warning
Step 0, TP Rank: 0, Loss: 2.9389801025390625
Step 0, TP Rank: 1, Loss: 2.9389801025390625
Step 1, TP Rank: 1, Loss: 0.28367936611175537
Step 1, TP Rank: 0, Loss: 0.28367936611175537
Step 2, TP Rank: 1, Loss: 2.2101123332977295
Step 2, TP Rank: 0, Loss: 2.2101123332977295
Step 3, TP Rank: 1, Loss: 1.304605484008789
Step 3, TP Rank: 0, Loss: 1.304605484008789
Step 4, TP Rank: 0, Loss: 0.023114478215575218
Step 4, TP Rank: 1, Loss: 0.023114478215575218
Step 5, TP Rank: 0, Loss: 0.006346769165247679
Step 5, TP Rank: 1, Loss: 0.006346769165247679
Step 6, TP Rank: 1, Loss: 0.0032010390423238277
Step 6, TP Rank: 0, Loss: 0.0032010390423238277
Step 7, TP Rank: 1, Loss: 0.00309620494954288
Step 7, TP Rank: 0, Loss: 0.00309620494954288
Step 8, TP Rank: 1, Loss: 0.00309001631103456
Step 8, TP Rank: 0, Loss: 0.00309001631103456
Step 9, TP Rank: 0, Loss: 0.0029638134874403477
Step 9, TP Rank: 1, Loss: 0.0029638134874403477
```

### Suggestion

We can see that the loss is not properly aligned with the `mps` device, suggesting that the current fallback to `cpu` does not work as expected.

I do not think it makes a lot of sense to spend a lot of time on this since doing TP on the `mps` device is not really relevant, so I suggest to simply fail to notice the user that it does not work.

Changed files:
filename: src/transformers/integrations/tensor_parallel.py
patch: @@ -55,7 +55,7 @@ def initialize_tensor_parallelism(
         # Detect the accelerator on the machine. If no accelerator is available, it returns CPU.
         device_type = torch._C._get_accelerator().type
         if device_type == "mps":
-            device_type = "cpu"  # fallback
+            raise RuntimeError("Tensor parallelism is not supported on MPS devices.")
         current_device = getattr(torch, device_type)
         if not torch.distributed.is_initialized():
             try: