d9d 0.1.0__py3-none-any.whl

d9d/module/model/qwen3_moe/model.py

@@ -0,0 +1,373 @@
+from collections.abc import Mapping
+from typing import cast
+
+import torch
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from d9d.module.base import ModuleLateInit
+from d9d.module.block.embedding import SplitTokenEmbeddings
+from d9d.module.block.head import SplitLanguageModellingHead
+from d9d.module.block.hidden_states_aggregator import HiddenStatesAggregationMode, create_hidden_states_aggregator
+from d9d.module.block.positional import RotaryEmbeddingProvider
+from d9d.pipelining.api import (
+    ModuleSupportsPipelining,
+    PipelineStageInfo,
+    distribute_layers_for_pipeline_stage,
+)
+
+from .decoder_layer import Qwen3MoELayer
+from .params import Qwen3MoEForCausalLMParameters, Qwen3MoEParameters
+
+
+class Qwen3MoEModel(nn.Module, ModuleLateInit, ModuleSupportsPipelining):
+    """
+    The Qwen3 Mixture-of-Experts (MoE) Transformer Decoder backbone.
+
+    It is designed to be split across multiple pipeline stages.
+    """
+
+    def __init__(
+            self,
+            params: Qwen3MoEParameters,
+            stage: PipelineStageInfo,
+            hidden_states_snapshot_mode: HiddenStatesAggregationMode,
+            enable_checkpointing: bool
+    ):
+        """
+        Constructs the Qwen3MoEModel object.
+
+        Args:
+            params: Configuration parameters for the full model.
+            stage: Information about the pipeline stage this instance belongs to.
+            hidden_states_snapshot_mode: Configures intermediate hidden state aggregation & snapshotting mode
+            enable_checkpointing: If True, enables activation checkpointing for transformer layers to save memory.
+        """
+
+        super().__init__()
+
+        if stage.is_current_stage_first:
+            self.embed_tokens = SplitTokenEmbeddings(
+                hidden_size=params.layer.hidden_size,
+                split_vocab_size=params.split_vocab_size,
+                split_order=params.split_vocab_order
+            )
+
+        # we use ModuleDict here to properly handle pipelining and loading weights after the model
+        # was pipelined
+        layer_start, layer_end = distribute_layers_for_pipeline_stage(
+            num_layers=params.num_hidden_layers,
+            num_virtual_layers_pre=params.pipeline_num_virtual_layers_pre,  # embeddings
+            num_virtual_layers_post=params.pipeline_num_virtual_layers_post,  # LM head
+            stage=stage
+        )
+
+        self._num_layers_before = layer_start
+        self._layers_iter = list(map(str, range(layer_start, layer_end)))
+        layers = nn.ModuleDict({
+            str(layer_idx): Qwen3MoELayer(params=params.layer) for layer_idx in self._layers_iter
+        })
+        self.layers: Mapping[str, Qwen3MoELayer] = cast(Mapping[str, Qwen3MoELayer], layers)
+
+        self.rope_provider = RotaryEmbeddingProvider(
+            max_position_ids=params.max_position_ids,
+            rope_base=params.rope_base,
+            head_dim=params.layer.head_dim
+        )
+
+        if stage.is_current_stage_last:
+            self.norm = nn.RMSNorm(
+                normalized_shape=params.layer.hidden_size,
+                eps=params.layer.rms_norm_eps
+            )
+
+        self._stage = stage
+        self._hidden_states_snapshot_mode = hidden_states_snapshot_mode
+        self._hidden_size = params.layer.hidden_size
+        self._enable_checkpointing = enable_checkpointing
+
+    def output_dtype(self) -> torch.dtype:
+        """
+        Returns the data type of the model output hidden states.
+        """
+        return self.layers[self._layers_iter[0]].input_layernorm.weight.dtype
+
+    def forward(
+            self,
+            input_ids: torch.Tensor | None = None,
+            hidden_states: torch.Tensor | None = None,
+            position_ids: torch.Tensor | None = None,
+            hidden_states_snapshot: torch.Tensor | None = None,
+            hidden_states_agg_mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Executes the forward pass for the current pipeline stage.
+
+        Args:
+            input_ids: Indices of input sequence tokens. Required if this is the
+                first pipeline stage.
+            hidden_states: Hidden states from the previous pipeline stage. Required
+                if this is not the first pipeline stage.
+            position_ids: Indices of positions of each input sequence tokens in the
+                position embeddings.
+            hidden_states_snapshot: Accumulated tensor of aggregated hidden states
+                from previous stages. Used if snapshotting is enabled.
+            hidden_states_agg_mask: Mask used to aggregate hidden states for
+                snapshots.
+
+        Returns:
+            A dictionary containing:
+                *   'hidden_states': The output of the last layer in this stage.
+                *   'hidden_states_snapshot': (Optional) The updated snapshot tensor.
+        """
+        state_aggregator = create_hidden_states_aggregator(self._hidden_states_snapshot_mode, hidden_states_agg_mask)
+
+        if input_ids is not None:
+            last_hidden_states = self.embed_tokens(input_ids)
+            state_aggregator.add_hidden_states(last_hidden_states)
+        else:
+            last_hidden_states = hidden_states
+
+        rope_params = self.rope_provider(position_ids)
+
+        for decoder_layer_name in self._layers_iter:
+            decoder_layer = self.layers[decoder_layer_name]
+
+            if self._enable_checkpointing:
+                last_hidden_states = checkpoint(
+                    decoder_layer, last_hidden_states, rope_params,
+                    use_reentrant=False
+                )
+            else:
+                last_hidden_states = decoder_layer(last_hidden_states, rope_params)
+
+            state_aggregator.add_hidden_states(last_hidden_states)
+
+        if self._stage.is_current_stage_last:
+            last_hidden_states = self.norm(last_hidden_states)
+
+        return {
+            "hidden_states": last_hidden_states,
+            "hidden_states_snapshot": state_aggregator.pack_with_snapshot(hidden_states_snapshot)
+        }
+
+    def reset_moe_stats(self):
+        """
+        Resets routing statistics for all MoE layers in this stage.
+        """
+
+        for layer_name in self._layers_iter:
+            self.layers[layer_name].reset_moe_stats()
+
+    @property
+    def moe_tokens_per_expert(self) -> torch.Tensor:
+        """
+        Retrieves the number of tokens routed to each expert across all layers.
+
+        Returns:
+            A tensor of shape (num_local_layers, num_experts) containing counts.
+        """
+
+        return torch.stack(
+            [self.layers[layer_name].moe_tokens_per_expert for layer_name in self._layers_iter],
+            dim=0
+        )
+
+    def reset_parameters(self):
+        """Resets module parameters"""
+
+        if self._stage.is_current_stage_first:
+            self.embed_tokens.reset_parameters()
+
+        self.rope_provider.reset_parameters()
+
+        for decoder_layer_name in self._layers_iter:
+            decoder_layer = self.layers[decoder_layer_name]
+            decoder_layer.reset_parameters()
+
+        if self._stage.is_current_stage_last:
+            self.norm.reset_parameters()
+
+    def infer_stage_inputs_from_pipeline_inputs(
+            self, inputs: dict[str, torch.Tensor], n_microbatches: int
+    ) -> dict[str, torch.Tensor]:
+        input_ids = inputs["input_ids"]
+
+        pp_inputs = {}
+
+        # for calculation - input ids or prev hidden state
+        if self._stage.is_current_stage_first:
+            pp_inputs["input_ids"] = torch.empty(
+                (input_ids.shape[0] // n_microbatches, input_ids.shape[1]),
+                dtype=torch.long,
+                device=input_ids.device
+            )
+        else:
+            pp_inputs["hidden_states"] = torch.empty(
+                (input_ids.shape[0] // n_microbatches, input_ids.shape[1], self._hidden_size),
+                dtype=self.output_dtype(),
+                device=input_ids.device
+            )
+            if self._hidden_states_snapshot_mode != HiddenStatesAggregationMode.no:
+                num_layers_before = self._num_layers_before + 1  # 1 for embedding
+                pp_inputs["hidden_states_snapshot"] = torch.empty(
+                    (num_layers_before, input_ids.shape[0] // n_microbatches, self._hidden_size),
+                    dtype=self.output_dtype(),
+                    device=input_ids.device
+                )
+
+        return pp_inputs
+
+    def infer_stage_outputs_from_pipeline_inputs(
+            self, inputs: dict[str, torch.Tensor], n_microbatches: int
+    ) -> dict[str, torch.Tensor]:
+        input_ids = inputs["input_ids"]
+
+        # for calculation - last hidden state
+        pp_outputs = {
+            "hidden_states": torch.empty(
+                (input_ids.shape[0] // n_microbatches, input_ids.shape[1], self._hidden_size),
+                dtype=self.output_dtype(),
+                device=input_ids.device
+            )
+        }
+
+        # for state caching
+        if self._hidden_states_snapshot_mode != HiddenStatesAggregationMode.no:
+            num_layers_before = self._num_layers_before + 1
+            num_layers_current = len(self.layers)
+            num_layers_after = num_layers_before + num_layers_current
+            pp_outputs["hidden_states_snapshot"] = torch.empty(
+                (num_layers_after, input_ids.shape[0] // n_microbatches, self._hidden_size),
+                dtype=self.output_dtype(),
+                device=input_ids.device
+            )
+
+        return pp_outputs
+
+
+class Qwen3MoEForCausalLM(nn.Module, ModuleLateInit, ModuleSupportsPipelining):
+    """
+    A Qwen3 MoE model wrapped with a Causal Language Modeling head.
+
+    It is designed to be split across multiple pipeline stages.
+    """
+
+    def __init__(
+            self,
+            params: Qwen3MoEForCausalLMParameters,
+            stage: PipelineStageInfo,
+            hidden_states_snapshot_mode: HiddenStatesAggregationMode,
+            enable_checkpointing: bool
+    ):
+        """
+        Constructs the Qwen3MoEForCausalLM object.
+
+        Args:
+            params: Full model configuration parameters.
+            stage: Pipeline stage information for this instance.
+            hidden_states_snapshot_mode: Configures intermediate hidden state aggregation & snapshotting mode.
+            enable_checkpointing: Whether to enable activation checkpointing.
+        """
+
+        super().__init__()
+
+        self.model = Qwen3MoEModel(
+            params.model,
+            stage,
+            hidden_states_snapshot_mode=hidden_states_snapshot_mode,
+            enable_checkpointing=enable_checkpointing
+        )
+
+        if stage.is_current_stage_last:
+            self.lm_head = SplitLanguageModellingHead(
+                split_vocab_size=params.model.split_vocab_size,
+                split_order=params.model.split_vocab_order,
+                hidden_size=params.model.layer.hidden_size
+            )
+
+        self._stage = stage
+        self._hidden_size = params.model.layer.hidden_size
+
+    def forward(
+            self,
+            input_ids: torch.Tensor | None = None,
+            hidden_states: torch.Tensor | None = None,
+            position_ids: torch.Tensor | None = None,
+            hidden_states_snapshot: torch.Tensor | None = None,
+            hidden_states_agg_mask: torch.Tensor | None = None,
+            labels: torch.Tensor | None = None
+    ) -> dict[str, torch.Tensor]:
+        """
+        Executes the model forward pass.
+
+        If this is the last stage, it expects `labels` to be provided and computes
+        the cross-entropy loss (returned as 'logps' typically representing per-token loss).
+
+        Args:
+            input_ids: Input token IDS (for Stage 0).
+            hidden_states: Hidden states from previous stage (for Stage > 0).
+            position_ids: Positional indices for RoPE.
+            hidden_states_snapshot: Intermediate state collector.
+            hidden_states_agg_mask: Mask for state aggregation.
+            labels: Target tokens for loss computation (Last Stage).
+
+        Returns:
+            Dictionary containing 'hidden_states', optionally 'hidden_states_snapshot',
+            and per-token 'logps' if on the last stage.
+        """
+
+        model_outputs = self.model(
+            input_ids=input_ids,
+            hidden_states=hidden_states,
+            position_ids=position_ids,
+            hidden_states_snapshot=hidden_states_snapshot,
+            hidden_states_agg_mask=hidden_states_agg_mask
+        )
+        if self._stage.is_current_stage_last:
+            lm_out = self.lm_head(
+                hidden_states=model_outputs["hidden_states"],
+                labels=labels
+            )
+            model_outputs["logps"] = lm_out
+        return model_outputs
+
+    def reset_parameters(self):
+        """
+        Resets module parameters.
+        """
+
+        self.model.reset_parameters()
+
+        if self._stage.is_current_stage_last:
+            self.lm_head.reset_parameters()
+
+    def reset_moe_stats(self):
+        """
+        Resets MoE routing statistics in the backbone.
+        """
+
+        self.model.reset_moe_stats()
+
+    @property
+    def moe_tokens_per_expert(self) -> torch.Tensor:
+        """
+        Accesses MoE routing statistics from the backbone.
+        """
+
+        return self.model.moe_tokens_per_expert
+
+    def infer_stage_inputs_from_pipeline_inputs(
+            self, inputs: dict[str, torch.Tensor], n_microbatches: int
+    ) -> dict[str, torch.Tensor]:
+        return self.model.infer_stage_inputs_from_pipeline_inputs(inputs, n_microbatches)
+
+    def infer_stage_outputs_from_pipeline_inputs(
+            self, inputs: dict[str, torch.Tensor], n_microbatches: int
+    ) -> dict[str, torch.Tensor]:
+        pp_outputs = self.model.infer_stage_outputs_from_pipeline_inputs(inputs, n_microbatches)
+
+        if self._stage.is_current_stage_last:
+            pp_outputs["logps"] = torch.empty(inputs["input_ids"].shape, dtype=torch.float32)
+
+        return pp_outputs

d9d/module/model/qwen3_moe/params.py

@@ -0,0 +1,69 @@
+from pydantic import BaseModel
+
+
+class Qwen3MoELayerParameters(BaseModel):
+    """
+    Configuration parameters for a single Qwen3 MoE layer.
+
+    Attributes:
+        hidden_size: Dimension of the model's hidden states.
+        intermediate_size: Dimension of the feed-forward hidden state.
+        num_experts: Total number of experts in the MoE layer.
+        experts_top_k: Number of experts to route tokens to.
+        num_attention_heads: Number of attention heads for the query.
+        num_key_value_heads: Number of attention heads for key and value.
+        rms_norm_eps: Epsilon value found in the RMSNorm layers.
+        head_dim: Dimension of a single attention head.
+    """
+
+    hidden_size: int
+    intermediate_size: int
+    num_experts: int
+    experts_top_k: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    head_dim: int
+
+
+class Qwen3MoEParameters(BaseModel):
+    """
+    Configuration parameters for the Qwen3 Mixture-of-Experts model backbone.
+
+    Attributes:
+        layer: Configuration shared across all transformer layers.
+        num_hidden_layers: The total number of transformer layers.
+        rope_base: Base value for RoPE frequency calculation.
+        max_position_ids: Maximum sequence length.
+        split_vocab_size: A dictionary mapping vocabulary segment names to their sizes.
+        split_vocab_order: The sequence in which vocabulary splits are correctly ordered.
+        pipeline_num_virtual_layers_pre: The number of 'virtual' layers representing the
+            computational cost of modules on the *first* stage, before the main
+            layers (e.g., token and positional embeddings).
+        pipeline_num_virtual_layers_post: The number of 'virtual' layers representing the
+            computational cost of modules on the *last* stage, after the main
+            layers (e.g., the final layer normalization and LM head).
+    """
+
+    layer: Qwen3MoELayerParameters
+
+    num_hidden_layers: int
+    rope_base: int
+    max_position_ids: int
+
+    split_vocab_size: dict[str, int]
+    split_vocab_order: list[str]
+
+    pipeline_num_virtual_layers_pre: int = 0
+    pipeline_num_virtual_layers_post: int = 0
+
+
+class Qwen3MoEForCausalLMParameters(BaseModel):
+    """
+    Configuration parameters for Qwen3 Mixture-of-Experts model with a Causal Language Modeling head.
+
+    Attributes:
+        model: The configuration for the underlying Qwen3 MoE model.
+    """
+
+    model: Qwen3MoEParameters

d9d/module/parallelism/api/__init__.py

@@ -0,0 +1,18 @@
+"""
+Horizontal parallelism strategies and utilities for d9d modules.
+
+This package provides high-level helper functions to apply specific distributed
+parallelism strategies to PyTorch modules compatible with the d9d ecosystem.
+"""
+
+from .expert_parallel import parallelize_expert_parallel
+from .fully_sharded import parallelize_fsdp
+from .hybrid_sharded import parallelize_hsdp
+from .replicate_parallel import parallelize_replicate
+
+__all__ = [
+    "parallelize_expert_parallel",
+    "parallelize_fsdp",
+    "parallelize_hsdp",
+    "parallelize_replicate"
+]

d9d/module/parallelism/api/expert_parallel.py

@@ -0,0 +1,36 @@
+from torch.distributed import DeviceMesh
+from torch.distributed.tensor import Replicate
+from torch.distributed.tensor.parallel import parallelize_module
+
+from d9d.module.block.moe import MoELayer
+from d9d.module.parallelism.style import ShardMoESparseExpertsParallel, ToLocalParallel
+
+
+def parallelize_expert_parallel(
+        module: MoELayer,
+        mesh_experts: DeviceMesh,
+        expert_shard_dim: str = "ep_shard"
+):
+    """
+    Applies Expert Parallelism to a MoE layer.
+
+    This function configures the provided Mixture of Experts layer for distributed
+    execution.
+
+    It partitions the sparse experts across the specified dimension
+    of the device mesh (Expert Parallelism) and replicates along other dims.
+
+    Simultaneously, it configures the router to be fully replicated across
+    the mesh.
+
+    Args:
+        module: The MoE layer instance to parallelize.
+        mesh_experts: The device mesh containing the expert parallel resources.
+        expert_shard_dim: The name of the mesh dimension where experts should be sharded.
+    """
+
+    parallelize_module(module, mesh_experts, ShardMoESparseExpertsParallel(shard_dim_name=expert_shard_dim))
+    parallelize_module(module.router, mesh_experts, ToLocalParallel(
+        param_placement=tuple(Replicate() for _ in range(mesh_experts.ndim)),
+        grad_placement=tuple(Replicate() for _ in range(mesh_experts.ndim))
+    ))

d9d/module/parallelism/api/fully_sharded.py

@@ -0,0 +1,43 @@
+from typing import Any
+
+from torch import nn
+from torch.distributed import DeviceMesh
+from torch.distributed.fsdp import FSDPModule, fully_shard
+
+
+def _force_fsdp_grad_reduction_policy(module: FSDPModule):
+    module.set_force_sum_reduction_for_comms(enable=True)
+    module.set_gradient_divide_factor(1.0)
+    module.set_requires_all_reduce(False)
+
+
+def parallelize_fsdp(
+        module: nn.Module,
+        mesh: DeviceMesh,
+        *args: Any,
+        **kwargs: Any
+):
+    """
+    Applies Fully Sharded Data Parallel (FSDP) with forced gradient summation.
+
+    This function wraps the provided module with PyTorch's ``fully_shard`` API using
+    the specified device mesh. Unlike standard FSDP usage, this function explicitly
+    configures the module to sum gradients across the mesh
+    instead of averaging them and disables internal all-sum-reduce hooks.
+    This is intended for d9d to handle gradient normalization and reduction across replicas externally.
+
+    Args:
+        module: The module to shard.
+        mesh: The device mesh over which to shard the module.
+        *args: Additional positional arguments passed to ``fully_shard``.
+        **kwargs: Additional keyword arguments passed to ``fully_shard``.
+    """
+
+    if mesh.ndim != 1:
+        raise ValueError("FSDP mesh should contain exactly one dimension - for HSDP, please apply "
+                         "parallelize_replicate(...) first!")
+
+    fully_shard(module, *args, mesh=mesh, **kwargs)
+    if not isinstance(module, FSDPModule):
+        raise RuntimeError("Torch FSDP did not convert the module into FSDPModule")
+    _force_fsdp_grad_reduction_policy(module)

d9d/module/parallelism/api/hybrid_sharded.py

@@ -0,0 +1,49 @@
+from typing import Any
+
+from torch import nn
+from torch.distributed import DeviceMesh
+
+from .fully_sharded import parallelize_fsdp
+from .replicate_parallel import parallelize_replicate
+
+
+def parallelize_hsdp(
+        module: nn.Module,
+        mesh: DeviceMesh,
+        shard_dim: str = "dp_cp_shard",
+        *fsdp_args: Any,
+        **fsdp_kwargs: Any
+):
+    """
+    Applies Hybrid Sharded Data Parallelism (HSDP) to a module.
+
+    This function decomposes the provided device mesh into sharding dimensions
+    and replication dimensions. It applies replication parallelism
+    across the replication dimensions and Fully Sharded Data Parallelism (FSDP)
+    across the specified shard dimension.
+
+    Args:
+        module: The module to parallelize.
+        mesh: The device mesh over which to distribute the module.
+        shard_dim: The name of the mesh dimension used for FSDP sharding. Any
+            dimension in the mesh not matching this name will be treated as a
+            replication dimension.
+        *fsdp_args: Positional arguments passed to the underlying FSDP parallelizer.
+        **fsdp_kwargs: Keyword arguments passed to the underlying FSDP parallelizer.
+
+    Raises:
+        ValueError: If the device mesh does not have named dimensions.
+    """
+
+    replicate_dims = mesh.mesh_dim_names
+
+    if replicate_dims is None:
+        raise ValueError("Cannot use with unnamed device meshes")
+
+    replicate_dims = tuple(x for x in replicate_dims if x != shard_dim and mesh[x].size() > 1)
+
+    if len(replicate_dims) > 0:
+        parallelize_replicate(module, mesh[replicate_dims])
+
+    if mesh[shard_dim].size() != 1:
+        parallelize_fsdp(module, mesh[shard_dim], *fsdp_args, **fsdp_kwargs)

d9d/module/parallelism/api/replicate_parallel.py

@@ -0,0 +1,33 @@
+from torch import nn
+from torch.distributed import DeviceMesh
+from torch.distributed.tensor import Replicate
+from torch.distributed.tensor.parallel import parallelize_module
+
+from d9d.module.parallelism.style import ToLocalParallel
+
+
+def parallelize_replicate(
+    module: nn.Module,
+    mesh: DeviceMesh,
+):
+    """
+    Applies replicated parallelism to the module.
+
+    This function configures the provided module to be fully replicated across the
+    given device mesh. It utilizes the ``ToLocalParallel`` style, which manages
+    ``DTensor`` wrapping for parameters and gradients (via ``Replicate`` placements)
+    while ensuring that the underlying computation sees standard local tensors during the forward pass.
+
+    This approach is effectively Data Parallelism managed via the DTensor
+    APIs, allowing seamless integration of modules that require local tensor inputs
+    into a broader distributed mesh context.
+
+    Args:
+     module: The module to parallelize.
+     mesh: The device mesh over which to replicate the module.
+    """
+
+    parallelize_module(module, mesh, ToLocalParallel(
+        param_placement=tuple(Replicate() for _ in range(mesh.ndim)),
+        grad_placement=tuple(Replicate() for _ in range(mesh.ndim))
+    ))