d9d 0.1.0__py3-none-any.whl

d9d/pipelining/infra/stage/stage.py

@@ -0,0 +1,321 @@
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from torch import nn
+
+from d9d.pipelining.api import ModuleSupportsPipelining, PipelineStageInfo
+
+from .communications import StageCommunicationHandler
+from .computations import BackwardComputeHandler, ForwardComputeHandler
+
+
+class PipelineStage:
+    """
+    Represents a single structural stage in a Pipelined Model.
+
+    This class acts as an orchestrator that combines `StageCommunicationHandler` (for I/O)
+    and `Forward/BackwardComputeHandler` (for execution). It abstracts away the complexity
+    of buffer management, distributed communication, and gradient calculation from the scheduler.
+    """
+
+    def __init__(
+            self,
+            info: PipelineStageInfo,
+            module: nn.Module,
+            group: dist.ProcessGroup,
+            stage_to_host_topology: dict[int, int]
+    ):
+        """
+        Constructs a PipelineStage object.
+
+        Args:
+            info: Metadata about the stage (index, total stages).
+            module: The PyTorch module executed by this stage.
+            group: The distributed process group for pipeline communications.
+            stage_to_host_topology: Dict mapping stage ID to PP rank hosting it.
+        """
+
+        self._info = info
+        self._module = module
+        self._group = group
+        self._stage_to_host_topology = stage_to_host_topology
+
+        self._has_backward = False
+
+        self._forward_comm: StageCommunicationHandler | None = None
+        self._backward_comm: StageCommunicationHandler | None = None
+
+        self._forward_comp = ForwardComputeHandler(
+            stage_index=info.current_stage,
+            module=module
+        )
+        self._backward_comp = BackwardComputeHandler(
+            stage_index=info.current_stage,
+            module=module
+        )
+
+    @property
+    def info(self) -> PipelineStageInfo:
+        return self._info
+
+    def configure_buffers(
+            self,
+            num_microbatches: int,
+            has_backward: bool,
+            pipeline_inputs: dict[str, torch.Tensor]
+    ):
+        """
+        Initializes the communication handlers and buffers for the stage.
+
+        This must be called before execution to establish P2P buffer sizes and directions.
+
+        Args:
+            num_microbatches: Total number of microbatches to process.
+            has_backward: Does this pipeline stage should store info for a backward pass
+            pipeline_inputs: Pipeline input data.
+        """
+
+        self._has_backward = has_backward
+
+        prev_stage_idx = None if self._info.is_current_stage_first else self._info.current_stage - 1
+        next_stage_idx = None if self._info.is_current_stage_last else self._info.current_stage + 1
+
+        with torch.device("meta"):
+            if not isinstance(self._module, ModuleSupportsPipelining):
+                raise TypeError("Module does not implement ModuleSupportsPipelining protocol")
+            inputs_meta = self._module.infer_stage_inputs_from_pipeline_inputs(
+                inputs=pipeline_inputs,
+                n_microbatches=num_microbatches
+            )
+            outputs_meta = self._module.infer_stage_outputs_from_pipeline_inputs(
+                inputs=pipeline_inputs,
+                n_microbatches=num_microbatches
+            )
+
+        self._forward_comm = StageCommunicationHandler(
+            name="fwd",
+            stage_index=self._info.current_stage,
+            num_microbatches=num_microbatches,
+            input_stage_index=prev_stage_idx,
+            input_args=inputs_meta,
+            output_stage_index=next_stage_idx,
+            output_args=outputs_meta,
+            group=self._group,
+            stage_idx_to_host_rank=self._stage_to_host_topology
+        )
+        self._forward_comm.set_input_requires_grad_(requires_grad=has_backward)
+
+        if has_backward:
+            # for grad - current stage receives OUTPUTS as inputs and sends INPUTS as outputs
+            # because it is reversed forward
+            self._backward_comm = StageCommunicationHandler(
+                name="bwd",
+                stage_index=self._info.current_stage,
+                num_microbatches=num_microbatches,
+                input_stage_index=next_stage_idx,
+                input_args=outputs_meta,
+                output_stage_index=prev_stage_idx,
+                output_args=inputs_meta,
+                group=self._group,
+                stage_idx_to_host_rank=self._stage_to_host_topology
+            )
+        else:
+            self._backward_comm = None
+
+    def set_local_fwd_input(self, inputs: dict[str, torch.Tensor], microbatch_index: int):
+        """
+        Sets local forward inputs manually.
+
+        Used for the V-shape schedulers.
+        """
+
+        if self._forward_comm is None:
+            raise ValueError("You must configure stage buffers first")
+
+        self._forward_comm.set_inputs_local(inputs, microbatch_index)
+
+    def get_local_fwd_output(self, microbatch_index: int) -> dict[str, torch.Tensor]:
+        return self._forward_comp.get_outputs(microbatch_index)
+
+    def pop_local_bwd_output(self, microbatch_index: int) -> dict[str, torch.Tensor]:
+        """
+        Retrieves local backward outputs (gradients).
+        """
+
+        if not self._has_backward:
+            raise ValueError()
+
+        return self._backward_comp.pop_for_sending(microbatch_index)
+
+    def set_local_bwd_input(self, inputs: dict[str, torch.Tensor], microbatch_index: int):
+        """
+        Sets local backward inputs (output gradients) manually.
+        """
+
+        if not self._has_backward:
+            raise ValueError()
+
+        if self._backward_comm is None:
+            raise ValueError("You must configure stage buffers first")
+
+        self._backward_comm.set_inputs_local(inputs, microbatch_index)
+
+    def get_fwd_recv_ops(self, microbatch_index: int) -> list[dist.P2POp]:
+        """Returns P2P ops to receive forward inputs for the given microbatch."""
+
+        if self._forward_comm is None:
+            raise ValueError("You must configure stage buffers first")
+
+        return self._forward_comm.create_receive_ops(microbatch_index)
+
+    def get_fwd_send_ops(self, microbatch_index: int) -> list[dist.P2POp]:
+        """Returns P2P ops to send forward outputs for the given microbatch."""
+
+        if self._forward_comm is None:
+            raise ValueError("You must configure stage buffers first")
+
+        fwd_result = self._forward_comp.get_outputs(microbatch_index)
+        return self._forward_comm.create_send_ops(fwd_result)
+
+    def get_bwd_recv_ops(self, microbatch_index: int) -> list[dist.P2POp]:
+        """Returns P2P ops to receive backward gradients for the given microbatch."""
+
+        if not self._has_backward:
+            return []
+
+        if self._backward_comm is None:
+            raise ValueError("You must configure stage buffers first")
+
+        return self._backward_comm.create_receive_ops(microbatch_index)
+
+    def get_bwd_send_ops(self, microbatch_index: int) -> list[dist.P2POp]:
+        """Returns P2P ops to send backward gradients for the given microbatch."""
+
+        if not self._has_backward:
+            return []
+
+        if self._backward_comm is None:
+            raise ValueError("You must configure stage buffers first")
+
+        bwd_result = self._backward_comp.pop_for_sending(microbatch_index)
+        return self._backward_comm.create_send_ops(bwd_result)
+
+    def forward_one_chunk(
+            self,
+            microbatch_index: int,
+            pipeline_inputs: dict[str, torch.Tensor],
+            pipeline_kwargs: dict[str, Any] | None = None,
+    ):
+        """
+        Executes a forward pass for a single microbatch chunk.
+
+        Fetches inputs from the communication buffer (or `pipeline_inputs` if first stage),
+        runs the computation, and caches the result.
+
+        Args:
+            microbatch_index: The microbatch index.
+            pipeline_inputs: Inputs provided locally (only used if this is the first stage).
+            pipeline_kwargs: Additional arguments for the module.
+
+        Returns:
+            The output tensors of the forward pass.
+        """
+
+        if self._forward_comm is None:
+            raise ValueError("You must configure stage buffers first")
+
+        if self._info.is_current_stage_first:
+            inputs = pipeline_inputs
+        else:
+            inputs = self._forward_comm.get_inputs(microbatch_index)
+
+        kwargs = pipeline_kwargs or {}
+
+        self._forward_comp.run(
+            microbatch_index=microbatch_index,
+            inputs=inputs,
+            kwargs=kwargs
+        )
+
+    def backward_one_chunk(
+            self,
+            microbatch_index: int,
+            loss: torch.Tensor | None = None,
+            full_backward: bool = True
+    ):
+        """
+        Executes a backward pass for a single microbatch chunk.
+
+        Can perform either a full backward or just the input gradients (if `full_backward=False`).
+        It fetches required data from forward cache and communication buffers.
+
+        Args:
+            microbatch_index: The microbatch index.
+            loss: The loss tensor (only used if this is the last stage).
+            full_backward: If True, computes grads for inputs and weights. If False, only for inputs.
+        """
+
+        if not self._has_backward:
+            raise ValueError()
+
+        if self._backward_comm is None:
+            raise ValueError("You must configure stage buffers first")
+
+        inputs, fwd_outputs = self._forward_comp.pop_inputs_outputs(microbatch_index)
+
+        outputs: dict[str, torch.Tensor]
+        outputs_grad: dict[str, torch.Tensor] | None
+
+        if self._info.is_current_stage_last:
+            if loss is None:
+                raise ValueError("Cannot perform backward on last stage without loss specified")
+            outputs = {"loss": loss}
+            outputs_grad = None
+        else:
+            outputs = fwd_outputs
+            outputs_grad = self._backward_comm.get_inputs(microbatch_index)
+
+        if full_backward:
+            self._backward_comp.backward_full(
+                microbatch_index=microbatch_index,
+                inputs=inputs,
+                outputs=outputs,
+                outputs_grad=outputs_grad
+            )
+        else:
+            self._backward_comp.backward_input(
+                microbatch_index=microbatch_index,
+                inputs=inputs,
+                outputs=outputs,
+                outputs_grad=outputs_grad
+            )
+
+        if self._info.is_current_stage_last and not self._info.is_current_stage_first:
+            for t in fwd_outputs.values():
+                if not t._is_view():  # noqa: SLF001
+                    t.detach_()
+
+    def backward_weight_one_chunk(self, microbatch_index: int):
+        """
+        Executes the weight gradient accumulation part of the backward pass.
+
+        This assumes `backward_one_chunk(..., full_backward=False)` was already called
+        for this microbatch.
+
+        Args:
+            microbatch_index: The microbatch index.
+        """
+
+        if not self._has_backward:
+            raise ValueError()
+
+        self._backward_comp.backward_weight(microbatch_index=microbatch_index)
+
+    def reset(self):
+        """Resets the internal state of communication handlers, clearing gradients on buffers."""
+
+        if self._forward_comm is not None:
+            self._forward_comm.reset()
+        if self._backward_comm is not None:
+            self._backward_comm.reset()

