d9d 0.1.0__py3-none-any.whl

d9d/core/dist_ops/object.py

@@ -0,0 +1,68 @@
+from typing import TypeVar, cast
+
+import torch.distributed as dist
+
+T = TypeVar("T")
+
+
+def gather_object(
+        obj: T,
+        group: dist.ProcessGroup,
+        group_dst: int
+) -> list[T] | None:
+    """
+    Gathers picklable objects from the whole process group to a specific destination rank.
+
+    This acts as a wrapper around torch.distributed.gather_object that automatically
+    initializes the output buffer list on the destination rank.
+
+    Args:
+        obj: The local object to send. Must be picklable.
+        group: The process group to work on.
+        group_dst: The rank within the group that will receive the objects.
+
+    Returns:
+        A list of objects from all ranks on the destination rank; None on other ranks.
+    """
+
+    if group.rank() == group_dst:
+        # We initialize with None, but we cast to list[T] because we know
+        # dist.gather_object will populate these slots with actual objects.
+        save_list = cast(list[T], [None for _ in range(group.size())])
+    else:
+        save_list = None
+    dist.gather_object(
+        obj,
+        save_list,
+        group=group,
+        group_dst=group_dst
+    )
+    return save_list
+
+
+def all_gather_object(
+        obj: T,
+        group: dist.ProcessGroup
+) -> list[T]:
+    """
+    Gathers picklable objects from the whole process group to all ranks.
+
+    This acts as a wrapper around torch.distributed.all_gather_object that automatically
+    initializes the output buffer list on all ranks.
+
+    Args:
+        obj: The local object to send. Must be picklable.
+        group: The process group to work on.
+
+    Returns:
+        A list of objects containing the data gathered from all ranks.
+    """
+    # We initialize with None, but we cast to list[T] because we know
+    # dist.gather_object will populate these slots with actual objects.
+    save_list = cast(list[T], [None for _ in range(group.size())])
+    dist.all_gather_object(
+        save_list,
+        obj,
+        group=group
+    )
+    return save_list

d9d/core/dist_ops/tensor.py

