d9d 0.1.0__py3-none-any.whl

d9d/internals/grad_sync/synchronizer.py

@@ -0,0 +1,257 @@
+import dataclasses
+from collections import defaultdict
+from typing import cast
+
+import torch
+from torch import nn
+from torch.distributed import DeviceMesh
+from torch.distributed.tensor import DTensor, Replicate, Shard
+
+from .bucket import AbstractGradientBucket, LocalGradientBucket, SyncGradientBucket
+
+
+def _find_reduce_mesh(data: DTensor) -> DeviceMesh | None:
+    """
+    Identifies the sub-mesh required for gradient reduction based on tensor placements.
+
+    Args:
+        data: The parameter tensor.
+
+    Returns:
+        The DeviceMesh subset needed for reduction, or None if no reduction is needed.
+    """
+
+    reduce_dims: set[int] = set()
+
+    for dim_i, dim_placement in enumerate(data.placements):
+        match dim_placement:
+            case Replicate():
+                reduce_dims.add(dim_i)
+            case Shard():
+                pass
+            case _:
+                raise ValueError(f"Unknown grad placement: {dim_placement}")
+
+    if len(reduce_dims) == 0:
+        return None
+
+    device_mesh: DeviceMesh = data.device_mesh
+
+    # we are sure that device mesh contain dim names so we cast(...)
+    mesh_dim_names = cast(tuple[str, ...], device_mesh.mesh_dim_names)
+    reduce_mesh = device_mesh[tuple(
+        mesh_dim_names[dim_i] for dim_i in reduce_dims
+    )]
+
+    return reduce_mesh
+
+
+@dataclasses.dataclass(frozen=True)
+class _ParameterGroupMarker:
+    """
+    Identifier for grouping compatible parameters into buckets.
+    """
+
+    group_i: int
+    reduce_mesh: DeviceMesh | None
+    device: torch.device
+    grad_dtype: torch.dtype | None
+
+
+def _group_params_for_buckets(
+        param_groups: list[list[nn.Parameter]]
+) -> dict[_ParameterGroupMarker, list[nn.Parameter]]:
+    """
+    Sorts parameters into groups based on their synchronization requirements.
+
+    Args:
+        param_groups: List of parameter groups (from optimizer).
+
+    Returns:
+        Dictionary mapping group markers to lists of parameters.
+    """
+
+    regrouped_params = defaultdict(list)
+    for param_group_i, param_group in enumerate(param_groups):
+        # iterate in reverse order to maximize overlap
+        for param in param_group[::-1]:
+            if not param.requires_grad:
+                continue
+
+            if not isinstance(param.data, DTensor):
+                raise TypeError("All params should be DTensors in a distributed setup")
+
+            reduce_mesh = _find_reduce_mesh(param.data)
+
+            group = _ParameterGroupMarker(
+                group_i=param_group_i,
+                reduce_mesh=reduce_mesh,
+                device=param.data.device,
+                grad_dtype=param.grad_dtype
+            )
+
+            regrouped_params[group].append(param)
+
+    return regrouped_params
+
+
+def _make_bucket(
+        require_accumulations: int,
+        group_marker: _ParameterGroupMarker,
+        parameters: list[nn.Parameter],
+        communicate_stream: torch.cuda.Stream
+) -> AbstractGradientBucket:
+    """
+    Factory function to create the appropriate bucket type.
+    """
+
+    if group_marker.reduce_mesh is None:
+        return LocalGradientBucket(parameters)
+    else:
+        if group_marker.grad_dtype is None:
+            raise ValueError("Gradient dtype could not be None")
+
+        return SyncGradientBucket(
+            parameters=parameters,
+            require_accumulations=require_accumulations,
+            device=group_marker.device,
+            grad_dtype=group_marker.grad_dtype,
+            reduce_mesh=group_marker.reduce_mesh,
+            communicate_stream=communicate_stream
+        )
+
+
+def _fill_buckets(
+        param_groups: dict[_ParameterGroupMarker, list[nn.Parameter]],
+        bucket_size_mb: int,
+        require_accumulations: int,
+        communicate_stream: torch.cuda.Stream
+) -> list[AbstractGradientBucket]:
+    """
+    Splits grouped parameters into buckets based on size constraints.
+
+    Args:
+        param_groups: Parameters grouped by sync requirements.
+        bucket_size_mb: Max size for each bucket in megabytes.
+        require_accumulations: Number of gradient accumulations required before syncing gradients.
+
+    Returns:
+        List of configured gradient buckets.
+    """
+
+    # TODO: Better grouping - probably we could trace autograd graph and use some topological clustering here
+    # TODO: to maximize overlap even better - current implementation just iterates over parameters in reverse order
+    buckets = []
+
+    bucket_size = bucket_size_mb * 1024 * 1024
+
+    for param_group_marker, param_group in param_groups.items():
+        current_bucket_size = 0
+        unfinished_bucket: list[nn.Parameter] = []
+        for param in param_group:
+            param_bytes = param.numel() * param.element_size()
+            if current_bucket_size + param_bytes >= bucket_size and unfinished_bucket:
+                buckets.append(_make_bucket(
+                    require_accumulations=require_accumulations,
+                    group_marker=param_group_marker,
+                    parameters=unfinished_bucket,
+                    communicate_stream=communicate_stream
+                ))
+                unfinished_bucket = []
+                current_bucket_size = 0
+
+            unfinished_bucket.append(param)
+            current_bucket_size += param_bytes
+
+        if unfinished_bucket:
+            buckets.append(_make_bucket(
+                require_accumulations=require_accumulations,
+                group_marker=param_group_marker,
+                parameters=unfinished_bucket,
+                communicate_stream=communicate_stream
+            ))
+    return buckets
+
+
+class GradientSynchronizer:
+    """
+    Manages gradient synchronization for distributed training.
+
+    This class handles the bucketing of parameters, memory allocation for flat
+    gradient buffers, and the orchestration of asynchronous all-reduce operations
+    during the backward pass.
+    """
+
+    def __init__(
+            self,
+            param_groups: list[list[nn.Parameter]],
+            bucket_size_mb: int,
+            require_accumulations: int
+    ):
+        """
+        Constructs a GradientSynchronizer.
+
+        Args:
+            param_groups: List of parameter groups.
+            bucket_size_mb: Maximal size of a single gradient bucket in MB.
+            require_accumulations: Number of micro-batches to accumulate before reducing.
+        """
+
+        self._param_groups = param_groups
+        self._bucket_size_mb = bucket_size_mb
+        self._require_accumulations = require_accumulations
+
+        self._communicate_stream: torch.cuda.Stream | None = None
+        self._can_sync: bool
+        self._buckets: list[AbstractGradientBucket] = []
+
+    def bind(self):
+        """
+        Initializes the synchronizer for training.
+
+        Groups parameters, creates buckets, allocates memory, and registers hooks.
+        Must be called before the backward pass.
+        """
+
+        stream = torch.cuda.Stream()
+        self._communicate_stream = stream
+        self._buckets = _fill_buckets(
+            _group_params_for_buckets(self._param_groups),
+            bucket_size_mb=self._bucket_size_mb,
+            require_accumulations=self._require_accumulations,
+            communicate_stream=stream
+        )
+
+        for bucket in self._buckets:
+            bucket.bind()
+
+    def unbind(self):
+        """
+        Releases resources.
+
+        Destroys buckets, frees memory buffers, and removes hooks.
+        """
+
+        for bucket in self._buckets:
+            bucket.unbind()
+
+        self._buckets = []
+        self._communicate_stream = None
+
+    def wait(self):
+        """
+        Waits for all bucket operations (async reductions) to complete.
+        """
+
+        torch.cuda.current_stream().wait_stream(self._communicate_stream)
+
+        for bucket in self._buckets:
+            bucket.mark_sync()
+
+    def zero_grad(self):
+        """
+        Resets gradients and accumulation counters for all managed parameters.
+        """
+
+        for bucket in self._buckets:
+            bucket.zero_grad()

