pytorch
diff --git a/‎torch/distributed/tensor/parallel/fsdp.py
Copy file name to clipboardExpand all lines: torch/distributed/tensor/parallel/fsdp.py
+3-1Lines changed: 3 additions & 1 deletion b/‎torch/distributed/tensor/parallel/fsdp.py
Copy file name to clipboardExpand all lines: torch/distributed/tensor/parallel/fsdp.py
+3-1Lines changed: 3 additions & 1 deletion
@@ -304,7 +304,9 @@ def _all_gather_dtensor(
 
     placements = list(copy.deepcopy(tensor.placements))
     # FSDP + TP: [Shard(0), tp_placement] -> [Replicate(), tp_placement]
-    placements[0] = Replicate()
+    # HSDP + TP: [Replicate(), Shard(0), tp_placement] -> [Replicate(), Replicate(), tp_placement]
+    for i in range(0, len(placements) - 1):
+        placements[i] = Replicate()
     tensor = tensor.redistribute(
         device_mesh=tensor.device_mesh,
         placements=placements,