d9d 0.1.0__py3-none-any.whl

d9d/module/parallelism/model/qwen3_moe.py

@@ -0,0 +1,99 @@
+from d9d.core.dist_context import DENSE_DOMAIN, EXPERT_DOMAIN, DistributedContext
+from d9d.module.model.qwen3_moe import Qwen3MoEForCausalLM, Qwen3MoEModel
+from d9d.module.parallelism.api import parallelize_expert_parallel, parallelize_hsdp
+from d9d.pipelining.api import PipelineStageInfo
+
+
+def parallelize_qwen3_moe_model(
+        dist_context: DistributedContext,
+        model: Qwen3MoEModel,
+        stage: PipelineStageInfo
+):
+    """
+    Parallelizes the base Qwen3 MoE model components.
+
+    This function configures the model layers for distributed execution within a pipeline
+    stage. It applies Hybrid Sharded Data Parallelism (HSDP) to dense components (embeddings,
+    norms, attention) and Expert Parallelism (EP) to the Mixture-of-Experts (MLP) layers.
+
+    Current usage constraints:
+    *   Tensor Parallelism is not supported (we may implement it later).
+    *   Context Parallelism is not supported (we will implement it later).
+
+    Args:
+        dist_context: The distributed context.
+        model: The Qwen3 MoE base model to parallelize.
+        stage: Information about the current pipeline stage.
+
+    Raises:
+        ValueError: If Tensor Parallel or Context Parallel is enabled in the context.
+    """
+
+    dims = dist_context.mesh_params
+    dense_mesh = dist_context.mesh_for(DENSE_DOMAIN)
+    expert_mesh = dist_context.mesh_for(EXPERT_DOMAIN)
+
+    if dims.has_tensor_parallel:
+        raise ValueError("Tensor Parallel currently is not supported for this model.")
+    if dims.has_context_parallel_replicate or dims.has_context_parallel_shard:
+        raise ValueError("Context Parallel currently is not supported for this model.")
+
+    if stage.is_current_stage_first:
+        parallelize_hsdp(
+            model.embed_tokens,
+            mesh=dense_mesh["dp_replicate", "dp_cp_shard", "cp_replicate"]
+        )
+
+    if stage.is_current_stage_last:
+        parallelize_hsdp(
+            model.norm,
+            mesh=dense_mesh["dp_replicate", "dp_cp_shard", "cp_replicate"],
+        )
+
+    for layer in model.layers.values():
+        parallelize_expert_parallel(
+            layer.mlp,
+            mesh_experts=expert_mesh["ep_replicate", "ep_shard"]
+        )
+
+        parallelize_hsdp(
+            layer.self_attn,
+            mesh=dense_mesh["dp_replicate", "dp_cp_shard", "cp_replicate"],
+        )
+        parallelize_hsdp(
+            layer.input_layernorm,
+            mesh=dense_mesh["dp_replicate", "dp_cp_shard", "cp_replicate"],
+        )
+        parallelize_hsdp(
+            layer.post_attention_layernorm,
+            mesh=dense_mesh["dp_replicate", "dp_cp_shard", "cp_replicate"],
+        )
+
+
+def parallelize_qwen3_moe_for_causal_lm(
+        dist_context: DistributedContext,
+        model: Qwen3MoEForCausalLM,
+        stage: PipelineStageInfo
+):
+    """
+    Parallelizes the Qwen3 MoE Causal LM model.
+
+    This function delegates backbone parallelization to ``parallelize_qwen3_moe_model``
+    and additionally configures the language model head with Hybrid Sharded Data
+    Parallelism (HSDP).
+
+    Args:
+        dist_context: The distributed context containing device meshes and topology info.
+        model: The Qwen3 MoE Causal LM model to parallelize.
+        stage: Information about the current pipeline stage.
+    """
+
+    dense_mesh = dist_context.mesh_for(DENSE_DOMAIN)
+
+    parallelize_qwen3_moe_model(dist_context, model.model, stage)
+
+    if stage.is_current_stage_last:
+        parallelize_hsdp(
+            model.lm_head,
+            mesh=dense_mesh["dp_replicate", "dp_cp_shard", "cp_replicate"],
+        )

d9d/module/parallelism/style/__init__.py

@@ -0,0 +1,7 @@
+from .shard_experts import ShardMoESparseExpertsParallel
+from .to_local import ToLocalParallel
+
+__all__ = [
+    "ShardMoESparseExpertsParallel",
+    "ToLocalParallel"
+]

d9d/module/parallelism/style/shard_experts.py

