d9d 0.1.0__py3-none-any.whl

d9d/loop/component/stepper.py

@@ -0,0 +1,52 @@
+from typing import Any
+
+from torch.distributed.checkpoint.stateful import Stateful
+
+from d9d.loop.config import StepActionPeriod, StepActionSpecial
+
+
+class Stepper(Stateful):
+    def __init__(self, initial_step: int, total_steps: int):
+        self._current_step = initial_step
+        self._total_steps = total_steps
+
+    def step(self):
+        self._current_step += 1
+
+    @property
+    def current_step(self) -> int:
+        return self._current_step
+
+    @property
+    def total_steps(self) -> int:
+        return self._total_steps
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "current_step": self._current_step,
+            "total_steps": self._total_steps
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        if state_dict["total_steps"] != self._total_steps:
+            raise ValueError(f'Step count differs: saved {state_dict["total_steps"]}, '
+                             f'current {self._total_steps}. Perhaps project configuration changed?')
+
+        self._current_step = state_dict["current_step"]
+
+    def should_do_action(self, action: StepActionPeriod, enable_on_last_step_if_periodic: bool = False) -> bool:
+        match action:
+            case StepActionSpecial.disable:
+                return False
+            case StepActionSpecial.last_step:
+                return self._current_step == self._total_steps
+            case int():
+                if action <= 0:
+                    raise ValueError()
+
+                will_do_periodic = self._current_step % action == 0
+                will_do_last = enable_on_last_step_if_periodic and self._current_step == self._total_steps
+
+                return will_do_periodic or will_do_last
+            case _:
+                raise ValueError("Invalid step configuration")

d9d/loop/component/timeout_manager.py