d9d/internals/pipeline_state/__init__.py

@@ -0,0 +1,14 @@
+"""
+Pipeline State management package.
+
+This package provides mechanisms to store, retrieve, and synchronize state
+across different stages of a distributed pipeline, providing global and sharded view for these states.
+"""
+
+from .api import PipelineState
+from .handler import PipelineStateHandler
+
+__all__ = [
+    "PipelineState",
+    "PipelineStateHandler"
+]

d9d/internals/pipeline_state/api.py

@@ -0,0 +1,45 @@
+import abc
+from typing import Any
+
+
+class PipelineState(abc.ABC):
+    """
+    Object representing the state of a pipeline.
+
+    This class defines the interface for accessing state variables like a dictionary,
+    abstracting away whether the underlying storage is local, sharded, or global.
+    """
+
+    @abc.abstractmethod
+    def __setitem__(self, key: str, value: Any):
+        """
+        Sets a state value for a given key.
+
+        Args:
+            key: The identifier for the state variable.
+            value: The value to store.
+        """
+
+    @abc.abstractmethod
+    def __getitem__(self, item: str) -> Any:
+        """
+        Retrieves a state value for a given key.
+
+        Args:
+            item: The identifier for the state variable.
+
+        Returns:
+            The value associated with the key.
+        """
+
+    @abc.abstractmethod
+    def __contains__(self, item: str) -> bool:
+        """
+        Checks if a key exists in the state.
+
+        Args:
+            item: The identifier to check.
+
+        Returns:
+            True if the key exists, False otherwise.
+        """