@@ -0,0 +1,60 @@
+from torch import nn
+from torch.distributed import DeviceMesh
+from torch.distributed.tensor import (
+    Replicate,
+    Shard,
+    distribute_module,
+    distribute_tensor,
+)
+from torch.distributed.tensor.parallel import ParallelStyle
+
+from d9d.module.block.moe import GroupedLinear, MoELayer
+
+
+class ShardMoESparseExpertsParallel(ParallelStyle):
+    """
+    Parallel style that shards MoE experts across a specific mesh dimension.
+
+    This style is designed for ``MoELayer`` instances using ``GroupedLinear`` for experts.
+    It splits the experts across the specified
+    dimension of the device mesh (Expert Parallelism). Other dimensions in the
+    mesh treat the parameters as Replicated.
+
+    It also initializes the necessary distributed communication groups within the
+    MoE layer to handle token dispatching.
+    """
+
+    def __init__(self, shard_dim_name: str):
+        self._shard_dim_name = shard_dim_name
+
+    def _partition_experts(self, module_name: str, mod: nn.Module, device_mesh: DeviceMesh):
+        if not isinstance(mod, GroupedLinear):
+            raise TypeError("This plan should be applied only on GroupedLinear")
+
+        mesh_dim_names = device_mesh.mesh_dim_names
+
+        if mesh_dim_names is None:
+            raise ValueError("This plan should be applied only on named DeviceMeshes")
+
+        placements = [
+            Shard(0) if dim_name == self._shard_dim_name else Replicate()
+            for dim_name
+            in mesh_dim_names
+        ]
+        weight = nn.Parameter(
+            distribute_tensor(mod.weight, device_mesh, placements),
+            requires_grad=mod.weight.requires_grad
+        )
+        mod.weight = weight
+
+    def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
+        if not isinstance(module, MoELayer):
+            raise TypeError("This plan should be applied only on MoELayer")
+
+        module.enable_distributed_communicator(device_mesh.get_group(self._shard_dim_name))
+
+        for submod in module.modules():
+            if isinstance(submod, GroupedLinear):
+                distribute_module(submod, device_mesh, self._partition_experts)
+
+        return module

d9d/module/parallelism/style/to_local.py

@@ -0,0 +1,86 @@
+from typing import Any
+
+from torch import nn
+from torch.distributed import DeviceMesh
+from torch.distributed.tensor import Placement, distribute_module, distribute_tensor
+from torch.distributed.tensor.parallel import ParallelStyle
+
+
+def _build_to_local_patched_class(
+        module: nn.Module,
+        grad_placement: tuple[Placement, ...],
+        param_names: list[str]
+) -> type:
+    param_name_to_property = {
+        param_name: property(
+            lambda self, pn=param_name: self._parameters[pn].to_local(grad_placements=grad_placement)  # type: ignore
+        )
+        for param_name in param_names
+    }
+    return type(
+        f"Replicate{module.__class__.__name__}",
+        (module.__class__,),
+        param_name_to_property,
+    )
+
+
+class _ModulePatch:
+    def __init__(self, class_mapper: dict[str, type]):
+        self._class_mapper = class_mapper
+
+    def __call__(self, mod: nn.Module, *args: Any, **kwargs: Any):
+        for submod_name, submod in mod.named_modules():
+            submod.__class__ = self._class_mapper[submod_name]
+
+
+class ToLocalParallel(ParallelStyle):
+    """
+    Parallel style that distributes parameters and gradients but executes with local tensors.
+
+    This style wraps standard tensor distribution (via ``DTensor``) but injects
+    runtime hooks to temporarily unwrap ``DTensor`` parameters into local ``torch.Tensor``
+    during the forward pass.
+
+    This is useful for parallel strategies (like Replicate)
+    where the underlying calculation logic is not DTensor-aware, but the parameters must remain
+    distributed for gradient synchronization and for distributed checkpointing.
+    """
+
+    def __init__(self, param_placement: tuple[Placement, ...], grad_placement: tuple[Placement, ...]):
+        """
+        Constructs ToLocalParallel object.
+
+        Args:
+            param_placement: Tuple of placements defining how parameters are distributed.
+            grad_placement: Tuple of placements defining how gradients are synchronized.
+        """
+
+        self._grad_placement = grad_placement
+        self._param_placement = param_placement
+
+    def _distribute_params(self, name: str, module: nn.Module, device_mesh: DeviceMesh):
+        for param_name, param in module.named_parameters(recurse=False):
+            new_param = nn.Parameter(
+                distribute_tensor(param.data, device_mesh, self._param_placement),
+                requires_grad=param.requires_grad
+            )
+
+            module.register_parameter(param_name, new_param)
+
+    def _apply(self, master_module: nn.Module, device_mesh: DeviceMesh):
+        patched_classes = {}
+        original_classes = {}
+
+        for submod_name, submod in master_module.named_modules():
+            param_names = [name for name, p in submod.named_parameters(recurse=False)]
+            patched_classes[submod_name] = _build_to_local_patched_class(submod, self._grad_placement, param_names)
+            original_classes[submod_name] = submod.__class__
+
+            distribute_module(
+                submod,
+                device_mesh,
+                self._distribute_params
+            )
+
+        master_module.register_forward_pre_hook(_ModulePatch(patched_classes))
+        master_module.register_forward_hook(_ModulePatch(original_classes))