@@ -0,0 +1,54 @@
+from enum import StrEnum
+
+from d9d.core.dist_context import DistributedContext
+from d9d.loop.config.config import TimeoutConfig
+
+
+class TimeoutState(StrEnum):
+    none = "none"
+    set_initial = "set_initial"
+    set_regular = "set_regular"
+
+
+class TimeoutManager:
+    """
+    Manages the dynamic adjustment of distributed timeouts during the job loop.
+
+    This manager handles the transition from initialization timeouts (which may need
+    to be longer due to JIT compilation, caching, or startup overhead) to regular
+    step execution timeouts.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            config: TimeoutConfig
+    ):
+        """
+        Constructs the timeout manager.
+
+        Args:
+            dist_context: The distributed context where timeouts are applied.
+            config: Configuration containing initialization and step timeout values.
+        """
+
+        self._dist_context = dist_context
+        self._config = config
+        self._state = TimeoutState.none
+
+    def step(self):
+        """
+        Updates the distributed backend timeout based on the current phase.
+        """
+
+        match self._state:
+            case TimeoutState.none:
+                self._dist_context.set_timeout(self._config.init_timeout)
+                self._state = TimeoutState.set_initial
+            case TimeoutState.set_initial:
+                self._dist_context.set_timeout(self._config.step_timeout)
+                self._state = TimeoutState.set_regular
+            case TimeoutState.set_regular:
+                pass  # do nothing
+            case _:
+                raise ValueError("Unknown timeout state")

d9d/loop/component/train_task_operator.py

@@ -0,0 +1,152 @@
+import dataclasses
+
+import torch
+
+from d9d.core.dist_context import DistributedContext
+from d9d.core.types import PyTree
+from d9d.internals.pipeline_state import PipelineStateHandler
+from d9d.loop.control import BuildForwardInputsContext, BuildForwardInputsResult, TrainTask, UpdateMetricsContext
+from d9d.metric.impl import ComposeMetric
+from d9d.pipelining.factory.factory import PipelineScheduleInfo
+
+from .loss_computer import STATE_LOSS, STATE_LOSS_WEIGHT, LossComputer
+from .model_stage_factory import TrackedModules
+
+
+@dataclasses.dataclass(kw_only=True)
+class ForwardResult:
+    """
+    Encapsulates the scalar results of a forward pass step.
+
+    Attributes:
+        loss: The computed loss tensor.
+        loss_weight: The weight associated with this loss (usually batch size).
+    """
+
+    loss: torch.Tensor
+    loss_weight: torch.Tensor
+
+
+class TrainTaskOperator:
+    """
+    Orchestrates the execution of the forward and backward passes for a specific training task.
+
+    This class abstracts the difference between standard execution
+    and pipeline-parallel execution. It manages input construction, schedule execution,
+    loss computation, and metric updates within the lifecycle of a single step.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            task: TrainTask,
+            pp_schedule: PipelineScheduleInfo | None,
+            tracked_modules: TrackedModules,
+            loss_computer: LossComputer,
+            pipeline_state: PipelineStateHandler,
+            metrics: ComposeMetric
+    ):
+        """
+        Constructs the TrainTaskOperator.
+
+        Args:
+            dist_context: The distributed context.
+            task: The user-defined training task logic.
+            pp_schedule: Information about the pipeline schedule.
+            tracked_modules: The model modules being trained.
+            loss_computer: Component responsible for calculating loss from outputs.
+            pipeline_state: Handler for transient state storage during the step.
+            metrics: Metric collection to update after the pass.
+        """
+
+        self._dist_context = dist_context
+        self._task = task
+        self._pp_schedule = pp_schedule
+        self._tracked_modules = tracked_modules
+        self._loss_computer = loss_computer
+        self._pipeline_state = pipeline_state
+        self._metrics = metrics
+
+    def _forward_backward_pipelining(self, model_inputs: BuildForwardInputsResult):
+        if self._pp_schedule is None:
+            raise ValueError("Cannot run pipelined pass if pipelining is disabled")
+
+        self._pp_schedule.schedule.configure_buffers(
+            inputs=model_inputs.inputs,
+            kwargs=model_inputs.kwargs,
+            sharding_spec=model_inputs.pipeline_sharding_spec
+        )
+        self._pp_schedule.schedule.step(
+            inputs=model_inputs.inputs,
+            kwargs=model_inputs.kwargs
+        )
+
+    def _forward_backward_regular(self, model_inputs: BuildForwardInputsResult):
+        pipeline_outputs = self._tracked_modules(
+            **model_inputs.inputs,
+            **model_inputs.kwargs
+        )
+        loss = self._loss_computer.compute_loss_mul_weight(
+            pipeline_outputs=pipeline_outputs,
+            microbatch_idx=None
+        )
+        # free to avoid bwd peaking memory
+        del pipeline_outputs
+        loss.backward()
+
+    def forward_backward(self, batch: PyTree) -> ForwardResult | None:
+        """
+        Executes the forward and backward passes for a single batch.
+
+        This method handles:
+
+        1. Context preparation and input building via the `TrainTask`.
+        2. Execution via Pipeline Parallel schedule or standard Forward/Backward.
+        3. Metric updates based on the results.
+        4. Reliable cleanup of the pipeline state.
+
+        Args:
+            batch: The input batch data.
+
+        Returns:
+            A `ForwardResult` containing the loss and weight if this rank is responsible
+            for loss calculation (e.g., the last pipeline stage or in standard DP).
+            Returns `None` if this rank is an intermediate pipeline stage that does
+            not compute loss.
+        """
+
+        try:
+            # Do forward and backward pass
+            model_inputs = self._task.build_forward_inputs(
+                BuildForwardInputsContext(
+                    batch=batch,
+                    state=self._pipeline_state.global_state()
+                )
+            )
+
+            if self._dist_context.mesh_params.has_pipeline_parallel:
+                self._forward_backward_pipelining(model_inputs)
+            else:
+                self._forward_backward_regular(model_inputs)
+
+            # Update metrics if possible
+
+            pipeline_state = self._pipeline_state.global_state()
+
+            if (
+                    self._dist_context.mesh_params.has_pipeline_parallel and
+                    self._pp_schedule is not None and
+                    not self._pp_schedule.has_last_stage
+            ):
+                return None
+
+            self._task.update_metrics(UpdateMetricsContext(
+                state=pipeline_state,
+                metrics=self._metrics.children
+            ))
+            return ForwardResult(
+                loss=pipeline_state[STATE_LOSS],
+                loss_weight=pipeline_state[STATE_LOSS_WEIGHT]
+            )
+        finally:
+            self._pipeline_state.reset()

d9d/loop/config/__init__.py

