d9d 0.1.0__py3-none-any.whl

d9d/core/autograd/__init__.py

@@ -0,0 +1,7 @@
+from .grad_context import GLOBAL_GRAD_CONTEXT, GlobalGradContext, GradDirection
+
+__all__ = [
+    "GLOBAL_GRAD_CONTEXT",
+    "GlobalGradContext",
+    "GradDirection"
+]

d9d/core/autograd/grad_context.py

@@ -0,0 +1,85 @@
+from contextlib import contextmanager
+from enum import StrEnum
+
+
+class GradDirection(StrEnum):
+    """
+    Enum representing the specific gradient edges to compute.
+
+    This is used to manually control gradient flow in custom autograd functions
+    during split backward passes.
+
+    Attributes:
+        inputs: Mark gradient edge as pointing to the module's inputs (activations).
+        weight: Mark gradient edge as pointing to the module's parameters (weights).
+    """
+
+    inputs = "inputs"
+    weight = "weights"
+
+
+class GlobalGradContext:
+    """
+    Global state manager for controlling gradient computation in custom autograd functions.
+
+    This context addresses a limitation in PyTorch where custom `torch.autograd.Function`
+    implementations set `ctx.needs_input_grad` to True for all edges requiring grad,
+    even during partial backward passes (e.g., `torch.autograd.backward(inputs=...)`).
+
+    For additional information on this limitation, please refer to a
+        [related issue](https://github.com/pytorch/pytorch/issues/174017).
+
+    This class allows:
+
+    1. For the training code - to explicitly signal which gradient edges (inputs vs weights)
+        should currently be computed, allowing custom ops to skip unnecessary computations.
+    2. For module code - to check whether it's required to compute a gradient edge.
+    """
+
+    def __init__(self):
+        """Constructs a GlobalGradContext object with all directions enabled by default."""
+
+        # both directions by default
+        self._enabled_directions: set[GradDirection] = {GradDirection.inputs, GradDirection.weight}
+
+    def check_direction(self, direction: GradDirection | None) -> bool:
+        """
+        Checks if the gradient calculation for the given direction is currently enabled.
+
+        Args:
+            direction: The direction to check (inputs or weights). If None,
+                returns True.
+
+        Returns:
+            True if the direction is enabled or None is passed, False otherwise.
+        """
+
+        if direction is None:
+            return True
+
+        return direction in self._enabled_directions
+
+    @contextmanager
+    def with_directions(self, *directions: GradDirection):
+        """
+        Context manager that sets the enabled gradient directions.
+
+        This overrides the current state for the duration of the context
+        and restores the previous state afterwards.
+
+        Args:
+            *directions: The gradient directions to enable.
+        """
+        prev_directions = self._enabled_directions
+        self._enabled_directions = set(directions)
+        yield
+        self._enabled_directions = prev_directions
+
+
+GLOBAL_GRAD_CONTEXT = GlobalGradContext()
+"""
+The singleton instance of GlobalGradContext.
+
+This should be used by custom autograd functions to check `GLOBAL_GRAD_CONTEXT.check_direction()`
+during their backward pass.
+"""

d9d/core/dist_context/__init__.py

@@ -0,0 +1,19 @@
+"""
+This package configures the distributed environment and device meshes.
+"""
+
+from .configured import DistributedContext
+from .device_mesh_domains import BATCH_DOMAIN, DENSE_DOMAIN, EXPERT_DOMAIN, FLAT_DOMAIN, REGULAR_DOMAIN
+from .log import build_dist_logger
+from .params import DeviceMeshParameters
+
+__all__ = [
+    "BATCH_DOMAIN",
+    "DENSE_DOMAIN",
+    "EXPERT_DOMAIN",
+    "FLAT_DOMAIN",
+    "REGULAR_DOMAIN",
+    "DeviceMeshParameters",
+    "DistributedContext",
+    "build_dist_logger"
+]

d9d/core/dist_context/configured.py

