pytorch / pytorch/pytorch

[DTensor] clip_grad_norm_ fails when per-parameter norms come from different DeviceMesh objects

Open
#180,346 1 comment 0 reactions 0 assignees View on GitHub
bot-triaged enhancement module: dtensor oncall: distributed oncall: distributed parallelisms ptd-bot-triaged triaged
Dominant language
Python
Stars
103k
Forks
29.7k
PR merge metrics
PR metrics pending

Description

### 🚀 The feature, motivation and pitch

`torch.nn.utils.clip_grad._get_total_norm()` groups tensors by (device, dtype) only, computes per-tensor norms, and then stacks those norms together. For DTensor inputs, the norms themselves remain DTensors and retain mesh metadata, even if such tensors are scalars ultimately. If the norms come from different DeviceMesh objects, the final `torch.stack(...)` fails even though the tensors share the same device and dtype.

Expected: `clip_grad_norm_` handles mixed meshes correctly. I considered extracting DTensor scalar norms via `item()`, but that only seems safe when each norm is known to be a replicated scalar. For general DTensor layouts, `item()` reads the local value and can be numerically wrong, so it does not look like a sound generic fix.

Actual: the generic `torch.stack` path fails because the norm DTensors carry different mesh metadata. This follows directly from the current `_get_total_norm()` implementation.

### Alternatives

_No response_

### Additional context

This seems distinct from the already reported performance issue in the same code path (https://github.com/pytorch/pytorch/issues/169445): this report is about a runtime failure with mixed meshes, not about stack complexity.

Related https://github.com/pytorch/pytorch/issues/121020.

Reproduction
```python
import torch
import torch.distributed as dist
from torch.distributed._tensor import DTensor, Replicate
from torch.distributed.device_mesh import DeviceMesh
from torch.testing._internal.common_distributed import spawn_threads_and_init_comms

WORLD_SIZE = 2

@spawn_threads_and_init_comms(world_size=WORLD_SIZE)
def repro_mixed_mesh_clip_grad_norm_failure(world_size: int) -> None:
rank = dist.get_rank()

mesh_a = DeviceMesh("cpu", torch.arange(world_size))
mesh_b = DeviceMesh("cpu", torch.arange(world_size - 1, -1, -1))

placements = [Replicate()]

local_a = torch.tensor([3.0 + rank], dtype=torch.float32)
local_b = torch.tensor([4.0 + rank], dtype=torch.float32)

dt_a = DTensor.from_local(local_a, mesh_a, placements)
dt_b = DTensor.from_local(local_b, mesh_b, placements)

norm_type = 2.0
norms = [torch.linalg.vector_norm(g, norm_type) for g in [dt_a, dt_b]]

first_device = dt_a.device
total_norm = torch.linalg.vector_norm(
torch.stack([norm.to(first_device) for norm in norms]),
norm_type,
)
print(rank, total_norm)

if __name__ == "__main__":
repro_mixed_mesh_clip_grad_norm_failure(WORLD_SIZE)
```

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy @tianyu-l @XilunWu @SherlockNoMad @ppwwyyxx @xmfan

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.