d9d/internals/pipeline_state/handler.py

@@ -0,0 +1,111 @@
+from typing import Any
+
+from d9d.core.sharding import ShardingSpecLeaf
+
+from .api import PipelineState
+from .storage import PipelineStateStorage
+
+
+class PipelineStateGlobal(PipelineState):
+    """
+    Represents the global (unsharded) view of the pipeline state.
+    """
+
+    def __init__(self, storage: PipelineStateStorage):
+        """
+        Constructs a PipelineStateGlobal object.
+
+        Args:
+            storage: The underlying storage backend.
+        """
+
+        self._storage = storage
+
+    def __setitem__(self, key: str, value: Any):
+        self._storage.store_global((key,), value)
+
+    def __getitem__(self, item: str) -> Any:
+        return self._storage.acquire_global((item,))
+
+    def __contains__(self, item: str) -> bool:
+        return self._storage.contains((item,))
+
+
+class PipelineStateShard(PipelineState):
+    """
+    Represents a sharded view of the pipeline state for a specific shard ID.
+    """
+
+    def __init__(self, storage: PipelineStateStorage, current_shard: int):
+        """
+        Constructs a PipelineStateShard object.
+
+        Args:
+            storage: The underlying storage backend.
+            current_shard: The index of the partial shard this view represents.
+        """
+
+        self._storage = storage
+        self._current_shard = current_shard
+
+    def __setitem__(self, key: str, value: Any):
+        self._storage.store_shard((key,), value, self._current_shard)
+
+    def __getitem__(self, item: str) -> Any:
+        return self._storage.acquire_shard((item,), self._current_shard)
+
+    def __contains__(self, item: str) -> bool:
+        return self._storage.contains((item,))
+
+
+class PipelineStateHandler:
+    """
+    Manages the lifecycle and access patterns of pipeline states.
+
+    This handler initializes the underlying storage and provides specific views
+    (global or sharded) into that storage.
+    """
+
+    def __init__(self, sharding_spec: dict[str, ShardingSpecLeaf], num_shards: int):
+        """
+        Constructs a PipelineStateHandler object.
+
+        Args:
+            sharding_spec: A definition of how specific keys should be sharded.
+            num_shards: The total number of shards in the pipeline.
+        """
+
+        self._storage = PipelineStateStorage(
+            sharding_spec={(k,): v for k, v in sharding_spec.items()},
+            num_shards=num_shards
+        )
+
+    def global_state(self) -> PipelineState:
+        """
+        Returns a view interface for accessing global state.
+
+        Returns:
+            A PipelineState interface that accesses the full, aggregated data.
+        """
+
+        return PipelineStateGlobal(self._storage)
+
+    def sharded_state(self, shard_id: int) -> PipelineState:
+        """
+        Returns a view interface for accessing state specific to a shard ID.
+
+        Args:
+            shard_id: The index of the shard to access.
+
+        Returns:
+            A PipelineState interface that accesses partial data for the given shard.
+        """
+
+        return PipelineStateShard(self._storage, shard_id)
+
+    def reset(self):
+        """
+        Resets the underlying storage, clearing all state.
+        """
+
+        self._storage.reset()