@@ -0,0 +1,192 @@
+from collections.abc import Sequence
+from typing import cast
+
+import torch
+import torch.distributed as dist
+
+
+def gather(
+        tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+        group_dst: int,
+        async_op: bool = False
+) -> list[torch.Tensor] | tuple[list[torch.Tensor] | None, dist.Work] | None:
+    """
+    Gathers tensors from the process group to a specific destination rank.
+
+    This function assumes that tensors on all ranks have the same shape and dtype
+    as the tensor on the current rank. It automatically allocates the output
+    buffer list on the destination.
+
+    Args:
+        tensor: The local tensor to send.
+        group: The process group to work on.
+        group_dst: The rank within the group that will receive the tensors.
+        async_op: Whether the operation should be asynchronous.
+
+    Returns:
+        If async_op is False: A list of tensors on the destination rank, None elsewhere.
+        If async_op is True: A tuple containing (buffer_list, work_handle).
+    """
+
+    if group.rank() == group_dst:
+        save_list = [torch.empty_like(tensor) for _ in range(group.size())]
+    else:
+        save_list = None
+
+    work = dist.gather(
+        tensor,
+        save_list,
+        group=group,
+        group_dst=group_dst,
+        async_op=async_op
+    )
+
+    if async_op:
+        return save_list, work
+    else:
+        return save_list
+
+
+def all_gather(
+        tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+        async_op: bool = False
+) -> list[torch.Tensor] | tuple[list[torch.Tensor], dist.Work]:
+    """
+    Gathers tensors from the whole process group to all ranks.
+
+    This function assumes that tensors on all ranks have the same shape and dtype
+    as the tensor on the current rank. It automatically allocates the output
+    buffer list.
+
+    Args:
+        tensor: The local tensor to send.
+        group: The process group to work on.
+        async_op: Whether the operation should be asynchronous.
+
+    Returns:
+        If async_op is False: A list of gathered tensors.
+        If async_op is True: A tuple containing (buffer_list, work_handle).
+    """
+
+    save_list = [torch.empty_like(tensor) for _ in range(group.size())]
+    work = dist.all_gather(
+        save_list,
+        tensor,
+        group=group,
+        async_op=async_op
+    )
+    if async_op:
+        return save_list, work
+    else:
+        return save_list
+
+
+def _all_gather_shapes(
+        tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+) -> Sequence[torch.Tensor]:
+    all_ndim = [torch.empty((), dtype=torch.long, device=tensor.device) for _ in range(group.size())]
+    all_ndim_wait = dist.all_gather(
+        all_ndim,
+        torch.tensor(tensor.ndim, dtype=torch.long, device=tensor.device),
+        group=group,
+        async_op=True
+    )
+    all_ndim_wait.wait()
+
+    all_shape = [torch.empty(cast(int, ndim.item()), dtype=torch.long, device=tensor.device) for ndim in all_ndim]
+    all_shape_wait = dist.all_gather(
+        all_shape,
+        torch.tensor(tensor.shape, dtype=torch.long, device=tensor.device),
+        group=group,
+        async_op=True
+    )
+    all_shape_wait.wait()
+
+    return all_shape
+
+
+def all_gather_variadic_shape(
+        tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+        async_op: bool = False
+) -> list[torch.Tensor] | tuple[list[torch.Tensor], dist.Work]:
+    """
+    Gathers tensors of different shapes from the whole process group to all ranks.
+
+    Unlike standard all_gather, this function first communicates the shape of the
+    tensor on every rank allowing for dynamic sizing.
+
+    Args:
+        tensor: The local tensor to send.
+        group: The process group to work on.
+        async_op: Whether the final data gathering should be asynchronous.
+                  Note that shape gathering is always synchronous.
+
+    Returns:
+        If async_op is False: A list of gathered tensors of varying shapes.
+        If async_op is True: A tuple containing (buffer_list, work_handle).
+    """
+
+    all_shape = _all_gather_shapes(tensor, group)
+
+    all_result = [torch.empty(tuple(shape), dtype=tensor.dtype, device=tensor.device) for shape in all_shape]
+    all_result_wait = dist.all_gather(
+        all_result,
+        tensor,
+        group=group,
+        async_op=async_op
+    )
+    if async_op:
+        return all_result, all_result_wait
+    else:
+        return all_result
+
+
+def gather_variadic_shape(
+        tensor: torch.Tensor,
+        group: dist.ProcessGroup,
+        group_dst: int
+) -> list[torch.Tensor] | None:
+    """
+    Gathers tensors of different shapes from the process group to a specific rank.
+
+    This function coordinates shape exchange and uses point-to-point communication
+    (isend/irecv) to gather tensors that may differ in shape across ranks.
+
+    Currently, does not support async_op.
+
+    Args:
+        tensor: The local tensor to send.
+        group: The process group to work on.
+        group_dst: The rank within the group that will receive the tensors.
+
+    Returns:
+        A list of tensors of varying shapes on the destination rank; None on other ranks.
+    """
+
+    is_current_dst = group.rank() == group_dst
+
+    all_shape = _all_gather_shapes(tensor, group)
+
+    if is_current_dst:
+        all_recv_futures: list[dist.Work] = []
+        all_result: list[torch.Tensor] = cast(list[torch.Tensor], [None for _ in range(group.size())])
+        for group_src_i in range(group.size()):
+            if group_src_i == group_dst:
+                all_result[group_src_i] = tensor
+                continue
+            all_result[group_src_i] = torch.empty(
+                tuple(all_shape[group_src_i]), dtype=tensor.dtype, device=tensor.device
+            )
+            all_recv_future = dist.irecv(all_result[group_src_i], group=group, group_src=group_src_i)
+            all_recv_future = cast(dist.Work, all_recv_future)  # we know we are on dst rank
+            all_recv_futures.append(all_recv_future)
+        for recv_future in all_recv_futures:
+            recv_future.wait()
+        return all_result
+    else:
+        dist.isend(tensor=tensor, group=group, group_dst=group_dst)
+        return None

