d9d 0.1.0__tar.gz

d9d-0.1.0/PKG-INFO

@@ -0,0 +1,90 @@
+Metadata-Version: 2.4
+Name: d9d
+Version: 0.1.0
+Summary: d9d - d[istribute]d - distributed training framework based on PyTorch that tries to be efficient yet hackable
+License: Apache-2.0
+Author: Maksim Afanasyev
+Author-email: mr.applexz@gmail.com
+Requires-Python: >=3.11,<3.15
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Provides-Extra: aim
+Provides-Extra: cce
+Provides-Extra: moe
+Provides-Extra: visualization
+Requires-Dist: aim (>=3.0.0,<4.0.0) ; extra == "aim"
+Requires-Dist: cut-cross-entropy (>=25.9.3) ; extra == "cce"
+Requires-Dist: deep-ep (>=1.2.1) ; extra == "moe"
+Requires-Dist: nv-grouped-gemm (>=1.1.4) ; extra == "moe"
+Requires-Dist: plotly (>=6.0.0) ; extra == "visualization"
+Requires-Dist: pydantic (>=2.0.0)
+Requires-Dist: safetensors (>=0.7.0)
+Requires-Dist: setuptools (>=70.0.0) ; extra == "aim"
+Requires-Dist: torch (>=2.10.0)
+Requires-Dist: torchdata (>=0.11.0)
+Requires-Dist: tqdm (>=4.0.0)
+Requires-Dist: triton (>=3.6.0)
+Project-URL: Documentation, https://d9d-project.github.io/d9d
+Project-URL: Homepage, https://d9d-project.github.io/d9d
+Project-URL: Issues, https://github.com/d9d-project/d9d/issues
+Project-URL: Repository, https://github.com/d9d-project/d9d
+Description-Content-Type: text/markdown
+
+# The d9d Project
+
+**d9d** is a distributed training framework built on top of PyTorch 2.0. It aims to be hackable, modular, and efficient, designed to scale from single-GPU debugging to massive clusters running 6D-Parallelism.
+
+[LET'S START TRAINING 🚀](https://d9d-project.github.io/d9d/)
+
+## Why another framework?
+
+Distributed training frameworks such as **Megatron-LM** are monolithic in the way you run a script from the command line to train any of a set of *predefined* models, using *predefined* regimes. While powerful, these systems can be difficult to hack and integrate into novel research workflows. Their focus is often on providing a complete, end-to-end solution, which can limit flexibility for experimentally-driven research.
+
+Conversely, creating your own distributed training solution from scratch is tricky. You have to implement many low-level components (like distributed checkpoints and synchronization) that are identical across setups, and manually tackle common performance bottlenecks.
+
+**d9d** was designed to fill the gap between monolithic frameworks and homebrew setups, providing a modular yet effective solution for distributed training.
+
+## What d9d is and isn't
+
+In terms of **core concept**:
+
+*   **IS** a pluggable framework for implementing distributed training regimes for your deep learning models.
+*   **IS** built on clear interfaces and building blocks that may be composed and implemented in your own way.
+*   **IS NOT** an all-in-one CLI platform for setting up pre-training and post-training like **torchtitan**, **Megatron-LM**, or **torchforge**.
+
+In terms of **codebase & engineering**:
+
+*   **IS** built on a **strong engineering foundation**: We enforce strict type-checking and rigorous linting to catch errors before execution.
+*   **IS** reliable: The framework is backed by a suite of **over 450 tests**, covering unit logic, integration flows, and End-to-End distributed scenarios.
+*   **IS** eager to use performance hacks (like **DeepEp** or custom kernels) if they improve MFU, even if they aren't PyTorch-native.
+*   **IS NOT** for legacy setups: We do not maintain backward compatibility with older PyTorch versions or hardware. We prioritize simplicity and modern APIs (like `DTensor`).
+
+## Key Philosophies
+
+To achieve the balance between hackability and performance, d9d adheres to specific design principles:
+
+*   **Composition over Monoliths**: We avoid "God Classes" like `DistributedDataParallel` or `ParallelDims` that assume ownership of the entire execution loop. Instead, we provide composable and extendable APIs. For instance, specific horizontal parallelism strategies for specific layers (`parallelize_replicate`, `parallelize_expert_parallel`, ...).
+*   **White-Box Modelling**: We encourage standard PyTorch code. Models are not wrapped in obscure metadata specifications; they are standard `nn.Module`s that implement lightweight protocols.
+*   **Pragmatic Efficiency**: While we prefer native PyTorch, we are eager to integrate non-native solutions if they improve MFU. For example, we implement MoE using **DeepEp** communications, reindexing kernels from **Megatron-LM**, and efficient grouped-GEMM implementations.
+*   **Graph-Based State Management**: Our IO system treats model checkpoints as directed acyclic graphs. This allows you to transform architectures (e.g., merging `q`, `k`, `v` into `qkv`) on-the-fly while streaming from disk, without massive memory overhead.
+*   **DTensors**: We mandate that distributed parameters be represented as `torch.distributed.tensor.DTensor`. This simplifies checkpointing by making them topology-aware automatically. We leverage modern PyTorch 2.0 APIs (`DeviceMesh`) as much as possible.
+
+---
+
+## Examples
+
+### Qwen3-MoE Pretraining
+An example showing causal LM pretraing for the Qwen3-MoE model.
+
+WIP: MoE load balancing is currently work in progress.
+
+[Link](https://github.com/d9d-project/d9d/blob/main/example/qwen3_moe/pretrain.py).
+

d9d-0.1.0/README.md

@@ -0,0 +1,49 @@
+# The d9d Project
+
+**d9d** is a distributed training framework built on top of PyTorch 2.0. It aims to be hackable, modular, and efficient, designed to scale from single-GPU debugging to massive clusters running 6D-Parallelism.
+
+[LET'S START TRAINING 🚀](https://d9d-project.github.io/d9d/)
+
+## Why another framework?
+
+Distributed training frameworks such as **Megatron-LM** are monolithic in the way you run a script from the command line to train any of a set of *predefined* models, using *predefined* regimes. While powerful, these systems can be difficult to hack and integrate into novel research workflows. Their focus is often on providing a complete, end-to-end solution, which can limit flexibility for experimentally-driven research.
+
+Conversely, creating your own distributed training solution from scratch is tricky. You have to implement many low-level components (like distributed checkpoints and synchronization) that are identical across setups, and manually tackle common performance bottlenecks.
+
+**d9d** was designed to fill the gap between monolithic frameworks and homebrew setups, providing a modular yet effective solution for distributed training.
+
+## What d9d is and isn't
+
+In terms of **core concept**:
+
+*   **IS** a pluggable framework for implementing distributed training regimes for your deep learning models.
+*   **IS** built on clear interfaces and building blocks that may be composed and implemented in your own way.
+*   **IS NOT** an all-in-one CLI platform for setting up pre-training and post-training like **torchtitan**, **Megatron-LM**, or **torchforge**.
+
+In terms of **codebase & engineering**:
+
+*   **IS** built on a **strong engineering foundation**: We enforce strict type-checking and rigorous linting to catch errors before execution.
+*   **IS** reliable: The framework is backed by a suite of **over 450 tests**, covering unit logic, integration flows, and End-to-End distributed scenarios.
+*   **IS** eager to use performance hacks (like **DeepEp** or custom kernels) if they improve MFU, even if they aren't PyTorch-native.
+*   **IS NOT** for legacy setups: We do not maintain backward compatibility with older PyTorch versions or hardware. We prioritize simplicity and modern APIs (like `DTensor`).
+
+## Key Philosophies
+
+To achieve the balance between hackability and performance, d9d adheres to specific design principles:
+
+*   **Composition over Monoliths**: We avoid "God Classes" like `DistributedDataParallel` or `ParallelDims` that assume ownership of the entire execution loop. Instead, we provide composable and extendable APIs. For instance, specific horizontal parallelism strategies for specific layers (`parallelize_replicate`, `parallelize_expert_parallel`, ...).
+*   **White-Box Modelling**: We encourage standard PyTorch code. Models are not wrapped in obscure metadata specifications; they are standard `nn.Module`s that implement lightweight protocols.
+*   **Pragmatic Efficiency**: While we prefer native PyTorch, we are eager to integrate non-native solutions if they improve MFU. For example, we implement MoE using **DeepEp** communications, reindexing kernels from **Megatron-LM**, and efficient grouped-GEMM implementations.
+*   **Graph-Based State Management**: Our IO system treats model checkpoints as directed acyclic graphs. This allows you to transform architectures (e.g., merging `q`, `k`, `v` into `qkv`) on-the-fly while streaming from disk, without massive memory overhead.
+*   **DTensors**: We mandate that distributed parameters be represented as `torch.distributed.tensor.DTensor`. This simplifies checkpointing by making them topology-aware automatically. We leverage modern PyTorch 2.0 APIs (`DeviceMesh`) as much as possible.
+
+---
+
+## Examples
+
+### Qwen3-MoE Pretraining
+An example showing causal LM pretraing for the Qwen3-MoE model.
+
+WIP: MoE load balancing is currently work in progress.
+
+[Link](https://github.com/d9d-project/d9d/blob/main/example/qwen3_moe/pretrain.py).

d9d-0.1.0/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-0.1.0/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-0.1.0/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-0.1.0/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-0.1.0/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()
+]