[DTensor] Fix CI: Python 3.10 compat, lint, and pyi stubs

- Fix _deduplicate_src_specs to use hashable (start, stop) tuples
  instead of slice objects as dict keys (slice is unhashable on
  Python < 3.12)
- Add dtensor_send/dtensor_recv stubs to tensorclass.pyi
- Fix lint: remove unused imports (TensorDictPipe TYPE_CHECKING,
  Sequence in megatron, _TransferPlan/ParameterPlan in tests),
  unused variable (sends in ModelTransferPlan.execute), loop
  variable naming (key -> _key in base.py)
- Fix docstring formatting (D205/D415 in _chunk_slice,
  ModelTransferPlan)
- Add examples/*.py to T201 (print) lint ignore in setup.cfg
- Remove unused variable in example file
- Auto-format with ufmt/black

Made-with: Cursor
ghstack-source-id: 4fd0dc9
Pull-Request: #1652
Made-with: Cursor

Author: vmoens

examples/dtensor_transfer_distributed_test.py

@@ -22,11 +22,10 @@
 
 import torch
 import torch.distributed as dist
-from torch.distributed.device_mesh import DeviceMesh
-from torch.distributed.tensor import Shard
-from torch.distributed.tensor import distribute_tensor
 
 from tensordict import TensorDict
+from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.tensor import distribute_tensor, Shard
 
 
 def log(msg: str):
@@ -54,9 +53,7 @@ def test_strategy_a_materialize():
     mesh = DeviceMesh("cuda", torch.arange(world_size))
 
     torch.manual_seed(42)
-    full_a = torch.arange(
-        world_size * 10, dtype=torch.float32, device="cuda"
-    )
+    full_a = torch.arange(world_size * 10, dtype=torch.float32, device="cuda")
     full_b = torch.randn(4, world_size * 8, dtype=torch.float32, device="cuda")
 
     dt_a = distribute_tensor(full_a, mesh, [Shard(0)])
@@ -120,9 +117,7 @@ def test_strategy_b_redistribute():
 
     mesh = DeviceMesh("cuda", torch.arange(world_size))
 
-    full_tensor = torch.arange(
-        world_size * 12, dtype=torch.float32, device="cuda"
-    )
+    full_tensor = torch.arange(world_size * 12, dtype=torch.float32, device="cuda")
     dt = distribute_tensor(full_tensor, mesh, [Shard(0)])
     td_src = TensorDict(weight=dt)
 
@@ -151,9 +146,9 @@ def test_strategy_b_redistribute():
 
         expected_local = list(full_tensor.chunk(world_size))[0]
         received = td_recv["weight"]
-        assert torch.allclose(received, expected_local), (
-            f"weight mismatch: got {received}, expected {expected_local}"
-        )
+        assert torch.allclose(
+            received, expected_local
+        ), f"weight mismatch: got {received}, expected {expected_local}"
         log("  Verification PASSED!")
 
     dist.barrier()
@@ -332,8 +327,10 @@ def main():
 
     torch.cuda.set_device(rank % torch.cuda.device_count())
 
-    log(f"Initialized: world_size={world_size}, "
-        f"device=cuda:{rank % torch.cuda.device_count()}")
+    log(
+        f"Initialized: world_size={world_size}, "
+        f"device=cuda:{rank % torch.cuda.device_count()}"
+    )
 
     test_plain_tensor()
     test_strategy_a_materialize()

examples/dtensor_transfer_plan_test.py

@@ -12,10 +12,7 @@
 
 import torch
 
-from tensordict._dtensor import (
-    _compute_all_local_slices,
-    _compute_transfer_plan,
-)
+from tensordict._dtensor import _compute_all_local_slices, _compute_transfer_plan
 from torch.distributed.tensor.placement_types import Replicate, Shard
 
 
@@ -52,9 +49,9 @@ def test_shard4_to_shard2():
 
     expected = list(full.chunk(2))
     for i in range(2):
-        assert torch.equal(dst_shards[i], expected[i]), (
-            f"rank {i}: expected {expected[i]}, got {dst_shards[i]}"
-        )
+        assert torch.equal(
+            dst_shards[i], expected[i]
+        ), f"rank {i}: expected {expected[i]}, got {dst_shards[i]}"
     print("  PASSED\n")
 
 
@@ -78,8 +75,7 @@ def test_2d_to_1d():
     print(f"  Number of transfers: {len(plan.transfers)}")
     for t in plan.transfers:
         print(
-            f"    rank {t.src_rank} -> rank {t.dst_rank}: "
-            f"global {t.global_slices}"
+            f"    rank {t.src_rank} -> rank {t.dst_rank}: " f"global {t.global_slices}"
         )
 
     top = full[:5]
@@ -134,8 +130,7 @@ def test_dp_tp_to_tp_only():
     print("\n  Transfers:")
     for t in plan.transfers:
         print(
-            f"    rank {t.src_rank} -> rank {t.dst_rank}: "
-            f"global {t.global_slices}"
+            f"    rank {t.src_rank} -> rank {t.dst_rank}: " f"global {t.global_slices}"
         )
 
     src_shards = {s.rank: full[s.slices].clone() for s in src_specs}
@@ -148,9 +143,9 @@ def test_dp_tp_to_tp_only():
     for s in dst_specs:
         expected = full[s.slices]
         actual = dst_buffers[s.rank]
-        assert torch.equal(actual, expected), (
-            f"rank {s.rank}: mismatch!\n  expected: {expected}\n  got: {actual}"
-        )
+        assert torch.equal(
+            actual, expected
+        ), f"rank {s.rank}: mismatch!\n  expected: {expected}\n  got: {actual}"
     print("  PASSED\n")
 
 
@@ -160,8 +155,6 @@ def test_replicate_to_shard():
     print("Test: Replicate on 4 ranks -> Shard(0) on 2 ranks")
     print("=" * 60)
 
-    full = torch.arange(100, dtype=torch.float32)
-
     plan = _compute_transfer_plan(
         global_shape=[100],
         src_mesh_shape=[4],
@@ -173,12 +166,11 @@ def test_replicate_to_shard():
     print(f"  Number of transfers: {len(plan.transfers)}")
     for t in plan.transfers:
         print(
-            f"    rank {t.src_rank} -> rank {t.dst_rank}: "
-            f"global {t.global_slices}"
+            f"    rank {t.src_rank} -> rank {t.dst_rank}: " f"global {t.global_slices}"
         )
-    assert all(t.src_rank == 0 for t in plan.transfers), (
-        "Expected deduplication to use only rank 0 as source"
-    )
+    assert all(
+        t.src_rank == 0 for t in plan.transfers
+    ), "Expected deduplication to use only rank 0 as source"
     print("  PASSED\n")
 
 

examples/minimal_p2p_test.py

@@ -2,6 +2,7 @@
 """Minimal NCCL P2P test to verify send/recv works."""
 
 import json
+
 import torch
 import torch.distributed as dist
 

setup.cfg

@@ -18,6 +18,7 @@ per-file-ignores =
     packaging/*/**.py: T201
     .github/scripts/*.py: T201
     benchmarks/*.py: T201
+    examples/*.py: T201
 exclude = venv
 extend-select = B901, C401, C408, C409
 

tensordict/_dtensor.py

@@ -17,15 +17,11 @@
 import json
 import struct
 from dataclasses import dataclass, field
-from typing import Any, Callable, Protocol, runtime_checkable, Sequence, TYPE_CHECKING
+from typing import Any, Callable, Protocol, runtime_checkable, Sequence
 
-import numpy as np
 import torch
 from torch import Tensor
 
-if TYPE_CHECKING:
-    from tensordict._ucxx import TensorDictPipe
-
 _has_ucxx = importlib.util.find_spec("ucxx") is not None
 
 
@@ -194,8 +190,9 @@ def sharded(
 
 
 def _chunk_slice(total_size: int, num_chunks: int, chunk_idx: int) -> slice:
-    """Return the slice for ``chunk_idx`` when splitting *total_size* into
-    *num_chunks* using ``torch.chunk`` semantics (last chunk may be smaller).
+    """Return the slice for ``chunk_idx`` when splitting into chunks.
+
+    Uses ``torch.chunk`` semantics (last chunk may be smaller).
     """
     chunk_size = (total_size + num_chunks - 1) // num_chunks
     start = chunk_idx * chunk_size
@@ -313,9 +310,9 @@ def _deduplicate_src_specs(
     if not has_replica:
         return src_specs
 
-    seen: dict[tuple[slice, ...], _ShardSpec] = {}
+    seen: dict[tuple[tuple[int, int], ...], _ShardSpec] = {}
     for spec in src_specs:
-        key = spec.slices
+        key = tuple((s.start, s.stop) for s in spec.slices)
         if key not in seen or spec.rank < seen[key].rank:
             seen[key] = spec
     return list(seen.values())
@@ -356,9 +353,7 @@ def _compute_transfer_plan(
     )
 
     # Deduplicate replicated src specs
-    src_specs_dedup = _deduplicate_src_specs(
-        src_specs, src_placements, src_mesh_shape
-    )
+    src_specs_dedup = _deduplicate_src_specs(src_specs, src_placements, src_mesh_shape)
 
     plan = _TransferPlan(global_shape=global_shape)
 
@@ -422,9 +417,7 @@ def execute_transfer_plan(
 
     if dst_buffer is not None:
         for transfer in recvs:
-            chunk_shape = tuple(
-                s.stop - s.start for s in transfer.global_slices
-            )
+            chunk_shape = tuple(s.stop - s.start for s in transfer.global_slices)
             buf = torch.empty(
                 chunk_shape, dtype=dst_buffer.dtype, device=dst_buffer.device
             )
@@ -486,7 +479,7 @@ def recv_object(self, src: int) -> Any:
         length = int(length_t.item())
         data_t = torch.empty(length, dtype=torch.uint8, device="cuda")
         dist.recv(data_t, src=src, group=self.group)
-        return json.loads(bytes(data_t.cpu().numpy()))
+        return json.loads(bytes(data_t.cpu().tolist()))
 
 
 class _UCXXBackend:
@@ -510,10 +503,10 @@ class _UCXXBackend:
     def __init__(self, endpoint):
         self._endpoint = endpoint
 
-    def _tensor_to_numpy(self, t: Tensor) -> np.ndarray:
-        return np.frombuffer(
-            t.contiguous().view(torch.uint8).numpy(), dtype=np.uint8
-        )
+    def _tensor_to_numpy(self, t: Tensor):
+        import numpy as np
+
+        return np.frombuffer(t.contiguous().view(torch.uint8).numpy(), dtype=np.uint8)
 
     def send_tensor(self, tensor: Tensor, dst: int, *, tag: int = 0) -> None:
         import asyncio
@@ -550,12 +543,16 @@ def recv_object(self, src: int) -> Any:
         return asyncio.run(self._arecv_object())
 
     async def _asend_object(self, obj: Any) -> None:
+        import numpy as np
+
         data = json.dumps(obj).encode("utf-8")
         length = struct.pack("<Q", len(data))
         await self._endpoint.send(np.frombuffer(length, dtype=np.uint8))
         await self._endpoint.send(np.frombuffer(data, dtype=np.uint8).copy())
 
     async def _arecv_object(self) -> Any:
+        import numpy as np
+
         len_buf = np.empty(8, dtype=np.uint8)
         await self._endpoint.recv(len_buf)
         length = struct.unpack("<Q", len_buf.tobytes())[0]
@@ -642,9 +639,9 @@ class ParameterPlan:
 
 
 class ModelTransferPlan:
-    """Precomputed plan for transferring an entire model's parameters
-    between two differently-sharded meshes.
+    """Precomputed plan for transferring an entire model's parameters.
 
+    Transfers between two differently-sharded meshes.
     Designed for LLM post-training: compute once at setup, execute
     every training iteration with near-zero overhead.
 
@@ -662,7 +659,9 @@ class ModelTransferPlan:
         )
     """
 
-    def __init__(self, param_plans: list[ParameterPlan], batches: list[list[ParameterPlan]]):
+    def __init__(
+        self, param_plans: list[ParameterPlan], batches: list[list[ParameterPlan]]
+    ):
         self._param_plans = param_plans
         self._batches = batches
 
@@ -803,7 +802,6 @@ def execute(
             if src_tensor is not None and pp.transform is not None:
                 src_tensor = pp.transform(src_tensor)
 
-            sends = pp.plan.sends_for_rank(rank)
             recvs = pp.plan.recvs_for_rank(rank)
 
             # Allocate dst buffer if this rank receives data
@@ -832,9 +830,7 @@ def execute(
         return result
 
     @staticmethod
-    def _rank_coords_for(
-        rank: int, desc: ShardingDescriptor
-    ) -> tuple[int, ...]:
+    def _rank_coords_for(rank: int, desc: ShardingDescriptor) -> tuple[int, ...]:
         """Find the mesh coordinates for *rank* in the descriptor's mesh."""
         if desc.rank_map is not None:
             for coords, r in desc.rank_map.items():
@@ -889,7 +885,9 @@ def summary(self) -> str:
         if n_optimal:
             lines.append(f"  Strategy C (optimal P2P): {n_optimal} params")
         if n_materialize:
-            lines.append(f"  Strategy A (materialize): {n_materialize} params (have transforms)")
+            lines.append(
+                f"  Strategy A (materialize): {n_materialize} params (have transforms)"
+            )
         if n_direct:
             lines.append(f"  Direct copy: {n_direct} params (same sharding)")
         lines.append(f"  Total transfer: {self.total_bytes / 1024**2:.1f} MB (float32)")

tensordict/base.py

@@ -9595,6 +9595,7 @@ def _recv(
             )
 
         return _tag
+
     def dtensor_send(
         self,
         dst,
@@ -9762,7 +9763,7 @@ def _dtensor_send_materialize(self, dst, *, backend) -> None:
 
         dst_int = dst if isinstance(dst, int) else 0
         backend.send_object(metadata, dst_int)
-        for key, tensor in tensors:
+        for _key, tensor in tensors:
             backend.send_tensor(tensor.contiguous(), dst_int)
 
     def _dtensor_recv_materialize(self, src, *, backend) -> None:
@@ -9833,7 +9834,7 @@ def _dtensor_send_redistribute(self, dst, *, backend) -> None:
 
         dst_int = dst if isinstance(dst, int) else 0
         backend.send_object(metadata, dst_int)
-        for key, tensor in tensors:
+        for _key, tensor in tensors:
             backend.send_tensor(tensor.contiguous(), dst_int)
 
     def _dtensor_recv_redistribute(self, src, *, backend) -> None:
@@ -9878,13 +9879,12 @@ def _dtensor_send_optimal(self, dst, *, backend, dst_mesh, dst_placements) -> No
         Computes which slices of each local shard need to go to which dst
         rank, then issues targeted P2P sends for just those slices.
         """
-        from torch import distributed as dist
-
         from tensordict._dtensor import (
             _compute_transfer_plan,
             _mesh_all_ranks,
             _mesh_to_rank_map,
         )
+        from torch import distributed as dist
 
         my_rank = dist.get_rank()
 
@@ -9946,13 +9946,12 @@ def _dtensor_recv_optimal(self, src, *, backend, src_mesh, src_placements) -> No
         Computes which slices this rank needs and from which src ranks,
         then issues targeted P2P recvs and assembles the local shard.
         """
-        from torch import distributed as dist
-
         from tensordict._dtensor import (
             _compute_transfer_plan,
             _mesh_all_ranks,
             _mesh_to_rank_map,
         )
+        from torch import distributed as dist
 
         my_rank = dist.get_rank()
 
@@ -10011,9 +10010,7 @@ def _dtensor_recv_optimal(self, src, *, backend, src_mesh, src_placements) -> No
                     local_tensor[transfer.dst_slices] = buf
                     tag += 1
 
-                self._set_str(
-                    key, local_tensor, inplace=True, validated=True
-                )
+                self._set_str(key, local_tensor, inplace=True, validated=True)
             else:
                 # Non-DTensor: recv from the first src rank
                 first_src = _mesh_all_ranks(src_mesh)[0]