d9d 0.1.0__py3-none-any.whl

d9d/loop/auto/auto_optimizer.py

@@ -0,0 +1,196 @@
+import abc
+from abc import ABC
+from collections.abc import Iterable
+from typing import Annotated, Literal
+
+import torch
+from pydantic import BaseModel, Field
+from torch import nn
+from torch.optim import SGD, Adam, AdamW, Optimizer
+
+from d9d.loop.control import InitializeOptimizerStageContext, OptimizerProvider
+from d9d.optim.stochastic import StochasticAdamW
+
+
+class BaseAutoOptimizerConfig(BaseModel, ABC):
+    """
+    Abstract base class for optimizer configurations.
+    """
+
+    @abc.abstractmethod
+    def build(self, params: Iterable[nn.Parameter]) -> Optimizer:
+        """
+        Creates the PyTorch optimizer instance.
+
+        Args:
+            params: An iterable of model parameters to optimize.
+
+        Returns:
+            The instantiated optimizer.
+        """
+        ...
+
+
+class StochasticAdamWOptimizerConfig(BaseAutoOptimizerConfig):
+    """
+    Configuration for the Stochastic AdamW optimizer.
+
+    Attributes:
+        name: Discriminator tag.
+        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.
+        state_dtype: Data Type to use for the optimizer states.
+    """
+    name: Literal["stochastic_adamw"] = "stochastic_adamw"
+
+    lr: float
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 1e-2
+    state_dtype: str
+
+    def build(self, params: Iterable[nn.Parameter]) -> Optimizer:
+        """Builds StochasticAdamW with the configured parameters."""
+        return StochasticAdamW(
+            params=params,
+            lr=self.lr,
+            betas=self.betas,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            state_dtype=getattr(torch, self.state_dtype)
+        )
+
+
+class AdamWOptimizerConfig(BaseAutoOptimizerConfig):
+    """
+    Configuration for the PyTorch AdamW optimizer.
+
+    Attributes:
+        name: Discriminator tag.
+        lr: The learning rate.
+        betas: Coefficients for computing running averages of gradient and its square.
+        eps: Term added to the denominator to improve numerical stability.
+        weight_decay: Weight decay coefficient.
+        amsgrad: Whether to use the AMSGrad variant.
+        maximize: Whether to maximize the params based on the objective (as opposed to minimizing).
+    """
+    name: Literal["adamw"] = "adamw"
+
+    lr: float
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 1e-2
+    amsgrad: bool = False
+    maximize: bool = False
+
+    def build(self, params: Iterable[nn.Parameter]) -> Optimizer:
+        """Builds fused AdamW with the configured parameters."""
+        return AdamW(
+            params=params,
+            lr=self.lr,
+            betas=self.betas,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            amsgrad=self.amsgrad,
+            maximize=self.maximize,
+            fused=True
+        )
+
+
+class AdamOptimizerConfig(BaseAutoOptimizerConfig):
+    """
+    Configuration for the PyTorch Adam optimizer.
+
+    Attributes:
+        name: Discriminator tag.
+        lr: The learning rate.
+        betas: Coefficients for computing running averages of gradient and its square.
+        eps: Term added to the denominator to improve numerical stability.
+        weight_decay: Weight decay coefficient.
+        decoupled_weight_decay: Whether to apply decoupled weight decay.
+        amsgrad: Whether to use the AMSGrad variant.
+        maximize: Whether to maximize the params based on the objective.
+    """
+    name: Literal["adam"] = "adam"
+
+    lr: float
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 1e-2
+    decoupled_weight_decay: bool = False
+    amsgrad: bool = False
+    maximize: bool = False
+
+    def build(self, params: Iterable[nn.Parameter]) -> Optimizer:
+        """Builds fused Adam with the configured parameters."""
+        return Adam(
+            params=params,
+            lr=self.lr,
+            betas=self.betas,
+            eps=self.eps,
+            weight_decay=self.weight_decay,
+            decoupled_weight_decay=self.decoupled_weight_decay,
+            amsgrad=self.amsgrad,
+            maximize=self.maximize,
+            fused=True
+        )
+
+
+class SGDOptimizerConfig(BaseAutoOptimizerConfig):
+    """
+    Configuration for the PyTorch SGD optimizer.
+
+    Attributes:
+        name: Discriminator tag.
+        lr: The learning rate.
+        momentum: Momentum factor.
+        dampening: Dampening for momentum.
+        weight_decay: Weight decay (L2 penalty).
+        nesterov: Enables Nesterov momentum.
+        maximize: Whether to maximize the params based on the objective.
+    """
+    name: Literal["sgd"] = "sgd"
+
+    lr: float
+    momentum: float = 0
+    dampening: float = 0
+    weight_decay: float = 0
+    nesterov: bool = False
+    maximize: bool = False
+
+    def build(self, params: Iterable[nn.Parameter]) -> Optimizer:
+        """Builds fused SGD with the configured parameters."""
+        return SGD(
+            params,
+            lr=self.lr,
+            momentum=self.momentum,
+            dampening=self.dampening,
+            weight_decay=self.weight_decay,
+            nesterov=self.nesterov,
+            maximize=self.maximize,
+            fused=True
+        )
+
+
+AutoOptimizerConfig = Annotated[
+    StochasticAdamWOptimizerConfig |
+    AdamWOptimizerConfig |
+    AdamOptimizerConfig |
+    SGDOptimizerConfig,
+    Field(discriminator="name")
+]
+
+
+class AutoOptimizerProvider(OptimizerProvider):
+    """
+    OptimizerProvider that builds a PyTorch optimizer based on a configuration object.
+    """
+
+    def __init__(self, config: AutoOptimizerConfig):
+        """Constructs the provider with the given configuration."""
+        self._config = config
+
+    def __call__(self, context: InitializeOptimizerStageContext) -> Optimizer:
+        return self._config.build(context.model.parameters())