d9d/core/protocol/__init__.py

@@ -0,0 +1,8 @@
+"""Package providing protocol definitions for standard PyTorch objects."""
+
+from .training import LRSchedulerProtocol, OptimizerProtocol
+
+__all__ = [
+    "LRSchedulerProtocol",
+    "OptimizerProtocol"
+]

d9d/core/protocol/training.py

@@ -0,0 +1,38 @@
+from typing import Protocol, runtime_checkable
+
+from torch.distributed.checkpoint.stateful import Stateful
+
+
+@runtime_checkable
+class OptimizerProtocol(Protocol, Stateful):
+    """
+    Protocol defining an interface for standard PyTorch Optimizer object.
+
+    This protocol ensures that the wrapped optimizer supports standard
+    API and state checkpointing via the Stateful interface.
+    """
+
+    def step(self):
+        """Performs a single optimization step."""
+
+        ...
+
+    def zero_grad(self):
+        """Sets the gradients of all optimized tensors to zero."""
+
+        ...
+
+
+@runtime_checkable
+class LRSchedulerProtocol(Protocol, Stateful):
+    """
+    Protocol defining an interface for a Learning Rate Scheduler.
+
+    This protocol ensures that the wrapped scheduler supports stepping
+    and state checkpointing via the Stateful interface.
+    """
+
+    def step(self):
+        """Performs a single learning rate scheduling step."""
+
+        ...

d9d/core/sharding/__init__.py

@@ -0,0 +1,15 @@
+from .auto_spec import shard_spec_nothing, shard_spec_on_dim
+from .shard import shard_tree
+from .spec import ShardingSpec, ShardingSpecLeaf, SpecReplicate, SpecShard
+from .unshard import unshard_tree
+
+__all__ = [
+    "ShardingSpec",
+    "ShardingSpecLeaf",
+    "SpecReplicate",
+    "SpecShard",
+    "shard_spec_nothing",
+    "shard_spec_on_dim",
+    "shard_tree",
+    "unshard_tree"
+]

d9d/core/sharding/auto_spec.py

@@ -0,0 +1,66 @@
+from typing import Any
+
+import torch
+import torch.utils._pytree as pytree  # noqa: PLC2701
+
+from d9d.core.types import PyTree
+
+from .spec import ShardingSpec, ShardingSpecLeaf, SpecReplicate, SpecShard
+
+
+def _tree_item_to_shard(item: Any, shard_on_dim: int) -> ShardingSpecLeaf:
+    if isinstance(item, list):
+        if shard_on_dim != 0:
+            raise ValueError(f"Cannot shard list on dim {shard_on_dim}. Lists behave as 1D sequences.")
+        return SpecShard(0)
+    elif isinstance(item, torch.Tensor):
+        if item.ndim == 0:
+            return SpecReplicate()
+        if item.ndim <= shard_on_dim:
+            raise ValueError(f"Cannot shard {item.ndim}-dimensional tensor on dim {shard_on_dim}")
+        return SpecShard(shard_on_dim)
+    else:
+        return SpecReplicate()
+
+
+def shard_spec_on_dim(tree: PyTree[Any], dim: int) -> ShardingSpec:
+    """
+    Creates a sharding specification to split all tensors in the tree on a specific dimension.
+
+    Iterates over the input tree:
+    *   If a leaf is a Tensor with enough dimensions, it is mapped to a SpecShard(dim) object.
+    *   If a leaf is a list, it is mapped to a SpecShard(0) object (only dim=0 is allowed for lists).
+    *   Other types and 0-dim tensors are mapped to SpecReplicate.
+
+    Args:
+        tree: The input PyTree structure.
+        dim: The dimension index to shard eligible tensors on.
+
+    Returns:
+        A new PyTree matching the input structure, containing SpecShard or SpecReplicate objects.
+
+    Raises:
+        ValueError: If a tensor exists in the tree with rank less than or equal to 'dim'.
+    """
+
+    return pytree.tree_map(
+        lambda x: _tree_item_to_shard(x, dim),
+        tree,
+        is_leaf=lambda x: isinstance(x, (torch.Tensor, list))
+    )
+
+
+def shard_spec_nothing(tree: PyTree[Any]) -> ShardingSpec:
+    """
+    Creates a sharding specification where no sharding is performed.
+
+    This effectively clones the tree structure but replaces every leaf with SpecReplicate.
+
+    Args:
+        tree: The input PyTree structure.
+
+    Returns:
+        A new PyTree matching the input structure, containing strictly SpecReplicate for all leaves.
+    """
+
+    return pytree.tree_map(lambda _: SpecReplicate(), tree, is_leaf=lambda x: isinstance(x, (torch.Tensor, list)))