d9d/pipelining/infra/stage/struct_helper.py

@@ -0,0 +1,46 @@
+from collections.abc import Iterable, Sequence
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+class DictFlattener:
+    """
+    Helper class to flatten and unflatten dictionaries into sequences deterministically.
+    """
+
+    def __init__(self, keys: Iterable[str]):
+        """
+        Constructs a DictFlattener object.
+
+        Args:
+            keys: The collection of dictionary keys to manage. They will be sorted internally.
+        """
+
+        self._order_to_key = {i: x for i, x in enumerate(sorted(keys))}
+
+    def flatten(self, inputs: dict[str, T]) -> list[T]:
+        """
+        Converts a dictionary into a list based on the sorted internal key order.
+
+        Args:
+            inputs: The dictionary to flatten. Must contain all keys provided at init.
+
+        Returns:
+            A list of values sorted by their corresponding keys.
+        """
+
+        return [inputs[self._order_to_key[i]] for i in range(len(inputs))]
+
+    def unflatten(self, outputs: Sequence[T]) -> dict[str, T]:
+        """
+        Reconstructs a dictionary from a sequence of values.
+
+        Args:
+            outputs: A sequence of values corresponding to the sorted internal key order.
+
+        Returns:
+            A dictionary mapping original keys to the provided values.
+        """
+
+        return {self._order_to_key[i]: out for i, out in enumerate(outputs)}