d9d/loop/component/__init__.py

@@ -0,0 +1,35 @@
+from .batch_maths import BatchMaths
+from .checkpointer import StateCheckpointer
+from .data_loader_factory import DataLoaderFactory
+from .garbage_collector import ManualGarbageCollector
+from .gradient_clipper import GradientClipper
+from .gradient_manager import GradientManager
+from .job_logger import JobLogger
+from .job_profiler import JobProfiler
+from .loss_computer import LossComputer
+from .model_stage_exporter import ModelStageExporter
+from .model_stage_factory import ModelStageFactory, TrackedModules
+from .optimizer_factory import OptimizerFactory
+from .stepper import Stepper
+from .timeout_manager import TimeoutManager
+from .train_task_operator import ForwardResult, TrainTaskOperator
+
+__all__ = [
+    "BatchMaths",
+    "DataLoaderFactory",
+    "ForwardResult",
+    "GradientClipper",
+    "GradientManager",
+    "JobLogger",
+    "JobProfiler",
+    "LossComputer",
+    "ManualGarbageCollector",
+    "ModelStageExporter",
+    "ModelStageFactory",
+    "OptimizerFactory",
+    "StateCheckpointer",
+    "Stepper",
+    "TimeoutManager",
+    "TrackedModules",
+    "TrainTaskOperator"
+]

d9d/loop/component/batch_maths.py