@@ -0,0 +1,215 @@
+import datetime
+import logging
+import os
+import socket
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+import torch
+from torch.distributed import DeviceMesh
+
+from .device_mesh_domains import ALL_DOMAIN_PROVIDERS, REGULAR_DOMAIN
+from .log import build_dist_logger
+
+if TYPE_CHECKING:
+    from .params import DeviceMeshParameters
+
+
+def _resolve_master_addr() -> str:
+    if "MASTER_ADDR" not in os.environ:
+        return "127.0.0.1"
+
+    master_addr = os.environ["MASTER_ADDR"]
+
+    try:
+        return socket.gethostbyname(master_addr)
+    except OSError:
+        return master_addr
+
+
+def _build_mesh_domains(params: "DeviceMeshParameters") -> dict[str, DeviceMesh]:
+    return {
+        provider.name: provider.build_mesh(params)
+        for provider in ALL_DOMAIN_PROVIDERS
+    }
+
+
+class DistributedContext:
+    """
+    Acts as the single source of truth for the distributed execution environment.
+
+    It acts as the central repository for the distributed configuration, managing the creation
+    and synchronization of PyTorch DeviceMeshes for different domains (Regular domain, Expert Parallel domain, ...).
+
+    All assertions regarding rank placement, group memberships, and parallel topology
+    must be derived from this context to ensure consistency.
+    """
+
+    def __init__(self, params: "DeviceMeshParameters", log_level: int):
+        self._params = params
+
+        if params.is_distributed:
+            meshes = _build_mesh_domains(params)
+            regular_mesh = meshes[REGULAR_DOMAIN]
+
+            self._meshes = meshes
+            self._num_nodes = regular_mesh.size() // torch.cuda.device_count()
+            self._logger = build_dist_logger(
+                f'pp:{regular_mesh.get_local_rank("pp")}-'
+                f'dpr:{regular_mesh.get_local_rank("dp_replicate")}-'
+                f'dps:{regular_mesh.get_local_rank("dp_shard")}-'
+                f'cps:{regular_mesh.get_local_rank("cp_shard")}-'
+                f'cpr:{regular_mesh.get_local_rank("cp_replicate")}-'
+                f'tp:{regular_mesh.get_local_rank("tp")}',
+                level=log_level
+            )
+        else:
+            self._meshes = {}
+            self._num_nodes = 1
+            self._logger = build_dist_logger("local", level=log_level)
+
+        self._local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+        self._global_rank = int(os.environ.get("RANK", "0"))
+
+        self._node_rank = self._global_rank // torch.cuda.device_count()
+
+        self._master_addr = _resolve_master_addr()
+        self._current_device = torch.device("cuda")
+
+        torch.cuda.set_device(self._local_rank)
+
+    @property
+    def logger(self) -> logging.Logger:
+        """Returns the logger instance configured for distributed logging."""
+
+        return self._logger
+
+    def mesh_for(self, domain: str) -> DeviceMesh:
+        """
+        Returns the device mesh view associated with a specific logical domain.
+
+        Available Domains and Dimensions:
+            *   `regular` (`REGULAR_DOMAIN`): The most granular mesh for fully decomposed parallelism.
+                Dimensions: ``('pp', 'dp_replicate', 'dp_shard', 'cp_shard', 'cp_replicate', 'tp')``
+            *   `expert` (`EXPERT_DOMAIN`): Mesh optimized for distributing MoE (Mixture of Experts) layers.
+                Dimensions: ``('pp', 'replicate', 'ep')``
+            *   `dense` (`DENSE_DOMAIN`): Mesh optimized for distributing dense layers.
+                Dimensions: ``('pp', 'dp_replicate', 'dp_cp_shard', 'cp_replicate', 'tp')``
+            *   `batch` (`BATCH_DOMAIN`): Mesh optimized for distributing input data.
+                Dimensions: ``('pp', 'dp', 'cp', 'tp')``
+            *   `flat` (`FLAT_DOMAIN`): Mesh containing a single dimension with all the processes.
+                Dimensions: ``('world')``
+
+        Args:
+            domain: The name of the domain to retrieve.
+
+        Returns:
+            The PyTorch DeviceMesh configured for the requested domain.
+
+        Raises:
+            ValueError: If the specified domain does not exist.
+        """
+
+        if domain not in self._meshes:
+            raise ValueError(f"Domain {domain} does not exist")
+        return self._meshes[domain]
+
+    @property
+    def is_main_process(self) -> bool:
+        """Checks if the current process is the global rank 0."""
+
+        return self._global_rank == 0
+
+    @property
+    def is_local_main_process(self) -> bool:
+        """Checks if the current process is the rank 0 on the specific node."""
+
+        return self._local_rank == 0
+
+    def wait_world(self):
+        """Blocks process execution until all ranks reach this point."""
+
+        torch.distributed.barrier(device_ids=[torch.cuda.current_device()])
+        torch.cuda.synchronize()
+
+    def set_timeout(self, timeout_seconds: float):
+        """
+        Updates the NCCL/process group timeout for all underlying meshes.
+
+        Args:
+            timeout_seconds: New timeout duration in seconds.
+        """
+
+        self.logger.info(f"Setting global timeout to {timeout_seconds} seconds")
+        self.wait_world()
+
+        groups: list[torch.distributed.ProcessGroup | None] = [None]
+        for mesh in self._meshes.values():
+            for dim in range(mesh.ndim):
+                groups.append(mesh.get_group(dim))
+
+        for group in groups:
+            torch.distributed.distributed_c10d._set_pg_timeout(datetime.timedelta(seconds=timeout_seconds), group)  # noqa: SLF001
+
+    @contextmanager
+    def local_main_process_first(self):
+        """
+        Context manager that executes the block on the local main process first.
+
+        Other local ranks wait at the entrance. The local main process waits at the
+        exit to synchronize before continuing.
+        """
+        if not self.is_local_main_process:
+            self.wait_world()
+
+        yield
+
+        if self.is_local_main_process:
+            self.wait_world()
+
+    @contextmanager
+    def main_process_first(self):
+        """
+        Context manager that executes the block on the global main process first.
+
+        All other ranks wait at the entrance. The global main process waits at the
+        exit to synchronize before continuing.
+        """
+
+        if not self.is_main_process:
+            self.wait_world()
+
+        yield
+
+        if self.is_main_process:
+            self.wait_world()
+
+    @property
+    def current_device(self) -> torch.device:
+        """Returns the CUDA device associated with this rank."""
+
+        return self._current_device
+
+    @property
+    def mesh_params(self) -> "DeviceMeshParameters":
+        """Returns the parameters used to initialize this context."""
+
+        return self._params
+
+    @property
+    def master_addr(self) -> str:
+        """Returns the IP address or domain name of the master node."""
+
+        return self._master_addr
+
+    @property
+    def node_rank(self) -> int:
+        """Returns the index of the node this process is running on."""
+
+        return self._node_rank
+
+    @property
+    def num_nodes(self) -> int:
+        """Returns the total number of nodes in the cluster."""
+
+        return self._num_nodes

