d9d 0.1.0__py3-none-any.whl

d9d/loop/control/lr_scheduler_provider.py

@@ -0,0 +1,47 @@
+import abc
+import dataclasses
+import typing
+from typing import Protocol
+
+from torch.optim import Optimizer
+
+from d9d.core.dist_context import DistributedContext
+from d9d.core.protocol import LRSchedulerProtocol
+
+
+@dataclasses.dataclass(kw_only=True)
+class InitializeLRSchedulerContext:
+    """
+    Context data required to initialize an LR scheduler.
+
+    Attributes:
+        dist_context: The distributed context.
+        total_steps: The total number of training steps.
+        optimizer: The optimizer instance that the scheduler will control.
+    """
+
+    dist_context: DistributedContext
+    total_steps: int
+    optimizer: Optimizer
+
+
+@typing.runtime_checkable
+class LRSchedulerProvider(Protocol):
+    """
+    Protocol for defining how Learning Rate schedulers are created.
+    """
+
+    @abc.abstractmethod
+    def __call__(
+            self,
+            context: InitializeLRSchedulerContext
+    ) -> LRSchedulerProtocol:
+        """
+        Initializes the LR scheduler for a specific model pipeline stage.
+
+        Args:
+            context: Context for this operation.
+
+        Returns:
+            The instantiated LR scheduler adhering to the protocol.
+        """

d9d/loop/control/model_provider.py

@@ -0,0 +1,162 @@
+import abc
+import dataclasses
+from typing import Generic, TypeVar
+
+from torch import nn
+
+from d9d.core.dist_context import DistributedContext
+from d9d.core.types import ScalarTree
+from d9d.model_state.mapper import ModelStateMapper
+from d9d.pipelining.api import PipelineStageInfo
+
+
+@dataclasses.dataclass(kw_only=True)
+class InitializeModelStageContext:
+    """
+    Context data required for initializing a specific model pipeline stage.
+
+    Attributes:
+        dist_context: The distributed execution context.
+        stage: Metadata describing the current pipeline stage being initialized.
+    """
+
+    dist_context: DistributedContext
+    stage: PipelineStageInfo
+
+
+TModel = TypeVar("TModel", bound=nn.Module)
+
+
+@dataclasses.dataclass(kw_only=True)
+class InitializeModelStageResult(Generic[TModel]):
+    """
+    The result of initializing a model stage.
+
+    Attributes:
+        model: The PyTorch module.
+        state_mapper: The mapper defining how to load weights into this module.
+    """
+
+    model: TModel
+    state_mapper: ModelStateMapper
+
+
+@dataclasses.dataclass(kw_only=True)
+class ParallelizeModelStageContext(Generic[TModel]):
+    """
+    Context data required for horizontally parallelizing a model stage.
+
+    Attributes:
+        dist_context: The distributed execution context.
+        stage: Metadata describing the current pipeline stage.
+        model: The PyTorch module to be parallelized.
+    """
+
+    dist_context: DistributedContext
+    stage: PipelineStageInfo
+    model: TModel
+
+
+@dataclasses.dataclass(kw_only=True)
+class PrepareExportModelStageContext(Generic[TModel]):
+    """
+    Context data required for preparing a model stage for export.
+
+    Attributes:
+        dist_context: The distributed execution context.
+        model: The PyTorch module to be exported.
+    """
+
+    dist_context: DistributedContext
+    model: TModel
+
+
+@dataclasses.dataclass(kw_only=True)
+class PrepareExportModelStageResult:
+    """
+    The result of preparing a model stage for export.
+
+    Attributes:
+        state_mapper: The mapper defining how model parameters map to disk storage.
+    """
+
+    state_mapper: ModelStateMapper
+
+
+class ModelProvider(abc.ABC, Generic[TModel]):
+    """
+    Abstract interface for defining the lifecycle of a distributed model.
+
+    This provider handles initialization, parallelization (sharding/replication/etc), and export preparation
+    for models within the d9d framework.
+    """
+
+    @abc.abstractmethod
+    def initialize_model_stage(
+            self,
+            context: InitializeModelStageContext
+    ) -> InitializeModelStageResult[TModel]:
+        """
+        Initializes the model architecture for a specific pipeline stage.
+
+        This method is responsible for constructing the `nn.Module` for the requested stage.
+
+        Construction occurs within a meta-device context; therefore, weights
+        should not be loaded directly here. Instead, a `ModelStateMapper` must be returned
+        to define how weights from a checkpoint map to the newly created module parameters.
+
+        This allows for architecture modifications, such as injecting LoRA adapters,
+        provided that the returned mapper reflects the new structure.
+
+        Args:
+            context: Context for this operation.
+
+        Returns:
+            Result of this operation.
+        """
+
+        ...
+
+    @abc.abstractmethod
+    def parallelize_model_stage(
+            self,
+            context: ParallelizeModelStageContext[TModel]
+    ):
+        """
+        Converts the model parameters into distributed tensors (DTensors).
+
+        Implementations should modify the model in-place. This involves converting
+        standard parameters into DTensors by replicating or sharding them according
+        to the desired parallelism strategies.
+
+        Args:
+            context: Context for this operation.
+        """
+
+    @abc.abstractmethod
+    def prepare_export_model_stage(
+            self,
+            context: PrepareExportModelStageContext[TModel]
+    ) -> PrepareExportModelStageResult:
+        """
+        Prepares the state mapper required for saving the model to disk.
+
+        This methods defines how the current in-memory model structure maps back to the
+        serialized checkpoint format.
+
+        Args:
+            context: Context for this operation.
+
+        Returns:
+            Result of this operation.
+        """
+
+    def dump_hparams(self) -> ScalarTree:
+        """
+        Exports hyperparameters associated with this model for logging.
+
+        Returns:
+            A dictionary of hyperparameter names and values.
+        """
+
+        return {}