@@ -0,0 +1,106 @@
+from d9d.core.dist_context import BATCH_DOMAIN, DistributedContext
+from d9d.loop.config import BatchingConfig, PipeliningConfig
+
+
+class BatchMaths:
+    """
+    Calculates derived batching dimensions and iteration counts for distributed training loops.
+
+    This class bridges the gap between global configuration (Global Batch Size) and
+    local execution constraints (Microbatch Size, Data Parallel World Size).
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            config_batching: BatchingConfig,
+            config_pipelining: PipeliningConfig | None
+    ):
+        """
+        Constructs the batch mathematics calculator.
+
+        Validates that the Global Batch Size is perfectly divisible by the
+        effective parallel microbatch capacity (DP size * Microbatch size).
+
+        Args:
+            dist_context: The distributed context containing mesh layout information.
+            config_batching: Configuration detailing batch sizes.
+            config_pipelining: Optional configuration for pipeline parallelism capabilities.
+
+        Raises:
+            ValueError: If global batch size is not divisible by the product of
+                Data Parallel size and Microbatch size.
+        """
+
+        self._dist_context = dist_context
+        self._config_batching = config_batching
+        self._config_pipelining = config_pipelining
+
+        global_batch = self._config_batching.global_batch_size
+        dp_size = self._dist_context.mesh_for(BATCH_DOMAIN)["dp"].size()
+        microbatch_size = self._config_batching.microbatch_size
+
+        global_microbatch = dp_size * microbatch_size
+
+        if global_batch % global_microbatch != 0:
+            raise ValueError("Global Batch Size must be divisible by (Data Parallel cardinality * Microbatch Size)")
+
+        self._global_microbatch_size = global_microbatch
+
+    @property
+    def global_batch_size(self) -> int:
+        """
+        Returns the global batch size across the world.
+        """
+
+        return self._config_batching.global_batch_size
+
+    @property
+    def num_microbatches_pipelining(self) -> int:
+        """
+        Returns the number of microbatches handled by the pipeline scheduler per step.
+
+        If pipeline parallelism is enabled, this is the total number of microbatches
+        processed to form one global batch. If disabled, this returns 1.
+        """
+
+        if not self._dist_context.mesh_params.has_pipeline_parallel:
+            return 1
+
+        return self._config_batching.global_batch_size // self._global_microbatch_size
+
+    @property
+    def num_microbatches_gradient_accumulation(self) -> int:
+        """
+        Returns the number of gradient accumulation iterations for non-pipelined training.
+
+        If pipeline parallelism is enabled, this returns 1 (as accumulation is handled
+        internally by the pipeline schedule). If disabled, this is the number of
+        forward/backward passes the training loop must execute before an optimizer step.
+        """
+
+        if self._dist_context.mesh_params.has_pipeline_parallel:
+            return 1
+
+        return self._config_batching.global_batch_size // self._global_microbatch_size
+
+    @property
+    def data_loader_batch_size(self) -> int:
+        """
+        Returns the quantity of samples this local rank needs to fetch for one optimizer step.
+
+        This is calculated as `microbatch_size * total_microbatches_per_step`.
+        """
+
+        return self._config_batching.microbatch_size * self.num_microbatches_pipelining
+
+    @property
+    def num_backward_calls(self) -> int:
+        """
+        Returns the total number of backward passes executed per optimizer step.
+
+        This represents the total gradient accumulation factor, regardless of whether
+        it is handled by a pipeline schedule or a simple loop.
+        """
+
+        return self.num_microbatches_pipelining * self.num_microbatches_gradient_accumulation

d9d/loop/component/checkpointer.py

@@ -0,0 +1,172 @@
+import re
+import shutil
+from pathlib import Path
+
+import torch
+import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint.stateful import Stateful
+
+from d9d.core.dist_context import DistributedContext
+from d9d.loop.config import CheckpointingConfig
+
+from .garbage_collector import ManualGarbageCollector
+from .stepper import Stepper
+
+# TODO feat(max): async checkpointing may break everything up, but I guess we still have to support it
+
+_SAVE_RE = re.compile(r"^save-(\d+)$")
+
+
+def _save_iter_predicate(x: Path) -> int:
+    match = _SAVE_RE.fullmatch(x.stem)
+    if match is None:
+        raise ValueError("Malformed checkpoint name")
+    return int(match.group(1))
+
+
+class StateCheckpointer:
+    """
+    Manages the lifecycle of distributed training checkpoints.
+
+    This class handles saving and loading the training state (JobState object)
+    using PyTorch Distributed Checkpoint (DCP). It manages checkpoint versioning,
+    storage rotation (keeping only N latest), and synchronization across distributed ranks.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            stepper: Stepper,
+            config: CheckpointingConfig,
+            gc: ManualGarbageCollector,
+            run_name: str | None
+    ):
+        """
+        Constructs the StateCheckpoint object.
+
+        Args:
+            dist_context: The distributed context.
+            stepper: The training stepper tracking the current iteration/step.
+            config: Configuration object containing checkpointing parameters.
+            gc: Garbage collector for manual memory management during IO.
+            run_name: Optional specific run name to append to the save directory.
+        """
+        self._dist_context = dist_context
+        self._stepper = stepper
+        self._gc = gc
+
+        if run_name:
+            self._save_dir = config.save_dir / run_name
+        else:
+            self._save_dir = config.save_dir
+
+        self._config = config
+
+    def _free_memory(self):
+        self._gc.collect_forced()
+        torch.cuda.empty_cache()
+
+    def _get_sorted_checkpoint_dirs(self) -> list[Path]:
+        if not self._save_dir:
+            return []
+
+        if not self._save_dir.is_dir():
+            return []
+
+        checkpoint_dirs = [x for x in self._save_dir.iterdir() if x.is_dir() and _SAVE_RE.fullmatch(x.stem)]
+        checkpoint_dirs = sorted(checkpoint_dirs, key=_save_iter_predicate)
+        return checkpoint_dirs
+
+    def _next_checkpoint_id(self) -> Path:
+        next_name = f"save-{self._stepper.current_step}"
+        return self._save_dir / next_name
+
+    def _purge_old_checkpoints(self):
+        if not self._dist_context.is_main_process:
+            return
+        if not self._config.num_to_keep:
+            return
+
+        to_delete = self._get_sorted_checkpoint_dirs()[:-self._config.num_to_keep]
+
+        for delete_dir in to_delete:
+            self._dist_context.logger.info(f"Purging checkpoint {delete_dir}")
+            shutil.rmtree(delete_dir)
+
+    def _checkpoint(self, state: Stateful):
+        next_checkpoint_id = self._next_checkpoint_id()
+
+        self._dist_context.logger.info("Freeing up memory before checkpointing")
+        self._free_memory()
+        self._dist_context.logger.info("Waiting for world before saving checkpoint")
+        self._dist_context.wait_world()
+        self._dist_context.logger.info(f"Saving checkpoint {next_checkpoint_id}")
+
+        save_from = {"state": state}
+        dcp.save(
+            state_dict=save_from,
+            checkpoint_id=next_checkpoint_id
+        )
+
+        self._purge_old_checkpoints()
+        self._free_memory()
+
+        self._dist_context.logger.info("Waiting for world after saving checkpoint")
+        self._dist_context.wait_world()
+        self._dist_context.logger.info("Checkpoint successfully saved across the world")
+
+    def checkpoint_if_needed(self, state: Stateful):
+        """
+        Checks if a checkpoint is due based on the configuration and saves if necessary.
+
+        This checks the stepper to see if the current step matches the configured
+        saving period (or if it is the final step).
+
+        Args:
+            state: The Stateful object to save.
+        """
+
+        if self._stepper.should_do_action(self._config.period_steps, enable_on_last_step_if_periodic=True):
+            self._checkpoint(state)
+
+    def _last_checkpoint_id(self) -> Path | None:
+        checkpoints = self._get_sorted_checkpoint_dirs()
+        if len(checkpoints) == 0:
+            return None
+        return checkpoints[-1]
+
+    def _load(self, state: Stateful):
+        last_checkpoint = self._last_checkpoint_id()
+
+        if last_checkpoint is None:
+            self._dist_context.logger.info("Starting job from scratch")
+            return
+
+        self._dist_context.logger.info("Waiting for world before loading checkpoint")
+        self._dist_context.wait_world()
+        self._dist_context.logger.info(f"Loading checkpoint {last_checkpoint}")
+
+        load_into = {
+            "state": state
+        }
+        dcp.load(
+            state_dict=load_into,
+            checkpoint_id=last_checkpoint
+        )
+        self._free_memory()
+
+        self._dist_context.logger.info("Waiting for world after loading checkpoint")
+        self._dist_context.wait_world()
+        self._dist_context.logger.info("Checkpoint successfully loaded across the world")
+
+    def load_last_checkpoint(self, state: Stateful):
+        """
+        Attempts to load the most recent checkpoint available in the save directory.
+
+        If no checkpoint is found, the state remains unchanged (starting from scratch).
+
+        Args:
+            state: The stateful object to which loaded parameters will be applied.
+        """
+
+        self._load(state)