d9d/core/dist_context/device_mesh_domains.py

@@ -0,0 +1,185 @@
+import abc
+from typing import TYPE_CHECKING
+
+from torch.distributed import DeviceMesh, init_device_mesh
+
+if TYPE_CHECKING:
+    from .params import DeviceMeshParameters
+
+
+class DeviceMeshDomain(abc.ABC):
+    """
+    Abstract base class for a Device Mesh provider.
+
+    A Domain defines a specific strategy for organizing available GPUs into a
+    multidimensional grid (Mesh) to support specific parallelism techniques.
+    """
+
+    @property
+    @abc.abstractmethod
+    def name(self) -> str:
+        """Returns the unique identifier for this mesh domain."""
+
+        ...
+
+    @abc.abstractmethod
+    def build_mesh(self, params: "DeviceMeshParameters") -> DeviceMesh:
+        """
+        Constructs the device mesh configuration.
+
+        Args:
+            params: Global configuration parameters for the distributed environment.
+
+        Returns:
+            The initialized PyTorch DeviceMesh for this specific domain.
+        """
+
+        ...
+
+
+REGULAR_DOMAIN = "regular"
+
+
+class RegularDomain(DeviceMeshDomain):
+    @property
+    def name(self) -> str:
+        return "regular"
+
+    def build_mesh(self, params: "DeviceMeshParameters") -> DeviceMesh:
+        return init_device_mesh(
+            device_type="cuda",
+            mesh_shape=(
+                params.pipeline_parallel,
+                params.data_parallel_replicate,
+                params.data_parallel_shard,
+                params.context_parallel_shard,
+                params.context_parallel_replicate,
+                params.tensor_parallel
+            ),
+            mesh_dim_names=(
+                "pp",
+                "dp_replicate",
+                "dp_shard",
+                "cp_shard",
+                "cp_replicate",
+                "tp"
+            )
+        )
+
+
+EXPERT_DOMAIN = "expert"
+
+
+class ExpertDomain(DeviceMeshDomain):
+    @property
+    def name(self) -> str:
+        return EXPERT_DOMAIN
+
+    def build_mesh(self, params: "DeviceMeshParameters") -> DeviceMesh:
+        replicate_degree = (
+                params.data_parallel_replicate *
+                params.context_parallel_replicate *
+                params.data_parallel_shard *
+                params.context_parallel_shard
+        )
+        return init_device_mesh(
+            device_type="cuda",
+            mesh_shape=(
+                params.pipeline_parallel,
+                replicate_degree // params.expert_parallel,
+                params.expert_parallel
+            ),
+            mesh_dim_names=(
+                "pp",
+                "ep_replicate",
+                "ep_shard"
+            )
+        )
+
+
+DENSE_DOMAIN = "dense"
+
+
+class DenseDomain(DeviceMeshDomain):
+    @property
+    def name(self) -> str:
+        return DENSE_DOMAIN
+
+    def build_mesh(self, params: "DeviceMeshParameters") -> DeviceMesh:
+        return init_device_mesh(
+            device_type="cuda",
+            mesh_shape=(
+                params.pipeline_parallel,
+                params.data_parallel_replicate,
+                params.data_parallel_shard * params.context_parallel_shard,
+                params.context_parallel_replicate,
+                params.tensor_parallel
+            ),
+            mesh_dim_names=(
+                "pp",
+                "dp_replicate",
+                "dp_cp_shard",
+                "cp_replicate",
+                "tp"
+            )
+        )
+
+
+BATCH_DOMAIN = "batch"
+
+
+class BatchDomain(DeviceMeshDomain):
+    @property
+    def name(self) -> str:
+        return BATCH_DOMAIN
+
+    def build_mesh(self, params: "DeviceMeshParameters") -> DeviceMesh:
+        return init_device_mesh(
+            device_type="cuda",
+            mesh_shape=(
+                params.pipeline_parallel,
+                params.data_parallel_replicate * params.data_parallel_shard,
+                params.context_parallel_replicate * params.context_parallel_shard,
+                params.tensor_parallel
+            ),
+            mesh_dim_names=(
+                "pp",
+                "dp",
+                "cp",
+                "tp"
+            )
+        )
+
+
+FLAT_DOMAIN = "flat"
+
+
+class FlatDomain(DeviceMeshDomain):
+    @property
+    def name(self) -> str:
+        return FLAT_DOMAIN
+
+    def build_mesh(self, params: "DeviceMeshParameters") -> DeviceMesh:
+        mesh_shape = (
+            params.pipeline_parallel *
+            params.data_parallel_replicate *
+            params.data_parallel_shard *
+            params.context_parallel_replicate *
+            params.context_parallel_shard *
+            params.tensor_parallel
+        )
+        return init_device_mesh(
+            device_type="cuda",
+            mesh_shape=(
+                mesh_shape,
+            ),
+            mesh_dim_names=(
+                "world",
+            )
+        )
+
+
+ALL_DOMAIN_PROVIDERS: list[DeviceMeshDomain] = [
+    RegularDomain(), DenseDomain(), ExpertDomain(), BatchDomain(),
+    FlatDomain()
+]