d9d/core/sharding/shard.py

@@ -0,0 +1,154 @@
+from collections.abc import Sequence
+from typing import TypeVar, cast
+
+import torch
+import torch.utils._pytree as pytree  # noqa: PLC2701
+
+from d9d.core.types import PyTree
+
+from .spec import ShardingSpec, SpecReplicate, SpecShard
+
+TLeaf = TypeVar("TLeaf")
+TSameTree = TypeVar("TSameTree", bound=PyTree)
+
+
+def _shard_list(
+        item: list[TLeaf],
+        spec: SpecShard,
+        num_shards: int,
+        enforce_even_split: bool
+) -> Sequence[list[TLeaf] | TLeaf]:
+    if spec.dim != 0:
+        raise ValueError(f"Lists can only be sharded on dim 0, got {spec.dim}")
+
+    if spec.do_stack:
+        if len(item) != num_shards:
+            raise ValueError(
+                f"do_stack=True requires list length ({len(item)}) to match num_shards ({num_shards})"
+            )
+        return item
+
+    if enforce_even_split and len(item) % num_shards != 0:
+        raise ValueError(
+            f"Tried to shard a list with length {len(item)} "
+            f"to {num_shards} shards, but the length is not perfectly divisible."
+        )
+
+    shard_size, shard_extra = divmod(len(item), num_shards)
+    return [
+        item[
+            shard_id * shard_size + min(shard_id, shard_extra):
+            (shard_id + 1) * shard_size + min(shard_id + 1, shard_extra)
+        ]
+        for shard_id in range(num_shards)
+    ]
+
+
+def _shard_tensor(
+        item: torch.Tensor,
+        spec: SpecShard,
+        num_shards: int,
+        enforce_even_split: bool
+) -> Sequence[torch.Tensor]:
+    if item.ndim == 0:
+        raise ValueError("Found a 0-dim Tensor for sharding")
+
+    if spec.do_stack:
+        if item.shape[spec.dim] != num_shards:
+            raise ValueError(
+                f"do_stack=True requires tensor shape[{spec.dim}] ({item.shape[spec.dim]}) "
+                f"to match num_shards ({num_shards})"
+            )
+        return torch.unbind(item, dim=spec.dim)
+
+    if enforce_even_split and item.shape[spec.dim] % num_shards != 0:
+        raise ValueError(
+            f"Tried to shard a tensor with shape {item.shape} on dim {spec.dim} "
+            f"to {num_shards} shards, but the dimension is not perfectly divisible."
+        )
+
+    return torch.tensor_split(item, sections=num_shards, dim=spec.dim)
+
+
+def _shard_leaf_to_list(
+        item: TLeaf,
+        spec: SpecShard | SpecReplicate,
+        num_shards: int,
+        enforce_even_split: bool
+) -> Sequence[TLeaf]:
+    """Helper to split an item into a list of items for each rank."""
+    if isinstance(spec, SpecReplicate):
+        # Replicated: strict copy reference for all shards
+        return [item] * num_shards
+
+    if not isinstance(spec, SpecShard):
+        raise TypeError(f"Unknown sharding spec object type: {type(spec)}")
+
+    if isinstance(item, torch.Tensor):
+        return cast(Sequence[TLeaf], _shard_tensor(
+            item=item,
+            num_shards=num_shards,
+            enforce_even_split=enforce_even_split,
+            spec=spec
+        ))
+    elif isinstance(item, list):
+        return cast(Sequence[TLeaf], _shard_list(
+            item=item,
+            num_shards=num_shards,
+            enforce_even_split=enforce_even_split,
+            spec=spec
+        ))
+    else:
+        raise TypeError(
+            f"Sharding spec found a SpecShard object, but the item was not a Tensor and not a list (got {type(item)})"
+        )
+
+
+def shard_tree(
+        tree: TSameTree,
+        sharding_spec: ShardingSpec,
+        num_shards: int,
+        enforce_even_split: bool
+) -> tuple[TSameTree, ...]:
+    """
+    Shards a PyTree into a tuple of PyTrees, one for each shard rank.
+
+    This function takes a single global data structure and splits it into `num_shards`
+    structures.
+
+    *   If a spec leaf is a ``SpecShard(dim)``, the tensor or list is split along that dimension,
+        and the ``i``-th slice goes to the ``i``-th output tree.
+    *   If a spec leaf is ``SpecReplicate``, the item is replicated (reference copy) to all
+        output trees.
+
+    Args:
+        tree: The structure containing tensors to be sharded.
+        sharding_spec: A structure matching 'tree' containing ``SpecShard`` or ``SpecReplicate`` objects.
+        num_shards: The total number of shards to split the tensors into.
+        enforce_even_split: If True, raises a ValueError if a tensor's dimension
+            size is not perfectly divisible by ``num_shards``.
+
+    Returns:
+        A tuple of length ``num_shards``. Each element is a PyTree matching
+        the structure of the input ``tree``, containing the local data for
+        that specific rank.
+
+    Raises:
+        ValueError: If tree structures do not match, or valid sharding conditions
+            are not met.
+    """
+    flat_spec, spec_struct = pytree.tree_flatten(sharding_spec)
+
+    try:
+        flat_tree = spec_struct.flatten_up_to(tree)
+    except (ValueError, TypeError) as e:
+        raise ValueError("Tree structure does not match sharding spec") from e
+
+    sharded_leaves_per_node = [
+        _shard_leaf_to_list(item, spec, num_shards, enforce_even_split)
+        for item, spec in zip(flat_tree, flat_spec, strict=True)
+    ]
+
+    rank_leaves = list(zip(*sharded_leaves_per_node, strict=True))
+
+    return tuple(spec_struct.unflatten(leaves) for leaves in rank_leaves)

