d9d 0.1.0__py3-none-any.whl

d9d/peft/full_tune/config.py

@@ -0,0 +1,20 @@
+from re import Pattern
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class FullTuneConfig(BaseModel):
+    """
+    Configuration for Full Fine-Tuning.
+
+    Allows specifying which modules should be fully fine-tuned using regex patterns.
+
+    Attributes:
+        kind: Discriminator field, always "full_tune".
+        module_name_pattern: Regular expression matching module names to unfreeze.
+    """
+
+    kind: Literal["full_tune"] = "full_tune"
+
+    module_name_pattern: Pattern

d9d/peft/full_tune/method.py

@@ -0,0 +1,46 @@
+from typing import Self
+
+from torch import nn
+
+from ..base import PeftInjectionResult, PeftMethod
+from .config import FullTuneConfig
+
+
+class FullTune(PeftMethod[FullTuneConfig]):
+    """
+    Implements Full Fine-Tuning as a 'PEFT' method.
+
+    Instead of injecting adapters, this method simply identifies existing parameters
+    that match the configuration pattern and marks them for training.
+    """
+
+    def __init__(self, config: FullTuneConfig):
+        """
+        Constructs a FullTune object.
+
+        Args:
+            config: Configuration defining the module name patterns to fine-tune.
+        """
+
+        self._config = config
+
+    def inject(self, module: nn.Module) -> PeftInjectionResult:
+        params_to_train = []
+
+        for mod_name, mod in module.named_modules():
+            is_applicable = self._config.module_name_pattern.fullmatch(mod_name)
+
+            if is_applicable:
+                params_to_train.extend(mod.parameters())
+
+        return PeftInjectionResult(
+            parameters_to_train=params_to_train,
+            load_state_mappers=[]
+        )
+
+    def merge(self, module: nn.Module):
+        pass  # do nothing here
+
+    @classmethod
+    def from_config(cls, config: FullTuneConfig) -> Self:
+        return cls(config)

d9d/peft/lora/__init__.py

@@ -0,0 +1,15 @@
+"""
+Package for Low-Rank Adaptation (LoRA) implementation.
+"""
+
+from .config import LoRAConfig, LoRAParameters
+from .layer import LoRAGroupedLinear, LoRALinear
+from .method import LoRA
+
+__all__ = [
+    "LoRA",
+    "LoRAConfig",
+    "LoRAGroupedLinear",
+    "LoRALinear",
+    "LoRAParameters"
+]

d9d/peft/lora/config.py

@@ -0,0 +1,35 @@
+from re import Pattern
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class LoRAParameters(BaseModel):
+    """
+    Hyperparameters for LoRA layers.
+
+    Attributes:
+        r: Rank of the low-rank adaptation matrices.
+        alpha: Scaling factor for the learned weights.
+        dropout: Dropout probability for the input to LoRA layers.
+    """
+
+    r: int
+    alpha: int
+    dropout: float
+
+
+class LoRAConfig(BaseModel):
+    """
+    Configuration for LoRA application.
+
+    Attributes:
+        kind: Discriminator field, always "lora".
+        module_name_pattern: Regular expression matching module names to wrap with LoRA.
+        params: Hyperparameters for the LoRA layers.
+    """
+
+    kind: Literal["lora"] = "lora"
+
+    module_name_pattern: Pattern
+    params: LoRAParameters

d9d/peft/lora/layer.py