d9d/internals/pipeline_state/storage.py

@@ -0,0 +1,236 @@
+import copy
+from collections import UserDict
+from typing import Any, TypeVar, cast
+
+import torch
+import torch.utils._pytree as pytree  # noqa: PLC2701
+
+from d9d.core.sharding import ShardingSpecLeaf, SpecReplicate, SpecShard, shard_tree, unshard_tree
+
+StateKey = tuple[str, ...]
+
+
+TMap = TypeVar("TMap")
+
+
+def _detach_leaf(x: TMap) -> TMap:
+    """
+    Detaches a tensor from the computation graph if the input is a tensor.
+
+    Args:
+        x: The input object.
+
+    Returns:
+        The detached tensor or original object.
+    """
+
+    if isinstance(x, torch.Tensor):
+        return cast(TMap, x.detach())
+    return x
+
+
+class ShardedState(UserDict):
+    """
+    Container for holding state broken down by shard indices.
+    """
+
+
+class PipelineStateStorage:
+    """
+    Low-level storage backend handling sharding and aggregation of state data.
+
+    This class manages the transition between sharded data
+    and global data. It uses sharding specifications to determine
+    how to split or join data.
+    """
+
+    def __init__(
+            self,
+            sharding_spec: dict[StateKey, ShardingSpecLeaf],
+            num_shards: int
+    ):
+        """
+        Constructs a PipelineStateStorage object.
+
+        Args:
+            sharding_spec: Dictionary mapping state keys to their sharding behaviors.
+            num_shards: Total number of shards involved in the storage.
+        """
+
+        self._sharding_spec_orig = copy.deepcopy(sharding_spec)
+
+        self._state: dict[StateKey, Any] = {}
+        self._state_sharding_spec: dict[StateKey, ShardingSpecLeaf] = {}
+
+        self._num_shards = num_shards
+
+    def _guess_sharding_spec_for_shard(self, key: StateKey, shard: Any) -> ShardingSpecLeaf:
+        # Stack if scalar (tensor or item), cat otherwise
+
+        if key in self._sharding_spec_orig:
+            return self._sharding_spec_orig[key]
+
+        if isinstance(shard, torch.Tensor):
+            do_stack = shard.ndim == 0
+            return SpecShard(dim=0, do_stack=do_stack)
+        elif isinstance(shard, list):
+            return SpecShard(dim=0)
+        else:
+            return SpecShard(dim=0, do_stack=True)
+
+    def _guess_sharding_spec_for_global(self, key: StateKey, state: Any) -> ShardingSpecLeaf:
+        if key in self._sharding_spec_orig:
+            return self._sharding_spec_orig[key]
+
+        if isinstance(state, torch.Tensor):
+            if state.ndim == 0:
+                return SpecReplicate()
+            else:
+                return SpecShard(dim=0)
+        elif isinstance(state, list):
+            return SpecShard(dim=0)
+        else:
+            return SpecReplicate()
+
+    def store_global(self, key: StateKey, state: Any):
+        """
+        Stores a value in the global scope.
+
+        If the key does not have a sharding spec, one will be inferred. Detaches tensors.
+
+        Args:
+            key: The identifier key.
+            state: The unified value to store.
+        """
+
+        state = pytree.tree_map(_detach_leaf, state)
+
+        if key not in self._state_sharding_spec:
+            self._state_sharding_spec[key] = self._guess_sharding_spec_for_global(key, state)
+
+        self._state[key] = state
+
+    def store_shard(self, key: StateKey, state: Any, shard_id: int):
+        """
+        Stores a value for a specific shard index.
+
+        Raises error if attempting to shard an already global key without conversion.
+
+        Args:
+            key: The identifier key.
+            state: The partial value for the shard.
+            shard_id: The index of the shard.
+
+        Raises:
+            ValueError: If trying to store sharded state into an unsharded container.
+        """
+
+        if key not in self._state:
+            self._state[key] = ShardedState()
+
+        container = self._state[key]
+
+        if not isinstance(container, ShardedState):
+            raise ValueError(f"Trying to store sharded state into an unsharded one: {key}")
+
+        state = pytree.tree_map(_detach_leaf, state)
+
+        # dynamically populate sharding spec to know whether it is stacking or not
+        if key not in self._state_sharding_spec:
+            self._state_sharding_spec[key] = self._guess_sharding_spec_for_shard(key, state)
+
+        self._state[key][shard_id] = state
+
+    def _ensure_global(self, key: StateKey):
+        if key not in self._state:
+            raise ValueError(f"Cannot access non-existing state {key}")
+
+        state = self._state[key]
+
+        if not isinstance(state, ShardedState):
+            return
+
+        # here we know we are in ShardedState
+
+        shards = [state[shard_id] for shard_id in range(self._num_shards)]
+        resharded = unshard_tree(shards, self._state_sharding_spec[key])
+
+        self._state[key] = resharded
+
+    def _ensure_sharded(self, key: StateKey):
+        if key not in self._state:
+            raise ValueError(f"Cannot access non-existing state {key}")
+
+        state = self._state[key]
+
+        if isinstance(state, ShardedState):
+            return
+
+        # here we know we are in global state
+
+        sharded = shard_tree(
+            state,
+            self._state_sharding_spec[key],
+            num_shards=self._num_shards,
+            enforce_even_split=True
+        )
+
+        sharded_state = ShardedState({
+            shard_idx: shard for shard_idx, shard in enumerate(sharded)
+        })
+
+        self._state[key] = sharded_state
+
+    def acquire_global(self, key: StateKey) -> Any:
+        """
+        Retrieves data for a key in its global form.
+
+        Args:
+            key: The state key.
+
+        Returns:
+            The aggregated global data.
+        """
+
+        self._ensure_global(key)
+        return self._state[key]
+
+    def acquire_shard(self, key: StateKey, shard: int) -> Any:
+        """
+        Retrieves data for a key specific to a shard index.
+
+        Args:
+            key: The state key.
+            shard: The shard index.
+
+        Returns:
+            The data slice corresponding to the shard.
+        """
+
+        self._ensure_sharded(key)
+        state = self._state[key]
+
+        if isinstance(state, ShardedState):
+            return state[shard]
+        else:
+            return state
+
+    def contains(self, key: StateKey) -> bool:
+        """
+        Checks if a key exists in storage.
+
+        Args:
+            key: The state key.
+
+        Returns:
+            True if present.
+        """
+
+        return key in self._state
+
+    def reset(self):
+        """
+        Clears all stored state.
+        """
+
+        self._state.clear()

d9d/internals/profiling/__init__.py

@@ -0,0 +1,7 @@
+"""Exposes the internal distributed profiler."""
+
+from .profile import Profiler
+
+__all__ = [
+    "Profiler"
+]