d9d/optim/stochastic/__init__.py

@@ -0,0 +1,5 @@
+from .adamw import StochasticAdamW
+
+__all__ = [
+    "StochasticAdamW"
+]

d9d/optim/stochastic/adamw.py

@@ -0,0 +1,158 @@
+from typing import cast
+
+import torch
+from torch.distributed.tensor import DTensor
+from torch.optim import Optimizer
+from torch.optim.optimizer import ParamsT, StateDict
+
+from d9d.kernel.stochastic import adamw_stochastic_bf16_
+
+_GENERATOR_STATE_KEY = "_d9d_generator_state"
+
+
+def _new_buffer(p: torch.Tensor, dtype_override: torch.dtype) -> torch.Tensor:
+    if isinstance(p, DTensor):
+        local_p = p.to_local()
+    else:
+        local_p = p
+
+    out = torch.zeros_like(local_p, dtype=dtype_override).contiguous()
+
+    if isinstance(p, DTensor):
+        out = DTensor.from_local(
+            local_tensor=out,
+            device_mesh=p.device_mesh,
+            placements=p.placements,
+            run_check=False,
+            shape=p.shape,
+            stride=p.stride(),
+        )
+
+    return out
+
+
+def _tensor_to_local(tensor: torch.Tensor) -> torch.Tensor:
+    if isinstance(tensor, DTensor):
+        return tensor.to_local()
+    return tensor
+
+
+class StochasticAdamW(Optimizer):
+    """Implements the AdamW algorithm with Stochastic Rounding.
+
+    This optimizer is designed to handle stochastic rounding primarily for BF16 training,
+    leveraging a custom kernel.
+
+    Parameters must be in BF16. Gradients could be both in BF16 and FP32.
+
+    It natively supports PyTorch distributed ``DTensor`` parameters.
+
+    It maintains its own random number generator state to ensure reproducibility.
+    """
+
+    def __init__(
+            self,
+            params: ParamsT,
+            lr: float,
+            betas: tuple[float, float] = (0.9, 0.999),
+            eps: float = 1e-8,
+            weight_decay: float = 1e-2,
+            generator: torch.Generator | None = None,
+            state_dtype: torch.dtype = torch.float32,
+    ):
+        """Constructs a new StochasticAdamW optimizer.
+
+         Args:
+             params: Iterable of parameters to optimize or dicts defining parameter groups.
+             lr: Learning rate.
+             betas: Coefficients used for computing running averages of gradient and its square.
+             eps: Term added to the denominator to improve numerical stability.
+             weight_decay: Weight decay coefficient.
+             generator: Pseudorandom number generator for stochastic rounding. If None,
+                 a new generator is created and seeded from the main PyTorch generator.
+             state_dtype: Data Type to use for the optimizer states.
+         """
+
+        if lr <= 0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if eps <= 0:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if weight_decay <= 0:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        if generator is None:
+            generator = torch.Generator(device="cpu")
+            # make the generator fork from pytorch's main generator
+            seed = cast(int, torch.randint(0, 2**32, (1,)).item())
+            generator.manual_seed(seed)
+
+        self._generator = generator
+
+        defaults = {
+            "lr": lr,
+            "betas": betas,
+            "eps": eps,
+            "weight_decay": weight_decay,
+            "state_dtype": state_dtype
+        }
+        super().__init__(params, defaults)
+
+    def state_dict(self) -> StateDict:
+        state_dict = super().state_dict()
+        state_dict[_GENERATOR_STATE_KEY] = self._generator.get_state()
+        return state_dict
+
+    def load_state_dict(self, state_dict: StateDict) -> None:
+        if _GENERATOR_STATE_KEY in state_dict:
+            self._generator.set_state(state_dict.pop(_GENERATOR_STATE_KEY))
+        super().load_state_dict(state_dict)
+
+    @torch.no_grad()
+    def step(self, closure: None = None) -> None:  # type: ignore[override]
+        if closure is not None:
+            raise ValueError("Closure is not supported")
+
+        for group in self.param_groups:
+            lr = group["lr"]
+            beta1, beta2 = group["betas"]
+            eps = group["eps"]
+            weight_decay = group["weight_decay"]
+            state_dtype = group["state_dtype"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError("StochasticAdamW does not support sparse gradients")
+
+                state = self.state[p]
+
+                # State Initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = _new_buffer(p, dtype_override=state_dtype)
+                    state["exp_avg_sq"] = _new_buffer(p, dtype_override=state_dtype)
+
+                state["step"] += 1
+                exp_avg = state["exp_avg"]
+                exp_avg_sq = state["exp_avg_sq"]
+
+                adamw_stochastic_bf16_(
+                    params=_tensor_to_local(p),
+                    grads=_tensor_to_local(grad),
+                    exp_avg=_tensor_to_local(exp_avg),
+                    exp_avg_sq=_tensor_to_local(exp_avg_sq),
+                    lr=lr,
+                    beta1=beta1,
+                    beta2=beta2,
+                    eps=eps,
+                    weight_decay=weight_decay,
+                    step=state["step"],
+                    generator=self._generator
+                )

d9d/peft/__init__.py

@@ -0,0 +1,13 @@
+"""
+Provides core logic for PEFT (Parameter-Efficient Fine-Tuning) application and base definitions.
+"""
+
+from .applicator import inject_peft_and_freeze, merge_peft
+from .base import PeftInjectionResult, PeftMethod
+
+__all__ = [
+    "PeftInjectionResult",
+    "PeftMethod",
+    "inject_peft_and_freeze",
+    "merge_peft"
+]

d9d/peft/all/__init__.py

@@ -0,0 +1,12 @@
+"""
+Package for composing multiple PEFT methods into a stack.
+"""
+
+from .config import PeftStackConfig
+from .method import PeftStack, peft_method_from_config
+
+__all__ = [
+    "PeftStack",
+    "PeftStackConfig",
+    "peft_method_from_config"
+]

d9d/peft/all/config.py

@@ -0,0 +1,31 @@
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from d9d.peft.full_tune.config import FullTuneConfig
+from d9d.peft.lora.config import LoRAConfig
+
+
+class PeftStackConfig(BaseModel):
+    """
+    Configuration for applying a stack of multiple PEFT methods sequentially.
+
+    Attributes:
+        kind: Discriminator field, always "stack".
+        methods: A list of specific PEFT configurations (e.g., LoRA, FullTune) to apply in order.
+    """
+
+    kind: Literal["stack"] = "stack"
+
+    methods: list["AnyPeftConfig"]
+
+
+AnyPeftConfig = Annotated[
+    LoRAConfig
+    | FullTuneConfig
+    | PeftStackConfig,
+    Field(discriminator="kind"),
+]
+"""
+Union type representing any valid PEFT configuration, discriminated by the 'kind' field.
+"""

d9d/peft/all/method.py

@@ -0,0 +1,76 @@
+from typing import Self, cast
+
+from pydantic import BaseModel
+from torch import nn
+
+from ..all.config import PeftStackConfig
+from ..base import PeftInjectionResult, PeftMethod, TConfig
+from ..full_tune.config import FullTuneConfig
+from ..full_tune.method import FullTune
+from ..lora.config import LoRAConfig
+from ..lora.method import LoRA
+
+
+class PeftStack(PeftMethod[PeftStackConfig]):
+    """
+    A composite PEFT method that applies a list of methods sequentially.
+    """
+
+    def __init__(self, methods: list[PeftMethod]):
+        """
+        Constructs a PeftStack object.
+
+        Args:
+            methods: A list of instantiated PEFT methods to apply in order.
+        """
+
+        self._methods = methods
+
+    def inject(self, module: nn.Module) -> PeftInjectionResult:
+        params_to_train = []
+        state_mappers = []
+
+        for method in self._methods:
+            result = method.inject(module)
+            params_to_train.extend(result.parameters_to_train)
+            state_mappers.extend(result.load_state_mappers)
+
+        return PeftInjectionResult(
+            parameters_to_train=params_to_train,
+            load_state_mappers=state_mappers
+        )
+
+    def merge(self, module: nn.Module):
+        for method in self._methods[::-1]:
+            method.merge(module)
+
+    @classmethod
+    def from_config(cls, config: PeftStackConfig) -> Self:
+        methods = []
+
+        for method in config.methods:
+            methods.append(peft_method_from_config(method))
+
+        return cls(methods)
+
+
+_PEFT_CONFIG_MAP: dict[type[BaseModel], type[PeftMethod]] = {
+    LoRAConfig: LoRA,
+    FullTuneConfig: FullTune,
+    PeftStackConfig: PeftStack
+}
+
+
+def peft_method_from_config(config: TConfig) -> PeftMethod[TConfig]:
+    """
+    Factory function to instantiate the correct PeftMethod based on the configuration type.
+
+    Args:
+        config: A specific PEFT configuration object (e.g., LoRAConfig).
+
+    Returns:
+        The corresponding method instance.
+    """
+
+    method_cls = cast(type[PeftMethod[TConfig]], _PEFT_CONFIG_MAP[type(config)])
+    return method_cls.from_config(config)

d9d/peft/applicator.py

@@ -0,0 +1,47 @@
+from torch import nn
+
+from d9d.model_state.mapper import ModelStateMapper
+from d9d.model_state.mapper.compose import ModelStateMapperParallel
+
+from .base import PeftMethod
+
+
+def inject_peft_and_freeze(method: PeftMethod, module: nn.Module) -> ModelStateMapper:
+    """
+    Applies a PEFT method to a module, freezes non-trained parameters, and prepares state mapping.
+
+    This function performs three main steps:
+
+    1. Sets `requires_grad=False` for all parameters in the module.
+    2. Calls the method's `inject` to modify the model structure.
+    3. Sets `requires_grad=True` for the parameters returned by the injection result.
+
+    Args:
+        method: The PEFT method strategy to apply.
+        module: The PyTorch module to modify.
+
+    Returns:
+        A ModelStateMapper capable of loading checkpoint weights into the modified structure.
+    """
+
+    for param in module.parameters():
+        param.requires_grad = False
+
+    result = method.inject(module)
+
+    for param in result.parameters_to_train:
+        param.requires_grad = True
+
+    return ModelStateMapperParallel(result.load_state_mappers)
+
+
+def merge_peft(method: PeftMethod, module: nn.Module):
+    """
+    Merges PEFT adaptations back into the base model weights.
+
+    Args:
+        method: The PEFT method strategy originally applied.
+        module: The PyTorch module to merge.
+    """
+
+    method.merge(module)

d9d/peft/base.py

@@ -0,0 +1,70 @@
+import abc
+import dataclasses
+from typing import Generic, Self, TypeVar
+
+from pydantic import BaseModel
+from torch import nn
+
+from d9d.model_state.mapper import ModelStateMapper
+
+
+@dataclasses.dataclass(slots=True)
+class PeftInjectionResult:
+    """
+    Encapsulates the result of injecting a PEFT method into a model.
+
+    Attributes:
+        parameters_to_train: A list of parameters that should remain trainable.
+        load_state_mappers: A list of mappers required to load pre-trained weights into the modified structure.
+    """
+
+    parameters_to_train: list[nn.Parameter]
+    load_state_mappers: list[ModelStateMapper]
+
+
+TConfig = TypeVar("TConfig", bound=BaseModel)
+
+
+class PeftMethod(abc.ABC, Generic[TConfig]):
+    """
+    Abstract base class for all Parameter-Efficient Fine-Tuning methods.
+    """
+
+    @abc.abstractmethod
+    def inject(self, module: nn.Module) -> PeftInjectionResult:
+        """
+        Modifies the module in-place to apply the PEFT strategy.
+
+        Args:
+            module: The PyTorch module to modify.
+
+        Returns:
+            Result object containing trainable parameters and structure mappers.
+        """
+        ...
+
+    @abc.abstractmethod
+    def merge(self, module: nn.Module):
+        """
+        Merges the trained adapters back into the base model parameters.
+
+        Args:
+            module: The PyTorch module to update.
+        """
+
+        ...
+
+    @classmethod
+    @abc.abstractmethod
+    def from_config(cls, config: TConfig) -> Self:
+        """
+        Creates an instance of the method from a configuration object.
+
+        Args:
+            config: The configuration object.
+
+        Returns:
+            An instance of the PeftMethod.
+        """
+
+        ...

d9d/peft/full_tune/__init__.py

@@ -0,0 +1,11 @@
+"""
+Package for Full Fine-Tuning functionality within the PEFT framework.
+"""
+
+from .config import FullTuneConfig
+from .method import FullTune
+
+__all__ = [
+    "FullTune",
+    "FullTuneConfig"
+]