@@ -0,0 +1,177 @@
+import torch
+from torch import nn
+
+from d9d.module.block.moe import GroupedLinear
+
+from .config import LoRAParameters
+
+
+class LoRALinear(nn.Module):
+    """
+    A LoRA wrapper around a standard PyTorch Linear layer.
+
+    Wraps a base linear layer and adds low-rank adaptation matrices A and B.
+
+    Attributes:
+        lora_A: The A matrix (in_features -> r).
+        lora_B: The B matrix (r -> out_features).
+        base: The original base Linear layer.
+        dropout: Scaling dropout layer.
+    """
+
+    def __init__(
+            self,
+            base_layer: nn.Linear,
+            params: LoRAParameters
+    ):
+        """
+        Constructs a LoRALinear layer.
+
+        Args:
+            base_layer: The original Linear layer to wrap.
+            params: LoRA hyperparameters (r, alpha, dropout).
+
+        Raises:
+            ValueError: If the base layer has a bias (currently unsupported).
+        """
+
+        super().__init__()
+        self.lora_A = nn.Linear(
+            base_layer.in_features, params.r, bias=False,
+            device=base_layer.weight.device,
+            dtype=base_layer.weight.dtype
+        )
+        self.lora_B = nn.Linear(
+            params.r, base_layer.out_features, bias=False,
+            device=base_layer.weight.device,
+            dtype=base_layer.weight.dtype
+        )
+        self.base = base_layer
+
+        if base_layer.bias is not None:
+            raise ValueError("LoRA is unsupported with biased linear layers")
+
+        self.dropout: nn.Dropout = nn.Dropout(params.dropout)
+
+        self._scale: float = params.alpha / params.r
+
+        self.reset_parameters()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Takes input tensor, computes base output and LoRA adaptation, and returns the sum.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            The output of base(x) + scale * (B @ A @ dropout(x)).
+        """
+
+        base_x = self.base(x)
+        adapt_x = self._scale * self.lora_B(self.lora_A(self.dropout(x)))
+        return base_x + adapt_x
+
+    @torch.no_grad()
+    def merge_with_base_(self) -> nn.Linear:
+        """
+        Collapse the LoRA weights into the base linear layer.
+
+        Returns:
+            The modified base linear layer with updated weights.
+        """
+
+        mod = self.base
+        mod.weight.data += (self.lora_B.weight.data @ self.lora_A.weight.data) * self._scale
+        return mod
+
+    def reset_parameters(self):
+        """
+        Resets LoRA parameters. A is random, B is zeroed.
+        """
+
+        self.lora_A.reset_parameters()
+        nn.init.zeros_(self.lora_B.weight)
+
+
+class LoRAGroupedLinear(nn.Module):
+    """
+    A LoRA wrapper around a GroupedLinear layer (commonly used in MoE or grouped query attention).
+
+    Attributes:
+        lora_A: The A matrix (grouped linear).
+        lora_B: The B matrix (grouped linear).
+        base: The original base GroupedLinear layer.
+        dropout: Scaling dropout layer.
+    """
+
+    def __init__(
+            self,
+            base_layer: GroupedLinear,
+            params: LoRAParameters
+    ):
+        """
+        Constructs a LoRAGroupedLinear layer.
+
+        Args:
+            base_layer: The original GroupedLinear layer to wrap.
+            params: LoRA hyperparameters.
+        """
+
+        super().__init__()
+        self.lora_A = GroupedLinear(
+            base_layer.n_groups, base_layer.in_features, params.r,
+            device=base_layer.weight.device,
+            dtype=base_layer.weight.dtype
+        )
+        self.lora_B = GroupedLinear(
+            base_layer.n_groups,
+            params.r,
+            base_layer.out_features,
+            device=base_layer.weight.device,
+            dtype=base_layer.weight.dtype
+        )
+        self.base = base_layer
+
+        self.dropout = nn.Dropout(params.dropout)
+
+        self._scale = params.alpha / params.r
+
+        self.reset_parameters()
+
+    def forward(self, x: torch.Tensor, x_groups: torch.Tensor) -> torch.Tensor:
+        """
+        Computes forward pass for grouped inputs.
+
+        Args:
+            x: Input tensor.
+            x_groups: A tensor indicating group indices for each input.
+
+        Returns:
+            Combined output of base and LoRA path.
+        """
+
+        base_x = self.base(x, x_groups)
+        adapt_x = self._scale * self.lora_B(self.lora_A(self.dropout(x), x_groups), x_groups)
+        return base_x + adapt_x
+
+    @torch.no_grad()
+    def merge_with_base_(self) -> GroupedLinear:
+        """
+        Collapse the LoRA weights into the base GroupedLinear layer.
+
+        Returns:
+            The modified GroupedLinear layer.
+        """
+
+        mod = self.base
+        mod.weight.data += (torch.bmm(self.lora_A.weight.data, self.lora_B.weight.data)) * self._scale
+        return mod
+
+    def reset_parameters(self):
+        """
+        Resets LoRA parameters. A is random, B is zeroed.
+        """
+
+        self.lora_A.reset_parameters()
+        nn.init.zeros_(self.lora_B.weight)

d9d/peft/lora/method.py

@@ -0,0 +1,132 @@
+from typing import Self
+
+import torch
+from torch import nn
+
+from d9d.model_state.mapper import ModelStateMapper
+from d9d.model_state.mapper.leaf import ModelStateMapperRename
+from d9d.module.block.moe import GroupedLinear
+
+from ..base import PeftInjectionResult, PeftMethod
+from .config import LoRAConfig
+from .layer import LoRAGroupedLinear, LoRALinear
+
+_CAN_APPLY_MODULES = (nn.Linear, GroupedLinear)
+_LORA_MODULES = (LoRALinear, LoRAGroupedLinear)
+
+
+def named_modules_without_lora(
+        module: nn.Module,
+        memo: set[nn.Module] | None = None,
+        prefix: str = "",
+        remove_duplicate: bool = True
+):
+    """
+    Yields named modules, skipping submodules that are already LoRA layers.
+
+    This prevents recursively re-injecting LoRA into an already wrapped layer during
+    traversal.
+
+    Args:
+        module: The root module to traverse.
+        memo: Set of processed modules to avoid duplicates.
+        prefix: Current namespace prefix.
+        remove_duplicate: Whether to skip modules seen in memo.
+
+    Yields:
+        Tuple of (name, module).
+    """
+
+    if isinstance(module, _LORA_MODULES):
+        return
+
+    if memo is None:
+        memo = set()
+    if module in memo:
+        return
+
+    if remove_duplicate:
+        memo.add(module)
+
+    yield prefix, module
+
+    for name, submodule in module.named_children():
+        if submodule is None:
+            continue
+
+        submodule_prefix = prefix + ("." if prefix else "") + name
+        yield from named_modules_without_lora(submodule, memo, submodule_prefix, remove_duplicate)
+
+
+class LoRA(PeftMethod[LoRAConfig]):
+    """
+    Implements the Low-Rank Adaptation (LoRA) injection strategy.
+
+    It scans the module structure for `nn.Linear` or `GroupedLinear` layers matching
+    the configured name pattern. Matched layers are replaced with LoRA wrappers.
+
+    It also generates `ModelStateMapperRename` objects. Since the original weight
+    `layer.weight` is now at `layer.base.weight` inside the wrapper, the mapper
+    ensures that loading a standard checkpoint still works by redirecting the key.
+    """
+
+    def __init__(self, config: LoRAConfig):
+        """
+         Constructs a LoRA method.
+
+         Args:
+             config: LoRA configuration containing patterns and hyperparameters.
+         """
+
+        self._config = config
+
+    def inject(self, module: nn.Module) -> PeftInjectionResult:
+        params_to_train: list[nn.Parameter] = []
+        state_mappers: list[ModelStateMapper] = []
+
+        for mod_name, mod in named_modules_without_lora(module):
+            if not isinstance(mod, _CAN_APPLY_MODULES):
+                continue
+
+            if not self._config.module_name_pattern.fullmatch(mod_name):
+                continue
+
+            lora_mod: LoRALinear | LoRAGroupedLinear
+            if isinstance(mod, nn.Linear):
+                lora_mod = LoRALinear(mod, self._config.params)
+            elif isinstance(mod, GroupedLinear):
+                lora_mod = LoRAGroupedLinear(mod, self._config.params)
+            else:
+                raise ValueError(f"Unknown layer {type(mod)} for LoRA")
+
+            params_to_train.extend(lora_mod.lora_A.parameters())
+            params_to_train.extend(lora_mod.lora_B.parameters())
+
+            state_mappers.append(ModelStateMapperRename(
+                name_from=f"{mod_name}.weight",
+                name_to=f"{mod_name}.base.weight"
+            ))
+
+            module.set_submodule(mod_name, lora_mod)
+
+        return PeftInjectionResult(
+            parameters_to_train=params_to_train,
+            load_state_mappers=state_mappers
+        )
+
+    def merge(self, module: nn.Module):
+        for mod_name, mod in module.named_modules():
+            if not isinstance(mod, _LORA_MODULES):
+                continue
+
+            if not self._config.module_name_pattern.fullmatch(mod_name):
+                continue
+
+            with torch.no_grad():
+                orig_mod = mod.merge_with_base_()
+
+            module.set_submodule(mod_name, orig_mod)
+
+    @classmethod
+    def from_config(cls, config: LoRAConfig) -> Self:
+        return cls(config)

d9d/pipelining/api/__init__.py

@@ -0,0 +1,19 @@
+"""
+Pipelining API that is intended to be accessible by end user.
+"""
+
+from .module import (
+    ModuleSupportsPipelining,
+    PipelineStageInfo,
+    distribute_layers_for_pipeline_stage,
+)
+from .schedule import PipelineSchedule
+from .sharding import PipelineShardingSpec
+
+__all__ = [
+    "ModuleSupportsPipelining",
+    "PipelineSchedule",
+    "PipelineShardingSpec",
+    "PipelineStageInfo",
+    "distribute_layers_for_pipeline_stage"
+]

d9d/pipelining/api/module.py

@@ -0,0 +1,149 @@
+import dataclasses
+import typing
+
+import torch
+
+
+@dataclasses.dataclass
+class PipelineStageInfo:
+    """
+    Holds information about the current position within the distributed pipeline.
+
+    Attributes:
+        current_stage: The 0-based index of the current pipeline stage.
+        num_stages: The total number of stages in the pipeline.
+    """
+
+    current_stage: int
+    num_stages: int
+
+    @property
+    def is_current_stage_first(self) -> bool:
+        """
+        Determines if this is the first stage in the pipeline.
+
+        Returns:
+            True if current_stage is 0.
+        """
+
+        return self.current_stage == 0
+
+    @property
+    def is_current_stage_last(self) -> bool:
+        """
+        Determines if this is the last stage in the pipeline.
+
+        Returns:
+            True if current_stage is the last index.
+        """
+
+        return self.current_stage == self.num_stages - 1
+
+
+def distribute_layers_for_pipeline_stage(
+        num_layers: int,
+        num_virtual_layers_pre: int,
+        num_virtual_layers_post: int,
+        stage: PipelineStageInfo
+) -> tuple[int, int]:
+    """
+    Calculates the layer index range for a specific pipeline stage.
+
+    This function distributes a given number of layers across multiple pipeline
+    stages as evenly as possible. It accounts for additional, non-layer
+    computational load on the first and last stages (e.g., embeddings and the
+    LM head) by using the concept of 'virtual layers' to reserve capacity.
+
+    Args:
+        num_layers: The total number of primary model layers to be distributed
+            (e.g., the transformer blocks).
+        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).
+        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).
+        stage: An object containing total stages and current stage index.
+
+    Returns:
+        A tuple (start_index, end_index), representing the slice of layers for
+            the given stage. The start_index is inclusive and the end_index is
+            exclusive.
+
+    Raises:
+        ValueError: If the pipeline configuration results in a stage having zero
+            or negative layers assigned (pipeline too long for the model size).
+    """
+
+    num_layers_virtual = num_layers + num_virtual_layers_pre + num_virtual_layers_post
+
+    base_layers_per_stage = num_layers_virtual // stage.num_stages
+    extra_layers = num_layers_virtual % stage.num_stages
+
+    layer_count_per_stage = []
+
+    for proposed_stage_i in range(stage.num_stages):
+        proposed_stage = PipelineStageInfo(num_stages=stage.num_stages, current_stage=proposed_stage_i)
+        layers = base_layers_per_stage + 1 if proposed_stage_i < extra_layers else base_layers_per_stage
+
+        adjustment = 0
+        if proposed_stage.is_current_stage_first:
+            adjustment += num_virtual_layers_pre
+        if proposed_stage.is_current_stage_last:
+            adjustment += num_virtual_layers_post
+
+        actual_layers = layers - adjustment
+
+        if actual_layers <= 0:
+            raise ValueError(f"Tried to distribute layers, but got {actual_layers} on "
+                             f"stage {proposed_stage.current_stage}. Perhaps the pipeline is too long for this model?")
+
+        layer_count_per_stage.append(actual_layers)
+
+    start_layer_id = sum(layer_count_per_stage[:stage.current_stage])
+    num_layers_in_stage = layer_count_per_stage[stage.current_stage]
+
+    return start_layer_id, start_layer_id + num_layers_in_stage
+
+
+@typing.runtime_checkable
+class ModuleSupportsPipelining(typing.Protocol):
+    """
+    Protocol for modules that support pipeline parallelism metadata inference.
+
+    Classes implementing this protocol enable the framework to pre-calculate
+    tensor shapes and types required for inter-stage communication (p2p)
+    without executing the full forward pass.
+    """
+
+    def infer_stage_inputs_from_pipeline_inputs(
+            self, inputs: dict[str, torch.Tensor], n_microbatches: int
+    ) -> dict[str, torch.Tensor]:
+        """
+        Infers the input tensors metadata for the current pipeline stage based on global batch inputs.
+
+        Args:
+            inputs: Global inputs for the pipeline.
+            n_microbatches: Number of microbatches the global batch is split into.
+
+        Returns:
+            Dictionary of input tensors expected by this specific stage locally.
+        """
+
+        ...
+
+    def infer_stage_outputs_from_pipeline_inputs(
+            self, inputs: dict[str, torch.Tensor], n_microbatches: int
+    ) -> dict[str, torch.Tensor]:
+        """
+        Infers the output tensors metadata for the current pipeline stage based on global batch inputs.
+
+        Args:
+            inputs: Global inputs for the pipeline (typically a batch).
+            n_microbatches: Number of microbatches the global batch is split into.
+
+        Returns:
+            Dictionary of output tensors produced by this specific stage locally.
+        """
+
+        ...

d9d/pipelining/api/schedule.py

@@ -0,0 +1,50 @@
+import abc
+from typing import Any
+
+import torch
+
+from .sharding import PipelineShardingSpec
+
+# TODO: feature - support any PyTrees as pipeline parameters
+
+
+class PipelineSchedule(abc.ABC):
+    """Abstract base class defining the interface for pipeline execution schedules."""
+
+    @abc.abstractmethod
+    def configure_buffers(
+            self,
+            inputs: dict[str, torch.Tensor],
+            kwargs: dict[str, Any],
+            sharding_spec: PipelineShardingSpec | None
+    ):
+        """
+        Configures internal state and buffers based on input shapes.
+
+        This method allows the schedule to pre-allocate memory or setup sharding
+        specifications based on the structure of the input data before execution begins.
+
+        Args:
+            inputs: A dictionary of input tensors.
+            kwargs: A dictionary of keyword arguments.
+            sharding_spec: A specification defining how inputs and kwargs should be split
+                into micro-batches. If None, assumes standard split-by-zero-dim behavior.
+        """
+
+        ...
+
+    @abc.abstractmethod
+    def step(self, inputs: dict[str, torch.Tensor], kwargs: dict[str, Any]):
+        """
+        Executes a single pipeline step using the provided inputs.
+
+         This typically involves distributing inputs across microbatches,
+         executing forward and backward passes according to the specific schedule logic,
+         and handling communications between stages.
+
+         Args:
+             inputs: A dictionary of global input tensors.
+             kwargs: A dictionary of global keyword arguments.
+         """
+
+        ...

d9d/pipelining/api/sharding.py

@@ -0,0 +1,9 @@
+import dataclasses
+
+from d9d.core.sharding import ShardingSpec
+
+
+@dataclasses.dataclass
+class PipelineShardingSpec:
+    input_data: ShardingSpec | None = None
+    input_kwargs: ShardingSpec | None = None

d9d/pipelining/factory/__init__.py

@@ -0,0 +1,21 @@
+from .config import (
+    AnyPipelineScheduleConfig,
+    PipelineSchedule1F1BConfig,
+    PipelineScheduleDualPipeVConfig,
+    PipelineScheduleGPipeConfig,
+    PipelineScheduleInferenceConfig,
+    PipelineScheduleLoopedBFSConfig,
+    PipelineScheduleZeroBubbleVConfig,
+)
+from .factory import build_schedule
+
+__all__ = [
+    "AnyPipelineScheduleConfig",
+    "PipelineSchedule1F1BConfig",
+    "PipelineScheduleDualPipeVConfig",
+    "PipelineScheduleGPipeConfig",
+    "PipelineScheduleInferenceConfig",
+    "PipelineScheduleLoopedBFSConfig",
+    "PipelineScheduleZeroBubbleVConfig",
+    "build_schedule"
+]