d9d/loop/control/optimizer_provider.py

@@ -0,0 +1,45 @@
+import abc
+import dataclasses
+import typing
+from typing import Protocol
+
+from torch import nn
+from torch.optim import Optimizer
+
+from d9d.core.dist_context import DistributedContext
+
+
+@dataclasses.dataclass(kw_only=True)
+class InitializeOptimizerStageContext:
+    """
+    Context data required to initialize an optimizer.
+
+    Attributes:
+        dist_context: The distributed context.
+        model: The model instance for which parameters will be optimized.
+    """
+
+    dist_context: DistributedContext
+    model: nn.Module
+
+
+@typing.runtime_checkable
+class OptimizerProvider(Protocol):
+    """
+    Protocol for defining how optimizers are created for model pipeline stages.
+    """
+
+    @abc.abstractmethod
+    def __call__(
+            self,
+            context: InitializeOptimizerStageContext
+    ) -> Optimizer:
+        """
+        Initializes the optimizer for a specific training stage.
+
+        Args:
+            context: Context for this operation.
+
+        Returns:
+            The instantiated PyTorch optimizer.
+        """

d9d/loop/control/task.py

@@ -0,0 +1,304 @@
+import abc
+import dataclasses
+import typing
+from collections.abc import Mapping
+from typing import Any, Protocol
+
+import torch
+from torch.distributed.checkpoint.stateful import Stateful
+
+from d9d.core.dist_context import DistributedContext
+from d9d.core.types import PyTree, ScalarTree
+from d9d.pipelining.api import PipelineShardingSpec
+
+if typing.TYPE_CHECKING:
+    from d9d.internals.pipeline_state import PipelineState
+    from d9d.loop.component import Stepper
+    from d9d.metric import Metric
+
+
+TBatch = typing.TypeVar("TBatch", bound=PyTree)
+
+
+@dataclasses.dataclass(kw_only=True)
+class BuildForwardInputsContext(typing.Generic[TBatch]):
+    """
+    Context data to prepare inputs for the model forward pass.
+
+    Attributes:
+        batch: The raw batch data loaded from the DataLoader object.
+        state: The current state of the pipeline. You can assign any data to this state object, and it will be
+            accessible during this pipeline step (e.g. when computing loss)
+    """
+
+    batch: TBatch
+    state: "PipelineState"
+
+
+@dataclasses.dataclass(kw_only=True)
+class BuildForwardInputsResult:
+    """
+    The result of processing the raw batch into model inputs.
+
+    Attributes:
+        inputs: A dictionary of inputs that are passed to model pipeline as input data
+            (first stage only if using pipeline parallelism).
+        kwargs: A dictionary of keyword arguments passed to each pipeline stage.
+        pipeline_sharding_spec: A specification defining how inputs and kwargs should be split
+            into micro-batches for pipeline parallelism. If None, the framework assumes
+            standard behavior where all the non-scalar Tensors and lists are split by 0 dimension.
+    """
+
+    inputs: dict[str, torch.Tensor]
+    kwargs: dict[str, Any]
+    pipeline_sharding_spec: PipelineShardingSpec | None = None
+
+
+@dataclasses.dataclass(kw_only=True)
+class FinalizeContext:
+    """Context data provided when the task is being finalized."""
+
+
+class BaseTask(abc.ABC, Stateful, typing.Generic[TBatch]):
+    """Abstract base class representing a unit of work (Task) in the training/inference loop."""
+
+    @abc.abstractmethod
+    def build_forward_inputs(self, ctx: BuildForwardInputsContext[TBatch]) -> BuildForwardInputsResult:
+        """
+        Transforms raw data loaded from the DataLoader into arguments for the model.
+
+        Args:
+            ctx: Context object.
+
+        Returns:
+            Result object.
+        """
+
+        ...
+
+    def state_dict(self) -> dict[str, Any]:
+        """
+        Returns the state dictionary for checkpointing this task.
+
+        Returns:
+            A dictionary containing the task's state.
+        """
+
+        return {}
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """
+        Restores the task's state from the provided dictionary.
+
+        Args:
+            state_dict: The state dictionary to load.
+        """
+        # do nothing by default
+
+    def finalize(self, ctx: FinalizeContext) -> None:
+        """
+        Performs cleanup or final actions when the task execution finishes.
+
+        Args:
+             ctx: Context object.
+        """
+
+
+@dataclasses.dataclass(kw_only=True)
+class ComputeLossContext:
+    """
+    Context data provided to calculate the loss during training.
+
+    Attributes:
+        pipeline_results: The outputs returned by the model's forward pass.
+        state: The current state of the pipeline. You can assign any data to this state object, and it will be
+            accessible during this pipeline step (e.g. when calculating metrics)
+        stepper: Component tracking the current step.
+    """
+
+    pipeline_results: Mapping[str, torch.Tensor]
+    state: "PipelineState"
+    stepper: "Stepper"
+
+
+@dataclasses.dataclass(kw_only=True)
+class ComputeLossResult:
+    """
+    The result of the loss computation.
+
+    Attributes:
+        loss: The scalar tensor representing the loss to be backpropagated.
+        loss_weight: The weight to apply to the loss (for synchronizing gradients using weighted mean).
+            None for 1.0.
+    """
+
+    loss: torch.Tensor
+    loss_weight: torch.Tensor | None
+
+
+@dataclasses.dataclass(kw_only=True)
+class CreateMetricsContext:
+    """Context data provided to initialize metrics."""
+
+
+@dataclasses.dataclass(kw_only=True)
+class CreateMetricsResult:
+    """
+    Result of metric initialization.
+
+    Attributes:
+        metrics: A dictionary mapping metric names to Metric instances.
+    """
+
+    metrics: dict[str, "Metric"]
+
+
+@dataclasses.dataclass(kw_only=True)
+class UpdateMetricsContext:
+    """
+    Context data provided to update metrics after a step.
+
+    Attributes:
+        state: The current state of the pipeline.
+        metrics: The dictionary of metrics to be updated.
+    """
+
+    state: "PipelineState"
+    metrics: Mapping[str, "Metric"]
+
+
+class TrainTask(BaseTask, abc.ABC, typing.Generic[TBatch]):
+    """Abstract base class for defining training-specific logic."""
+
+    @abc.abstractmethod
+    def compute_loss(self, ctx: ComputeLossContext) -> ComputeLossResult:
+        """
+        Calculates the loss based on model outputs.
+
+        Args:
+            ctx: Context object.
+
+        Returns:
+            Result object.
+        """
+
+        ...
+
+    def create_metrics(self, ctx: CreateMetricsContext) -> CreateMetricsResult:
+        """
+        Initializes metrics to be tracked during training.
+
+        Args:
+             ctx: Context object.
+
+        Returns:
+            Result object.
+        """
+
+        return CreateMetricsResult(metrics={})
+
+    def update_metrics(self, ctx: UpdateMetricsContext):
+        """
+        Updates the state of the metrics at the end of training step.
+
+        Args:
+            ctx: Context object.
+        """
+
+    def dump_hparams(self) -> ScalarTree:
+        """
+        Exports hyperparameters associated with this task for logging.
+
+        Returns:
+            A dictionary of hyperparameter names and values.
+        """
+
+        return {}
+
+
+@dataclasses.dataclass(kw_only=True)
+class TrainTaskProviderContext:
+    """
+    Context data provided to the factory creating a TrainTask.
+
+    Attributes:
+        dist_context: Information about the distributed environment.
+    """
+
+    dist_context: DistributedContext
+
+
+@typing.runtime_checkable
+class TrainTaskProvider(Protocol):
+    """Protocol that creates a TrainTask instance."""
+
+    def __call__(self, ctx: TrainTaskProviderContext) -> TrainTask:
+        """
+        Creates and returns a new TrainTask.
+
+        Args:
+            ctx: Context object.
+
+        Returns:
+            An instantiated TrainTask.
+        """
+
+        ...
+
+
+@dataclasses.dataclass(kw_only=True)
+class ProcessOutputsContext:
+    """
+    Context data provided to process outputs during inference.
+
+    Attributes:
+        outputs: The outputs returned by the model's forward pass.
+        state: The current state of the pipeline.
+    """
+
+    outputs: dict[str, torch.Tensor]
+    state: "PipelineState"
+
+
+class InferenceTask(BaseTask, abc.ABC, typing.Generic[TBatch]):
+    """Abstract base class for defining inference-specific logic."""
+
+    @abc.abstractmethod
+    def process_outputs(self, ctx: ProcessOutputsContext):
+        """
+        Processes the model outputs (e.g. saving to disk, decoding tokens).
+
+        Args:
+            ctx: Context containing the model outputs and pipeline state.
+        """
+
+        ...
+
+
+@dataclasses.dataclass(kw_only=True)
+class InferenceTaskProviderContext:
+    """
+    Context data provided to the factory creating an InferenceTask.
+
+    Attributes:
+        dist_context: Information about the distributed environment.
+    """
+
+    dist_context: DistributedContext
+
+
+@typing.runtime_checkable
+class InferenceTaskProvider(Protocol):
+    """Protocol for a callable that creates an InferenceTask instance."""
+
+    def __call__(self, ctx: InferenceTaskProviderContext) -> InferenceTask:
+        """
+        Creates and returns a new InferenceTask.
+
+        Args:
+            ctx: Context providing distributed environment information.
+
+        Returns:
+            An instantiated InferenceTask.
+        """
+        ...

d9d/loop/run/__init__.py

@@ -0,0 +1,6 @@
+from .train import Trainer, TrainingConfigurator
+
+__all__ = [
+    "Trainer",
+    "TrainingConfigurator"
+]