Give each tensor parallel model in a process its own device mesh - #48383

What does this PR do?

Two tensor parallel models in the same process silently share a communicator, which deadlocks as soon as they are used at the same time.

initialize_tensor_parallelism built its mesh with init_device_mesh(device_type, (tp_size,)). That mesh has no dimension names, and when tp_size == world_size torch's 1-D whole-world shortcut returns the default process group outright, so the second model gets the same communicator as the first.

Naming them apart is needed too, and this is the part that is easy to miss. A DeviceMesh is identified by its ranks, layout, device type, dimension names and creating thread. The process group behind it is not part of that identity, so two meshes over the same ranks compare and hash equal. DTensor's sharding propagation is an LRU cache keyed on OpSchema, which hashes the mesh that way, so the second model hits the first model's entry and its parameters come back holding the first model's mesh. distribute_tensor keeps the mesh it is given; torch.nn.Parameter(dtensor) is what loses it, because aten.detach goes through the cache.

Nothing errors either way. The models look correctly sharded, they just share a communicator.

_build_tp_mesh names the first mesh tp and gives every later one its own process group and its own name (tp_1, tp_2, ...). DistributedHelper.extract_tp_mesh matches the tp_ prefix so continuous batching still finds the TP mesh.

Single-model runs are unchanged apart from the mesh now carrying the name tp.

Repro

torchrun --nproc-per-node 2 repro.py, Qwen3-0.6B, two models, TP 2:

The underlying sharp edge is torch's, and is worth its own issue there: a mesh's identity arguably should include the process group it is backed by. Here is that part on its own, no transformers involved, torchrun --nproc-per-node 2:

pg_a, pg_b = dist.new_group(ranks), dist.new_group(ranks)
mesh_a = DeviceMesh.from_group(pg_a, "cuda")
mesh_b = DeviceMesh.from_group(pg_b, "cuda")
assert mesh_a == mesh_b                                    # they are the same mesh to torch
b = distribute_tensor(torch.randn(8, 8, device="cuda"), mesh_b, [Shard(0)], src_data_rank=None)
assert b.device_mesh is mesh_b                             # still fine here
p = torch.nn.Parameter(b)                                  # aten.detach, through the cache
assert p.device_mesh is mesh_a                             # and now it is on the other mesh
assert p.device_mesh.get_group() is pg_a                   # so its collectives go to pg_a

Naming the two meshes apart (mesh_dim_names=("tp",) and ("tp_gen",)) is enough to keep them distinct, which is what this PR does.

Who can review?

@SunMarc / @3outeille (tensor parallelism), @remi-or (continuous batching)

"""
torchrun --nproc-per-node 2 repro_tp_two_models_share_group.py
"""

import os

import torch
import torch.distributed as dist
from transformers import AutoModelForCausalLM
from transformers.distributed.configuration_utils import DistributedConfig

rank = int(os.environ["RANK"])
world = int(os.environ["WORLD_SIZE"])
dist.init_process_group("nccl")
torch.cuda.set_device(rank)

first = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", dtype=torch.bfloat16, distributed_config=DistributedConfig(tp_size=world))
second = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", dtype=torch.bfloat16, distributed_config=DistributedConfig(tp_size=world))
w1 = first.model.layers[0].self_attn.q_proj.weight
w2 = second.model.layers[0].self_attn.q_proj.weight

if rank == 0:
    print(f"mesh names:         {w1.device_mesh.mesh_dim_names} and {w2.device_mesh.mesh_dim_names}", flush=True)
    print(f"same mesh object:   {w1.device_mesh is w2.device_mesh}", flush=True)
    print(f"same process group: {w1.device_mesh.get_group() is w2.device_mesh.get_group()}", flush=True)

assert w1.device_mesh.get_group() is not w2.device_mesh.get_group(), (
    "the two models share a communicator, so using them concurrently will deadlock"
)
if rank == 0:
    print("REPRO_OK: each model has its own communicator", flush=True)
dist.destroy_process_group()