d9d/core/sharding/spec.py

@@ -0,0 +1,28 @@
+import dataclasses
+
+from d9d.core.types import PyTree
+
+
+@dataclasses.dataclass(slots=True, frozen=True)
+class SpecReplicate:
+    """
+    Specifies that a leaf node should be replicated across all shards.
+    """
+
+
+@dataclasses.dataclass(slots=True, frozen=True)
+class SpecShard:
+    """
+    Specifies that a leaf node should be split along a specific dimension.
+
+    Attributes:
+        dim: The dimension to split.
+        do_stack: If True, sharding will squeeze the sharded dimension (it should be exactly the num_shards length)
+    """
+
+    dim: int
+    do_stack: bool = False
+
+
+ShardingSpecLeaf = SpecReplicate | SpecShard
+ShardingSpec = PyTree[ShardingSpecLeaf]

d9d/core/sharding/unshard.py

@@ -0,0 +1,117 @@
+from collections.abc import Sequence
+from typing import TypeVar, cast
+
+import torch
+import torch.utils._pytree as pytree  # noqa: PLC2701
+
+from d9d.core.types import PyTree
+
+from .spec import ShardingSpec, ShardingSpecLeaf, SpecReplicate, SpecShard
+
+TLeaf = TypeVar("TLeaf")
+TSameTree = TypeVar("TSameTree", bound=PyTree)
+
+
+def _unshard_list(
+        group: Sequence[list[TLeaf] | TLeaf],
+        spec: SpecShard
+) -> list[TLeaf]:
+    if spec.dim != 0:
+        raise ValueError(f"Lists can only be unsharded on dim 0, got {spec.dim}")
+
+    if spec.do_stack:
+        return cast(list[TLeaf], list(group))
+
+    merged_list: list[TLeaf] = []
+    for x in group:
+        merged_list.extend(cast(list[TLeaf], x))
+    return merged_list
+
+
+def _unshard_tensor(
+        group: list[torch.Tensor],
+        spec: SpecShard
+) -> torch.Tensor:
+    if spec.do_stack:
+        return torch.stack(group, dim=spec.dim)
+
+    return torch.cat(group, dim=spec.dim)
+
+
+def _unshard_leaf_from_group(
+        group: Sequence[TLeaf],
+        spec: ShardingSpecLeaf
+) -> TLeaf:
+    """Helper to merge a group of items from different ranks into one."""
+    if isinstance(spec, SpecReplicate):
+        return group[0]
+
+    if not isinstance(spec, SpecShard):
+        raise TypeError(f"Unknown sharding spec object type: {type(spec)}")
+
+    first_item = group[0]
+
+    if isinstance(first_item, torch.Tensor):
+        return cast(TLeaf, _unshard_tensor(
+            cast(list[torch.Tensor], group),
+            spec
+        ))
+    elif spec.do_stack or isinstance(first_item, list):
+        return cast(TLeaf, _unshard_list(group, spec))
+    else:
+        raise TypeError(f"Expected Tensor or list instances, got {type(group[0])}")
+
+
+def unshard_tree(
+        sharded_trees: Sequence[TSameTree],
+        sharding_spec: ShardingSpec
+) -> TSameTree:
+    """
+    Combines a sequence of PyTrees (one per rank) into a single global PyTree.
+
+    This is the inverse of ``shard_tree``. It iterates over the provided trees,
+    gathering corresponding leaves from each rank.
+
+    *   If the spec for a leaf is ``SpecShard(dim)``, the tensors from all ranks are
+        concatenated (or stacked if ``do_stack=True``) along that dimension.
+    *   If the spec is ``SpecReplicate``, it assumes the data is replicated
+        and takes the item from the first rank.
+
+    Args:
+        sharded_trees: A sequence (list or tuple) of PyTrees. There must be
+            one tree for each shard rank, and they must all share the same
+            structure as ``sharding_spec``.
+        sharding_spec: A structure matching the input trees containing
+            ``SpecShard`` or ``SpecReplicate`` objects.
+
+    Returns:
+        A single PyTree where distinct shards have been merged into full tensors.
+
+    Raises:
+        ValueError: If ``sharded_trees`` is empty, or if unit structures do
+            not match the spec.
+    """
+    if not sharded_trees:
+        raise ValueError("sharded_trees sequence cannot be empty")
+
+    flat_spec, spec_struct = pytree.tree_flatten(sharding_spec)
+
+    flat_shards_per_rank = []
+    for i, tree in enumerate(sharded_trees):
+        try:
+            leaves = spec_struct.flatten_up_to(tree)
+        except (ValueError, TypeError) as e:
+            raise ValueError(
+                f"Structure mismatch at shard {i}: tree does not match sharding spec structure"
+            ) from e
+
+        flat_shards_per_rank.append(leaves)
+
+    grouped_leaves = list(zip(*flat_shards_per_rank, strict=True))
+
+    reconstructed_leaves = [
+        _unshard_leaf_from_group(group, spec)
+        for group, spec in zip(grouped_leaves, flat_spec, strict=True)
+    ]
+
+    return spec_struct.unflatten(reconstructed_leaves)