d9d/core/dist_context/log.py

@@ -0,0 +1,30 @@
+import logging
+import sys
+
+
+def build_dist_logger(qualifier: str, level: int) -> logging.Logger:
+    """
+    Configures and returns a logger instance for d9d.
+
+    The logger is configured to write to stdout with a formatter that includes
+    the provided rank qualifier, allowing for easier debugging in distributed logs.
+
+    Args:
+        qualifier: A string identifying the current rank's position in the mesh.
+        level: Log level to set by default
+
+    Returns:
+        A configured logging.Logger instance.
+    """
+
+    dist_logger = logging.getLogger("d9d")
+    dist_logger.setLevel(level)
+    dist_logger.handlers.clear()
+    ch = logging.StreamHandler(sys.stdout)
+    ch.setLevel(level)
+    formatter = logging.Formatter(
+        f"[d9d] [{qualifier}] %(asctime)s - %(levelname)s - %(message)s"
+    )
+    ch.setFormatter(formatter)
+    dist_logger.addHandler(ch)
+    return dist_logger

d9d/core/dist_context/params.py

@@ -0,0 +1,113 @@
+import logging
+from typing import Self
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+from .configured import DistributedContext
+
+
+class DeviceMeshParameters(BaseModel):
+    """
+    Configuration parameters for initializing Distributed Device Meshes.
+
+    Attributes:
+        pipeline_parallel: Degree of pipeline parallelism (PP).
+        data_parallel_replicate: Degree of data parallel replication (DDP).
+        data_parallel_shard: Degree of data parallel sharding (FSDP).
+        context_parallel_replicate: Degree of context parallel (CP) replication.
+        context_parallel_shard: Degree of context parallel (FSCP) sharding.
+        tensor_parallel: Degree of tensor parallelism (TP).
+        expert_parallel: Degree of expert parallelism (EP/MoE).
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    pipeline_parallel: int = 1
+
+    data_parallel_replicate: int = 1
+    data_parallel_shard: int = 1
+
+    context_parallel_replicate: int = 1
+    context_parallel_shard: int = 1
+
+    tensor_parallel: int = 1
+
+    expert_parallel: int = 1
+
+    @property
+    def has_pipeline_parallel(self) -> bool:
+        """Checks if pipeline parallelism is enabled (degree > 1)."""
+
+        return self.pipeline_parallel > 1
+
+    @property
+    def has_data_parallel_replicate(self) -> bool:
+        """Checks if data parallel replication is enabled (degree > 1)."""
+
+        return self.data_parallel_replicate > 1
+
+    @property
+    def has_data_parallel_shard(self) -> bool:
+        """Checks if data parallel sharding is enabled (degree > 1)."""
+
+        return self.data_parallel_shard > 1
+
+    @property
+    def has_context_parallel_replicate(self) -> bool:
+        return self.context_parallel_replicate > 1
+
+    @property
+    def has_context_parallel_shard(self) -> bool:
+        return self.context_parallel_shard > 1
+
+    @property
+    def has_tensor_parallel(self) -> bool:
+        return self.tensor_parallel > 1
+
+    @property
+    def has_expert_parallel(self) -> bool:
+        """Checks if expert parallelism is enabled (degree > 1)."""
+        return self.expert_parallel > 1
+
+    @property
+    def is_distributed(self) -> bool:
+        """Checks if any form of parallelism is enabled."""
+
+        return (
+                self.has_pipeline_parallel or
+                self.has_data_parallel_replicate or
+                self.has_data_parallel_shard or
+                self.has_context_parallel_shard or
+                self.has_context_parallel_replicate or
+                self.has_expert_parallel or
+                self.has_tensor_parallel
+        )
+
+    @model_validator(mode="after")
+    def _check_ep_divisibility(self) -> Self:
+        """Validates that DP/CP/TP dimensions can support the requested EP/ETP degrees."""
+        dp_cp_tp_degree = (
+            self.data_parallel_shard *
+            self.data_parallel_replicate *
+            self.context_parallel_shard *
+            self.context_parallel_replicate *
+            self.tensor_parallel
+        )
+        ep_degree = self.expert_parallel
+
+        if dp_cp_tp_degree % ep_degree != 0:
+            raise ValueError(
+                f"Total data/context/tensor parallelism degree ({dp_cp_tp_degree}) must be divisible by "
+                f"total expert parallelism degree ({ep_degree})."
+            )
+        return self
+
+    def build(self, log_level: int = logging.INFO) -> "DistributedContext":
+        """
+        Initializes the DistributedContext using these parameters.
+
+        Returns:
+            A new DistributedContext instance containing the initialized device meshes.
+        """
+
+        return DistributedContext(self, log_level)

d9d/core/dist_ops/__init__.py

@@ -0,0 +1,16 @@
+"""
+This module provides high-level wrappers around `torch.distributed` collective operations.
+"""
+
+
+from .object import all_gather_object, gather_object
+from .tensor import all_gather, all_gather_variadic_shape, gather, gather_variadic_shape
+
+__all__ = [
+    "all_gather",
+    "all_gather_object",
+    "all_gather_variadic_shape",
+    "gather",
+    "gather_object",
+    "gather_variadic_shape"
+]