@@ -0,0 +1,36 @@
+from .config import (
+    BatchingConfig,
+    CheckpointingConfig,
+    DataLoadingConfig,
+    DeterminismConfig,
+    GarbageCollectionConfig,
+    GradientClippingConfig,
+    GradientManagerConfig,
+    InferenceConfig,
+    JobLoggerConfig,
+    ModelStageFactoryConfig,
+    PipeliningConfig,
+    ProfilingConfig,
+    TimeoutConfig,
+    TrainerConfig,
+)
+from .types import StepActionPeriod, StepActionSpecial
+
+__all__ = [
+    "BatchingConfig",
+    "CheckpointingConfig",
+    "DataLoadingConfig",
+    "DeterminismConfig",
+    "GarbageCollectionConfig",
+    "GradientClippingConfig",
+    "GradientManagerConfig",
+    "InferenceConfig",
+    "JobLoggerConfig",
+    "ModelStageFactoryConfig",
+    "PipeliningConfig",
+    "ProfilingConfig",
+    "StepActionPeriod",
+    "StepActionSpecial",
+    "TimeoutConfig",
+    "TrainerConfig"
+]

d9d/loop/config/config.py

@@ -0,0 +1,225 @@
+from pathlib import Path
+
+from pydantic import BaseModel
+
+from d9d.pipelining.factory import AnyPipelineScheduleConfig
+from d9d.tracker import AnyTrackerConfig, RunConfig
+
+from .types import StepActionPeriod
+
+
+class BatchingConfig(BaseModel):
+    """
+    Configuration for batch sizing logic.
+
+    Attributes:
+        global_batch_size: The total effective batch size across all distributed
+            replicas and gradient accumulation steps.
+        microbatch_size: The distinct batch size fed into the model during a single
+            forward pass on a single device.
+    """
+    global_batch_size: int
+    microbatch_size: int
+
+
+class DeterminismConfig(BaseModel):
+    """
+    Configuration for reproducibility and random number generation.
+
+    Attributes:
+        base_seed: The base integer seed used to initialize random number
+            generators (Python, NumPy, PyTorch) across all ranks.
+    """
+    base_seed: int
+
+
+class PipeliningConfig(BaseModel):
+    """
+    Configuration for pipeline parallelism orchestration.
+
+    Attributes:
+        schedule: The specific scheduling strategy configuration used to manage pipeline execution.
+    """
+    schedule: AnyPipelineScheduleConfig
+
+
+class GarbageCollectionConfig(BaseModel):
+    """
+    Configuration for manual Python garbage collection control.
+
+    Attributes:
+        period_steps: How frequently to manually trigger the Python garbage collector.
+    """
+    period_steps: StepActionPeriod
+
+
+class DataLoadingConfig(BaseModel):
+    """
+    Configuration for PyTorch DataLoaders.
+
+    Attributes:
+        num_workers: The number of subprocesses to use for data loading.
+        pin_memory: Whether to copy tensors into CUDA pinned memory before returning them.
+        persistent_workers: If True, the data loader will not shutdown the worker processes
+            after a dataset has been consumed once.
+    """
+    num_workers: int
+    pin_memory: bool
+    persistent_workers: bool
+
+
+class CheckpointingConfig(BaseModel):
+    """
+    Configuration for saving model snapshots.
+
+    Attributes:
+        save_dir: The root directory where checkpoints will be stored.
+        period_steps: How frequently to save a checkpoint.
+        num_to_keep: The maximum number of recent checkpoints to retain. If None,
+            all checkpoints are kept.
+    """
+    save_dir: Path
+    period_steps: StepActionPeriod
+    num_to_keep: int | None
+
+
+class ModelStageFactoryConfig(BaseModel):
+    """
+    Configuration for initializing model weights.
+
+    Attributes:
+        source_checkpoint: Path to an initial checkpoint to load into the model
+            before training starts. If None, random initialization is used.
+        checkpoint_only_trainable_parameters: If True, only parameters with
+            requires_grad=True will be saved in checkpoints. Useful for PEFT/LoRA.
+    """
+    source_checkpoint: Path | None
+    checkpoint_only_trainable_parameters: bool
+
+
+class GradientClippingConfig(BaseModel):
+    """
+    Configuration for gradient norm clipping.
+
+    Attributes:
+        max_norm: The maximum norm value for gradient clipping. If None,
+            no clipping is performed.
+        log_total_steps: Frequency at which to log the total gradient norm.
+    """
+    max_norm: float | None
+    log_total_steps: StepActionPeriod
+
+
+class ProfilingConfig(BaseModel):
+    """
+    Configuration for the PyTorch Profiler.
+
+    Attributes:
+        enabled: Whether to enable the profiler.
+        traces_dir: Directory where trace files will be saved.
+        period_steps: Total length of a profiling cycle (wait + warmup + active).
+        warmup_steps: Number of steps to ignore before recording to allow for warming-up.
+        active_steps: Number of steps to actively record traces.
+    """
+    enabled: bool
+
+    traces_dir: Path
+
+    period_steps: int
+    warmup_steps: int
+    active_steps: int
+
+
+class JobLoggerConfig(BaseModel):
+    """
+    Configuration for experiment tracking and logging.
+
+    Attributes:
+        period_steps: How frequently metrics are flushed to the logger.
+        tracker: Logic for the specific tracking backend (e.g., WandB, MLflow, stdout).
+    """
+    period_steps: StepActionPeriod
+    tracker: AnyTrackerConfig
+
+
+class GradientManagerConfig(BaseModel):
+    """
+    Configuration for gradient synchronization.
+
+    Attributes:
+        grad_dtype: The data type to use for storing the gradient. If None, follows the model's dtype.
+        bucket_size_mb: The size of gradient buckets in Megabytes for communication.
+    """
+    grad_dtype: str | None
+    bucket_size_mb: int
+
+
+class TimeoutConfig(BaseModel):
+    """
+    Configuration for distributed process group timeouts.
+
+    Attributes:
+        init_timeout: Timeout in seconds for initializing the process group.
+        step_timeout: Timeout in seconds for individual step communications.
+    """
+    init_timeout: int = 10000
+    step_timeout: int = 100
+
+
+class TrainerConfig(BaseModel):
+    """
+    Top-level configuration object defining a complete training job.
+
+    Attributes:
+        run: Meta-information about the run (name, ID, tags).
+        batching: Batch sizing strategy.
+        data_loading: DataLoader settings.
+        logging: Experiment tracking settings.
+        pipelining: Pipeline Parallelism schedule and settings. If None,
+            pipeline parallelism is disabled.
+        model_stage_factory: Model initialization and additional checkpointing logic.
+        determinism: Random seed settings.
+        gc: Garbage collection settings.
+        checkpointing: Checkpoint saving settings.
+        gradient_clipping: Gradient clipping settings.
+        profiling: Profiler settings.
+        gradient_manager: Gradient Synchronization Settings.
+        timeout: Distributed timeout settings.
+    """
+    run: RunConfig
+    batching: BatchingConfig
+    data_loading: DataLoadingConfig
+    logging: JobLoggerConfig
+    pipelining: PipeliningConfig | None
+    model_stage_factory: ModelStageFactoryConfig
+    determinism: DeterminismConfig
+    gc: GarbageCollectionConfig
+    checkpointing: CheckpointingConfig
+    gradient_clipping: GradientClippingConfig
+    profiling: ProfilingConfig | None
+    gradient_manager: GradientManagerConfig
+    timeout: TimeoutConfig
+
+
+class InferenceConfig(BaseModel):
+    """
+    Top-level configuration object defining an inference/evaluation job.
+
+    Attributes:
+        batching: Batch sizing strategy.
+        data_loading: DataLoader settings.
+        model_stage_factory: Model initialization logic.
+        determinism: Random seed settings.
+        gc: Garbage collection settings.
+        checkpointing: Checkpointing settings.
+        profiling: Profiler settings.
+        timeout: Distributed timeout settings.
+    """
+    batching: BatchingConfig
+    data_loading: DataLoadingConfig
+    model_stage_factory: ModelStageFactoryConfig
+    determinism: DeterminismConfig
+    gc: GarbageCollectionConfig
+    checkpointing: CheckpointingConfig
+    profiling: ProfilingConfig | None
+    timeout: TimeoutConfig