d9d/pipelining/training/__init__.py

@@ -0,0 +1,7 @@
+from .optimizer import PipelinedOptimizer
+from .scheduler import PipelinedLRScheduler
+
+__all__ = [
+    "PipelinedLRScheduler",
+    "PipelinedOptimizer"
+]

d9d/pipelining/training/optimizer.py

@@ -0,0 +1,41 @@
+from typing import Any
+
+from torch.distributed import DeviceMesh
+
+from d9d.core.protocol import OptimizerProtocol
+
+
+class PipelinedOptimizer(OptimizerProtocol):
+    """
+    Wrapper that manages multiple optimizers for a pipeline parallel rank.
+
+    In a pipeline parallel setup, a single rank might host multiple stages, each having its own parameters
+    and optimizer.
+    This class aggregates them into a single interface.
+    """
+
+    def __init__(self, mesh_pp: DeviceMesh, optimizers: list[OptimizerProtocol]):
+        super().__init__()
+
+        self._mesh_pp = mesh_pp
+        self._optimizers = optimizers
+
+    def state_dict(self) -> dict[str, Any]:
+        pp_rank = self._mesh_pp.get_local_rank()
+        return {
+            f"pp_{pp_rank}_stage_{i}": optimizer.state_dict()
+            for i, optimizer in enumerate(self._optimizers)
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        pp_rank = self._mesh_pp.get_local_rank()
+        for i, optimizer in enumerate(self._optimizers):
+            optimizer.load_state_dict(state_dict[f"pp_{pp_rank}_stage_{i}"])
+
+    def step(self) -> None:
+        for optimizer in self._optimizers:
+            optimizer.step()
+
+    def zero_grad(self) -> None:
+        for optimizer in self._optimizers:
+            optimizer.zero_grad()

d9d/pipelining/training/scheduler.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from torch.distributed import DeviceMesh
+
+from d9d.core.protocol import LRSchedulerProtocol
+
+
+class PipelinedLRScheduler(LRSchedulerProtocol):
+    """
+    Wrapper that manages multiple LR schedulers for a pipeline parallel rank.
+
+    Similar to `PipelinedOptimizer`, this aggregates schedulers corresponding to
+    multiple model stages hosted on the current rank.
+    """
+
+    def __init__(self, mesh_pp: DeviceMesh, schedulers: list[LRSchedulerProtocol]):
+        self._mesh_pp = mesh_pp
+        self._schedulers = schedulers
+
+    def state_dict(self) -> dict[str, Any]:
+        pp_rank = self._mesh_pp.get_local_rank()
+        return {
+            f"pp_{pp_rank}_stage_{i}": scheduler.state_dict()
+            for i, scheduler in enumerate(self._schedulers)
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        pp_rank = self._mesh_pp.get_local_rank()
+        for i, scheduler in enumerate(self._schedulers):
+            scheduler.load_state_dict(state_dict[f"pp_{pp_rank}_stage_{i}"])
+
+    def step(self) -> None:
+        for scheduler in self._schedulers:
+            scheduler.step()

d9d/tracker/__init__.py

@@ -0,0 +1,14 @@
+"""
+Package providing a unified interface for experiment tracking and logging.
+"""
+
+from .base import BaseTracker, BaseTrackerRun, RunConfig
+from .factory import AnyTrackerConfig, tracker_from_config
+
+__all__ = [
+    "AnyTrackerConfig",
+    "BaseTracker",
+    "BaseTrackerRun",
+    "RunConfig",
+    "tracker_from_config"
+]

d9d/tracker/base.py

@@ -0,0 +1,124 @@
+import abc
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any, Generic, Self, TypeVar
+
+import torch
+from pydantic import BaseModel, Field
+from torch.distributed.checkpoint.stateful import Stateful
+
+
+class BaseTrackerRun(abc.ABC):
+    """
+    Abstract base class representing an active tracking session (run).
+
+    This object is responsible for the actual logging of metrics, parameters,
+    during train or inference run.
+    """
+
+    @abc.abstractmethod
+    def set_step(self, step: int):
+        """
+        Updates the global step counter for subsequent logs.
+
+        Args:
+            step: The current step index (e.g., iteration number).
+        """
+        ...
+
+    @abc.abstractmethod
+    def set_context(self, context: dict[str, str]):
+        """
+        Sets a persistent context dictionary for subsequent logs.
+
+        These context values (tags) will be attached to every metric logged
+        until changed.
+
+        Args:
+            context: A dictionary of tag names and values.
+        """
+        ...
+
+    @abc.abstractmethod
+    def scalar(self, name: str, value: float, context: dict[str, str] | None = None):
+        """
+        Logs a scalar value.
+
+        Args:
+            name: The name of the metric.
+            value: The scalar value to log.
+            context: Optional ephemeral context specific to this metric event.
+                Merged with global context if present.
+        """
+        ...
+
+    @abc.abstractmethod
+    def bins(self, name: str, values: torch.Tensor, context: dict[str, str] | None = None):
+        """
+        Logs a distribution/histogram of values.
+
+        Args:
+            name: The name of the metric.
+            values: A tensor containing the population of values to bin.
+            context: Optional ephemeral context specific to this metric event.
+                Merged with global context if present.
+        """
+        ...
+
+
+class RunConfig(BaseModel):
+    """
+    Configuration for initializing a specific logged run.
+
+    Attributes:
+        name: The display name of the experiment.
+        description: An optional description of the experiment.
+        hparams: A dictionary of hyperparameters to log at the start of the run.
+    """
+
+    name: str
+    description: str | None
+    hparams: dict[str, Any] = Field(default_factory=dict)
+
+
+TConfig = TypeVar("TConfig", bound=BaseModel)
+
+
+class BaseTracker(abc.ABC, Stateful, Generic[TConfig]):
+    """
+    Abstract base class for a tracker backend factory.
+
+    This class manages the lifecycle of runs and integration with the
+    distributed checkpointing system to ensure experiment continuity
+    (e.g., resuming the same run hash after a restart).
+    """
+
+    @contextmanager
+    @abc.abstractmethod
+    def open(self, properties: RunConfig) -> Generator[BaseTrackerRun, None, None]:
+        """
+        Context manager that initiates and manages an experiment run.
+
+        Args:
+            properties: Configuration metadata for the run.
+
+        Yields:
+            An active BaseTrackerRun instance for logging metrics.
+        """
+
+        ...
+
+    @classmethod
+    @abc.abstractmethod
+    def from_config(cls, config: TConfig) -> Self:
+        """
+        Factory method to create a tracker instance from a configuration object.
+
+        Args:
+            config: The backend-specific configuration object.
+
+        Returns:
+            An initialized instance of the tracker.
+        """
+
+        ...

d9d/tracker/factory.py

@@ -0,0 +1,57 @@
+import dataclasses
+from typing import Annotated
+
+from pydantic import Field
+
+from .base import BaseTracker
+from .provider.aim.config import AimConfig
+from .provider.null import NullTracker, NullTrackerConfig
+
+AnyTrackerConfig = Annotated[AimConfig | NullTrackerConfig, Field(discriminator="provider")]
+
+
+@dataclasses.dataclass
+class _TrackerImportFailed:
+    dependency: str
+    exception: ImportError
+
+
+_MAP: dict[type[AnyTrackerConfig], type[BaseTracker] | _TrackerImportFailed] = {
+    NullTrackerConfig: NullTracker
+}
+
+try:
+    from .provider.aim.tracker import AimTracker
+
+    _MAP[AimConfig] = AimTracker
+except ImportError as e:
+    _MAP[AimConfig] = _TrackerImportFailed(dependency="aim", exception=e)
+
+
+def tracker_from_config(config: AnyTrackerConfig) -> BaseTracker:
+    """
+    Instantiates a specific tracker implementation based on the configuration.
+
+    Based on the 'provider' field in the config, this function selects the
+    appropriate backend (e.g., Aim, Null). It handles checking for missing
+    dependencies for optional backends.
+
+    Args:
+        config: A specific tracker configuration object.
+
+    Returns:
+        An initialized BaseTracker instance.
+
+    Raises:
+        ImportError: If the dependencies for the requested provider are not installed.
+    """
+
+    tracker_type = _MAP[type(config)]
+
+    if isinstance(tracker_type, _TrackerImportFailed):
+        raise ImportError(
+            f"The tracker configuration {config.provider} could not be loaded - "
+            f"ensure these dependencies are installed: {tracker_type.dependency}"
+        ) from tracker_type.exception
+
+    return tracker_type.from_config(config)

d9d/tracker/provider/aim/config.py

@@ -0,0 +1,23 @@
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class AimConfig(BaseModel):
+    """
+    Configuration for the Aim tracker backend.
+
+    Attributes:
+        provider: Discriminator field, must be 'aim'.
+        repo: Path to the Aim repository directory or URL.
+        log_system_params: Whether to log system resource usage (CPU/GPU/Memory).
+        capture_terminal_logs: Whether to capture stdout/stderr.
+        system_tracking_interval: Interval in seconds for system monitoring.
+    """
+
+    provider: Literal["aim"] = "aim"
+
+    repo: str
+    log_system_params: bool = True
+    capture_terminal_logs: bool = True
+    system_tracking_interval: int = 10