d9d/loop/config/types.py

@@ -0,0 +1,24 @@
+from enum import StrEnum
+
+
+class StepActionSpecial(StrEnum):
+    """
+    Special flag values for configuring periodic actions.
+
+    Attributes:
+        last_step: Indicates the action should occur exactly once at the
+            very end of the training run.
+        disable: Indicates the action should never occur.
+    """
+    last_step = "last_step"
+    disable = "disable"
+
+
+StepActionPeriod = int | StepActionSpecial
+"""
+Union type representing a configuration for periodic events.
+
+Values:
+    int: The period in steps (frequency) at which the event occurs.
+    StepActionSpecial: A special flag indicating end-of-run execution or disabling.
+"""

d9d/loop/control/__init__.py

@@ -0,0 +1,61 @@
+from .dataset_provider import DatasetProvider, InitializeDatasetContext, InitializeDatasetResult
+from .lr_scheduler_provider import InitializeLRSchedulerContext, LRSchedulerProvider
+from .model_provider import (
+    InitializeModelStageContext,
+    InitializeModelStageResult,
+    ModelProvider,
+    ParallelizeModelStageContext,
+    PrepareExportModelStageContext,
+    PrepareExportModelStageResult,
+)
+from .optimizer_provider import InitializeOptimizerStageContext, OptimizerProvider
+from .task import (
+    BaseTask,
+    BuildForwardInputsContext,
+    BuildForwardInputsResult,
+    ComputeLossContext,
+    ComputeLossResult,
+    CreateMetricsContext,
+    CreateMetricsResult,
+    FinalizeContext,
+    InferenceTask,
+    InferenceTaskProvider,
+    InferenceTaskProviderContext,
+    ProcessOutputsContext,
+    TrainTask,
+    TrainTaskProvider,
+    TrainTaskProviderContext,
+    UpdateMetricsContext,
+)
+
+__all__ = [
+    "BaseTask",
+    "BuildForwardInputsContext",
+    "BuildForwardInputsResult",
+    "ComputeLossContext",
+    "ComputeLossResult",
+    "CreateMetricsContext",
+    "CreateMetricsResult",
+    "DatasetProvider",
+    "FinalizeContext",
+    "InferenceTask",
+    "InferenceTaskProvider",
+    "InferenceTaskProviderContext",
+    "InitializeDatasetContext",
+    "InitializeDatasetResult",
+    "InitializeLRSchedulerContext",
+    "InitializeModelStageContext",
+    "InitializeModelStageResult",
+    "InitializeOptimizerStageContext",
+    "LRSchedulerProvider",
+    "ModelProvider",
+    "OptimizerProvider",
+    "ParallelizeModelStageContext",
+    "PrepareExportModelStageContext",
+    "PrepareExportModelStageResult",
+    "ProcessOutputsContext",
+    "TrainTask",
+    "TrainTaskProvider",
+    "TrainTaskProviderContext",
+    "UpdateMetricsContext"
+]

d9d/loop/control/dataset_provider.py

@@ -0,0 +1,58 @@
+import dataclasses
+import typing
+from typing import Protocol
+
+from torch.utils.data import Dataset
+
+from d9d.core.dist_context import DistributedContext
+from d9d.core.types import CollateFn
+
+if typing.TYPE_CHECKING:
+    from d9d.loop.component import BatchMaths
+
+
+@dataclasses.dataclass(kw_only=True)
+class InitializeDatasetContext:
+    """Context data required to initialize a dataset provider.
+
+    Attributes:
+        dist_context: The distributed context containing rank and world size information.
+        batch_maths: The batch maths component handling global batch size calculations.
+    """
+
+    dist_context: DistributedContext
+    batch_maths: "BatchMaths"
+
+
+@dataclasses.dataclass(kw_only=True)
+class InitializeDatasetResult:
+    """The result of initializing a dataset provider.
+
+    Attributes:
+        dataset: The instantiated PyTorch Dataset.
+        collator: The function used to collate individual samples into a batch.
+    """
+
+    dataset: Dataset
+    collator: CollateFn
+
+
+@typing.runtime_checkable
+class DatasetProvider(Protocol):
+    """Protocol that allows users to define how datasets are loaded and collated.
+
+    Users should subclass this to provide custom data loading logic.
+    """
+
+    def __call__(self, context: InitializeDatasetContext) -> InitializeDatasetResult:
+        """
+        Initializes the dataset components.
+
+        It is important that the user must shard the dataset manually, perhaps using `d9d.dataset.ShardedDataset`.
+
+        Args:
+            context: Context for this operation.
+
+        Returns:
+            Result of this operation.
+        """